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/litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql new file mode 100644 index 00000000000..81d32345794 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql @@ -0,0 +1,2 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS "idx_background_interaction_settlement_settled_at" +ON "LiteLLM_BackgroundInteractionSettlement" ("settled_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d40bac380aa..83a0e6f9e05 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1981,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement { outcome String? @@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at") + @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } model LiteLLM_Lens { 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 ff901ede512..df52808643d 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -207,7 +207,7 @@ _async_failure_callback: List[Union[str, Callable[..., object], "CustomLogger"]] [] ) # internal variable - async custom callbacks are routed here. pre_call_rules: List[Callable] = [] -post_call_rules: List[Callable] = [] +post_call_rules: List[Callable[..., object]] = [] turn_off_message_logging: Optional[bool] = False standard_logging_payload_excluded_fields: Optional[List[str]] = ( None # Fields to exclude from StandardLoggingPayload before callbacks receive it @@ -393,6 +393,7 @@ key_alias_pattern: str | None = None enable_gemini_default_thinking_level_low: bool = ( False # opt-in: force thinkingLevel low/minimal for Gemini 3 thinking param mapping ) +gemini_include_thoughts_default: bool = True #################### logging: bool = True enable_loadbalancing_on_batch_endpoints: Optional[bool] = None @@ -1176,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 @@ -1184,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, @@ -1415,6 +1418,7 @@ from .exceptions import ( InvalidRequestError, BadRequestError, ImageFetchError, + MalformedRequestFieldError as MalformedRequestFieldError, VectorStoreSearchError, NotFoundError, PermissionDeniedError, 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/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.py b/litellm/caching/caching.py index 655131e3c69..5f5b8210096 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -25,7 +25,7 @@ from litellm._internal_context import current_service_target, service_target from litellm._logging import verbose_logger from litellm.constants import CACHED_STREAMING_CHUNK_DELAY from litellm.integrations.otel.runtime import phase_span -from litellm.litellm_core_utils.model_param_helper import ModelParamHelper +from litellm.litellm_core_utils.model_param_helper import ModelParamHelper, peek_reasoning_summary_aliases from litellm.types.caching import * from litellm.types.utils import EmbeddingResponse, is_litellm_owned_kwarg @@ -376,7 +376,7 @@ class Cache: ) return "".join(f"{field}: {value}" for field, value in scope_values if value is not None) - def get_cache_key(self, **kwargs) -> str: + def get_cache_key(self, **kwargs: object) -> str: """ Get the cache key for the given arguments. @@ -410,13 +410,16 @@ class Cache: if litellm.enable_caching_on_provider_specific_optional_params is True: # feature flagged for now if kwargs[param] is None: continue # ignore None params - param_value = kwargs[param] - cache_key += f"{param}: {param_value}" + cache_key += f"{param}: {kwargs[param]}" if is_semantic_cache: cache_key += self._get_semantic_cache_tenant_scope(kwargs) - hashed_cache_key = Cache._get_hashed_cache_key(cache_key) + reasoning_summary: Final = peek_reasoning_summary_aliases(kwargs) + reasoning_summary_material: Final = ( + "" if reasoning_summary is None else f"reasoning_summary: {reasoning_summary}" + ) + hashed_cache_key = Cache._get_hashed_cache_key(cache_key + reasoning_summary_material) hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs) verbose_logger.debug( "\nCreated cache key: %s (source material length: %d)", 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 9f29baac562..f448fc351c3 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -57,6 +57,7 @@ RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset( "model_group_retry_policy", "model_group_alias", "enable_weighted_failover", + "enable_mid_stream_fallback_continuation", "enable_tag_filtering", "tag_routing_prefix", "optional_pre_call_checks", @@ -146,6 +147,7 @@ HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0 # Streams exceeding this duration are terminated with a Timeout error. # None (default) = no limit. Set env var to a number of seconds to enable globally. LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = get_env_float_or_none("LITELLM_MAX_STREAMING_DURATION_SECONDS") +MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = 1.0 # Maximum number of base64 characters to keep in logging payloads. # Data URIs exceeding this are replaced with a size placeholder. @@ -389,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" @@ -1629,6 +1631,7 @@ INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response" SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted" +CALLER_SESSION_ID_METADATA_KEY: Final = "litellm_caller_session_id" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( "Truncation is a DB storage safeguard. " @@ -1705,6 +1708,7 @@ AZURE_SPEECH_FAST_TRANSCRIPTION_PATH: Final = "/speechtotext/transcriptions:tran AZURE_SPEECH_STT_DOMAIN: Final = "stt.speech.microsoft.com" AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN: Final = "api.cognitive.microsoft.com" AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER: Final = "Ocp-Apim-Subscription-Key" +AWS_ERROR_TYPE_HEADER: Final = "x-amzn-errortype" AZURE_SPEECH_SHORT_AUDIO_MODEL: Final = "short-audio" AZURE_SPEECH_BATCH_MODEL: Final = "batch-transcription" AZURE_SPEECH_FAST_TRANSCRIPTION_MODEL: Final = "fast-transcription" @@ -1855,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)) @@ -1894,6 +1899,7 @@ _background_interaction_cost_polling_env: Final = os.getenv( "BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true" ).lower() BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_polling_env == "true" +BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS: Final = 300.0 PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605) PROXY_BATCH_WRITE_AT: Final = get_env_int("PROXY_BATCH_WRITE_AT", 10) # in seconds, increased from 10 PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30) 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..d9ee6bdc270 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -86,6 +86,9 @@ from litellm.llms.openai.cost_calculation import ( from litellm.llms.openai.cost_calculation import ( video_output_cost_per_second, ) +from litellm.llms.openai.transcriptions.gpt_transformation import ( + OpenAIGPTAudioTranscriptionConfig, +) from litellm.llms.perplexity.cost_calculator import ( cost_per_token as perplexity_cost_per_token, ) @@ -325,6 +328,15 @@ def _transcription_usage_has_token_details( return (prompt_tokens_val > 0) or (completion_tokens_val > 0) +def _bills_provider_transcription_seconds(model: str | None, custom_llm_provider: str | None) -> bool: + if model is None or custom_llm_provider != LlmProviders.OPENAI.value: + return False + transcription_config: Final = ProviderConfigManager.get_provider_audio_transcription_config( + model=model.removeprefix(f"{LlmProviders.OPENAI.value}/"), provider=LlmProviders.OPENAI + ) + return isinstance(transcription_config, OpenAIGPTAudioTranscriptionConfig) + + OCRPricingField = Literal["ocr_cost_per_page", "ocr_cost_per_credit", "annotation_cost_per_page"] @@ -1739,6 +1751,12 @@ def completion_cost( audio_transcription_file_duration = _hidden.get( "audio_transcription_duration", getattr(completion_response, "duration", 0.0), + ) or ( + completion_response.usage.seconds + if isinstance(completion_response, TranscriptionResponse) + and isinstance(completion_response.usage, TranscriptionUsageDurationObject) + and _bills_provider_transcription_seconds(requested_model, custom_llm_provider) + else 0.0 ) elif call_type in _RERANK_CALL_TYPES: if completion_response is not None and isinstance(completion_response, RerankResponse): @@ -2063,13 +2081,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..040d80e1ac0 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -458,7 +458,7 @@ def prepare_call( ) from error try: - resolved_api_base, resolved_api_key = provider_config.resolve_connection( + resolved_api_base, static_api_key = provider_config.resolve_connection( dynamic_api_base or api_base, dynamic_api_key or api_key ) except ValueError as error: @@ -469,6 +469,7 @@ def prepare_call( model=model, llm_provider=provider, ) + resolved_api_key: Final = provider_config.resolve_credential(static_api_key, resolved_api_base, kwargs) if resolved_api_key is None and provider_config.api_key_required: raise litellm.AuthenticationError( message=f"Missing API key for Decisions provider '{provider}'", diff --git a/litellm/decisions/main.py b/litellm/decisions/main.py index 553fe4b884a..7d2ee03da0e 100644 --- a/litellm/decisions/main.py +++ b/litellm/decisions/main.py @@ -1,5 +1,7 @@ +import asyncio import warnings from collections.abc import Mapping +from functools import partial from typing import Final import httpx @@ -104,19 +106,22 @@ async def adecisions( videos: DecisionsVideoClips | None = None, **kwargs: object, ) -> OpenAIDecisionResponse | DecisionsResponse: - call: Final = _prepare_decisions_call( - model=model, - decision_input=input, - questions=questions, - safety_identifier=safety_identifier, - audio=audio, - videos=videos, - api_key=api_key, - api_base=api_base, - timeout=timeout, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - kwargs=kwargs, + call: Final = await asyncio.to_thread( + partial( + _prepare_decisions_call, + model=model, + decision_input=input, + questions=questions, + safety_identifier=safety_identifier, + audio=audio, + videos=videos, + api_key=api_key, + api_base=api_base, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + kwargs=kwargs, + ) ) return await asend(call) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 31388b00554..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, @@ -1217,6 +1221,7 @@ class MidStreamFallbackError(ServiceUnavailableError): num_retries: int | None = None, generated_content: str = "", is_pre_first_chunk: bool = False, + emitted_disqualifying_content: bool = False, ): original_status: Final = getattr(original_exception, "status_code", None) self.status_code = int(original_status) if original_status is not None else 503 @@ -1229,6 +1234,7 @@ class MidStreamFallbackError(ServiceUnavailableError): self.num_retries = num_retries self.generated_content = generated_content self.is_pre_first_chunk = is_pre_first_chunk + self.emitted_disqualifying_content = emitted_disqualifying_content # Create a response if one wasn't provided if response is None: 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/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 2a2a2a0326d..93dd3d0f5ee 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -695,9 +695,7 @@ class CustomGuardrail(CustomLogger): supported_event_hooks: list[GuardrailEventHooks], ) -> None: allowed_hooks: Final = frozenset(supported_event_hooks) | ( - frozenset((GuardrailEventHooks.logging_only,)) - if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks - else frozenset() + frozenset((GuardrailEventHooks.logging_only,)) if self.accepts_logging_only_mode() else frozenset() ) def _validate_event_hook_list_is_in_supported_event_hooks( @@ -869,8 +867,13 @@ class CustomGuardrail(CustomLogger): return True return False - def uses_apply_guardrail_interface(self) -> bool: - return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail + @classmethod + def uses_apply_guardrail_interface(cls) -> bool: + return cls.apply_guardrail is not CustomGuardrail.apply_guardrail + + @classmethod + def accepts_logging_only_mode(cls) -> bool: + return cls.uses_apply_guardrail_interface() and not cls.use_native_lifecycle_hooks @classmethod def supports_logging_only_scope(cls) -> bool: @@ -1029,7 +1032,7 @@ class CustomGuardrail(CustomLogger): ) return kwargs, result litellm_params: Final = kwargs.get("litellm_params") or {} - scratch_metadata: Final = { + scratch_metadata: Final[dict[str, object]] = { # mutable-ok: apply_guardrail records its verdict here key: value for key, value in (litellm_params.get("metadata") or {}).items() if key != "standard_logging_guardrail_information" @@ -1083,11 +1086,11 @@ class CustomGuardrail(CustomLogger): async def _scan_logged_call( self, - kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract + kwargs: dict[str, object], # mutable-ok: CustomLogger.async_logging_hook contract response: object | None, translation: "BaseTranslation", output_translation: "BaseTranslation", - scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata + scratch_metadata: dict[str, object], # mutable-ok: apply_guardrail records its verdict into request metadata ) -> None: scratch_fields: Final = self._copy_scratch_request_fields(kwargs) scratch_input, scratch_tools = scratch_fields or (None, None) diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index 85b91821f6c..adaacd615a8 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -219,7 +219,7 @@ class GitLabTemplateManager: elif value.replace(".", "").isdigit(): try: result[key] = float(value) - except Exception: + except ValueError: result[key] = value else: result[key] = value.strip("\"'") 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/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index 21d438d9621..275eb17787e 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -32,19 +32,24 @@ task to settle. A ``BackgroundSettlementStore`` makes the pending settlement durable across processes: the create registers the request context that billing needs (never provider credentials), the settlement is claimed exactly once through the store, and a delete on any replica rebuilds the -billing context from the store when the poll task is not local. Rows left -unclaimed by a process that died are resumed at startup. The default store is -in-memory, which keeps the SDK and single-process behavior unchanged; the -proxy installs a database-backed one. +billing context from the store when the poll task is not local. Every row +still unsettled when a process starts is resumed there, including one whose +settler died between its claim and its outcome: the database-backed store +treats a claim as a lease that its settler renews until it records the +outcome, so such a row is claimable again once the lease runs out. The +default store is in-memory, which keeps the SDK and single-process behavior +unchanged; the proxy installs a database-backed one. """ import asyncio +import itertools from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass, field from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias +from openai import APIStatusError from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError from pydantic_core import PydanticSerializationError, to_jsonable_python @@ -54,10 +59,12 @@ from litellm.constants import ( BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS, BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS, BACKGROUND_INTERACTION_COST_POLLING_ENABLED, + BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS, ) from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.types.interactions import InteractionsAPIResponse from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.router import CredentialLiteLLMParams from litellm.types.utils import CustomPricingLiteLLMParams if TYPE_CHECKING: @@ -101,6 +108,7 @@ class PendingBackgroundInteraction: custom_llm_provider: str create_context: BackgroundInteractionCreateContext created_at: datetime + claim_expires_at: datetime | None = None class BackgroundSettlementStore(Protocol): @@ -112,9 +120,11 @@ class BackgroundSettlementStore(Protocol): async def claim(self, interaction_id: str) -> bool: ... + async def renew_claim(self, interaction_id: str) -> None: ... + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: ... - async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: ... + async def unsettled(self) -> Sequence[PendingBackgroundInteraction]: ... @dataclass(frozen=True, slots=True) @@ -145,10 +155,13 @@ class InMemoryBackgroundSettlementStore: self._rows[interaction_id] = None return True + async def renew_claim(self, interaction_id: str) -> None: + return None + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: return None - async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: + async def unsettled(self) -> Sequence[PendingBackgroundInteraction]: return tuple(row for row in self._rows.values() if row is not None) @@ -169,11 +182,15 @@ class BackgroundInteractionPollContext: interaction_id: str custom_llm_provider: str logging_obj: "LiteLLMLoggingObj" - api_key: str | None = None - api_base: str | None = None + credentials: CredentialLiteLLMParams = field(default_factory=CredentialLiteLLMParams) + deployment_id: str | None = None + model_group: str | None = None initial_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS max_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS timeout_seconds: float = BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS + start_delay_seconds: float = 0.0 + outcome_write_timeout_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS / 2 + claim_renewal_interval_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS / 3 store: BackgroundSettlementStore = field(default_factory=InMemoryBackgroundSettlementStore) resumed: bool = False @@ -184,15 +201,41 @@ FetchInteraction: TypeAlias = Callable[[BackgroundInteractionPollContext], Await async def fetch_background_interaction(context: BackgroundInteractionPollContext) -> InteractionsAPIResponse: from litellm.interactions import aget + credentials: Final = context.credentials return await aget( interaction_id=context.interaction_id, custom_llm_provider=context.custom_llm_provider, - api_key=context.api_key, - api_base=context.api_base, + api_key=credentials.api_key, + api_base=credentials.api_base, + vertex_project=credentials.vertex_project, + vertex_location=credentials.vertex_location, + vertex_credentials=credentials.vertex_credentials, **{"no-log": True}, ) +_FETCH_CREDENTIAL_FIELDS: Final = ("api_key", "api_base", "vertex_project", "vertex_location", "vertex_credentials") + + +def _fetch_credentials(request_kwargs: Mapping[str, object]) -> CredentialLiteLLMParams: + try: + return CredentialLiteLLMParams.model_validate( + {field_name: request_kwargs.get(field_name) for field_name in _FETCH_CREDENTIAL_FIELDS} + ) + except ValidationError: + verbose_logger.debug( + "Fetching a background interaction with environment credentials: the request's are malformed" + ) + return CredentialLiteLLMParams() + + +_PERMANENT_FETCH_STATUS_CODES: Final = frozenset({400, 401, 403, 404}) + + +def _is_permanent_fetch_error(error: Exception) -> bool: + return isinstance(error, APIStatusError) and error.status_code in _PERMANENT_FETCH_STATUS_CODES + + def _poll_intervals(initial: float, maximum: float, timeout: float) -> Iterator[float]: elapsed = 0.0 # rebind-ok: the schedule accumulates the time it has already yielded interval = initial # rebind-ok: the schedule doubles the interval up to the cap @@ -231,6 +274,7 @@ _UNCARRIED_METADATA_KEY: Final = "user_api_key_auth" _JSON_VALUE: Final = TypeAdapter(JsonValue) _STRING: Final = TypeAdapter(str) _OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OBJECT_DICT: Final = TypeAdapter(dict[str, object]) def _carries(key: str) -> bool: @@ -256,10 +300,36 @@ def _as_datetime(start_time: datetime | float) -> datetime: return start_time if isinstance(start_time, datetime) else datetime.fromtimestamp(start_time, tz=timezone.utc) +def _spend_metadata(logging_obj: "LiteLLMLoggingObj") -> Mapping[str, object]: + """ + The metadata the creating worker's own spend row is built from, with the + request tags computed once, header-derived ones included, because a + rebuilt logging object has no request headers to derive them from. + """ + litellm_params: Final = _OBJECT_DICT.validate_python(logging_obj.model_call_details.get("litellm_params") or {}) + return { + **_OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details)), + **_request_tags(litellm_params), + } + + +def _request_tags(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + try: + return { + "tags": StandardLoggingPayloadSetup.get_request_tags( + litellm_params=_OBJECT_DICT.validate_python(litellm_params), + proxy_server_request=_OBJECT_DICT.validate_python(litellm_params.get("proxy_server_request") or {}), + ) + } + except AttributeError: + verbose_logger.debug("Keeping a background interaction's request tags as sent: they are not a list") + return {} + + def _create_context(logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str) -> BackgroundInteractionCreateContext: - metadata: Final = _OBJECT_MAPPING.validate_python( - get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) - ) + metadata: Final = _spend_metadata(logging_obj) litellm_params: Final = _OBJECT_MAPPING.validate_python(logging_obj.litellm_params) model: Final = logging_obj.model_call_details.get("model") return BackgroundInteractionCreateContext( @@ -313,6 +383,16 @@ async def _settled_elsewhere(context: BackgroundInteractionPollContext) -> bool: return False +async def _settled_elsewhere_before_the_window_ends( + context: BackgroundInteractionPollContext, intervals: Iterator[float] +) -> bool: + for interval in intervals: + await asyncio.sleep(interval) + if await _settled_elsewhere(context): + return True + return False + + async def _claim(context: BackgroundInteractionPollContext) -> bool | None: """ Exactly-once gate between every settler of one interaction, on every @@ -326,30 +406,118 @@ async def _claim(context: BackgroundInteractionPollContext) -> bool | None: return None -async def _record(context: BackgroundInteractionPollContext, outcome: SettlementOutcome) -> SettlementOutcome: +async def _wrote_outcome(context: BackgroundInteractionPollContext, outcome: SettlementOutcome) -> bool: try: await context.store.record_outcome(context.interaction_id, outcome) - except Exception: # noqa: BLE001 # the outcome is an audit trail; the claim already made the settlement exclusive - verbose_logger.exception("Could not record the settlement of background interaction %s", context.interaction_id) + except Exception as e: # noqa: BLE001 # a failed write is retried while the claim still shuts out takeovers + verbose_logger.warning( + "Could not record the settlement of background interaction %s: %s", context.interaction_id, e + ) + return False + return True + + +async def _recorded(context: BackgroundInteractionPollContext, outcome: SettlementOutcome) -> bool: + """ + Retries a failed write for half the claim lease, while no other settler + can take the row over, so a database blip does not leave the row to a + takeover or hold the budget reservation that only a recorded release frees. + """ + retries: Final = _poll_intervals( + initial=context.initial_interval_seconds, + maximum=context.max_interval_seconds, + timeout=context.outcome_write_timeout_seconds, + ) + for delay in itertools.chain((0.0,), retries): + await asyncio.sleep(delay) + if await _wrote_outcome(context, outcome): + return True + verbose_logger.error( + "Could not record the settlement of background interaction %s within %ss; the next process to take " + "over its lapsed claim settles it", + context.interaction_id, + context.outcome_write_timeout_seconds, + ) + return False + + +async def _bill(context: BackgroundInteractionPollContext, response: InteractionsAPIResponse) -> SettlementOutcome: + await _bill_settled_interaction(logging_obj=context.logging_obj, response=response) + await _recorded(context, "billed") + return "billed" + + +async def _release( + context: BackgroundInteractionPollContext, outcome: Literal["released", "unsettled"] +) -> SettlementOutcome: + """ + A release writes no spend row, so a takeover cannot tell whether the + settler before it released. The outcome is recorded first and the + reservation released only once it is, which leaves a takeover with + nothing to release twice. + """ + if await _recorded(context, outcome): + await _release_open_budget_reservation(logging_obj=context.logging_obj) return outcome +async def _renew_claim_until_cancelled(context: BackgroundInteractionPollContext) -> None: + while True: + await asyncio.sleep(context.claim_renewal_interval_seconds) + try: + await context.store.renew_claim(context.interaction_id) + except Exception as e: # noqa: BLE001 # a missed renewal is retried next interval, well inside the lease + verbose_logger.warning( + "Could not renew the settlement claim of background interaction %s: %s", context.interaction_id, e + ) + + +async def _settle( + context: BackgroundInteractionPollContext, response: InteractionsAPIResponse | None +) -> SettlementOutcome: + """ + Billing runs every success callback in turn, and a slow one can outlast + the claim lease, so the claim is renewed until the outcome is recorded. + A settler that stops renewing has died, and only then can another take + the row over. + """ + renewal: Final = asyncio.create_task(_renew_claim_until_cancelled(context)) + try: + if response is not None and response.status in _TERMINAL_STATUSES and response.usage is not None: + return await _bill(context, response) + return await _release(context, "released") + finally: + renewal.cancel() + + async def poll_and_log_background_interaction_cost( context: BackgroundInteractionPollContext, fetch_interaction: FetchInteraction = fetch_background_interaction, ) -> SettlementOutcome | None: last_response: InteractionsAPIResponse | None = None # rebind-ok: the give-up path settles from the last poll - for interval in _poll_intervals( + await asyncio.sleep(context.start_delay_seconds) + intervals: Final = _poll_intervals( initial=context.initial_interval_seconds, maximum=context.max_interval_seconds, timeout=context.timeout_seconds, - ): + ) + for interval in intervals: await asyncio.sleep(interval) if await _settled_elsewhere(context): return None try: response = await fetch_interaction(context) except Exception as e: # noqa: BLE001 # any fetch error must not kill the billing poll loop + if context.resumed and _is_permanent_fetch_error(e): + verbose_logger.warning( + "Stopped fetching resumed background interaction %s: the provider answered %s, so this " + "process cannot read it and leaves it to the replica that created it until the poll window ends", + context.interaction_id, + e, + ) + if await _settled_elsewhere_before_the_window_ends(context, intervals): + return None + break verbose_logger.debug( "Background interaction cost poll for %s failed, will retry: %s", context.interaction_id, @@ -363,11 +531,11 @@ async def poll_and_log_background_interaction_cost( continue if not claimed: return None - return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=response)) + return await _settle(context, response) if not await _claim(context): return None if last_response is not None and last_response.status in _TERMINAL_STATUSES: - return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=last_response)) + return await _settle(context, last_response) if last_response is not None and last_response.status not in _POLLABLE_STATUSES: verbose_logger.error( "Gave up cost polling for background interaction %s after %ss: its last status %r is in neither " @@ -379,20 +547,10 @@ async def poll_and_log_background_interaction_cost( ) else: verbose_logger.warning( - "Gave up cost polling for background interaction %s after %ss; its usage will not be tracked", + "Gave up cost polling for background interaction %s; its usage will not be tracked", context.interaction_id, - context.timeout_seconds, ) - await _release_open_budget_reservation(logging_obj=context.logging_obj) - return await _record(context, "unsettled") - - -async def _settle_terminal(logging_obj: "LiteLLMLoggingObj", response: InteractionsAPIResponse) -> SettlementOutcome: - if response.status in _TERMINAL_STATUSES and response.usage is not None: - await _bill_settled_interaction(logging_obj=logging_obj, response=response) - return "billed" - await _release_open_budget_reservation(logging_obj=logging_obj) - return "released" + return await _release(context, "unsettled") async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> None: @@ -405,7 +563,9 @@ async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> settlement must release the reservation here or the spend counters stay pinned at the estimated cost. """ - metadata: Final = get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) + metadata: Final = _OBJECT_MAPPING.validate_python( + get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) + ) budget_reservation: Final = metadata.get("user_api_key_budget_reservation") if not isinstance(budget_reservation, dict): return @@ -518,13 +678,18 @@ class _UnverifiedRegistrationStore: return False return await self.local.claim(interaction_id) + async def renew_claim(self, interaction_id: str) -> None: + if await self.local.is_claimed(interaction_id): + return + await self.durable.renew_claim(interaction_id) + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: if await self.local.is_claimed(interaction_id): return await self.durable.record_outcome(interaction_id, outcome) - async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: - return await self.durable.unclaimed() + async def unsettled(self) -> Sequence[PendingBackgroundInteraction]: + return await self.durable.unsettled() async def _registered_store( @@ -559,8 +724,6 @@ async def maybe_schedule_background_interaction_cost_polling( logging_obj: Final = create_kwargs.get("litellm_logging_obj") if not isinstance(logging_obj, Logging): return None - api_key: Final = create_kwargs.get("api_key") - api_base: Final = create_kwargs.get("api_base") pending: Final = PendingBackgroundInteraction( interaction_id=response.id, custom_llm_provider=custom_llm_provider, @@ -571,8 +734,7 @@ async def maybe_schedule_background_interaction_cost_polling( interaction_id=response.id, custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, - api_key=api_key if isinstance(api_key, str) else None, - api_base=api_base if isinstance(api_base, str) else None, + credentials=_fetch_credentials(create_kwargs), store=await _registered_store(store or _STORE.store, pending), ) return _track_poll(context, fetch_interaction) @@ -603,10 +765,7 @@ async def _settle_before_delete( ) -> SettlementOutcome | None: if not await _claim(context): return None - if response is None: - await _release_open_budget_reservation(logging_obj=context.logging_obj) - return await _record(context, "released") - return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=response)) + return await _settle(context, response) async def maybe_settle_background_interaction_before_delete( @@ -622,14 +781,11 @@ async def maybe_settle_background_interaction_before_delete( pending: Final = await _pending(settlement_store, interaction_id) if pending is None: return None - api_key: Final = delete_kwargs.get("api_key") - api_base: Final = delete_kwargs.get("api_base") context: Final = BackgroundInteractionPollContext( interaction_id=interaction_id, custom_llm_provider=pending.custom_llm_provider, logging_obj=_rebuild_logging_obj(pending.create_context), - api_key=api_key if isinstance(api_key, str) else None, - api_base=api_base if isinstance(api_base, str) else None, + credentials=_fetch_credentials(delete_kwargs), store=settlement_store, ) try: @@ -644,9 +800,9 @@ async def maybe_settle_background_interaction_before_delete( return await _settle_before_delete(context, response) -async def _unclaimed(store: BackgroundSettlementStore) -> Sequence[PendingBackgroundInteraction]: +async def _unsettled(store: BackgroundSettlementStore) -> Sequence[PendingBackgroundInteraction]: try: - return await store.unclaimed() + return await store.unsettled() except Exception: # noqa: BLE001 # an unreadable store at startup leaves its rows for the next boot verbose_logger.exception("Could not list the unsettled background interactions") return () @@ -662,17 +818,36 @@ class PollSchedule: DEFAULT_POLL_SCHEDULE: Final = PollSchedule() +def _deployment_id(create_context: BackgroundInteractionCreateContext) -> str | None: + model_info: Final = create_context.metadata.get("model_info") + deployment_id: Final = model_info.get("id") if isinstance(model_info, dict) else None + return deployment_id if isinstance(deployment_id, str) else None + + +def _model_group(create_context: BackgroundInteractionCreateContext) -> str | None: + model_group: Final = create_context.metadata.get("model_group") + return model_group if isinstance(model_group, str) else None + + +def _seconds_until(moment: datetime | None) -> float: + if moment is None: + return 0.0 + return max((moment - datetime.now(timezone.utc)).total_seconds(), 0.0) + + def _resumed_context( row: PendingBackgroundInteraction, store: BackgroundSettlementStore, schedule: PollSchedule ) -> BackgroundInteractionPollContext: - age_seconds: Final = (datetime.now(timezone.utc) - row.created_at).total_seconds() return BackgroundInteractionPollContext( interaction_id=row.interaction_id, custom_llm_provider=row.custom_llm_provider, logging_obj=_rebuild_logging_obj(row.create_context), + deployment_id=_deployment_id(row.create_context), + model_group=_model_group(row.create_context), initial_interval_seconds=schedule.initial_interval_seconds, max_interval_seconds=schedule.max_interval_seconds, - timeout_seconds=max(schedule.timeout_seconds - age_seconds, schedule.initial_interval_seconds), + timeout_seconds=schedule.timeout_seconds, + start_delay_seconds=_seconds_until(row.claim_expires_at), store=store, resumed=True, ) @@ -684,13 +859,15 @@ async def resume_unsettled_background_interactions( schedule: PollSchedule = DEFAULT_POLL_SCHEDULE, ) -> tuple["asyncio.Task[SettlementOutcome | None]", ...]: """ - Pick up every settlement no process has claimed, which is what a replica - that died mid-poll leaves behind. Each resumed poll keeps the remaining - share of the original timeout and gets at least one fetch, so a completed - interaction is still billed however late the resume comes. + Pick up every settlement no process has finished, which is what a replica + that died mid-poll or mid-settlement leaves behind. A row whose claim is + still within its lease waits until the lease runs out, then polls like any + other. Each resumed poll gets the full poll window however old its row is, + so a fetch that fails right after a late restart is retried rather than + ending the row as unsettled. """ return tuple( _track_poll(_resumed_context(row, store, schedule), fetch_interaction) - for row in await _unclaimed(store) + for row in await _unsettled(store) if row.interaction_id not in _ACTIVE_POLLS ) 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 e766a2be9c3..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,54 @@ 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 + self.model_call_details["continued_partial_cost"] = response_cost + self.model_call_details["continued_response_id"] = response_id + + def _add_continued_partial_spend(self, response: object) -> None: + partial_usage: Final = self.model_call_details.get("continued_partial_usage") + partial_cost: Final = self.model_call_details.get("continued_partial_cost") + own_usage: Final = getattr(response, "usage", None) + own_cost: Final = self.model_call_details.get("response_cost") + if not (isinstance(partial_usage, Usage) and isinstance(partial_cost, float) and isinstance(own_usage, Usage)): + return + from litellm.cost_calculator import BaseTokenUsageProcessor + + self.model_call_details["combined_usage_object"] = BaseTokenUsageProcessor.combine_usage_objects( + usage_objects=[own_usage, partial_usage] + ) + if not isinstance(own_cost, (int, float)): + return + self.model_call_details["response_cost"] = own_cost + partial_cost + if self.cost_breakdown is not None: + self.cost_breakdown = self.cost_breakdown | CostBreakdown( + total_cost=own_cost + partial_cost, + additional_costs={ + **self.cost_breakdown.get("additional_costs", {}), + "mid_stream_partial_cost": partial_cost, + }, + ) + def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None: """Bill a fully streamed response on the failure log when a post-call hook rejects it.""" usage: Final = getattr(assembled, "usage", None) @@ -2590,6 +2643,15 @@ class Logging(LiteLLMLoggingBaseClass): return response return response.model_copy(update={"model": model}) + def _with_continued_response_id( + self, + response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse, + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse: + response_id: Final = self.model_call_details.get("continued_response_id") + if not isinstance(response_id, str) or getattr(response, "id", None) == response_id: + return response + return response.model_copy(update={"id": response_id}) + def _success_handler_helper_fn( self, result: object = None, @@ -2623,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) @@ -2908,8 +2972,11 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = self.response_cost_calculator( result=complete_streaming_response ) + self._add_continued_partial_spend(complete_streaming_response) self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) - logged_streaming_response: Final = self._with_client_facing_stream_model(complete_streaming_response) + logged_streaming_response: Final = self._with_continued_response_id( + self._with_client_facing_stream_model(complete_streaming_response) + ) self.model_call_details["complete_streaming_response"] = logged_streaming_response ## STANDARDIZED LOGGING PAYLOAD self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( @@ -3445,12 +3512,15 @@ class Logging(LiteLLMLoggingBaseClass): self.model, ) self.model_call_details["response_cost"] = None + self._add_continued_partial_spend(complete_streaming_response) self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) await self._prepare_baseline_cache_estimate(complete_streaming_response) - logged_streaming_response: Final = self._with_client_facing_stream_model(complete_streaming_response) + logged_streaming_response: Final = self._with_continued_response_id( + self._with_client_facing_stream_model(complete_streaming_response) + ) self.model_call_details["async_complete_streaming_response"] = logged_streaming_response ## STANDARDIZED LOGGING PAYLOAD @@ -4282,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 @@ -4382,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) @@ -4406,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): @@ -6634,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, @@ -6796,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 6c7760d7da8..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 @@ -2,12 +2,13 @@ Helper utilities for tracking the cost of built-in tools. """ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Final, Literal, cast from pydantic import ValidationError import litellm +from litellm._logging import verbose_logger from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS from litellm.litellm_core_utils.llm_cost_calc.utils import ( get_web_search_requests_from_usage, @@ -30,20 +31,36 @@ from litellm.types.utils import ( ) +def _output_item_field(output_item: object, field: str) -> object: + if isinstance(output_item, dict): + fields: Final[Mapping[str, object]] = cast(Mapping[str, object], output_item) # cast-ok: narrowed by isinstance + return fields.get(field) + return getattr(output_item, field, None) + + def _output_item_type(output_item: object) -> str | None: - item_type: Final = output_item.get("type") if isinstance(output_item, dict) else getattr(output_item, "type", None) + item_type: Final = _output_item_field(output_item, "type") return item_type if isinstance(item_type, str) else None -def _reported_web_search_requests(response_object: ResponsesAPIResponse) -> int | None: - tool_usage: Final = getattr(response_object, "tool_usage", None) - if tool_usage is None: +def _responses_tool_usage(response_object: ResponsesAPIResponse) -> ResponsesToolUsage | None: + top_level: Final = getattr(response_object, "tool_usage", None) + raw: Final = ( + top_level if top_level is not None else getattr(getattr(response_object, "usage", None), "tool_usage", None) + ) + if raw is None: return None try: - web_search: Final = ResponsesToolUsage.model_validate(tool_usage).web_search + return ResponsesToolUsage.model_validate(raw) except ValidationError: return None - return None if web_search is None else web_search.num_requests + + +def _reported_web_search_requests(response_object: ResponsesAPIResponse) -> int | None: + tool_usage: Final = _responses_tool_usage(response_object) + if tool_usage is None: + return None + return None if tool_usage.web_search is None else tool_usage.web_search.num_requests def _usage_reports_server_side_web_search_calls(usage: Usage) -> bool: @@ -87,31 +104,48 @@ class StandardBuiltInToolCostTracking: usage=usage, ) + image_generation_cost: Final = StandardBuiltInToolCostTracking._handle_image_generation_cost( + response_object=response_object, + custom_llm_provider=custom_llm_provider, + ) + # Handle web search if StandardBuiltInToolCostTracking.response_object_includes_web_search_call( response_object=response_object, usage=usage ): - return google_maps_grounding_cost + StandardBuiltInToolCostTracking._handle_web_search_cost( - model=model, - custom_llm_provider=custom_llm_provider, - usage=usage, - standard_built_in_tools_params=standard_built_in_tools_params, - response_object=response_object, + return ( + google_maps_grounding_cost + + image_generation_cost + + StandardBuiltInToolCostTracking._handle_web_search_cost( + model=model, + custom_llm_provider=custom_llm_provider, + usage=usage, + standard_built_in_tools_params=standard_built_in_tools_params, + response_object=response_object, + ) ) # Handle file search if StandardBuiltInToolCostTracking.response_object_includes_file_search_call(response_object=response_object): - return google_maps_grounding_cost + StandardBuiltInToolCostTracking._handle_file_search_cost( + return ( + google_maps_grounding_cost + + image_generation_cost + + StandardBuiltInToolCostTracking._handle_file_search_cost( + model=model, + custom_llm_provider=custom_llm_provider, + standard_built_in_tools_params=standard_built_in_tools_params, + ) + ) + + # Handle Azure assistant features + return ( + google_maps_grounding_cost + + image_generation_cost + + StandardBuiltInToolCostTracking._handle_azure_assistant_costs( model=model, custom_llm_provider=custom_llm_provider, standard_built_in_tools_params=standard_built_in_tools_params, ) - - # Handle Azure assistant features - return google_maps_grounding_cost + StandardBuiltInToolCostTracking._handle_azure_assistant_costs( - model=model, - custom_llm_provider=custom_llm_provider, - standard_built_in_tools_params=standard_built_in_tools_params, ) @staticmethod @@ -210,6 +244,87 @@ class StandardBuiltInToolCostTracking: ) return max(count, 1) + @staticmethod + def _image_generation_tool_model(response_object: ResponsesAPIResponse) -> str | None: + tools: Final[Sequence[object]] = cast(Sequence[object], response_object.tools or ()) # cast-ok: dict tool union + for tool in tools: + if _output_item_field(tool, "type") != "image_generation": + continue + if isinstance(model := _output_item_field(tool, "model"), str) and model: + return model + return None + + @staticmethod + def _image_generation_call_cost(output_item: object, tool_model: str, custom_llm_provider: str | None) -> float: + from litellm.cost_calculator import ( + default_image_cost_calculator, # pyright: ignore[reportUnknownVariableType] # optional_params param is untyped + ) + + quality: Final = _output_item_field(output_item, "quality") + size: Final = _output_item_field(output_item, "size") + try: + return default_image_cost_calculator( + model=tool_model, + custom_llm_provider=custom_llm_provider, + quality=quality if isinstance(quality, str) and quality != "auto" else None, + n=1, + size=size if isinstance(size, str) and size != "auto" else None, + ) + except Exception as e: # noqa: BLE001 # pricing helpers raise bare Exception for unmapped models; bill 0.0 + verbose_logger.debug("Could not price Responses API image_generation_call item: %s", e) + return 0.0 + + @staticmethod + def _image_generation_token_cost( + response_object: ResponsesAPIResponse, tool_model: str, custom_llm_provider: str | None + ) -> float | None: + tool_usage: Final = _responses_tool_usage(response_object) + if tool_usage is None or tool_usage.image_gen is None or tool_usage.image_gen.total_tokens <= 0: + return None + try: + model_info: Final = litellm.get_model_info(model=tool_model, custom_llm_provider=custom_llm_provider) + except Exception as e: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models; fall back + verbose_logger.debug("Could not resolve pricing for image tool model %s: %s", tool_model, e) + return None + image_gen: Final = tool_usage.image_gen + input_details: Final = image_gen.input_tokens_details + output_details: Final = image_gen.output_tokens_details + token_cost: Final = ( + (input_details.text_tokens if input_details else 0) * (model_info.get("input_cost_per_token") or 0) + + (input_details.image_tokens if input_details else 0) * (model_info.get("input_cost_per_image_token") or 0) + + (output_details.image_tokens if output_details else 0) + * (model_info.get("output_cost_per_image_token") or 0) + + (output_details.text_tokens if output_details else 0) * (model_info.get("output_cost_per_token") or 0) + ) + return token_cost if token_cost > 0 else None + + @staticmethod + def _handle_image_generation_cost(response_object: object, custom_llm_provider: str | None) -> float: + if not isinstance(response_object, ResponsesAPIResponse): + return 0.0 + output: Final[Sequence[object]] = cast(Sequence[object], response_object.output) # cast-ok: isinstance-narrowed + completed_items: Final = tuple( + output_item + for output_item in output + if _output_item_type(output_item) == "image_generation_call" + and _output_item_field(output_item, "status") == "completed" + ) + if not completed_items: + return 0.0 + tool_model: Final = StandardBuiltInToolCostTracking._image_generation_tool_model(response_object) + if tool_model is None: + verbose_logger.debug("Responses API image_generation_call has no tools[].model to price, billing 0.0") + return 0.0 + token_cost: Final = StandardBuiltInToolCostTracking._image_generation_token_cost( + response_object, tool_model, custom_llm_provider + ) + if token_cost is not None: + return token_cost + return sum( + StandardBuiltInToolCostTracking._image_generation_call_cost(output_item, tool_model, custom_llm_provider) + for output_item in completed_items + ) + @staticmethod def _handle_file_search_cost( model: str, @@ -445,9 +560,12 @@ class StandardBuiltInToolCostTracking: return False elif isinstance(response_object, ResponsesAPIResponse): # response api explicitly includes web_search_call in the output - return StandardBuiltInToolCostTracking.response_includes_output_type( + if StandardBuiltInToolCostTracking.response_includes_output_type( response_object=response_object, output_type="web_search_call" - ) + ): + return True + reported: Final = _reported_web_search_requests(response_object) + return reported is not None and reported > 0 elif usage is not None: if get_web_search_requests_from_usage(usage) is not None or ( hasattr(usage, "prompt_tokens_details") @@ -740,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" @@ -758,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/llm_response_utils/get_headers.py b/litellm/litellm_core_utils/llm_response_utils/get_headers.py index b86cc954b0e..3c08f16fa80 100644 --- a/litellm/litellm_core_utils/llm_response_utils/get_headers.py +++ b/litellm/litellm_core_utils/llm_response_utils/get_headers.py @@ -1,8 +1,11 @@ from collections.abc import Mapping from typing import Final +import httpx from pydantic import TypeAdapter, ValidationError +from litellm.constants import PROVIDER_REQUEST_ID_HEADERS + _HEADERS: Final = TypeAdapter(Mapping[str, str]) _REQUEST_ID_HEADERS: Final = ("request-id", "x-request-id", "llm_provider-request-id", "llm_provider-x-request-id") @@ -16,6 +19,16 @@ def get_provider_request_id(headers: object) -> str | None: return next((normalized[name] for name in _REQUEST_ID_HEADERS if normalized.get(name)), None) +def provider_request_id_headers(headers: Mapping[str, object] | None) -> httpx.Headers | None: + """The provider request id headers in `headers` whose values are ASCII strings, or None when there are none.""" + request_id_headers: Final = tuple( + (name, value) + for name, value in (headers or {}).items() + if name.lower() in PROVIDER_REQUEST_ID_HEADERS and isinstance(value, str) and value.isascii() + ) + return httpx.Headers(request_id_headers) if request_id_headers else None + + def get_response_headers(_response_headers: Mapping[str, str] | None = None) -> dict: """ diff --git a/litellm/litellm_core_utils/model_param_helper.py b/litellm/litellm_core_utils/model_param_helper.py index 2dc52980875..c9c7a0e89ac 100644 --- a/litellm/litellm_core_utils/model_param_helper.py +++ b/litellm/litellm_core_utils/model_param_helper.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from functools import lru_cache from typing import ClassVar, Final @@ -22,6 +23,24 @@ from litellm.types.llms.anthropic import AnthropicMessagesRequest from litellm.types.rerank import RerankRequest +def peek_reasoning_summary_aliases(optional_params: Mapping[str, object]) -> object | None: + """Read AI-SDK-style reasoning summary from optional_params or nested extra_body. + + Uses key membership (not ``or`` chains) so falsy values like ``""`` are not skipped. + """ + if "reasoningSummary" in optional_params: + return optional_params["reasoningSummary"] + if "reasoning_summary" in optional_params: + return optional_params["reasoning_summary"] + extra_body: Final = optional_params.get("extra_body") + if isinstance(extra_body, dict): + if "reasoningSummary" in extra_body: + return extra_body["reasoningSummary"] + if "reasoning_summary" in extra_body: + return extra_body["reasoning_summary"] + return None + + class ModelParamHelper: # Cached at class level — deterministic set built from static OpenAI type annotations relevant_logging_args: ClassVar[frozenset[str]] = frozenset() diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 81fc865a5bc..044123d8729 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1894,6 +1894,11 @@ def strip_litellm_internal_message_fields(message: AllMessageValues) -> AllMessa return _message_without_fields(message, LITELLM_INTERNAL_MESSAGE_FIELDS) +def strip_prefill_prefix_field(message: AllMessageValues) -> AllMessageValues: + """Drop the ``prefix`` flag Mistral and DeepSeek read off a prefill, for endpoints that reject unknown keys.""" + return _message_without_fields(message, frozenset({"prefix"})) + + def filter_value_from_dict(dictionary: dict, key: str, depth: int = 0) -> Any: """ Filters a value from a dictionary diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0c18ce38b13..87e76fe39f7 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -9,10 +9,11 @@ import xml.etree.ElementTree as ET from collections.abc import Container, Iterable, Iterator, Mapping, Sequence from enum import Enum from types import MappingProxyType -from typing import Any, Final, TypeAlias, TypedDict, cast, overload +from typing import Annotated, Any, Final, TypeAlias, TypedDict, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment -from pydantic import BaseModel, JsonValue, TypeAdapter +from pydantic import BaseModel, BeforeValidator, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator +from pydantic_core import ErrorDetails, PydanticCustomError import litellm import litellm.types @@ -2582,6 +2583,48 @@ def _with_native_redacted_data( _AnthropicMessageList: TypeAlias = list[AllAnthropicPassThroughMessageValues] +def _client_tool_result(message: AllMessageValues, replaced_server_tool_answers: frozenset[str]) -> bool: + return ( + message_field(message, "role") in ("tool", "function") + and message_field(message, "tool_call_id") not in replaced_server_tool_answers + ) + + +def _tool_result_later_in_turn( + messages: Sequence[AllMessageValues], replaced_server_tool_answers: frozenset[str] +) -> tuple[bool, ...]: + from_the_end: Final = itertools.accumulate( + reversed(messages), + lambda later, message: ( + message_field(message, "role") in ("user", "tool", "function") + and (later or _client_tool_result(message, replaced_server_tool_answers)) + ), + initial=False, + ) + return tuple(from_the_end)[::-1][1:] + + +def _empty_text_placeholder(message: AllMessageValues) -> AnthropicMessagesTextParam: + placeholder: Final = AnthropicMessagesTextParam(type="text", text=_EMPTY_TEXT_PLACEHOLDER) + add_cache_control_to_content(anthropic_content_element=placeholder, original_content_element=message) + return placeholder + + +def _cache_breakpoint_placeholder( + message: AllMessageValues, tool_result_follows: bool +) -> tuple[AnthropicMessagesTextParam, ...]: + if tool_result_follows or not message_field(message, "cache_control"): + return () + return (_empty_text_placeholder(message),) + + +def _empty_user_turn_placeholder(span: Sequence[AllMessageValues]) -> AnthropicMessagesTextParam | None: + empty_user_message: Final = next( + (message for message in reversed(span) if message["role"] == "user" and not message.get("content")), None + ) + return None if empty_user_message is None else _empty_text_placeholder(empty_user_message) + + def anthropic_messages_pt( messages: list[AllMessageValues], model: str, @@ -2636,6 +2679,7 @@ def anthropic_messages_pt( is_vertex_ai = llm_provider.startswith("vertex_ai") if llm_provider else False force_base64 = is_bedrock_invoke or is_vertex_ai + tool_result_later_in_turn: Final = _tool_result_later_in_turn(messages, replaced_server_tool_answers) msg_i = 0 while msg_i < len(messages): user_content: list[AnthropicMessagesUserMessageValues] = [] @@ -2648,8 +2692,12 @@ def anthropic_messages_pt( ChatCompletionToolMessage | ChatCompletionUserMessage | ChatCompletionFunctionMessage ) = messages[msg_i] if user_message_types_block["role"] == "user": - if isinstance(user_message_types_block["content"], list): - for m in user_message_types_block["content"]: + if not (user_block_content := user_message_types_block.get("content")): + user_content.extend( + _cache_breakpoint_placeholder(user_message_types_block, tool_result_later_in_turn[msg_i]) + ) + elif isinstance(user_block_content, list): + for m in user_block_content: if m.get("type", "") == "image_url": m = cast(ChatCompletionImageObject, m) format = m["image_url"].get("format") if isinstance(m["image_url"], dict) else None @@ -2720,10 +2768,10 @@ def anthropic_messages_pt( _file_content_element, ) ) - elif isinstance(user_message_types_block["content"], str): + elif isinstance(user_block_content, str): _anthropic_content_text_element: AnthropicMessagesTextParam = { "type": "text", - "text": user_message_types_block["content"], + "text": user_block_content, } _content_element = add_cache_control_to_content( anthropic_content_element=_anthropic_content_text_element, @@ -2747,6 +2795,8 @@ def anthropic_messages_pt( if user_content: new_messages.append({"role": "user", "content": user_content}) + elif (placeholder := _empty_user_turn_placeholder(messages[init_msg_i:msg_i])) is not None: + new_messages.append({"role": "user", "content": [placeholder]}) # Track unique tool IDs in this merge block to avoid duplication unique_tool_ids: set[str] = set() @@ -4253,6 +4303,24 @@ def _with_text_when_document_only(message: BedrockMessageBlock) -> BedrockMessag return BedrockMessageBlock(role="user", content=[*blocks[:cut], placeholder, *blocks[cut:]]) +def _rstrip_text_block(block: BedrockContentBlock) -> BedrockContentBlock: + return BedrockContentBlock(text=block["text"].rstrip()) if "text" in block else block + + +def _assistant_turn_rstripped(message: BedrockMessageBlock) -> BedrockMessageBlock: + """Claude on Converse rejects a final assistant turn that ends in whitespace, as the Anthropic path strips it.""" + if message["role"] != "assistant": + return message + blocks: Final = message["content"] + last_text_index: Final = max((index for index, block in enumerate(blocks) if "text" in block), default=None) + return BedrockMessageBlock( + role="assistant", + content=[ + _rstrip_text_block(block) if index == last_text_index else block for index, block in enumerate(blocks) + ], + ) + + def _ensure_document_messages_have_text( contents: list[BedrockMessageBlock], ) -> list[BedrockMessageBlock]: @@ -4844,7 +4912,12 @@ class BedrockConverseMessagesProcessor: llm_provider=llm_provider, ) - return _ensure_document_messages_have_text(_rename_duplicate_bedrock_document_names(contents)) + converse_messages: Final = _ensure_document_messages_have_text( + _rename_duplicate_bedrock_document_names(contents) + ) + if not (to_claude and converse_messages): + return converse_messages + return [*converse_messages[:-1], _assistant_turn_rstripped(converse_messages[-1])] @staticmethod def translate_thinking_blocks_to_reasoning_content_blocks( @@ -5253,7 +5326,10 @@ def _bedrock_converse_messages_pt( llm_provider=llm_provider, ) - return _ensure_document_messages_have_text(_rename_duplicate_bedrock_document_names(contents)) + converse_messages: Final = _ensure_document_messages_have_text(_rename_duplicate_bedrock_document_names(contents)) + if not (to_claude and converse_messages): + return converse_messages + return [*converse_messages[:-1], _assistant_turn_rstripped(converse_messages[-1])] def make_valid_bedrock_tool_name(input_tool_name: str) -> str: @@ -5516,7 +5592,7 @@ def function_call_prompt( ] -def response_schema_prompt(model: str, response_schema: dict) -> str: +def response_schema_prompt(model: str, response_schema: dict, llm_provider: str) -> str: """ Decides if a user-defined custom prompt or default needs to be used @@ -5532,11 +5608,8 @@ def response_schema_prompt(model: str, response_schema: dict) -> str: custom_prompt_details = litellm.custom_prompt_dict["response_schema_prompt"] if custom_prompt_details is not None: - return custom_prompt( - role_dict=custom_prompt_details["roles"], - initial_prompt_value=custom_prompt_details["initial_prompt_value"], - final_prompt_value=custom_prompt_details["final_prompt_value"], - messages=response_schema_as_message, + return render_custom_prompt_template( + template=custom_prompt_details, messages=response_schema_as_message, model=model, llm_provider=llm_provider ) else: return default_response_schema_prompt(response_schema=response_schema) @@ -5556,8 +5629,95 @@ def default_response_schema_prompt(response_schema: dict) -> str: # Custom prompt template +def _mapping_or_none(value: object) -> object: + if value is None or issubclass(type(value), Mapping): + return value + verbose_logger.warning( + "Ignoring a custom prompt template roles value of type %s because it is not a mapping, " + "so those messages render without role markers", + type(value).__name__, + ) + return None + + +class CustomPromptRole(BaseModel): + model_config = ConfigDict(frozen=True) + + pre_message: str | None = None + post_message: str | None = None + + +class CustomPromptRoles(BaseModel): + model_config = ConfigDict(frozen=True, title="custom prompt template") + + roles: Mapping[str, Annotated[CustomPromptRole | None, BeforeValidator(_mapping_or_none)]] + + +class CustomPromptTemplate(BaseModel): + model_config = ConfigDict(frozen=True, title="custom prompt template") + + roles: Annotated[Mapping[object, object] | None, BeforeValidator(_mapping_or_none)] = None + initial_prompt_value: str | None = None + final_prompt_value: str | None = None + + @model_validator(mode="after") + def _sets_a_rendered_field(self) -> "CustomPromptTemplate": + if self.roles is None and self.initial_prompt_value is None and self.final_prompt_value is None: + raise PydanticCustomError( + "empty_custom_prompt_template", + "it sets none of roles, initial_prompt_value, or final_prompt_value", + ) + return self + + +def render_custom_prompt_template( + template: Mapping[str, object], + messages: Sequence[Mapping[str, object]], + model: str, + llm_provider: str, +) -> str: + try: + return _render_custom_prompt_template(template=template, messages=messages) + except ValidationError as error: + raise litellm.APIConnectionError( + message=f"The custom prompt template for {model} is invalid: {_describe_template_errors(error)}", + model=model, + llm_provider=llm_provider, + ) from error + + +def _describe_template_errors(error: ValidationError) -> str: + return "; ".join( + _describe_template_error(detail) for detail in error.errors(include_url=False, include_input=False) + ) + + +def _describe_template_error(detail: ErrorDetails) -> str: + location: Final = ".".join(str(part) for part in detail["loc"]) + return f"{location}: {detail['msg']}" if location else detail["msg"] + + +def _render_custom_prompt_template(template: Mapping[str, object], messages: Sequence[Mapping[str, object]]) -> str: + validated: Final = CustomPromptTemplate.model_validate(template) + message_roles: Final = frozenset(message.get("role") for message in messages) + template_roles: Final = validated.roles or {} + used_roles: Final = CustomPromptRoles.model_validate( + {"roles": {role: entry for role, entry in template_roles.items() if role in message_roles}} + ) + return custom_prompt( + role_dict={ + role: {"pre_message": entry.pre_message or "", "post_message": entry.post_message or ""} + for role, entry in used_roles.roles.items() + if entry is not None + }, + initial_prompt_value=validated.initial_prompt_value or "", + final_prompt_value=validated.final_prompt_value or "", + messages=list(messages), + ) + + def custom_prompt( - role_dict: dict, + role_dict: Mapping[str, Mapping[str, str]], messages: list, initial_prompt_value: str = "", final_prompt_value: str = "", diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py index 4d1e77126a9..fd0096f9f4d 100644 --- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py +++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py @@ -414,10 +414,11 @@ def _assistant_renders(message: object, answered: frozenset[str]) -> bool: def _renders(message: object, answered: frozenset[str]) -> bool: """Whether ``anthropic_messages_pt`` puts a block on the wire for this message. - A tool message always becomes a ``tool_result`` and a user message with string - content always becomes a text block (empty text gets a placeholder). A user list - renders only through parts of a type the converter emits; ``None``, an empty list, - and a list of other parts vanish. Assistant messages follow ``_assistant_renders``. + A tool message always becomes a ``tool_result``, and a user message with string or + missing content (``None``, an empty list) always puts its turn on the wire, an empty + turn getting a placeholder. A user list renders only through parts of a type the + converter emits, so a list of only other parts vanishes. Assistant messages follow + ``_assistant_renders``. """ role: Final = message_field(message, "role") if role in _TOOL_ROLES: @@ -425,8 +426,10 @@ def _renders(message: object, answered: frozenset[str]) -> bool: if role == "assistant": return _assistant_renders(message, answered) content: Final = message_field(message, "content") - return isinstance(content, str) or any( - message_field(part, "type") in _RENDERED_PART_TYPES for part in parts_of(content) + return ( + isinstance(content, str) + or not content + or any(message_field(part, "type") in _RENDERED_PART_TYPES for part in parts_of(content)) ) 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 46ad1b70abe..78fb82c98e7 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -242,6 +242,7 @@ class CustomStreamWrapper: self.completion_stream = completion_stream self.sent_first_chunk = False self.sent_last_chunk = False + self._failure_logging: asyncio.Task[None] | None = None self._closed = False self._stream_created_time: float = time.time() @@ -269,6 +270,7 @@ class CustomStreamWrapper: self.holding_chunk = "" self.complete_response = "" self.response_uptil_now = "" + self._emitted_disqualifying_content = False _model_info: Final[dict] = litellm_params.model_info or {} _api_base: Final = get_api_base( @@ -1126,31 +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: - 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, @@ -1865,9 +1860,7 @@ class CustomStreamWrapper: if response.choices: choice = response.choices[0] if isinstance(choice, StreamingChoices): - self.response_uptil_now += choice.delta.get("content", "") or "" - else: - self.response_uptil_now += "" + self._accumulate_streamed_delta(choice.delta) self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) # HANDLE STREAM OPTIONS self.chunks.append(response) @@ -2084,9 +2077,7 @@ class CustomStreamWrapper: if processed_chunk.choices: choice = processed_chunk.choices[0] if isinstance(choice, StreamingChoices): - self.response_uptil_now += choice.delta.get("content", "") or "" - else: - self.response_uptil_now += "" + self._accumulate_streamed_delta(choice.delta) self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) # Add mcp_list_tools to first chunk if present if not self.sent_first_chunk and processed_chunk.choices: @@ -2156,9 +2147,7 @@ class CustomStreamWrapper: choice = processed_chunk.choices[0] if isinstance(choice, StreamingChoices): - self.response_uptil_now += choice.delta.get("content", "") or "" - else: - self.response_uptil_now += "" + self._accumulate_streamed_delta(choice.delta) self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) # RETURN RESULT self.chunks.append(processed_chunk) @@ -2172,7 +2161,7 @@ class CustomStreamWrapper: if self.logging_obj is not None: self._record_partial_usage_for_failure() ## LOGGING - asyncio.create_task( + self._failure_logging = asyncio.create_task( self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) ) await self.aclose() @@ -2305,11 +2294,17 @@ class CustomStreamWrapper: if self.logging_obj is not None: self._record_partial_usage_for_failure() ## LOGGING - asyncio.create_task( + self._failure_logging = asyncio.create_task( self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) ) self._handle_stream_fallback_error(e) + async def wait_for_failure_logging(self, timeout: float) -> None: + """Let this stream's failure log read its own deployment before a fallback reuses the logging object.""" + if self._failure_logging is None: + return + await asyncio.wait((self._failure_logging,), timeout=timeout) + def _record_partial_usage_for_failure(self) -> None: """ A stream that breaks mid-flight still billed the provider for the chunks @@ -2346,6 +2341,35 @@ class CustomStreamWrapper: recover_error, ) + # Delta fields a text-only prefill continuation cannot carry, so a stream + # that emitted any of them is not eligible for mid-stream continuation. + _CONTINUATION_DISQUALIFYING_DELTA_FIELDS: Final = ( + "tool_calls", + "function_call", + "reasoning_content", + "thinking_blocks", + "reasoning_items", + "audio", + "images", + "annotations", + ) + + @classmethod + def _delta_disqualifies_continuation(cls, delta: object) -> bool: + get: Final = getattr(delta, "get", None) + if not callable(get): + return False + return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS) + + def _accumulate_streamed_delta(self, delta: object) -> None: + # Shared by the sync, async, and non-aiohttp iteration sites so answer + # text and the disqualifying-content latch stay in step across all three. + get: Final = getattr(delta, "get", None) + content: Final = get("content") if callable(get) else None + self.response_uptil_now += content if isinstance(content, str) else "" + if not self._emitted_disqualifying_content and self._delta_disqualifies_continuation(delta): + self._emitted_disqualifying_content = True + def _handle_stream_fallback_error(self, e: Exception) -> "NoReturn": """ Common error handling for both __next__ and __anext__. @@ -2377,6 +2401,7 @@ class CustomStreamWrapper: original_exception=e.original_exception, generated_content=self.response_uptil_now, is_pre_first_chunk=not self.sent_first_chunk, + emitted_disqualifying_content=self._emitted_disqualifying_content or e.emitted_disqualifying_content, ) # Map to OpenAI exception format. Some providers' mappers (e.g. @@ -2437,6 +2462,7 @@ class CustomStreamWrapper: original_exception=mapped_exception, generated_content=self.response_uptil_now, is_pre_first_chunk=not self.sent_first_chunk, + emitted_disqualifying_content=self._emitted_disqualifying_content, ) @staticmethod 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/completion/transformation.py b/litellm/llms/anthropic/completion/transformation.py index e3ad961ed76..73e94e85f04 100644 --- a/litellm/llms/anthropic/completion/transformation.py +++ b/litellm/llms/anthropic/completion/transformation.py @@ -14,8 +14,8 @@ import httpx import litellm from litellm.constants import DEFAULT_MAX_TOKENS from litellm.litellm_core_utils.prompt_templates.factory import ( - custom_prompt, prompt_factory, + render_custom_prompt_template, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import ( @@ -244,11 +244,8 @@ class AnthropicTextConfig(BaseConfig): if model in custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details: Final = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="anthropic" ) else: prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="anthropic") diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index d78de0037a0..ecf050d185e 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -6,11 +6,12 @@ from typing import Any, Final from pydantic import JsonValue +import litellm 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() @@ -36,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. @@ -61,14 +63,24 @@ class AnthropicTokenCounter(BaseTokenCounter): try: auth_header: Final = await AnthropicModelInfo.aget_auth_header( - api_key=litellm_params.get("api_key"), + api_key=litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key, api_base=api_base, litellm_params=litellm_params, allow_workload_identity=True, ) if auth_header is None: - verbose_logger.warning("No Anthropic credential found for token counting") - return None + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="anthropic_api", + error=True, + error_message=( + f"No Anthropic API key is set for model {request_model or model_to_use}, so tokens can't be " + "counted. Set api_key on the deployment or ANTHROPIC_API_KEY in the proxy environment." + ), + status_code=401, + ) result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request( model=model_to_use, diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 4c0613777c6..7b8c2550838 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -1,4 +1,5 @@ from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence +from dataclasses import dataclass from typing import ( TYPE_CHECKING, Final, @@ -62,6 +63,12 @@ class _CompletionKwargs(TypedDict, total=False, extra_items=object): custom_llm_provider: str +@dataclass(frozen=True, slots=True) +class _ChatTargetReasoning: + effort: str + summary: str | None = None + + def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool: """Return True when any message carries a ``compaction`` content block.""" for msg in messages: @@ -132,7 +139,7 @@ def request_safeguards_evaluator( ) -async def _prepare_context_managed_request( +async def prepare_context_managed_request( *, model: str, messages: _AnthropicMessages, @@ -145,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 ( @@ -186,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: @@ -259,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( @@ -284,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 ( @@ -310,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: @@ -338,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 @@ -441,15 +457,33 @@ class LiteLLMMessagesToCompletionTransformationHandler: completion_kwargs["reasoning_effort"] = updated_reasoning_effort @staticmethod - def _plain_effort_for_chat_target( + def _requested_reasoning_summary(thinking: Mapping[str, object]) -> str | None: + if thinking.get("type") == "disabled": + return None + summary: Final = thinking.get("summary") + if isinstance(summary, str): + return summary + display: Final = thinking.get("display") + return "auto" if display is not None and display != "omitted" else None + + @staticmethod + def _chat_target_reasoning( completion_kwargs: _CompletionKwargs, *, thinking: Mapping[str, object] | None, - ) -> str | None: - reasoning_effort: Final = completion_kwargs.get("reasoning_effort") - if not thinking or not isinstance(reasoning_effort, dict) or "summary" not in reasoning_effort: + ) -> _ChatTargetReasoning | None: + if not thinking: return None - effort: Final = reasoning_effort.get("effort") + reasoning_effort: Final = completion_kwargs.get("reasoning_effort") + wraps_summary: Final = isinstance(reasoning_effort, dict) and "summary" in reasoning_effort + requested_summary: Final = LiteLLMMessagesToCompletionTransformationHandler._requested_reasoning_summary( + thinking + ) + if not wraps_summary and requested_summary is None: + return None + effort: Final[object] = ( + reasoning_effort.get("effort") if isinstance(reasoning_effort, dict) else reasoning_effort + ) model: Final = completion_kwargs.get("model") if not isinstance(effort, str) or not isinstance(model, str) or not model: return None @@ -467,7 +501,12 @@ class LiteLLMMessagesToCompletionTransformationHandler: return None if resolved_provider == "litellm_proxy": return None - from litellm.responses.bridge_check import responses_api_bridge_check + from litellm.responses.bridge_check import chat_target_accepts_reasoning_summary, responses_api_bridge_check + + if requested_summary is not None and chat_target_accepts_reasoning_summary(local_model, resolved_provider): + return _ChatTargetReasoning(effort=effort, summary=requested_summary) + if not wraps_summary or not isinstance(reasoning_effort, dict): + return None web_search_options: Final = completion_kwargs.get("web_search_options") tools: Final = completion_kwargs.get("tools") @@ -481,7 +520,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: reasoning_effort=reasoning_effort, api_base=resolved_api_base, ) - return None if model_info.get("mode") == "responses" else effort + return None if model_info.get("mode") == "responses" else _ChatTargetReasoning(effort=effort) @staticmethod def _normalize_reasoning_effort( @@ -647,12 +686,14 @@ class LiteLLMMessagesToCompletionTransformationHandler: thinking=thinking, ) - plain_effort: Final = LiteLLMMessagesToCompletionTransformationHandler._plain_effort_for_chat_target( + chat_reasoning: Final = LiteLLMMessagesToCompletionTransformationHandler._chat_target_reasoning( completion_kwargs, thinking=thinking, ) - if plain_effort is not None: - completion_kwargs["reasoning_effort"] = plain_effort + if chat_reasoning is not None: + completion_kwargs["reasoning_effort"] = chat_reasoning.effort + if chat_reasoning is not None and chat_reasoning.summary is not None: + completion_kwargs["reasoning_summary"] = chat_reasoning.summary return completion_kwargs, tool_name_mapping @@ -677,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) @@ -684,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, @@ -696,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 @@ -809,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 @@ -834,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, @@ -846,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 98cd16c6403..3aad0be1913 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -465,7 +465,7 @@ class LiteLLMAnthropicMessagesAdapter: *, custom_llm_provider: str | None = None, preserve_midturn_system: bool = False, - ) -> list: + ) -> list[AllMessageValues]: new_messages: Final[list[AllMessageValues]] = [] replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) leading_count: Final = next( @@ -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/convertible_blocks.py b/litellm/llms/anthropic/pass_through/convertible_blocks.py new file mode 100644 index 00000000000..543a6efd4c4 --- /dev/null +++ b/litellm/llms/anthropic/pass_through/convertible_blocks.py @@ -0,0 +1,106 @@ +from collections.abc import Iterator, Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final + +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages +from litellm.types.utils import TokenCountMessagesFormat + +_CONVERTIBLE_BLOCK_TYPES: Final = MappingProxyType( + { + "user": frozenset({"text", "image", "document", "tool_result"}), + "assistant": frozenset({"text", "tool_use", "thinking", "redacted_thinking"}), + } +) +_SOURCED_BLOCK_TYPES: Final = frozenset({"image", "document"}) +_TOOL_RESULT_BLOCK_TYPES: Final = frozenset({"text", "image", "document"}) + + +_JSON_MESSAGES: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) + + +def _json_messages(messages: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, JsonValue], ...]: + try: + return _JSON_MESSAGES.validate_python(messages) + except ValidationError: + return () + + +def _source_label(block_type: str, source: JsonValue) -> str | None: + if not isinstance(source, dict): + return f"{block_type} without a source" + match source.get("type"): + case "base64" if source.get("data"): + return None + case "url" if source.get("url"): + return None + case "base64" | "url" as source_type: + return f"{block_type} with an empty {source_type} source" + case str() as source_type: + return f"{block_type} with a {source_type} source" + case _: + return f"{block_type} with an untyped source" + + +def _block_label(block: JsonValue, convertible: frozenset[str]) -> str | None: + if not isinstance(block, dict): + return None + block_type: Final = block.get("type") + if not isinstance(block_type, str): + return "a block without a type" + if block_type not in convertible: + return block_type + return _source_label(block_type, block.get("source")) if block_type in _SOURCED_BLOCK_TYPES else None + + +def _tool_result_labels(block: JsonValue) -> Iterator[str]: + if not isinstance(block, dict) or block.get("type") != "tool_result": + return iter(()) + inner: Final = block.get("content") + if not isinstance(inner, list): + return iter(()) + return ( + f"{label} inside a tool_result" + for item in inner + if (label := _block_label(item, _TOOL_RESULT_BLOCK_TYPES)) is not None + ) + + +def _content_labels(block: JsonValue, convertible: frozenset[str]) -> Iterator[str]: + label: Final = _block_label(block, convertible) + return chain(() if label is None else (label,), _tool_result_labels(block)) + + +def _message_labels(message: Mapping[str, JsonValue]) -> Iterator[str]: + role: Final = message.get("role") + convertible: Final = _CONVERTIBLE_BLOCK_TYPES.get(role) if isinstance(role, str) else None + content: Final = message.get("content") + if convertible is None or not isinstance(content, list): + return iter(()) + return chain.from_iterable(_content_labels(block, convertible) for block in content) + + +def unconvertible_anthropic_blocks(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + normalized: Final = _json_messages(normalize_native_messages(messages)) + return tuple(sorted(set(chain.from_iterable(_message_labels(message) for message in normalized)))) + + +def unconvertible_blocks_error( + messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat +) -> str | None: + match messages_format: + case "openai": + return None + case "anthropic": + labels: Final = unconvertible_anthropic_blocks(messages or ()) + return ( + "count_tokens can't convert these content blocks for this model, so the count would leave them out: " + + ", ".join(labels) + if labels + else None + ) + case _: + assert_never(messages_format) 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/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 0e0b68df449..53734a77f3c 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -30,7 +30,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): self, model: str, messages: list[dict[str, JsonValue]], - api_key: str, + api_key: str | None, api_base: str, litellm_params: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, @@ -80,6 +80,15 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): api_key=api_key, litellm_params=litellm_params, ) + if not {"x-api-key", "Authorization"} & headers.keys(): + raise AnthropicError( + status_code=401, + message=( + f"No Azure AI API key or Entra ID credential is set for model {model}, so tokens can't be " + "counted. Set api_key or an Entra ID credential on the deployment, or AZURE_AI_API_KEY or " + "AZURE_API_KEY in the proxy environment." + ), + ) # Use LiteLLM's async httpx client async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.AZURE_AI) 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 2707d24ee07..d1f45751d3e 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -7,17 +7,33 @@ from typing import Any, Final from pydantic import JsonValue +import litellm from litellm._logging import verbose_logger 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() +def _missing_api_base_response(request_model: str, model_to_use: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="azure_ai_anthropic_api", + error=True, + error_message=( + f"No Azure AI API base is set for model {request_model or model_to_use}, so tokens can't be counted. " + "Set api_base on the deployment or AZURE_AI_API_BASE or AZURE_API_BASE in the proxy environment." + ), + status_code=500, + ) + + class AzureAIAnthropicTokenCounter(BaseTokenCounter): """Token counter implementation for Azure AI Anthropic provider using the CountTokens API.""" @@ -38,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. @@ -65,18 +82,15 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): if not api_key: api_key = os.getenv("AZURE_AI_API_KEY") - # Get API base from deployment config or environment - api_base = litellm_params.get("api_base") - if not api_base: - api_base = os.getenv("AZURE_AI_API_BASE") - - if not api_key: - verbose_logger.warning("No Azure AI API key found for token counting") - return None + api_base: Final = ( + litellm_params.get("api_base") + or os.getenv("AZURE_AI_API_BASE") + or os.getenv("AZURE_API_BASE") + or litellm.api_base + ) if not api_base: - verbose_logger.warning("No Azure AI API base found for token counting") - return None + return _missing_api_base_response(request_model, model_to_use) try: result: Final = await azure_ai_anthropic_count_tokens_handler.handle_count_tokens_request( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index d7455439151..5946a1f4480 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -24,7 +24,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): def get_required_headers( self, - api_key: str, + api_key: str | None, litellm_params: dict[str, object] | None = None, ) -> dict[str, str]: """ @@ -45,12 +45,12 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): "Content-Type": "application/json", "anthropic-version": "2023-06-01", "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, - "x-api-key": api_key, # Azure AI Anthropic requires this header + **({"x-api-key": api_key} if api_key else {}), } # Also set up Azure auth headers for flexibility litellm_params = litellm_params or {} - if "api_key" not in litellm_params: + if api_key and "api_key" not in litellm_params: litellm_params["api_key"] = api_key litellm_params_obj: Final = GenericLiteLLMParams.model_validate(litellm_params) @@ -58,8 +58,9 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): # Get Azure auth headers (api-key or Authorization) azure_headers = BaseAzureLLM.base_validate_azure_environment(headers={}, litellm_params=litellm_params_obj) - # Merge Azure auth headers headers.update(azure_headers) + if "api-key" in azure_headers: + headers.setdefault("x-api-key", azure_headers["api-key"]) return headers 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/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 2b3c644f266..db4d5e0f820 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -488,6 +488,10 @@ class BaseConfig(ABC): """ return True + @property + def accepts_reasoning_summary_in_reasoning_effort(self) -> bool: + return False + def post_stream_processing(self, stream: object) -> object: """Hook for providers to post-process streaming responses. Default: pass-through.""" return stream diff --git a/litellm/llms/base_llm/containers/transformation.py b/litellm/llms/base_llm/containers/transformation.py index 49487cddb6f..7a04c4d9488 100644 --- a/litellm/llms/base_llm/containers/transformation.py +++ b/litellm/llms/base_llm/containers/transformation.py @@ -81,6 +81,7 @@ class BaseContainerConfig(ABC): self, headers: dict, api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, ) -> dict: return {} diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index 041473cb634..df7b7645840 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -495,6 +495,11 @@ class BaseDecisionsConfig(ABC): def resolve_connection(self, api_base: str | None, api_key: str | None) -> tuple[str | None, str | None]: return self.resolve_api_base(api_base), self.resolve_api_key(api_key) + def resolve_credential( + self, api_key: str | None, api_base: str, litellm_params: Mapping[str, object] + ) -> str | None: + return api_key + @staticmethod def _first_secret(names: tuple[str, ...]) -> str | None: return next((value for value in (get_secret_str(name) for name in names) if value), None) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 39ff2fa0528..55a4b3da0b3 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -4,7 +4,7 @@ import json from collections.abc import Callable, Iterator, Mapping, Sequence from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage @@ -292,6 +292,17 @@ def openai_tool_name(tool: object) -> str | None: return flat_name if isinstance(flat_name, str) else None +_FIELDS: Final = TypeAdapter(Mapping[object, object], config=ConfigDict(strict=True)) +_ITEMS: Final = TypeAdapter(Sequence[object], config=ConfigDict(strict=True)) + + +def list_field(data: object, key: str) -> Sequence[object]: + try: + return _ITEMS.validate_python(_FIELDS.validate_python(data).get(key)) + except ValidationError: + return () + + def anthropic_tool_names(tool: object) -> tuple[str, ...]: """Every name a /v1/messages tool dict can act under: the flat Anthropic ``name`` plus ``function.name`` for OpenAI-format tools the bridge forwards verbatim. Allowlist checks diff --git a/litellm/llms/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/base_llm/videos/async_task_utils.py b/litellm/llms/base_llm/videos/async_task_utils.py new file mode 100644 index 00000000000..d2fb3ed7606 --- /dev/null +++ b/litellm/llms/base_llm/videos/async_task_utils.py @@ -0,0 +1,123 @@ +import base64 +from collections.abc import Mapping +from io import BufferedReader, BytesIO +from itertools import accumulate, repeat +from pathlib import Path, PurePath +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm.images.utils import ImageEditRequestUtils +from litellm.types.videos.utils import decode_video_id_with_provider + +_FIELDS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_FILE_TUPLE_ADAPTER: Final = TypeAdapter(tuple[object, ...]) + + +def _fields(value: object) -> Mapping[str, object] | None: + try: + return _FIELDS_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _logged_litellm_params(logging_obj: object) -> Mapping[str, object]: + return _fields(getattr(logging_obj, "litellm_params", None)) or {} + + +def _id_safe_name(name: object) -> str | None: + return name if isinstance(name, str) and name and ";" not in name else None + + +def _stamped_models(metadata: object) -> tuple[str, str | None] | None: + fields: Final = _fields(metadata) + if fields is None: + return None + requested: Final = _id_safe_name(fields.get("model_group")) + served: Final = _id_safe_name(fields.get("deployment_model_name")) + if requested is None: + return None + return requested, None if served is None or served == requested or "*" in served else served + + +def routed_video_models(logging_obj: object) -> tuple[str, str | None] | None: + """ + A later status or content request checks the key's model access against the model encoded in the id, and + routes by it. One upstream model name can be served by several deployments (regions, accounts), so only the + group that served the task can send its polls back to the deployment that created it, while a key may be + scoped to the alias, routing group or wildcard name it asked for instead. The id carries both names: the one + the router was asked for on this hop, and the group that served it when that differs and is not a wildcard. + The router stamps both into litellm_metadata, and metadata carries a copy only when the caller sent none. + """ + params: Final = _logged_litellm_params(logging_obj) + stamped: Final = (_stamped_models(params.get(key)) for key in ("litellm_metadata", "metadata")) + return next((models for models in stamped if models), None) + + +def polled_video_models(logging_obj: object) -> tuple[str | None, str | None]: + """ + The served group survives into the returned id only when the router sent this poll to it, so an edited id + cannot write any name it likes back. + """ + video_id: Final = _logged_litellm_params(logging_obj).get("video_id") + if not isinstance(video_id, str): + return None, None + decoded: Final = decode_video_id_with_provider(video_id) + served: Final = decoded.get("served_model") + routed: Final = routed_video_models(logging_obj) + return decoded.get("model_id") or None, served if routed is None or served in routed else None + + +def _tuple_payload(file_obj: object) -> object: + if not isinstance(file_obj, tuple): + return file_obj + file_tuple: Final = _FILE_TUPLE_ADAPTER.validate_python(file_obj) + return file_tuple[1] if len(file_tuple) >= 2 else None + + +def _file_content(file_obj: object) -> object: + unwrapped: Final = accumulate(repeat(None), lambda current, _: _tuple_payload(current), initial=file_obj) + return next(item for item in unwrapped if not isinstance(item, tuple)) + + +def input_reference_bytes(file_obj: object) -> bytes: + content: Final = _file_content(file_obj) + if isinstance(content, PurePath): + return Path(content).read_bytes() + if isinstance(content, (BytesIO, BufferedReader)): + current_position: Final = content.tell() + _ = content.seek(0) + data: Final = content.read() + _ = content.seek(current_position) + return data + if isinstance(content, bytes): + return content + if isinstance(content, bytearray): + return bytes(content) + read: Final = getattr(content, "read", None) + if callable(read): + read_data: Final[object] = read() + if isinstance(read_data, bytes): + return read_data + raise ValueError("input_reference must be a URL string, bytes, a path, a file object, or a file tuple") + + +def _reference_object_url(reference: Mapping[str, object]) -> str: + image_url: Final = reference.get("image_url") + if isinstance(image_url, str) and image_url: + return image_url + raise ValueError("input_reference as an object must carry an image_url, a file_id is not supported") + + +def input_reference_url(image: object) -> str: + """A string is a URL the provider fetches itself and is never opened as a local path. OpenAI's object form + carries that URL, or a data URL, in ``image_url``.""" + if isinstance(image, str): + return image + reference: Final = _fields(image) + if reference is not None: + return _reference_object_url(reference) + content: Final = input_reference_bytes(image) + content_type: Final = ImageEditRequestUtils.get_image_content_type(content) + encoded: Final = base64.b64encode(content).decode("utf-8") + return f"data:{content_type};base64,{encoded}" diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 848120150cd..f1426aa9375 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -99,11 +99,16 @@ def _assume_role_params( ) -_SecureTransportBool = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]}) +_SecureTransportBool: Final = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]}) + +_CalledViaLastEquals: Final = TypedDict( + "_CalledViaLastEquals", {"aws:CalledViaLast": ReadOnly[Literal["aws-external-anthropic.amazonaws.com"]]} +) -class _SecureTransportCondition(TypedDict): +class _SessionPolicyCondition(TypedDict): Bool: ReadOnly[_SecureTransportBool] + StringEquals: NotRequired[ReadOnly[_CalledViaLastEquals]] class _SessionPolicyStatement(TypedDict): @@ -111,7 +116,7 @@ class _SessionPolicyStatement(TypedDict): Effect: ReadOnly[Literal["Allow"]] Action: ReadOnly[tuple[str, ...]] Resource: ReadOnly[Literal["*"]] - Condition: ReadOnly[_SecureTransportCondition] + Condition: ReadOnly[_SessionPolicyCondition] class WebIdentitySessionPolicy(TypedDict): @@ -147,11 +152,24 @@ _WEB_IDENTITY_SESSION_POLICY_ACTIONS: Final[Mapping[str, tuple[str, ...]]] = Map "aws-external-anthropic:Get*", "aws-external-anthropic:List*", ), + "ClaudePlatformOutboundIdentityLiteLLM": ( + "sts:GetWebIdentityToken", + "sts:TagGetWebIdentityToken", + ), "BedrockMantleLiteLLM": ("bedrock-mantle:CreateInference",), } ) -_SECURE_TRANSPORT_ONLY: Final = _SecureTransportCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"})) +_SECURE_TRANSPORT_ONLY: Final = _SessionPolicyCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"})) + +_WEB_IDENTITY_SESSION_POLICY_CONDITIONS: Final[Mapping[str, _SessionPolicyCondition]] = MappingProxyType( + { + "ClaudePlatformOutboundIdentityLiteLLM": _SessionPolicyCondition( + Bool=_SecureTransportBool({"aws:SecureTransport": "true"}), + StringEquals=_CalledViaLastEquals({"aws:CalledViaLast": "aws-external-anthropic.amazonaws.com"}), + ), + } +) def build_web_identity_session_policy() -> WebIdentitySessionPolicy: @@ -163,7 +181,7 @@ def build_web_identity_session_policy() -> WebIdentitySessionPolicy: Effect="Allow", Action=actions, Resource="*", - Condition=_SECURE_TRANSPORT_ONLY, + Condition=_WEB_IDENTITY_SESSION_POLICY_CONDITIONS.get(sid, _SECURE_TRANSPORT_ONLY), ) for sid, actions in _WEB_IDENTITY_SESSION_POLICY_ACTIONS.items() ), 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 a207bceb418..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 @@ -13,12 +13,14 @@ from litellm.caching.caching import InMemoryCache from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.litellm_core_utils.prompt_templates.common_utils import native_redacted_thinking_data 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, @@ -31,12 +33,8 @@ from litellm.types.llms.openai import ( ChatCompletionRedactedThinkingBlock, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk, - ChatCompletionToolCallFunctionChunk, - ChatCompletionUsageBlock, ) from litellm.types.utils import ( - ChatCompletionMessageToolCall, - Choices, Delta, ModelResponse, ModelResponseStream, @@ -227,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( @@ -263,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, @@ -323,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( @@ -363,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, @@ -448,6 +451,7 @@ class _EventStreamTally: f"x-amzn-requestid={_response_header(response_headers, 'x-amzn-requestid')!r}, " f"{self.bytes_received} bytes received" ) + request_id_headers: Final = provider_request_id_headers(response_headers) if not self.events: return BedrockError( status_code=502, @@ -455,6 +459,7 @@ class _EventStreamTally: "Bedrock answered the stream with HTTP 200 but its body decoded to no events " f"({detail}, first bytes={self.head!r})" ), + headers=request_id_headers, ) if not self.recognized_events: return BedrockError( @@ -464,10 +469,12 @@ class _EventStreamTally: f"event type (event types={sorted(self.unrecognized_event_types)}, {detail}, " f"first payload={self.unrecognized_head!r})" ), + headers=request_id_headers, ) return BedrockError( status_code=502, message=f"Bedrock stream ended with {undecoded} undecoded bytes after {self.events} events ({detail})", + headers=request_id_headers, ) @@ -837,11 +844,12 @@ class AWSEventStreamDecoder: event_stream_buffer: Final = EventStreamBuffer() tally: Final = _EventStreamTally(get_bedrock_stream_event_statuses()) + request_id_headers: Final = provider_request_id_headers(response_headers) for chunk in iterator: event_stream_buffer.add_data(chunk) tally.add_chunk(chunk) for event in event_stream_buffer: - message = self._decode_event(event, tally) + message = self._decode_event(event, tally, request_id_headers) if message: # sse_event = ServerSentEvent(data=message, event="completion") _data = json.loads(message) @@ -858,11 +866,12 @@ class AWSEventStreamDecoder: event_stream_buffer: Final = EventStreamBuffer() tally: Final = _EventStreamTally(get_bedrock_stream_event_statuses()) + request_id_headers: Final = provider_request_id_headers(response_headers) async for chunk in iterator: event_stream_buffer.add_data(chunk) tally.add_chunk(chunk) for event in event_stream_buffer: - message = self._decode_event(event, tally) + message = self._decode_event(event, tally, request_id_headers) if message: _data = json.loads(message) yield self.chunk_parser(chunk_data=_data) @@ -882,27 +891,34 @@ class AWSEventStreamDecoder: ) return response_stream_shape - def _decode_event(self, event: "EventStreamMessage", tally: _EventStreamTally) -> str | None: + def _decode_event( + self, event: "EventStreamMessage", tally: _EventStreamTally, request_id_headers: httpx.Headers | None + ) -> str | None: response_stream_shape: Final = self._response_stream_shape() response_dict: Final = bedrock_event_stream_response(event) tally.add_event(event, response_dict["headers"]) - return self._parse_message_from_response(response_dict, response_stream_shape) + return self._parse_message_from_response(response_dict, response_stream_shape, request_id_headers) def _parse_message_from_event(self, event: "EventStreamMessage") -> str | None: response_stream_shape: Final = self._response_stream_shape() return self._parse_message_from_response(bedrock_event_stream_response(event), response_stream_shape) def _parse_message_from_response( - self, response_dict: BedrockEventStreamResponseDict, response_stream_shape: "Shape" + self, + response_dict: BedrockEventStreamResponseDict, + response_stream_shape: "Shape", + request_id_headers: httpx.Headers | None = None, ) -> str | None: parsed_response: Final = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: - raise build_bedrock_stream_error(response_dict, response_stream_shape) + raise build_bedrock_stream_error(response_dict, response_stream_shape, headers=request_id_headers) event_type: Final = bedrock_event_stream_header(response_dict["headers"], ":event-type") event_error_status: Final = bedrock_stream_event_error_status(event_type) if event_type is not None and event_error_status is not None: - raise build_bedrock_stream_event_error(event_type, event_error_status, response_dict["body"]) + raise build_bedrock_stream_event_error( + event_type, event_error_status, response_dict["body"], headers=request_id_headers + ) if "chunk" in parsed_response: chunk = parsed_response.get("chunk") if not chunk: @@ -977,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/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 1b3b0c344df..f026e4f9a99 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -14,9 +14,9 @@ from litellm.litellm_core_utils.get_litellm_params import stored_control_options from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.prompt_templates.factory import ( cohere_message_pt, - custom_prompt, deepseek_r1_pt, prompt_factory, + render_custom_prompt_template, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call @@ -601,7 +601,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, provider) return None - def convert_messages_to_prompt(self, model, messages, provider, custom_prompt_dict) -> tuple[str, list | None]: + def convert_messages_to_prompt(self, model: str, messages, provider, custom_prompt_dict) -> tuple[str, list | None]: # handle anthropic prompts and amazon titan prompts prompt = "" chat_history: list | None = None @@ -609,11 +609,8 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): if model in custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details: Final = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), - final_prompt_value=model_prompt_details.get("final_prompt_value", ""), - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="bedrock" ) return prompt, None ## ELSE diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index e79865ef6bf..5862b05a619 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1913,13 +1913,18 @@ def bedrock_event_stream_header(headers: Mapping[str, object], name: str) -> str return value if isinstance(value, str) else None -def build_bedrock_stream_event_error(event_type: str, status_code: int, body: bytes) -> BedrockError: - return BedrockError(status_code=status_code, message=f"{event_type} {body.decode(errors='replace')}") +def build_bedrock_stream_event_error( + event_type: str, status_code: int, body: bytes, headers: httpx.Headers | None = None +) -> BedrockError: + return BedrockError( + status_code=status_code, message=f"{event_type} {body.decode(errors='replace')}", headers=headers + ) def build_bedrock_stream_error( response_dict: BedrockEventStreamResponseDict, response_stream_shape: Shape | None, + headers: httpx.Headers | None = None, ) -> BedrockError: """Build a BedrockError for a non-200 event-stream error event. @@ -1929,12 +1934,14 @@ def build_bedrock_stream_error( """ exception_type: Final = bedrock_event_stream_header(response_dict["headers"], ":exception-type") if exception_type is None: - return BedrockError(status_code=response_dict["status_code"], message=response_dict["body"].decode()) + return BedrockError( + status_code=response_dict["status_code"], message=response_dict["body"].decode(), headers=headers + ) member: Final = _structure_members(response_stream_shape).get(exception_type) modeled_status: Final = None if member is None else _modeled_error_status(member) status_code: Final = response_dict["status_code"] if modeled_status is None else modeled_status - return build_bedrock_stream_event_error(exception_type, status_code, response_dict["body"]) + return build_bedrock_stream_event_error(exception_type, status_code, response_dict["body"], headers=headers) class BedrockEventStreamDecoderBase: 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/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index f8486d3b274..08799c3a0a5 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -13,8 +13,8 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.prompt_templates.factory import ( - custom_prompt, prompt_factory, + render_custom_prompt_template, ) from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -261,11 +261,8 @@ class CodestralTextCompletion: if model in custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details: Final = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="codestral" ) else: prompt = prompt_factory(model=model, messages=messages) 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/dashscope/videos/transformation.py b/litellm/llms/dashscope/videos/transformation.py index 17b685a0e1a..1db91cbb127 100644 --- a/litellm/llms/dashscope/videos/transformation.py +++ b/litellm/llms/dashscope/videos/transformation.py @@ -3,13 +3,10 @@ DashScope (Alibaba Model Studio) async video task API: create, poll and download delete or remix endpoint, so those raise. """ -import base64 import json from collections.abc import Mapping from datetime import datetime -from io import BufferedReader, BytesIO from math import log -from pathlib import Path, PurePath from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final @@ -19,9 +16,13 @@ from pydantic import BaseModel, Field, StrictBool, StrictInt, StrictStr, TypeAda import litellm from litellm.exceptions import UnsupportedParamsError -from litellm.images.utils import ImageEditRequestUtils from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.videos.async_task_utils import ( + input_reference_url, + polled_video_models, + routed_video_models, +) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # house cached-client factory has no public alias and its stub leaves params untyped @@ -32,7 +33,6 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ModelInfo from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject from litellm.types.videos.utils import ( - decode_video_id_with_provider, encode_video_id_with_provider, extract_original_video_id, ) @@ -86,7 +86,6 @@ _EMPTY_PARAMS: Final[dict[str, object]] = {} # mutable-ok: BaseVideoConfig cont _TASK_RESPONSE_ADAPTER: Final = TypeAdapter(_TaskResponse) _PARAMETERS_ADAPTER: Final = TypeAdapter(dict[str, object]) -_FILE_TUPLE_ADAPTER: Final = TypeAdapter(tuple[object, ...]) _STATUS_MAP: Final = MappingProxyType( { @@ -383,39 +382,6 @@ def _is_priced_tier(resolution: object, priced_tiers: tuple[str, ...]) -> bool: return isinstance(resolution, str) and resolution.strip().upper() in priced_tiers -def _read_all_bytes(file_obj: object) -> bytes: - if isinstance(file_obj, tuple): - file_tuple: Final = _FILE_TUPLE_ADAPTER.validate_python(file_obj) - return _read_all_bytes(file_tuple[1] if len(file_tuple) >= 2 else None) - if isinstance(file_obj, PurePath): - return Path(file_obj).read_bytes() - if isinstance(file_obj, (BytesIO, BufferedReader)): - current_position: Final = file_obj.tell() - file_obj.seek(0) - content: Final = file_obj.read() - file_obj.seek(current_position) - return content - if isinstance(file_obj, bytes): - return file_obj - if isinstance(file_obj, bytearray): - return bytes(file_obj) - read: Final = getattr(file_obj, "read", None) - if callable(read): - data: Final = read() - if isinstance(data, bytes): - return data - raise ValueError("input_reference must be a URL string, bytes, a path, a file object, or a file tuple") - - -def _image_url(image: object) -> str: - if isinstance(image, str): - return image - content: Final = _read_all_bytes(image) - content_type: Final = ImageEditRequestUtils.get_image_content_type(content) - encoded: Final = base64.b64encode(content).decode("utf-8") - return f"data:{content_type};base64,{encoded}" - - def _media_reference_type(model: str) -> str | None: """Models without the field predate the media array and take the image on the flat ``img_url`` field.""" info: Final = _dashscope_model_info(model) @@ -558,58 +524,6 @@ def _video_url_from_task(task: _TaskResponse, raw_response: httpx.Response) -> s raise _video_error(raw_response, 502, "Video URL not found in task response. The task may not have succeeded yet.") -def _logged_litellm_params(logging_obj: object) -> Mapping[str, object]: - litellm_params: Final[object] = getattr(logging_obj, "litellm_params", None) - try: - return _PARAMETERS_ADAPTER.validate_python(litellm_params) - except ValidationError: - return _EMPTY_PARAMS - - -def _polled_models(logging_obj: object) -> tuple[str | None, str | None]: - """ - The proxy routes by the models encoded in the polled id, so they must survive into the returned id. The served - group survives only when the router sent this poll to it, so an edited id cannot write any name it likes back. - """ - video_id: Final = _logged_litellm_params(logging_obj).get("video_id") - if not isinstance(video_id, str): - return None, None - decoded: Final = decode_video_id_with_provider(video_id) - served: Final = decoded.get("served_model") - routed: Final = _poll_models(logging_obj) - return decoded.get("model_id") or None, served if routed is None or served in routed else None - - -def _poll_models(logging_obj: object) -> tuple[str, str | None] | None: - """ - A later status or content request checks the key's model access against the model encoded in the id, and - routes by it. DashScope serves the same model name from several regional hosts, so only the group that - served the task can send its polls back to the region that created it, while a key may be scoped to the - alias, routing group or wildcard name it asked for instead. The id carries both names: the one the router - was asked for on this hop, and the group that served it when that differs and is not a wildcard. The router - stamps both into litellm_metadata, and metadata carries a copy only when the caller sent none. - """ - params: Final = _logged_litellm_params(logging_obj) - stamped: Final = (_stamped_poll_models(params.get(key)) for key in ("litellm_metadata", "metadata")) - return next((models for models in stamped if models), None) - - -def _stamped_poll_models(metadata: object) -> tuple[str, str | None] | None: - try: - fields: Final = _PARAMETERS_ADAPTER.validate_python(metadata) - except ValidationError: - return None - requested: Final = _id_safe_name(fields.get("model_group")) - served: Final = _id_safe_name(fields.get("deployment_model_name")) - if requested is None: - return None - return requested, None if served is None or served == requested or "*" in served else served - - -def _id_safe_name(name: object) -> str | None: - return name if isinstance(name, str) and name and ";" not in name else None - - class DashScopeVideoConfig(BaseVideoConfig): def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseVideoConfig contract returns list return [ @@ -804,7 +718,7 @@ class DashScopeVideoConfig(BaseVideoConfig): input_reference: Final = video_create_optional_request_params.get("input_reference") if input_reference is None: return None, None - image_url: Final = _image_url(input_reference) + image_url: Final = input_reference_url(input_reference) media_type: Final = _media_reference_type(model) if media_type is None: return "img_url", image_url @@ -823,7 +737,7 @@ class DashScopeVideoConfig(BaseVideoConfig): requested: Final = self._requested_parameters(request_data) video_obj: Final = _video_object_from_task(task, raw_response, model=model, requested=requested) if custom_llm_provider and video_obj.id: - requested_model, served_model = _poll_models(logging_obj) or (model, None) + requested_model, served_model = routed_video_models(logging_obj) or (model, None) video_obj.id = encode_video_id_with_provider( video_obj.id, custom_llm_provider, requested_model, served_model ) @@ -869,7 +783,7 @@ class DashScopeVideoConfig(BaseVideoConfig): client: "HTTPHandler | None" = None, ) -> VideoObject: task: Final = _parse_task_response(raw_response) - model, served_model = _polled_models(logging_obj) + model, served_model = polled_video_models(logging_obj) video_obj: Final = _video_object_from_task(task, raw_response, model=model) if custom_llm_provider and video_obj.id: video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model, served_model) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 2145c19a71a..74bdd2cd733 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( merge_consecutive_system_messages, strip_litellm_internal_message_fields, strip_name_from_message, + strip_prefill_prefix_field, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.base_utils import ( @@ -487,6 +488,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): Databricks does not support: - 'name' in user message. - litellm's internal `thinking_blocks` / `reasoning_content` on assistant messages. + - the `prefix` prefill flag Mistral and DeepSeek read, which Claude does not need to continue an assistant turn. - non-Claude json_object "json" checks read only string content, so a single text block is sent as a string. """ new_messages = [] @@ -496,7 +498,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): else: _message = message _message = strip_name_from_message(_message, allowed_name_roles=["user"]) - _message = strip_litellm_internal_message_fields(_message) + _message = strip_prefill_prefix_field(strip_litellm_internal_message_fields(_message)) # Move message-level cache_control into a content block when content is a string. if "cache_control" in _message and isinstance(_message.get("content"), str): _message = self._move_cache_control_into_string_content_block(_message) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 57ee0f33879..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,32 +497,61 @@ 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, + provider_error_response, + unconvertible_messages_error, vertex_gemini_count_context, ) - count_context: Final = vertex_gemini_count_context(system=system, tools=tools) + 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) + blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format) + if blocks_error is not None: + return count_context_error_response(blocks_error, request_model, model_to_use) deployment = deployment or {} count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) + resolved_contents: Final = ( + contents + if contents is not None + else await gemini_count_contents( + request_view.messages, messages_format, model_to_use, "gemini", count_tokens_params_request + ) + ) + if isinstance(resolved_contents, VertexGeminiCountContextError): + return count_context_error_response(resolved_contents, request_model, model_to_use) count_tokens_params: Final = { "model": model_to_use, - "contents": contents, + "contents": resolved_contents, "system_instruction": count_context.system_instruction, "tools": count_context.tools, } count_tokens_params_request.update(count_tokens_params) - result: Final = await GoogleAIStudioTokenCounter().acount_tokens( - **count_tokens_params_request, - ) + try: + result: Final = await GoogleAIStudioTokenCounter().acount_tokens( + **count_tokens_params_request, + ) + except litellm.APIError as error: + return provider_error_response(error, request_model, model_to_use) if result is not None: return TokenCountResponse( 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/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 703d45acdf2..ac41310fb61 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -16,8 +16,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_ollama_image, - custom_prompt, ollama_pt, + render_custom_prompt_template, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_inline_remote_media, @@ -437,11 +437,8 @@ class OllamaConfig(BaseConfig): if model in custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details: Final = custom_prompt_dict[model] - ollama_prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=messages, + ollama_prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="ollama" ) elif text_completion_request: # handle `/completions` requests ollama_prompt = get_str_from_messages(messages=messages) 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..0c1eaaaae26 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -30,9 +30,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 +52,19 @@ from .chat.gpt_transformation import OpenAIGPTConfig, OpenAIUnknownModelConfig from .chat.o_series_transformation import OpenAIOSeriesConfig from .common_utils import ( BaseOpenAILLM, - OpenAIAsyncHTTPClient, OpenAIError, - OpenAIHTTPClient, build_output_token_limit_response, drop_params_from_unprocessable_entity_error, is_openai_backed_api_base, is_output_token_limit_error, model_dump_provider_response, ) -from .workload_identity import resolve_openai_workload_identity_config +from .workload_identity import ( + build_async_openai_client, + build_openai_client, + openai_static_api_key, + resolve_openai_workload_identity_config, +) openaiOSeriesConfig: Final = OpenAIOSeriesConfig() openAIGPT5Config: Final = OpenAIGPT5Config() @@ -426,7 +428,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) if workload_identity_config is not None else AsyncOpenAI( - api_key=api_key, + api_key=openai_static_api_key(api_key), base_url=api_base, http_client=async_http_client, timeout=timeout, @@ -448,7 +450,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) if workload_identity_config is not None else OpenAI( - api_key=api_key, + api_key=openai_static_api_key(api_key), base_url=api_base, http_client=sync_http_client, timeout=timeout, @@ -1456,10 +1458,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=None, organization: str | None = None, headers: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ): try: openai_aclient: Final = self._get_openai_client( is_async=True, + litellm_params=litellm_params, api_key=api_key, api_base=api_base, timeout=timeout, @@ -1521,6 +1525,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): aimg_generation=None, organization: str | None = None, headers: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> ImageResponse: data = {} try: @@ -1542,10 +1547,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=max_retries, organization=organization, headers=headers, + litellm_params=litellm_params, ) openai_client: Final[OpenAI] = self._get_openai_client( is_async=False, + litellm_params=litellm_params, api_key=api_key, api_base=api_base, timeout=timeout, @@ -1623,6 +1630,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): aspeech: bool | None = None, client=None, shared_session: Optional["ClientSession"] = None, + litellm_params: Mapping[str, object] | None = None, ) -> HttpxBinaryResponseContent: if aspeech is not None and aspeech is True: return self.async_audio_speech( @@ -1639,6 +1647,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj=logging_obj, client=client, shared_session=shared_session, + litellm_params=litellm_params, ) openai_client: Final = self._get_openai_client( @@ -1649,6 +1658,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=max_retries, client=client, shared_session=shared_session, + litellm_params=litellm_params, ) sync_client: Final = cast(OpenAI, openai_client) @@ -1686,6 +1696,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj: LiteLLMLoggingObj, client=None, shared_session: Optional["ClientSession"] = None, + litellm_params: Mapping[str, object] | None = None, ) -> HttpxBinaryResponseContent: openai_client: Final = cast( AsyncOpenAI, @@ -1697,6 +1708,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): max_retries=max_retries, client=client, shared_session=shared_session, + litellm_params=litellm_params, ), ) @@ -1743,26 +1755,27 @@ class OpenAIFilesAPI(BaseLLM): organization: str | None, client: OpenAI | AsyncOpenAI | None = None, _is_async: bool = False, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAI | AsyncOpenAI | None: - received_args: Final[Mapping[str, object]] = types.MappingProxyType(locals()) - openai_client: OpenAI | AsyncOpenAI | None = None - if client is None: - data: Final = {} - for k, v in received_args.items(): - if k == "self" or k == "client" or k == "_is_async": - pass - elif k == "api_base" and v is not None: - data["base_url"] = v - elif v is not None: - data[k] = v - if _is_async is True: - openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient()) - else: - openai_client = OpenAI(**data, http_client=OpenAIHTTPClient()) - else: - openai_client = client - - return openai_client + if client is not None: + return client + if _is_async: + return build_async_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) + return build_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) async def acreate_file( self, @@ -1782,6 +1795,7 @@ class OpenAIFilesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]: openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -1791,6 +1805,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -1824,6 +1839,7 @@ class OpenAIFilesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> HttpxBinaryResponseContent | Coroutine[None, None, HttpxBinaryResponseContent]: openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -1833,6 +1849,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -1889,6 +1906,7 @@ class OpenAIFilesAPI(BaseLLM): organization: str | None, chunk_size: int = 1024 * 1024, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> FileContentStreamingResult: openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -1898,6 +1916,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -1952,6 +1971,7 @@ class OpenAIFilesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ): openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -1961,6 +1981,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -1998,6 +2019,7 @@ class OpenAIFilesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ): openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -2007,6 +2029,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2047,6 +2070,7 @@ class OpenAIFilesAPI(BaseLLM): organization: str | None, purpose: str | None = None, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ): openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -2056,6 +2080,7 @@ class OpenAIFilesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2101,26 +2126,27 @@ class OpenAIBatchesAPI(BaseLLM): organization: str | None, client: OpenAI | AsyncOpenAI | None = None, _is_async: bool = False, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAI | AsyncOpenAI | None: - received_args: Final[Mapping[str, object]] = types.MappingProxyType(locals()) - openai_client: OpenAI | AsyncOpenAI | None = None - if client is None: - data: Final = {} - for k, v in received_args.items(): - if k == "self" or k == "client" or k == "_is_async": - pass - elif k == "api_base" and v is not None: - data["base_url"] = v - elif v is not None: - data[k] = v - if _is_async is True: - openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient()) - else: - openai_client = OpenAI(**data, http_client=OpenAIHTTPClient()) - else: - openai_client = client - - return openai_client + if client is not None: + return client + if _is_async: + return build_async_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) + return build_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) async def acreate_batch( self, @@ -2140,6 +2166,7 @@ class OpenAIBatchesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> LiteLLMBatch | Coroutine[None, None, LiteLLMBatch]: openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -2149,6 +2176,7 @@ class OpenAIBatchesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2184,6 +2212,7 @@ class OpenAIBatchesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ): openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -2193,6 +2222,7 @@ class OpenAIBatchesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2227,6 +2257,7 @@ class OpenAIBatchesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ): openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -2236,6 +2267,7 @@ class OpenAIBatchesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2276,6 +2308,7 @@ class OpenAIBatchesAPI(BaseLLM): after: str | None = None, limit: int | None = None, client: OpenAI | None = None, + litellm_params: Mapping[str, object] | None = None, ): openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, @@ -2285,6 +2318,7 @@ class OpenAIBatchesAPI(BaseLLM): organization=organization, client=client, _is_async=_is_async, + litellm_params=litellm_params, ) if openai_client is None: raise ValueError( @@ -2313,22 +2347,19 @@ class OpenAIAssistantsAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAI: - received_args: Final[Mapping[str, object]] = types.MappingProxyType(locals()) - if client is None: - data: Final = {} - for k, v in received_args.items(): - if k == "self" or k == "client": - pass - elif k == "api_base" and v is not None: - data["base_url"] = v - elif v is not None: - data[k] = v - openai_client = OpenAI(**data, http_client=OpenAIHTTPClient()) - else: - openai_client = client - - return openai_client + if client is not None: + return client + return build_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) def async_get_openai_client( self, @@ -2338,22 +2369,19 @@ class OpenAIAssistantsAPI(BaseLLM): max_retries: int | None, organization: str | None, client: AsyncOpenAI | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> AsyncOpenAI: - received_args: Final[Mapping[str, object]] = types.MappingProxyType(locals()) - if client is None: - data: Final = {} - for k, v in received_args.items(): - if k == "self" or k == "client": - pass - elif k == "api_base" and v is not None: - data["base_url"] = v - elif v is not None: - data[k] = v - openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient()) - else: - openai_client = client - - return openai_client + if client is not None: + return client + return build_async_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + litellm_params=litellm_params, + ) ### ASSISTANTS ### @@ -2369,6 +2397,8 @@ class OpenAIAssistantsAPI(BaseLLM): limit: int | None = 20, before: str | None = None, after: str | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> AsyncCursorPage[Assistant]: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2376,6 +2406,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) request_params: Final = { @@ -2403,6 +2434,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, aget_assistants: Literal[True], + *, + litellm_params: Mapping[str, object] | None = None, ) -> Coroutine[None, None, AsyncCursorPage[Assistant]]: ... @@ -2416,6 +2449,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: OpenAI | None, aget_assistants: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> SyncCursorPage[Assistant]: ... @@ -2434,7 +2469,9 @@ class OpenAIAssistantsAPI(BaseLLM): limit: int | None = 20, before: str | None = None, after: str | None = None, - ): + *, + litellm_params: Mapping[str, object] | None = None, + ) -> Coroutine[None, None, AsyncCursorPage[Assistant]] | SyncCursorPage[Assistant]: if aget_assistants is not None and aget_assistants is True: return self.async_get_assistants( api_key=api_key, @@ -2442,6 +2479,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) openai_client: Final = self.get_openai_client( @@ -2450,6 +2488,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2477,6 +2516,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, create_assistant_data: dict, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Assistant: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2484,6 +2525,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2501,6 +2543,8 @@ class OpenAIAssistantsAPI(BaseLLM): create_assistant_data: dict, client=None, async_create_assistants=None, + *, + litellm_params: Mapping[str, object] | None = None, ): if async_create_assistants is not None and async_create_assistants is True: return self.async_create_assistants( @@ -2509,6 +2553,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, create_assistant_data=create_assistant_data, ) @@ -2518,6 +2563,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2534,6 +2580,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, assistant_id: str, + *, + litellm_params: Mapping[str, object] | None = None, ) -> AssistantDeleted: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2541,6 +2589,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2558,6 +2607,8 @@ class OpenAIAssistantsAPI(BaseLLM): assistant_id: str, client=None, async_delete_assistants=None, + *, + litellm_params: Mapping[str, object] | None = None, ): if async_delete_assistants is not None and async_delete_assistants is True: return self.async_delete_assistant( @@ -2566,6 +2617,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, assistant_id=assistant_id, ) @@ -2575,6 +2627,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2593,6 +2646,8 @@ class OpenAIAssistantsAPI(BaseLLM): max_retries: int | None, organization: str | None, client: AsyncOpenAI | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAIMessage: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2600,6 +2655,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2630,6 +2686,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, a_add_message: Literal[True], + *, + litellm_params: Mapping[str, object] | None = None, ) -> Coroutine[None, None, OpenAIMessage]: ... @@ -2645,6 +2703,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: OpenAI | None, a_add_message: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> OpenAIMessage: ... @@ -2661,6 +2721,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client=None, a_add_message: bool | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ): if a_add_message is not None and a_add_message is True: return self.a_add_message( @@ -2671,6 +2733,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) openai_client: Final = self.get_openai_client( @@ -2679,6 +2742,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2704,6 +2768,8 @@ class OpenAIAssistantsAPI(BaseLLM): max_retries: int | None, organization: str | None, client: AsyncOpenAI | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> AsyncCursorPage[OpenAIMessage]: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2711,6 +2777,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2731,6 +2798,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, aget_messages: Literal[True], + *, + litellm_params: Mapping[str, object] | None = None, ) -> Coroutine[None, None, AsyncCursorPage[OpenAIMessage]]: ... @@ -2745,6 +2814,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: OpenAI | None, aget_messages: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> SyncCursorPage[OpenAIMessage]: ... @@ -2760,6 +2831,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client=None, aget_messages=None, + *, + litellm_params: Mapping[str, object] | None = None, ): if aget_messages is not None and aget_messages is True: return self.async_get_messages( @@ -2769,6 +2842,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) openai_client: Final = self.get_openai_client( @@ -2777,6 +2851,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2796,6 +2871,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, messages: Iterable[OpenAICreateThreadParamsMessage] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Thread: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2803,6 +2880,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2835,6 +2913,8 @@ class OpenAIAssistantsAPI(BaseLLM): messages: Iterable[OpenAICreateThreadParamsMessage] | None, client: AsyncOpenAI | None, acreate_thread: Literal[True], + *, + litellm_params: Mapping[str, object] | None = None, ) -> Coroutine[None, None, Thread]: ... @@ -2850,6 +2930,8 @@ class OpenAIAssistantsAPI(BaseLLM): messages: Iterable[OpenAICreateThreadParamsMessage] | None, client: OpenAI | None, acreate_thread: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Thread: ... @@ -2866,6 +2948,8 @@ class OpenAIAssistantsAPI(BaseLLM): messages: Iterable[OpenAICreateThreadParamsMessage] | None, client=None, acreate_thread=None, + *, + litellm_params: Mapping[str, object] | None = None, ): """ Here's an example: @@ -2885,6 +2969,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, messages=messages, ) @@ -2894,6 +2979,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2921,6 +3007,8 @@ class OpenAIAssistantsAPI(BaseLLM): max_retries: int | None, organization: str | None, client: AsyncOpenAI | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Thread: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -2928,6 +3016,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -2953,6 +3042,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None, aget_thread: Literal[True], + *, + litellm_params: Mapping[str, object] | None = None, ) -> Coroutine[None, None, Thread]: ... @@ -2967,6 +3058,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: OpenAI | None, aget_thread: Literal[False] | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Thread: ... @@ -2982,6 +3075,8 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client=None, aget_thread=None, + *, + litellm_params: Mapping[str, object] | None = None, ): if aget_thread is not None and aget_thread is True: return self.async_get_thread( @@ -2991,6 +3086,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) openai_client: Final = self.get_openai_client( @@ -2999,6 +3095,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -3032,6 +3129,8 @@ class OpenAIAssistantsAPI(BaseLLM): max_retries: int | None, organization: str | None, client: AsyncOpenAI | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Run: openai_client: Final = self.async_get_openai_client( api_key=api_key, @@ -3039,6 +3138,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) @@ -3143,6 +3243,8 @@ class OpenAIAssistantsAPI(BaseLLM): client, arun_thread: Literal[True], event_handler: AssistantEventHandler | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Coroutine[None, None, Run]: ... @@ -3165,6 +3267,8 @@ class OpenAIAssistantsAPI(BaseLLM): client, arun_thread: Literal[False] | None, event_handler: AssistantEventHandler | None, + *, + litellm_params: Mapping[str, object] | None = None, ) -> Run: ... @@ -3188,6 +3292,8 @@ class OpenAIAssistantsAPI(BaseLLM): client=None, arun_thread=None, event_handler: AssistantEventHandler | None = None, + *, + litellm_params: Mapping[str, object] | None = None, ): if arun_thread is not None and arun_thread is True: if stream is not None and stream is True: @@ -3197,6 +3303,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) return self.async_run_thread_stream( @@ -3224,6 +3331,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) openai_client: Final = self.get_openai_client( @@ -3232,6 +3340,7 @@ class OpenAIAssistantsAPI(BaseLLM): timeout=timeout, max_retries=max_retries, organization=organization, + litellm_params=litellm_params, client=client, ) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index b97e1329675..c7c997bc961 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -2,19 +2,21 @@ OpenAI Token Counter implementation using the Responses API /input_tokens endpoint. """ -import os from collections.abc import Mapping, Sequence from typing import Any, Final +import openai from pydantic import JsonValue from typing_extensions import assert_never from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify +from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.common_utils import OpenAIError, get_openai_credentials from litellm.llms.openai.responses.count_tokens.handler import ( OpenAICountTokensHandler, ) @@ -22,12 +24,32 @@ from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ResponsesCountRequestParts, ) -from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.llms.openai.workload_identity import ( + OpenAIWorkloadIdentityConfig, + get_workload_identity_bearer_token, + openai_static_api_key, + resolve_openai_workload_identity_config, +) +from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # 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 | OpenAIError | None: + api_base: Final = litellm_params.get("api_base") + 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): """Token counter implementation for OpenAI provider using the Responses API.""" @@ -79,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. @@ -91,30 +114,53 @@ class OpenAITokenCounter(BaseTokenCounter): if messages_format == "anthropic" else None ) - if malformed_error is not None: + request_error: Final = malformed_error or unconvertible_blocks_error(messages, messages_format) + if request_error is not None: return TokenCountResponse( total_tokens=0, request_model=request_model, model_used=model_to_use, tokenizer_type="openai_api", error=True, - error_message=malformed_error, + error_message=request_error, status_code=400, ) deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) - # Get OpenAI API key from deployment config or environment - api_key = litellm_params.get("api_key") - if not api_key: - api_key = os.getenv("OPENAI_API_KEY") + credentials: Final = get_openai_credentials( + api_base=litellm_params.get("api_base"), api_key=litellm_params.get("api_key") + ) + credential: Final = openai_static_api_key(credentials.api_key) or _workload_identity_config(litellm_params) - if not api_key: - verbose_logger.warning("No OpenAI API key found for token counting") - return None + 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, + ) - api_base: Final = litellm_params.get("api_base") + if credential is None: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=( + f"No OpenAI API key is set for model {request_model or model_to_use}, so tokens can't be " + "counted. Set api_key or workload identity on the deployment, or OPENAI_API_KEY in the proxy " + "environment." + ), + status_code=401, + ) + + api_base: Final = credentials.api_base input_items, instructions, counted_tools, text = self._responses_request_parts( model_to_use=model_to_use, @@ -130,6 +176,11 @@ class OpenAITokenCounter(BaseTokenCounter): return None try: + api_key: Final = ( + credential + if isinstance(credential, str) + else await asyncify(get_workload_identity_bearer_token)(credential) + ) result: Final = await openai_count_tokens_handler.handle_count_tokens_request( model=model_to_use, input=input_items, @@ -161,6 +212,17 @@ class OpenAITokenCounter(BaseTokenCounter): error_message=e.message, status_code=e.status_code, ) + except openai.AuthenticationError as e: + verbose_logger.warning("OpenAI workload identity token exchange failed: %s", e.message) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=e.message, + status_code=401, + ) except Exception as e: verbose_logger.warning("Error calling OpenAI CountTokens API: %s", e) 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/petals/completion/handler.py b/litellm/llms/petals/completion/handler.py index 425d3729d7c..a2ccd7eaf08 100644 --- a/litellm/llms/petals/completion/handler.py +++ b/litellm/llms/petals/completion/handler.py @@ -5,8 +5,8 @@ from typing import Final import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( - custom_prompt, prompt_factory, + render_custom_prompt_template, ) from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -43,11 +43,8 @@ def completion( if model in litellm.custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details: Final = litellm.custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="petals" ) else: prompt = prompt_factory(model=model, messages=messages) diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py index 7bd05c6e30b..0dd7e27722e 100644 --- a/litellm/llms/predibase/chat/transformation.py +++ b/litellm/llms/predibase/chat/transformation.py @@ -10,8 +10,8 @@ import litellm from litellm.constants import DEFAULT_MAX_TOKENS from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.prompt_templates.factory import ( - custom_prompt, prompt_factory, + render_custom_prompt_template, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.openai import AllMessageValues @@ -277,11 +277,8 @@ class PredibaseConfig(BaseConfig): custom_prompt_dict: Final = litellm_params.get("custom_prompt_dict", {}) if model in custom_prompt_dict: model_prompt_details: Final = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="predibase" ) else: prompt = prompt_factory(model=model, messages=messages) diff --git a/litellm/llms/sagemaker/completion/transformation.py b/litellm/llms/sagemaker/completion/transformation.py index 274f884de0b..a69846c816f 100644 --- a/litellm/llms/sagemaker/completion/transformation.py +++ b/litellm/llms/sagemaker/completion/transformation.py @@ -15,8 +15,8 @@ from httpx._models import Headers, Response import litellm from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.factory import ( - custom_prompt, prompt_factory, + render_custom_prompt_template, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.openai import AllMessageValues @@ -122,20 +122,14 @@ class SagemakerConfig(BaseConfig): if model in custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details.get("roles", None), - initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), - final_prompt_value=model_prompt_details.get("final_prompt_value", ""), - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="sagemaker" ) elif hf_model_name in custom_prompt_dict: # check if the base huggingface model has a registered custom prompt model_prompt_details = custom_prompt_dict[hf_model_name] - prompt = custom_prompt( - role_dict=model_prompt_details.get("roles", None), - initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), - final_prompt_value=model_prompt_details.get("final_prompt_value", ""), - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="sagemaker" ) else: if hf_model_name is None: 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/soniox/audio_transcription/handler.py b/litellm/llms/soniox/audio_transcription/handler.py index b53955ae75d..f86bc142898 100644 --- a/litellm/llms/soniox/audio_transcription/handler.py +++ b/litellm/llms/soniox/audio_transcription/handler.py @@ -19,11 +19,14 @@ import asyncio import math import time from collections.abc import Coroutine, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict +import litellm from litellm.litellm_core_utils.audio_utils.utils import ( get_audio_file_name, process_audio_file, @@ -35,7 +38,12 @@ from litellm.llms.custom_httpx.http_handler import ( get_httpx_client, ) from litellm.llms.soniox.audio_transcription.transformation import ( + SONIOX_HANDLER_ONLY_PARAMS, SonioxAudioTranscriptionConfig, + SonioxInvalidFormParam, + decode_soniox_form_params, + raise_soniox_form_error, + soniox_language_hints, ) from litellm.llms.soniox.common_utils import ( SONIOX_DEFAULT_CLEANUP, @@ -58,6 +66,39 @@ else: LiteLLMLoggingObj = Any +_CLEANUP_ENTRIES: Final = TypeAdapter(tuple[object, ...]) +_NOT_FORWARDED_TO_SONIOX: Final[frozenset[str]] = frozenset( + (*SONIOX_HANDLER_ONLY_PARAMS, "audio_url", "file_id", "response_format", "language") +) + + +def _optional_str(value: object) -> str | None: + return None if value is None else str(value) + + +def _float_or_default(value: object, default: float) -> float: + try: + parsed: Final = float(str(value)) + except ValueError: + return default + return parsed if math.isfinite(parsed) else default + + +def _int_or_default(value: object, default: int) -> int: + try: + return int(value) if isinstance(value, (int, float)) else int(str(value)) + except (ValueError, OverflowError): + return default + + +def _cleanup_targets(raw: object) -> tuple[str, ...]: + if raw is None: + return () + if isinstance(raw, str): + return (raw,) + return tuple(target for target in _CLEANUP_ENTRIES.validate_python(raw) if isinstance(target, str)) + + class _TranscriptionMeta(TypedDict, total=False): """Fields the handler reads from a Soniox transcription object.""" @@ -182,7 +223,7 @@ class SonioxAudioTranscriptionHandler: ) -> tuple[ dict[str, str], # auth headers str, # api_base (no trailing slash) - dict[str, object], # body for POST /v1/transcriptions (without file_id/audio_url) + Mapping[str, object], # body for POST /v1/transcriptions (without file_id/audio_url) _HandlerOptions, # handler-only options (poll interval, cleanup, ...) ]: # Validate env -> auth headers. @@ -198,52 +239,42 @@ class SonioxAudioTranscriptionHandler: base_url: Final = get_soniox_api_base(api_base) - # Operate on a local copy so we don't mutate the caller's dict - # (the caller may reuse `optional_params` for retries or logging). - params: Final = dict(optional_params) - - # Pull handler-only kwargs out of params so they aren't sent - # to Soniox. - poll_interval = float(params.pop("soniox_polling_interval", SONIOX_DEFAULT_POLL_INTERVAL)) - try: - max_attempts = int(params.pop("soniox_max_polling_attempts", SONIOX_DEFAULT_MAX_POLL_ATTEMPTS)) - except (ValueError, OverflowError): - max_attempts = SONIOX_DEFAULT_MAX_POLL_ATTEMPTS - cleanup_raw: Final = params.pop("soniox_cleanup", SONIOX_DEFAULT_CLEANUP) - if cleanup_raw is None: - cleanup: list[str] = [] - elif isinstance(cleanup_raw, str): - cleanup = [cleanup_raw] - else: - cleanup = list(cleanup_raw) - filename_override: Final = params.pop("filename", None) + decoded: Final = decode_soniox_form_params( + optional_params, drop_params=litellm_params.get("drop_params") is True or litellm.drop_params is True + ) + if isinstance(decoded, SonioxInvalidFormParam): + raise_soniox_form_error(decoded) # Server-side clamps. Caller-supplied poll settings (from request kwargs) # are bounded so an authenticated caller cannot force a worker into a # tight poll loop (zero interval) or pin it indefinitely (huge attempt # count). Total polling time is bounded by # SONIOX_MAX_POLL_ATTEMPTS * SONIOX_MAX_POLL_INTERVAL. - if not math.isfinite(poll_interval): - poll_interval = SONIOX_DEFAULT_POLL_INTERVAL - clamped_poll_interval: Final = max(SONIOX_MIN_POLL_INTERVAL, min(poll_interval, SONIOX_MAX_POLL_INTERVAL)) - clamped_max_attempts: Final = max(1, min(max_attempts, SONIOX_MAX_POLL_ATTEMPTS)) + poll_interval: Final = _float_or_default(decoded.get("soniox_polling_interval"), SONIOX_DEFAULT_POLL_INTERVAL) + max_attempts: Final = _int_or_default( + decoded.get("soniox_max_polling_attempts"), SONIOX_DEFAULT_MAX_POLL_ATTEMPTS + ) + cleanup: Final = _cleanup_targets(decoded.get("soniox_cleanup", SONIOX_DEFAULT_CLEANUP)) + language_hints: Final = soniox_language_hints(decoded.get("language"), decoded.get("language_hints")) # response_format is handled by LiteLLM post-processing, not Soniox. handler_opts: Final[_HandlerOptions] = { - "poll_interval": clamped_poll_interval, - "max_attempts": clamped_max_attempts, + "poll_interval": max(SONIOX_MIN_POLL_INTERVAL, min(poll_interval, SONIOX_MAX_POLL_INTERVAL)), + "max_attempts": max(1, min(max_attempts, SONIOX_MAX_POLL_ATTEMPTS)), "cleanup": cleanup, - "filename_override": filename_override, - "audio_url": params.pop("audio_url", None), - "file_id": params.pop("file_id", None), - "response_format": params.pop("response_format", None), + "filename_override": _optional_str(decoded.get("filename")), + "audio_url": _optional_str(decoded.get("audio_url")), + "file_id": _optional_str(decoded.get("file_id")), + "response_format": _optional_str(decoded.get("response_format")), } - # Soniox does not accept `language` directly; map_openai_params should - # already have translated it, but drop any leftover to be safe. - params.pop("language", None) - - return auth_headers, base_url, params, handler_opts + provider_params: Final = MappingProxyType( + { + **{key: value for key, value in decoded.items() if key not in _NOT_FORWARDED_TO_SONIOX}, + **({} if language_hints is None else {"language_hints": language_hints}), + } + ) + return auth_headers, base_url, provider_params, handler_opts def _build_create_body( self, diff --git a/litellm/llms/soniox/audio_transcription/transformation.py b/litellm/llms/soniox/audio_transcription/transformation.py index aa1ff239320..30d872f248f 100644 --- a/litellm/llms/soniox/audio_transcription/transformation.py +++ b/litellm/llms/soniox/audio_transcription/transformation.py @@ -9,9 +9,14 @@ async API requires multiple HTTP calls and does not fit the single-request contract of `base_llm_http_handler.audio_transcriptions`. """ -from typing import Any, Final +import re +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Final, NoReturn from httpx import Headers, Response +from pydantic import TypeAdapter, ValidationError from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, @@ -26,6 +31,7 @@ from litellm.llms.soniox.common_utils import ( render_soniox_tokens_as_srt, render_soniox_tokens_as_vtt, ) +from litellm.llms.soniox.types import SonioxContext from litellm.types.llms.openai import ( AllMessageValues, OpenAIAudioTranscriptionOptionalParams, @@ -58,6 +64,152 @@ SONIOX_HANDLER_ONLY_PARAMS: Final[list[str]] = [ ] +SONIOX_BOOL_PARAMS: Final[frozenset[str]] = frozenset( + {"enable_speaker_diarization", "enable_language_identification", "language_hints_strict"} +) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_STRUCTURED_CONTEXT: Final = TypeAdapter(SonioxContext) +_FORM_BOOL: Final = TypeAdapter(bool) +_ITEMS: Final = TypeAdapter(tuple[object, ...]) +_BRACKETED_FORM_FIELD: Final = re.compile( + r"(?Pcontext|translation)\[(?P\w+)\](?:\[\]\[(?P\w+)\])?" +) +_BracketedField = tuple[str, str, str | None, object] + + +@dataclass(frozen=True, slots=True) +class SonioxInvalidFormParam: + key: str + value: str + expected: str + + +def raise_soniox_form_error(error: SonioxInvalidFormParam) -> NoReturn: + raise SonioxException( + message=f"`{error.key}` must be {error.expected}, got {error.value!r}", status_code=400, headers=None + ) + + +def soniox_language_hints(language: object, hints: object) -> object: + """Soniox has no `language` field, so the OpenAI `language` becomes `language_hints` when none are set.""" + if hints is not None or not isinstance(language, str) or not language: + return hints + return (language,) + + +def _decode_bool(key: str, value: str) -> bool | SonioxInvalidFormParam: + try: + return _FORM_BOOL.validate_python(value.strip()) + except ValidationError: + return SonioxInvalidFormParam(key=key, value=value, expected="a boolean") + + +def _decode_polling_interval(value: str) -> str | SonioxInvalidFormParam: + try: + float(value) + except ValueError: + return SonioxInvalidFormParam(key="soniox_polling_interval", value=value, expected="a number") + return value + + +def _decode_context(value: str) -> SonioxContext | str: + if not value.lstrip().startswith("{"): + return value + try: + return _STRUCTURED_CONTEXT.validate_json(value) + except ValidationError: + return value + + +def _decode_translation(value: str) -> Mapping[str, object] | SonioxInvalidFormParam: + try: + return _JSON_OBJECT.validate_json(value) + except ValidationError: + return SonioxInvalidFormParam(key="translation", value=value, expected="a JSON object") + + +def _decode_language_hints(value: str) -> tuple[object, ...] | SonioxInvalidFormParam: + if not value.lstrip().startswith("["): + return tuple(hint.strip() for hint in value.split(",") if hint.strip()) + try: + return _ITEMS.validate_json(value) + except ValidationError: + return SonioxInvalidFormParam(key="language_hints", value=value, expected="a JSON list") + + +def _decode_form_value(key: str, value: object) -> object | SonioxInvalidFormParam: + if not isinstance(value, str): + return value + if key in SONIOX_BOOL_PARAMS: + return _decode_bool(key, value) + if key == "soniox_polling_interval": + return _decode_polling_interval(value) + if key == "context": + return _decode_context(value) + if key == "translation": + return _decode_translation(value) + if key == "language_hints": + return _decode_language_hints(value) + return value + + +def _nested_form_param(param: str, fields: tuple[_BracketedField, ...]) -> Mapping[str, object]: + owned: Final = tuple(field for field in fields if field[0] == param) + flat: Final = {name: value for _, name, entry_field, value in owned if entry_field is None} + listed: Final = {name for _, name, entry_field, _ in owned if entry_field is not None} + return { + **flat, + **{ + name: ( + { + entry_field: value + for _, owner, entry_field, value in owned + if owner == name and entry_field is not None + }, + ) + for name in listed + }, + } + + +def _is_soniox_context(value: object) -> bool: + try: + _STRUCTURED_CONTEXT.validate_python(value) + except ValidationError: + return False + return True + + +def _nest_bracketed_form_fields(params: Mapping[str, object]) -> Mapping[str, object]: + """The OpenAI SDK sends a dict in `extra_body` as `context[terms]` and `context[general][][key]` form keys.""" + matches: Final = {key: _BRACKETED_FORM_FIELD.fullmatch(key) for key in params} + fields: Final = tuple( + (match["param"], match["field"], match["entry_field"], params[key]) for key, match in matches.items() if match + ) + nested: Final = {param: _nested_form_param(param, fields) for param in {field[0] for field in fields}} + kept: Final = {param: value for param, value in nested.items() if param != "context" or _is_soniox_context(value)} + unbracketed: Final = { + key: value for key, value in params.items() if (match := matches[key]) is None or match["param"] not in kept + } + return {**kept, **unbracketed} + + +def decode_soniox_form_params( + optional_params: Mapping[str, object], + drop_params: bool = False, +) -> Mapping[str, object] | SonioxInvalidFormParam: + """Multipart form fields reach the proxy as strings; restore the JSON types Soniox expects.""" + decoded: Final = { + key: _decode_form_value(key, value) for key, value in _nest_bracketed_form_fields(optional_params).items() + } + invalid: Final = next((value for value in decoded.values() if isinstance(value, SonioxInvalidFormParam)), None) + if invalid is not None and not drop_params: + return invalid + return MappingProxyType( + {key: value for key, value in decoded.items() if not isinstance(value, SonioxInvalidFormParam)} + ) + + class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): """Configuration for Soniox async speech-to-text transcription.""" @@ -74,25 +226,26 @@ class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): model: str, drop_params: bool, ) -> dict: - # Translate the OpenAI `language` param into Soniox `language_hints`. - if "language" in non_default_params and non_default_params["language"]: - language: Final = non_default_params["language"] - existing_hints: Final = optional_params.get("language_hints") - if not existing_hints: - optional_params["language_hints"] = [language] - elif language not in existing_hints: - optional_params["language_hints"] = [language] + list(existing_hints) - - # Capture response_format for post-processing (not sent to Soniox API). - if "response_format" in non_default_params: - optional_params["response_format"] = non_default_params["response_format"] - - # Pass through Soniox-native kwargs unchanged. - for key in SONIOX_PASSTHROUGH_PARAMS + SONIOX_HANDLER_ONLY_PARAMS: - if key in non_default_params and non_default_params[key] is not None: - optional_params[key] = non_default_params[key] - - return optional_params + language: Final[object] = non_default_params.get("language") + language_params: Final[Mapping[str, object]] = ( + { + "language": language, + "language_hints": soniox_language_hints(language, optional_params.get("language_hints")), + } + if language + else {} + ) + response_format_params: Final[Mapping[str, object]] = ( + {"response_format": non_default_params["response_format"]} + if "response_format" in non_default_params + else {} + ) + passthrough_params: Final[Mapping[str, object]] = { + key: non_default_params[key] + for key in SONIOX_PASSTHROUGH_PARAMS + SONIOX_HANDLER_ONLY_PARAMS + if non_default_params.get(key) is not None + } + return {**optional_params, **language_params, **response_format_params, **passthrough_params} def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException: return SonioxException(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/soniox/types.py b/litellm/llms/soniox/types.py new file mode 100644 index 00000000000..806d2a8fb0e --- /dev/null +++ b/litellm/llms/soniox/types.py @@ -0,0 +1,24 @@ +from collections.abc import Sequence + +from pydantic import ConfigDict, with_config +from typing_extensions import ReadOnly, TypedDict + + +class SonioxContextGeneralEntry(TypedDict): + key: ReadOnly[str] + value: ReadOnly[str] + + +class SonioxTranslationTerm(TypedDict): + source: ReadOnly[str] + target: ReadOnly[str] + + +@with_config(ConfigDict(extra="forbid")) +class SonioxContext(TypedDict, total=False): + """Soniox `context` request field: https://soniox.com/docs/stt/concepts/context""" + + general: ReadOnly[Sequence[SonioxContextGeneralEntry]] + text: ReadOnly[str] + terms: ReadOnly[Sequence[str]] + translation_terms: ReadOnly[Sequence[SonioxTranslationTerm]] 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 0e4cb856632..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( @@ -738,14 +738,14 @@ def build_vertex_schema(parameters: dict, add_property_ordering: bool = False): # refs recursively and correctly detects/skips circular references. unpack_defs(parameters, defs) + _convert_schema_types(parameters) + # 5. Nullable fields: # * https://github.com/pydantic/pydantic/issues/1270 # * https://stackoverflow.com/a/58841311 # * https://github.com/pydantic/pydantic/discussions/4872 convert_anyof_null_to_nullable(parameters) - _convert_schema_types(parameters) - # Handle empty strings in enum values - Gemini doesn't accept empty strings in enums _fix_enum_empty_strings(parameters) @@ -926,34 +926,28 @@ def convert_anyof_null_to_nullable(schema, depth=0): ) """ Converts null objects within anyOf by removing them and adding nullable to all remaining objects """ anyof: Final = schema.get("anyOf", None) - if anyof is not None: - contains_null = False + if isinstance(anyof, list): + non_null_branches: Final = [ + atype for atype in anyof if not (isinstance(atype, dict) and atype.get("type") == "null") + ] + contains_null: Final = len(non_null_branches) < len(anyof) + anyof[:] = non_null_branches or anyof[:1] for atype in anyof: - if isinstance(atype, dict) and atype.get("type") == "null": - # remove null type - anyof.remove(atype) - contains_null = True - elif "type" not in atype and len(atype) == 0: - # Handle empty object case + if not isinstance(atype, dict): + continue + if not atype: atype["type"] = "object" - - if len(anyof) == 0: - # Edge case: response schema with only null type present is invalid in Vertex AI - raise ValueError( - "Invalid input: AnyOf schema with only null type is not supported. Please provide a non-null type." - ) - - if contains_null: - # set all types to nullable following guidance found here: https://cloud.google.com/vertex-ai/generative-ai/docs/samples/generativeaionvertexai-gemini-controlled-generation-response-schema-3#generativeaionvertexai_gemini_controlled_generation_response_schema_3-python - # Empty `items: {}` on array branches is left in place; downstream - # process_items() converts it to {"type": "object"}, which Vertex - # requires whenever type == "array" (even inside anyOf). - for atype in anyof: + if contains_null: + # set all types to nullable following guidance found here: https://cloud.google.com/vertex-ai/generative-ai/docs/samples/generativeaionvertexai-gemini-controlled-generation-response-schema-3#generativeaionvertexai_gemini_controlled_generation_response_schema_3-python + # Empty `items: {}` on array branches is left in place; downstream + # process_items() converts it to {"type": "object"}, which Vertex + # requires whenever type == "array" (even inside anyOf). atype["nullable"] = True + convert_anyof_null_to_nullable(atype, depth=depth + 1) properties: Final = schema.get("properties", None) if properties is not None: - for name, value in properties.items(): + for value in properties.values(): convert_anyof_null_to_nullable(value, depth=depth + 1) items: Final = schema.get("items", None) @@ -1280,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 @@ -1358,25 +1353,40 @@ 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, + provider_error_response, + unconvertible_messages_error, vertex_gemini_count_context, ) from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter - from litellm.llms.vertex_ai.gemini.transformation import ( - gemini_convert_messages_with_history, # pyright: ignore[reportPrivateUsage] # shared helper already used by gemini/chat, context_caching, and vertex_and_google_ai_studio_gemini - ) - count_context: Final = vertex_gemini_count_context(system=system, tools=tools) + 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) + blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format) + if blocks_error is not None: + return count_context_error_response(blocks_error, request_model, model_to_use) resolved_contents: Final = ( contents if contents is not None - else gemini_convert_messages_with_history( - messages=messages or [] # pyright: ignore[reportArgumentType] # preserve direct caller message handling + else await gemini_count_contents( + request_view.messages, messages_format, model_to_use, "vertex_ai", count_tokens_params_request ) ) + if isinstance(resolved_contents, VertexGeminiCountContextError): + return count_context_error_response(resolved_contents, request_model, model_to_use) count_tokens_params: Final = { "model": model_to_use, @@ -1390,15 +1400,7 @@ class VertexAITokenCounter(BaseTokenCounter): **count_tokens_params_request, ) except litellm.APIError as error: - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="gemini", - error=True, - error_message=error.message, - status_code=error.status_code, - ) + return provider_error_response(error, request_model, model_to_use) if "totalTokens" in result: return TokenCountResponse( diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index cb0bedf0238..871532b7af5 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -1,12 +1,19 @@ -from collections.abc import Mapping +import copy +from collections.abc import Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Final +from typing import TYPE_CHECKING, Final, Literal from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never from litellm.types.llms.anthropic import AnthropicMessagesTextParam -from litellm.types.utils import TokenCountResponse +from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.vertex_ai import ContentType +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse + +if TYPE_CHECKING: + from litellm.exceptions import APIError @dataclass(frozen=True, slots=True) @@ -18,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, ...]) @@ -31,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: @@ -41,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) @@ -130,6 +143,15 @@ def vertex_gemini_count_context( ) +def unconvertible_messages_error( + messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat +) -> VertexGeminiCountContextError | None: + from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error + + error_message: Final = unconvertible_blocks_error(messages, messages_format) + return None if error_message is None else VertexGeminiCountContextError(error_message=error_message) + + def count_context_error_response( error: VertexGeminiCountContextError, request_model: str, model_used: str ) -> TokenCountResponse: @@ -140,5 +162,189 @@ def count_context_error_response( tokenizer_type="gemini", error=True, error_message=error.error_message, - status_code=400, + status_code=error.status_code, ) + + +def provider_error_response(error: "APIError", request_model: str, model_used: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_used, + tokenizer_type="gemini", + error=True, + error_message=error.message, + status_code=error.status_code, + ) + + +@dataclass(frozen=True, slots=True) +class InferenceRequestView: + messages: Sequence[Mapping[str, object]] + system: object | 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, + ) + + 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) + + +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": + return await _anthropic_inference_view(messages, system, edits) + case _: + assert_never(messages_format) + + +def _chat_messages( + messages: Sequence[Mapping[str, object]], + messages_format: TokenCountMessagesFormat, + model: str, + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> Sequence[AllMessageValues]: + from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter + + match messages_format: + case "anthropic": + return LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=list(messages), # pyright: ignore[reportArgumentType] # raw Anthropic body, read the way /v1/messages inference reads it + model=model, + custom_llm_provider=custom_llm_provider, + ) + case "openai": + return copy.deepcopy(list(messages)) # pyright: ignore[reportReturnType] # chat messages from the request, copied because the AI Studio conversion rewrites them in place + case _: + assert_never(messages_format) + + +async def gemini_count_contents( + messages: Sequence[Mapping[str, object]], + messages_format: TokenCountMessagesFormat, + model: str, + custom_llm_provider: Literal["vertex_ai", "gemini"], + litellm_params: Mapping[str, object], +) -> Sequence[ContentType] | VertexGeminiCountContextError: + from litellm.exceptions import BadRequestError + from litellm.litellm_core_utils.asyncify import asyncify + from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media + from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig + from litellm.llms.vertex_ai.gemini.transformation import should_inline_remote_media + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + config: Final = GoogleAIStudioGeminiConfig() if custom_llm_provider == "gemini" else VertexGeminiConfig() + try: + chat_messages: Final = await async_inline_remote_media( + list(_chat_messages(messages, messages_format, model, custom_llm_provider)), + should_inline=should_inline_remote_media(custom_llm_provider), + ) + return await asyncify(config.transform_messages)( + messages=chat_messages, model=model, litellm_params=dict(litellm_params) + ) + except BadRequestError as error: + return VertexGeminiCountContextError(error_message=error.message) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 57327731863..7b832b321da 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -6,7 +6,7 @@ Why separate file? Make it easy to see how transformation works import os import re -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import quote @@ -1154,6 +1154,7 @@ def transform_request_body( user_response_schema_message: Final = response_schema_prompt( model=model, response_schema=optional_params.get("response_schema"), + llm_provider=custom_llm_provider, ) messages.append({"role": "user", "content": user_response_schema_message}) optional_params.pop("response_schema") @@ -1318,6 +1319,10 @@ def _vertex_inlines(media: RemoteMedia) -> bool: ) +def should_inline_remote_media(custom_llm_provider: str) -> Callable[[RemoteMedia], bool]: + return _ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + + async def async_transform_request_body( gemini_api_key: str | None, messages: list[AllMessageValues], @@ -1360,7 +1365,7 @@ async def async_transform_request_body( ) inlined_messages: Final = await async_inline_remote_media( - messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + messages, should_inline=should_inline_remote_media(custom_llm_provider) ) if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index bed7bc1bd74..257fafbb6b8 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -27,6 +27,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO, ) from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator from litellm.litellm_core_utils.prompt_templates.factory import ( encode_tool_call_id_with_signature, @@ -926,11 +927,38 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # - Requires propertyOrdering optional_params["response_schema"] = self._map_response_schema(value=schema) + @property + def accepts_reasoning_summary_in_reasoning_effort(self) -> bool: + return True + + @staticmethod + def _include_thoughts(requested: bool | None) -> bool: + if requested is not None: + return requested + return normalize_drop_params(litellm.gemini_include_thoughts_default) is not False + + @staticmethod + def _include_thoughts_requested_by_reasoning_effort(reasoning_effort: object) -> bool | None: + match reasoning_effort: + case {"summary": str()}: + return True + case _: + return None + + @staticmethod + def _map_reasoning_effort(effort: str, model: str, reasoning_effort: object) -> GeminiThinkingConfig: + include_thoughts: Final = VertexGeminiConfig._include_thoughts_requested_by_reasoning_effort(reasoning_effort) + if VertexGeminiConfig._is_gemini_3_or_newer(model): + return VertexGeminiConfig._map_reasoning_effort_to_thinking_level(effort, model, include_thoughts) + return VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(effort, model, include_thoughts) + @staticmethod def _map_reasoning_effort_to_thinking_budget( reasoning_effort: str, model: str | None = None, + include_thoughts: bool | None = None, ) -> GeminiThinkingConfig: + sends_thoughts: Final = VertexGeminiConfig._include_thoughts(include_thoughts) if reasoning_effort == "minimal": # Use model-specific minimum thinking budget or fallback # Check for exact matches first, then partial matches @@ -945,22 +973,22 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return { "thinkingBudget": budget, - "includeThoughts": True, + "includeThoughts": sends_thoughts, } elif reasoning_effort == "low": return { "thinkingBudget": DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, - "includeThoughts": True, + "includeThoughts": sends_thoughts, } elif reasoning_effort == "medium": return { "thinkingBudget": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, - "includeThoughts": True, + "includeThoughts": sends_thoughts, } elif reasoning_effort == "high": return { "thinkingBudget": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, - "includeThoughts": True, + "includeThoughts": sends_thoughts, } elif reasoning_effort == "disable": return { @@ -987,6 +1015,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): def _map_reasoning_effort_to_thinking_level( reasoning_effort: str, model: str | None = None, + include_thoughts: bool | None = None, ) -> GeminiThinkingConfig: """ Map reasoning_effort to thinking_level for Gemini 3+ models. @@ -998,17 +1027,18 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): GeminiThinkingConfig with thinkingLevel and includeThoughts """ supports_minimal: Final = bool(model) and VertexGeminiConfig._supports_minimal_thinking_level(model) + sends_thoughts: Final = VertexGeminiConfig._include_thoughts(include_thoughts) if reasoning_effort == "minimal": if supports_minimal: - return {"thinkingLevel": "minimal", "includeThoughts": True} + return {"thinkingLevel": "minimal", "includeThoughts": sends_thoughts} else: - return {"thinkingLevel": "low", "includeThoughts": True} + return {"thinkingLevel": "low", "includeThoughts": sends_thoughts} elif reasoning_effort == "low": - return {"thinkingLevel": "low", "includeThoughts": True} + return {"thinkingLevel": "low", "includeThoughts": sends_thoughts} elif reasoning_effort == "medium": - return {"thinkingLevel": "medium", "includeThoughts": True} + return {"thinkingLevel": "medium", "includeThoughts": sends_thoughts} elif reasoning_effort == "high": - return {"thinkingLevel": "high", "includeThoughts": True} + return {"thinkingLevel": "high", "includeThoughts": sends_thoughts} elif reasoning_effort in ("disable", "none"): return { "thinkingLevel": "minimal" if supports_minimal else "low", @@ -1061,6 +1091,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): status_code=400, ) + @staticmethod + def _thinking_asks_for_thoughts(thinking_param: Mapping[str, object]) -> bool: + display: Final = thinking_param.get("display") + return isinstance(thinking_param.get("summary"), str) or (display is not None and display != "omitted") + @staticmethod def _map_thinking_param( thinking_param: AnthropicThinkingParam, @@ -1068,6 +1103,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) -> GeminiThinkingConfig: thinking_enabled: Final = thinking_param.get("type") == "enabled" thinking_budget: Final = thinking_param.get("budget_tokens") + asks_for_thoughts: Final = VertexGeminiConfig._thinking_asks_for_thoughts(thinking_param) + sends_thoughts: Final = VertexGeminiConfig._include_thoughts(True if asks_for_thoughts else None) params: Final[GeminiThinkingConfig] = {} @@ -1077,7 +1114,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if thinking_budget == 0: params["includeThoughts"] = False else: - params["includeThoughts"] = True + params["includeThoughts"] = sends_thoughts # Follow provider defaults unless explicitly opted into legacy behavior. if litellm.enable_gemini_default_thinking_level_low is True: params["thinkingLevel"] = ( @@ -1089,7 +1126,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): else: # For older Gemini models, use thinkingBudget if thinking_enabled and not VertexGeminiConfig._is_thinking_budget_zero(thinking_budget): - params["includeThoughts"] = True + params["includeThoughts"] = sends_thoughts if thinking_budget is not None and isinstance(thinking_budget, int): params["thinkingBudget"] = thinking_budget @@ -1304,14 +1341,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): param_name="reasoning_effort", param_description="thinking_budget", ) - if VertexGeminiConfig._is_gemini_3_or_newer(model): - optional_params["thinkingConfig"] = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( - effort_value, model - ) - else: - optional_params["thinkingConfig"] = VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( - effort_value, model - ) + optional_params["thinkingConfig"] = VertexGeminiConfig._map_reasoning_effort( + effort=effort_value, model=model, reasoning_effort=value + ) elif param == "thinking": # Validate no conflict with thinking_level VertexGeminiConfig._validate_thinking_config_conflicts( 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/vllm/completion/handler.py b/litellm/llms/vllm/completion/handler.py index 78e6c74c2f7..f330f0928d0 100644 --- a/litellm/llms/vllm/completion/handler.py +++ b/litellm/llms/vllm/completion/handler.py @@ -5,8 +5,8 @@ from typing import Final import httpx from litellm.litellm_core_utils.prompt_templates.factory import ( - custom_prompt, prompt_factory, + render_custom_prompt_template, ) from litellm.utils import ModelResponse, Usage @@ -56,11 +56,8 @@ def completion( if model in custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details: Final = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="vllm" ) else: prompt = prompt_factory(model=model, messages=messages) @@ -144,11 +141,8 @@ def batch_completions(model: str, messages: list, optional_params=None, custom_p # check if the model has a registered custom prompt model_prompt_details: Final = custom_prompt_dict[model] for message in messages: - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=message, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=message, model=model, llm_provider="vllm" ) prompts.append(prompt) else: diff --git a/tests/unit/chat_completions/__init__.py b/litellm/llms/volcengine/videos/__init__.py similarity index 100% rename from tests/unit/chat_completions/__init__.py rename to litellm/llms/volcengine/videos/__init__.py diff --git a/litellm/llms/volcengine/videos/transformation.py b/litellm/llms/volcengine/videos/transformation.py new file mode 100644 index 00000000000..a75a4185186 --- /dev/null +++ b/litellm/llms/volcengine/videos/transformation.py @@ -0,0 +1,904 @@ +import json +from collections.abc import Mapping, Sequence +from math import log +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeVar + +import httpx +from httpx._types import RequestFiles +from pydantic import BaseModel, StrictBool, StrictInt, StrictStr, TypeAdapter, ValidationError + +import litellm +from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.videos.async_task_utils import ( + input_reference_url, + polled_video_models, + routed_video_models, +) +from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + _get_httpx_client, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # house cached-client factory has no public alias and its stub leaves params untyped + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # factory stub leaves params untyped +) +from litellm.llms.volcengine.common_utils import get_volcengine_base_url +from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ModelInfo +from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject +from litellm.types.videos.utils import encode_video_id_with_provider, extract_original_video_id +from litellm.utils import get_model_info + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + + +_ParsedT: Final = TypeVar("_ParsedT") + + +class VolcEngineVideoError(BaseLLMException): + pass + + +class _TaskError(BaseModel, frozen=True): + code: str | None = None + message: str | None = None + + +class _TaskContent(BaseModel, frozen=True): + video_url: str | None = None + + +class _TaskUsage(BaseModel, frozen=True): + total_tokens: int | None = None + + +class _Task(BaseModel, frozen=True): + id: str + status: str + model: str | None = None + error: _TaskError | None = None + content: _TaskContent | None = None + usage: _TaskUsage | None = None + created_at: int | None = None + updated_at: int | None = None + resolution: str | None = None + ratio: str | None = None + duration: int | float | None = None + + +class _TaskCreateResponse(BaseModel, frozen=True): + id: str | None = None + error: _TaskError | None = None + + +class _TaskList(BaseModel, frozen=True): + items: Sequence[_Task] = () + total: int | None = None + + +class _ContentItem(BaseModel, frozen=True): + type: str | None = None + role: str | None = None + + +class _KnownFields(BaseModel, frozen=True): + resolution: StrictStr | None = None + ratio: StrictStr | None = None + seed: StrictInt | None = None + camera_fixed: StrictBool | None = None + watermark: StrictBool | None = None + generate_audio: StrictBool | None = None + + +VOLCENGINE_VIDEO_TASKS_PATH: Final = "/api/v3/contents/generations/tasks" +VOLCENGINE_VIDEO_DEFAULT_DURATION_SECONDS: Final = 4 +VOLCENGINE_VIDEO_DEFAULT_RESOLUTION: Final = "720p" +_VOLCENGINE_API_MARKER: Final = "/api/" +_VOLCENGINE_COMPAT_SUFFIXES: Final = ("/api", "/v1") +_FRAMES_FIELD: Final = "volcengine_video_frames" + +_EMPTY_PARAMS: Final[dict[str, object]] = {} # mutable-ok: BaseVideoConfig contract empty params + +_TASK_ADAPTER: Final = TypeAdapter(_Task) +_TASK_CREATE_ADAPTER: Final = TypeAdapter(_TaskCreateResponse) +_TASK_LIST_ADAPTER: Final = TypeAdapter(_TaskList) +_FIELDS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_CONTENT_ITEMS_ADAPTER: Final = TypeAdapter(tuple[dict[str, object], ...]) +_FRAMES_ADAPTER: Final = TypeAdapter(dict[str, dict[str, str]]) +_EMPTY_FRAMES: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType({}) + +_STATUS_MAP: Final = MappingProxyType( + { + "queued": "queued", + "running": "in_progress", + "succeeded": "completed", + "failed": "failed", + "cancelled": "cancelled", + "expired": "failed", + } +) +_PENDING_STATUSES: Final = frozenset({"queued", "running"}) +_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "cancelled"}) + +_CREATE_PASSTHROUGH_KEYS: Final = ( + "resolution", + "ratio", + "duration", + "seed", + "camera_fixed", + "watermark", + "generate_audio", + "service_tier", + "priority", + "output_format", + "execution_expires_after", + "callback_url", + "safety_identifier", + "tools", + "omni_reference_task_type", +) + +_REQUEST_METADATA: Final = frozenset({"model", "prompt", "user", "extra_headers", "extra_query", "extra_body"}) +_UNSUPPORTED_BY_VOLCENGINE: Final = frozenset({"characters", "image"}) +_MAPPED_ELSEWHERE: Final = frozenset({*_REQUEST_METADATA, *_UNSUPPORTED_BY_VOLCENGINE, "seconds", "size", "parameters"}) +_READ_FROM_THE_REQUEST: Final = _MAPPED_ELSEWHERE - _UNSUPPORTED_BY_VOLCENGINE +_LITELLM_CONTROLS: Final = frozenset({"drop_params", "additional_drop_params", "organization"}) +_DROP_KEYS_ADAPTER: Final = TypeAdapter(frozenset[str]) + +# litellm/videos/utils.py merges the caller's raw ``extra_body`` back over the mapped params, which returns the keys +# this adapter maps itself a second time, as fields of their own. They are already accounted for, so they count as +# handled rather than unknown. +_MAPPED_EXTRA_BODY_COPIES: Final = frozenset({"size", "seconds", "user"}) +_CREATE_HANDLED_KEYS: Final = frozenset( + { + *_CREATE_PASSTHROUGH_KEYS, + "content", + "input_reference", + "frames", + *_MAPPED_EXTRA_BODY_COPIES, + *_LITELLM_CONTROLS, + } +) + +_VIDEO_INPUT_UNSUPPORTED: Final = ( + "{operation} is not supported through litellm: Ark bills a task with a video input at a different rate per " + "token, and litellm bills the video when the task is created, before that cost is known" +) +_INPUT_VIDEO_ROLE: Final = "reference_video" +_INPUT_VIDEO_TYPE: Final = "video_url" + +_RATIO_SET: Final = ("16:9", "9:16", "4:3", "3:4", "1:1", "21:9") +_RESOLUTION_SHORTEST_SIDE: Final = MappingProxyType({"480p": 480, "720p": 720, "1080p": 1080, "4k": 2160}) +_TIER_FRAME_ASPECT: Final = 16 / 9 + + +def _video_error(raw_response: httpx.Response, status_code: int, message: str) -> VolcEngineVideoError: + """The provider's own error status wins, and a body that failed on a 2xx answers the caller's status instead, + since a 2xx error has the proxy retry a call the provider already rejected.""" + return VolcEngineVideoError( + status_code=raw_response.status_code if raw_response.status_code >= 400 else status_code, + message=message, + headers=raw_response.headers, + ) + + +def _parse(adapter: TypeAdapter[_ParsedT], raw_response: httpx.Response) -> _ParsedT: + try: + return adapter.validate_python(raw_response.json()) + except (json.JSONDecodeError, ValidationError) as e: + raise _video_error(raw_response, 502, f"Unexpected Volcengine video response: {e}") from e + + +def _unsupported(model: str, message: str) -> UnsupportedParamsError: + return UnsupportedParamsError(message=message, model=model, llm_provider="volcengine") + + +def _droppable(model: str, message: str) -> UnsupportedParamsError: + return _unsupported(model, f"{message}, set drop_params=True to ignore it") + + +def _normalized_api_base(api_base: str) -> str: + """The Ark credential is shared with the chat, responses and Anthropic surfaces, so its api_base may carry any of + their paths, and the video paths carry their own version segment.""" + trimmed: Final = api_base.rstrip("/") + marker_index: Final = trimmed.rfind(_VOLCENGINE_API_MARKER) + if marker_index != -1: + return trimmed[:marker_index] + suffix: Final = next((suffix for suffix in _VOLCENGINE_COMPAT_SUFFIXES if trimmed.endswith(suffix)), None) + return trimmed[: -len(suffix)] if suffix is not None else trimmed + + +def _size_dimensions(size: str) -> tuple[int, int] | None: + width_str, separator, height_str = size.partition("x") + if not separator or not (width_str.isdigit() and height_str.isdigit()): + return None + if not int(width_str) or not int(height_str): + return None + return int(width_str), int(height_str) + + +def _ratio_value(ratio: str) -> float: + width_str, _, height_str = ratio.partition(":") + return int(width_str) / int(height_str) + + +def _ratio_from_size(size: str) -> str | None: + """Ark only takes its named ratios, and its own tiers do not reduce onto them (854x480 is 427:240), so the + closest ratio wins.""" + dimensions: Final = _size_dimensions(size) + if dimensions is None: + return None + width, height = dimensions + return min(_RATIO_SET, key=lambda candidate: abs(log(width / height / _ratio_value(candidate)))) + + +def _resolution_from_size(size: str) -> str | None: + """Ark keeps a tier's pixel count across its ratios (960x960 and 1280x720 are both 720p), so the tier whose 16:9 + frame has the closest pixel count wins.""" + dimensions: Final = _size_dimensions(size) + if dimensions is None: + return None + width, height = dimensions + return min( + _RESOLUTION_SHORTEST_SIDE, + key=lambda tier: abs(log(width * height / (_RESOLUTION_SHORTEST_SIDE[tier] ** 2 * _TIER_FRAME_ASPECT))), + ) + + +def _duration_param(seconds: object) -> int | None: + if isinstance(seconds, bool): + return None + if isinstance(seconds, int): + return seconds + if not isinstance(seconds, (str, float)): + return None + try: + parsed: Final = float(seconds) + except ValueError: + return None + return int(parsed) if parsed.is_integer() else None + + +def _duration_fields(seconds: object, model: str, drop_params: bool) -> Mapping[str, object]: + if seconds is None: + return _EMPTY_PARAMS + duration: Final = _duration_param(seconds) + if duration is not None: + return MappingProxyType({"duration": duration}) + if drop_params: + return _EMPTY_PARAMS + raise _droppable(model, f"seconds must be a whole number of seconds, got {seconds!r}") + + +def _size_fields(size: object, model: str, drop_params: bool) -> Mapping[str, object]: + if size is None: + return _EMPTY_PARAMS + ratio: Final = _ratio_from_size(size) if isinstance(size, str) else None + resolution: Final = _resolution_from_size(size) if isinstance(size, str) else None + if ratio is not None and resolution is not None: + return MappingProxyType({"ratio": ratio, "resolution": resolution}) + if drop_params: + return _EMPTY_PARAMS + raise _droppable(model, f"size must be WIDTHxHEIGHT, got {size!r}") + + +def _parameters_fields(parameters: object, model: str, drop_params: bool) -> Mapping[str, object]: + if parameters is None: + return _EMPTY_PARAMS + try: + fields: Final = _FIELDS_ADAPTER.validate_python(parameters) + except ValidationError: + if drop_params: + return _EMPTY_PARAMS + raise _droppable( + model, f"parameters must be an object of Ark request fields, got {type(parameters).__name__}" + ) from None + misplaced: Final = tuple(sorted(key for key in fields if key in _READ_FROM_THE_REQUEST)) + if misplaced and not drop_params: + raise _droppable(model, f"send {', '.join(misplaced)} as top-level request fields, not inside parameters") + return MappingProxyType({key: value for key, value in fields.items() if key not in _READ_FROM_THE_REQUEST}) + + +def _user_fields(user: object, model: str, drop_params: bool) -> Mapping[str, object]: + if isinstance(user, str) and user: + return MappingProxyType({"safety_identifier": user}) + if user is None or isinstance(user, str) or drop_params: + return _EMPTY_PARAMS + raise _droppable(model, f"user must be a string, got {type(user).__name__}") + + +def _additional_drop_keys(params: Mapping[str, object]) -> frozenset[str]: + try: + return _DROP_KEYS_ADAPTER.validate_python(params.get("additional_drop_params") or ()) + except ValidationError: + return frozenset() + + +def _without(params: Mapping[str, object], keys: frozenset[str]) -> Mapping[str, object]: + return MappingProxyType({key: value for key, value in params.items() if key not in keys}) + + +def _malformed_fields(fields: Mapping[str, object]) -> tuple[tuple[str, str], ...]: + try: + _ = _KnownFields.model_validate(dict(fields)) + except ValidationError as e: + return tuple((str(error["loc"][0]), error["msg"]) for error in e.errors() if error["loc"]) + return () + + +def _drops_params(litellm_params: GenericLiteLLMParams) -> bool: + return litellm_params.drop_params is True or litellm.drop_params is True + + +def _model_info(model: str) -> ModelInfo | None: + try: + return get_model_info(model=model, custom_llm_provider="volcengine") + except Exception: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models + return None + + +def _pricing_model(model: str, litellm_params: GenericLiteLLMParams) -> str: + """A deployment priced through ``model_info.base_model``, like an ``ep-`` endpoint id, bills at that model's + tiers.""" + model_info: Final = _FIELDS_ADAPTER.validate_python(litellm_params["model_info"] or _EMPTY_PARAMS) + base_model: Final = model_info.get("base_model") + return base_model if isinstance(base_model, str) and base_model else model + + +def _priced_tiers(model: str) -> tuple[str, ...] | None: + """A tier the model has no price for is one Ark does not offer it at, and litellm would bill it at the base rate. + A model with no tier price listed, missing from the cost map or registered unpriced by the router, has nothing + to check against.""" + info: Final = _model_info(model) + if info is None: + return None + return ( + tuple(tier for tier in _RESOLUTION_SHORTEST_SIDE if info.get(f"output_cost_per_second_{tier}") is not None) + or None + ) + + +def _delivered_frames(model: str | None) -> Mapping[str, Mapping[str, str]]: + info: Final = _model_info(model) if model else None + provider_specific: Final[object] = info.get("provider_specific_entry") if info is not None else None + frames: Final[object] = provider_specific.get(_FRAMES_FIELD) if isinstance(provider_specific, Mapping) else None + try: + return _FRAMES_ADAPTER.validate_python(frames) if frames is not None else _EMPTY_FRAMES + except ValidationError: + return _EMPTY_FRAMES + + +def _size_from_geometry(resolution: str | None, ratio: str | None, model: str | None) -> str | None: + """Ark echoes a tier and a ratio rather than pixels, and its rendered frame does not always follow the tier's + nominal size, so the frame listed for the model wins over scaling the tier's shortest side up the ratio.""" + if not resolution or not ratio or ":" not in ratio: + return None + tier: Final = resolution.strip().lower() + normalized_ratio: Final = ratio.strip() + listed: Final = _delivered_frames(model).get(tier, MappingProxyType({})).get(normalized_ratio) + if listed is not None: + return listed + shortest_side: Final = _RESOLUTION_SHORTEST_SIDE.get(tier) + width_str, _, height_str = normalized_ratio.partition(":") + if shortest_side is None or not (width_str.isdigit() and height_str.isdigit()): + return None + ratio_width: Final = int(width_str) + ratio_height: Final = int(height_str) + if not ratio_width or not ratio_height: + return None + shortest_ratio_side: Final = min(ratio_width, ratio_height) + return f"{-(-shortest_side * ratio_width // shortest_ratio_side)}x{-(-shortest_side * ratio_height // shortest_ratio_side)}" + + +def _is_input_video(item: _ContentItem) -> bool: + return item.type == _INPUT_VIDEO_TYPE or item.role == _INPUT_VIDEO_ROLE + + +def _supplied_items(supplied: object, model: str, drop_params: bool) -> tuple[Mapping[str, object], ...]: + try: + items: Final = _CONTENT_ITEMS_ADAPTER.validate_python(supplied) + entries: Final = tuple(_ContentItem.model_validate(item) for item in items) + except ValidationError: + if drop_params: + return () + raise _droppable(model, "content must be an array of Ark content items") from None + if any(entry.type == "text" for entry in entries) and not drop_params: + raise _droppable( + model, + "A text item in content is not supported, put the text in prompt since it is the only text input " + "guardrails scan before the request reaches Ark", + ) + if any(_is_input_video(entry) for entry in entries) and not drop_params: + raise _droppable(model, _VIDEO_INPUT_UNSUPPORTED.format(operation="Video input (a reference_video item)")) + kept: Final = (item for item, entry in zip(items, entries) if entry.type != "text" and not _is_input_video(entry)) + return tuple(MappingProxyType(item) for item in kept) + + +def _reference_items(params: Mapping[str, object], model: str, drop_params: bool) -> tuple[Mapping[str, object], ...]: + """A caller-supplied ``content`` array wins over ``input_reference``, since its roles are the only way to express + reference-to-video.""" + supplied: Final = params.get("content") + if supplied is not None: + return _supplied_items(supplied, model, drop_params) + reference: Final = params.get("input_reference") + if reference is None: + return () + try: + url: Final = input_reference_url(reference) + except ValueError as e: + if drop_params: + return () + raise _droppable(model, str(e)) from None + return (MappingProxyType({"type": "image_url", "image_url": {"url": url}, "role": "first_frame"}),) + + +def _sent_content( + prompt: str, params: Mapping[str, object], model: str, drop_params: bool +) -> tuple[Mapping[str, object], ...]: + """The prompt always leads, as it is the only text the input guardrails scan.""" + return (MappingProxyType({"type": "text", "text": prompt}), *_reference_items(params, model, drop_params)) + + +def _seconds_label(seconds: object) -> str | None: + if isinstance(seconds, float) and seconds.is_integer(): + return str(int(seconds)) + return None if seconds is None else str(seconds) + + +def _requested_usage(request_data: Mapping[str, object]) -> Mapping[str, object]: + """The create call is the only billed one and Ark's create response carries no usage, so the requested duration + and tier are what the video bills at.""" + duration: Final = _duration_param(request_data.get("duration")) + resolution: Final = request_data.get("resolution") + return MappingProxyType( + { + key: value + for key, value in ( + ("duration_seconds", float(duration) if duration is not None and duration > 0 else None), + ("video_resolution", resolution.strip().lower() if isinstance(resolution, str) else None), + ) + if value is not None + } + ) + + +def _task_usage(task: _Task) -> Mapping[str, object]: + resolution: Final = task.resolution.strip().lower() if task.resolution else None + return MappingProxyType( + { + key: value + for key, value in ( + ("duration_seconds", float(task.duration) if task.duration is not None else None), + ("video_resolution", resolution), + ("total_tokens", task.usage.total_tokens if task.usage is not None else None), + ) + if value is not None + } + ) + + +def _task_status(task: _Task, raw_response: httpx.Response) -> str: + status: Final = _STATUS_MAP.get(task.status.strip().lower()) + if status is None: # defaulting it would hand back a job the caller polls forever + raise _video_error(raw_response, 502, f"Unknown Volcengine task status: {task.status}") + return status + + +def _video_object_from_task(task: _Task, raw_response: httpx.Response, model: str | None) -> VideoObject: + status: Final = _task_status(task, raw_response) + usage: Final = _task_usage(task) + seconds: Final = usage.get("duration_seconds") + error: Final[dict[str, object] | None] = ( # mutable-ok: VideoObject.error is a dict field + {key: value for key, value in (("code", task.error.code), ("message", task.error.message)) if value is not None} + if task.error is not None + else None + ) + return VideoObject( + id=task.id, + object="video", + status=status, + created_at=task.created_at, + completed_at=task.updated_at if status in _TERMINAL_STATUSES else None, + error=error, + seconds=_seconds_label(seconds), + size=_size_from_geometry(task.resolution, task.ratio, task.model), + model=model or task.model, + usage=dict(usage) if usage else None, + ) + + +def _video_url_from_task(task: _Task, raw_response: httpx.Response) -> str: + """A download before the video is ready is the caller polling too early, not a fault of the deployment, so it + answers 400, which litellm neither retries nor counts toward a deployment cooldown. A 404 would cool the + deployment down, and a 409 or a 5xx would be retried against a job that is simply not finished.""" + if task.content is not None and task.content.video_url: + return task.content.video_url + status: Final = task.status.strip().lower() + if status in _PENDING_STATUSES: + raise _video_error( + raw_response, 400, f"Video is still processing (status: {status}). Please wait and try again." + ) + if task.error is not None and (task.error.message or task.error.code): + raise _video_error(raw_response, 400, f"Video generation failed: {task.error.message or task.error.code}") + if status == "expired": + raise _video_error(raw_response, 400, "Task expired. Ark keeps task records for 7 days after creation.") + if status in ("failed", "cancelled"): + raise _video_error(raw_response, 400, f"Video generation ended as {status} with no video URL.") + raise _video_error(raw_response, 502, "Video URL not found in task response.") + + +class VolcEngineVideoConfig(BaseVideoConfig): + """Volcengine Ark (Doubao-Seedance) on Ark's async task API: create, poll, download from the signed + ``content.video_url``, and list.""" + + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseVideoConfig contract returns list + return [ + "model", + "prompt", + "input_reference", + "seconds", + "size", + "user", + "extra_headers", + "content", + "parameters", + *_CREATE_PASSTHROUGH_KEYS, + ] + + def map_openai_params( + self, + video_create_optional_params: VideoCreateOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: BaseVideoConfig contract + requested: Final = _without(video_create_optional_params, _additional_drop_keys(video_create_optional_params)) + unsupported: Final = next( + (key for key, value in requested.items() if key in _UNSUPPORTED_BY_VOLCENGINE and value is not None), + None, + ) + if unsupported is not None and not drop_params: + raise _droppable(model, f"'{unsupported}' is not supported by Volcengine video generation") + passthrough: Final = MappingProxyType( + {key: value for key, value in requested.items() if value is not None and key not in _MAPPED_ELSEWHERE} + ) + return { + **_user_fields(requested.get("user"), model, drop_params), + **_size_fields(requested.get("size"), model, drop_params), + **_duration_fields(requested.get("seconds"), model, drop_params), + **passthrough, + **_parameters_fields(requested.get("parameters"), model, drop_params), + } + + def _resolve_api_key(self, api_key: str | None) -> str: + resolved_api_key: Final = api_key or get_secret_str("ARK_API_KEY") or get_secret_str("VOLCENGINE_API_KEY") + if resolved_api_key is None: + raise ValueError( + "Volcengine API key is required. Set ARK_API_KEY or VOLCENGINE_API_KEY environment variable or pass " + "api_key parameter." + ) + return resolved_api_key + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract; handler expects a mutable headers dict + model: str, + api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, + ) -> dict[str, str]: # mutable-ok: BaseVideoConfig contract + resolved_api_key: Final = self._resolve_api_key( + api_key + or (litellm_params.api_key if litellm_params is not None and litellm_params.api_key else None) + or litellm.api_key + ) + headers.update({"Authorization": f"Bearer {resolved_api_key}", "Content-Type": "application/json"}) + return headers + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict[str, object], # mutable-ok: BaseVideoConfig contract + ) -> str: + return _normalized_api_base( + get_volcengine_base_url(api_base or get_secret_str("VOLCENGINE_API_BASE") or get_secret_str("ARK_API_BASE")) + ) + + def transform_video_create_request( + self, + model: str, + prompt: str, + api_base: str, + video_create_optional_request_params: dict[str, object], # mutable-ok: BaseVideoConfig contract + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + ) -> tuple[dict[str, object], RequestFiles, str]: # mutable-ok: BaseVideoConfig contract + """ + Ark bills a video when the task finishes, while litellm bills it here at create, so every request has to + carry a length and a tier whose cost is known now. Ark-picked length (``duration: -1``) and a video input + (edits, extensions and reference videos, billed at a different rate per token) are rejected for that reason, + and an omitted length is sent as OpenAI's documented default of 4 seconds, since the 2.5 family's own default + is ``-1``. + """ + drop_params: Final = _drops_params(litellm_params) + dropped_by_config: Final = _additional_drop_keys(video_create_optional_request_params) + kept: Final = _without(video_create_optional_request_params, dropped_by_config) + params: Final = _without( + { + **{key: value for key, value in kept.items() if key != "parameters"}, + **_parameters_fields(kept.get("parameters"), model, drop_params), + }, + dropped_by_config, + ) + content: Final = _sent_content(prompt, params, model, drop_params) + if params.get("frames") is not None and not drop_params: + raise _droppable(model, "frames is not supported through litellm, pass seconds instead") + requested_duration: Final = params.get("duration") + parsed_duration: Final = None if requested_duration is None else _duration_param(requested_duration) + if requested_duration is not None and (parsed_duration is None or parsed_duration <= 0) and not drop_params: + raise _droppable( + model, + f"duration must be a positive whole number of seconds, got {requested_duration!r}. Ark-picked length " + "(-1) is not supported through litellm, since the video is billed when the task is created, before " + "Ark has picked its length", + ) + duration: Final = ( + parsed_duration + if parsed_duration is not None and parsed_duration > 0 + else VOLCENGINE_VIDEO_DEFAULT_DURATION_SECONDS + ) + + unknown_fields: Final = tuple( + sorted(key for key, value in params.items() if value is not None and key not in _CREATE_HANDLED_KEYS) + ) + if unknown_fields and not drop_params: + raise _droppable(model, f"litellm does not send {', '.join(unknown_fields)} to Ark video generation") + requested_fields: Final = MappingProxyType( + {key: params[key] for key in _CREATE_PASSTHROUGH_KEYS if params.get(key) is not None} + ) + malformed: Final = _malformed_fields(requested_fields) + if malformed and not drop_params: + raise _droppable(model, "; ".join(f"invalid {field}: {message}" for field, message in malformed)) + resolution: Final = requested_fields.get("resolution") + priced_tiers: Final = _priced_tiers(_pricing_model(model, litellm_params)) + drops_resolution: Final = ( + priced_tiers is not None and isinstance(resolution, str) and resolution.strip().lower() not in priced_tiers + ) + if drops_resolution and not drop_params: + raise _droppable( + model, f"{model} is not offered at resolution {resolution!r}, only at {', '.join(priced_tiers or ())}" + ) + dropped: Final = frozenset({*(field for field, _ in malformed), *(("resolution",) if drops_resolution else ())}) + kept_fields: Final = MappingProxyType( + {key: value for key, value in requested_fields.items() if key not in dropped} + ) + default_tier: Final = ( + MappingProxyType({"resolution": VOLCENGINE_VIDEO_DEFAULT_RESOLUTION}) + if "resolution" not in kept_fields + and priced_tiers is not None + and VOLCENGINE_VIDEO_DEFAULT_RESOLUTION in priced_tiers + else _EMPTY_PARAMS + ) + + request_data: Final[dict[str, object]] = { # mutable-ok: request body dict, JSON-serialized by the handler + "model": model, + "content": [dict(item) for item in content], + **default_tier, + **kept_fields, + "duration": duration, + } + return request_data, (), f"{api_base}{VOLCENGINE_VIDEO_TASKS_PATH}" + + def transform_video_create_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "Logging", + custom_llm_provider: str | None = None, + request_data: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract + ) -> VideoObject: + task: Final = _parse(_TASK_CREATE_ADAPTER, raw_response) + if not task.id: + message: Final = task.error.message if task.error is not None and task.error.message else "no task id" + raise _video_error(raw_response, 400, f"Volcengine did not return a video task id: {message}") + requested: Final = request_data or _EMPTY_PARAMS + usage: Final = _requested_usage(requested) + seconds: Final = usage.get("duration_seconds") + requested_model, served_model = routed_video_models(logging_obj) or (model, None) + resolution: Final = requested.get("resolution") + ratio: Final = requested.get("ratio") + return VideoObject( + id=encode_video_id_with_provider(task.id, custom_llm_provider, requested_model, served_model) + if custom_llm_provider + else task.id, + object="video", + status="queued", + model=model, + seconds=_seconds_label(seconds), + size=_size_from_geometry( + resolution if isinstance(resolution, str) else None, ratio if isinstance(ratio, str) else None, model + ), + usage=dict(usage) if usage else None, + ) + + @staticmethod + def _task_url(video_id: str, api_base: str) -> str: + original_task_id: Final = extract_original_video_id(video_id) + encoded_task_id: Final = encode_url_path_segment(original_task_id, field_name="video_id") + return f"{api_base}{VOLCENGINE_VIDEO_TASKS_PATH}/{encoded_task_id}" + + def transform_video_status_retrieve_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract + return self._task_url(video_id, api_base), _EMPTY_PARAMS + + def transform_video_status_retrieve_response( + self, + raw_response: httpx.Response, + logging_obj: "Logging", + custom_llm_provider: str | None = None, + client: "HTTPHandler | None" = None, + ) -> VideoObject: + task: Final = _parse(_TASK_ADAPTER, raw_response) + model, served_model = polled_video_models(logging_obj) + video_obj: Final = _video_object_from_task(task, raw_response, model) + if custom_llm_provider: + video_obj.id = encode_video_id_with_provider(task.id, custom_llm_provider, model, served_model) + return video_obj + + def transform_video_content_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + variant: str | None = None, + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract + return self._task_url(video_id, api_base), _EMPTY_PARAMS + + def transform_video_content_response( + self, + raw_response: httpx.Response, + logging_obj: "Logging", + client: "HTTPHandler | None" = None, + ) -> bytes: + video_url: Final = _video_url_from_task(_parse(_TASK_ADAPTER, raw_response), raw_response) + httpx_client: Final[HTTPHandler] = client if client is not None else _get_httpx_client() + video_response: Final = httpx_client.get(video_url) # pyright: ignore[reportUnknownMemberType] # HTTPHandler.get stub leaves params/headers untyped + video_response.raise_for_status() + return video_response.content + + async def async_transform_video_content_response( + self, + raw_response: httpx.Response, + logging_obj: "Logging", + client: "AsyncHTTPHandler | None" = None, + ) -> bytes: + video_url: Final = _video_url_from_task(_parse(_TASK_ADAPTER, raw_response), raw_response) + async_httpx_client: Final[AsyncHTTPHandler] = ( + client if client is not None else get_async_httpx_client(llm_provider=litellm.LlmProviders.VOLCENGINE) + ) + video_response: Final = await async_httpx_client.get(video_url) # pyright: ignore[reportUnknownMemberType] # HTTPHandler.get stub leaves params/headers untyped + video_response.raise_for_status() + return video_response.content + + def transform_video_remix_request( + self, + video_id: str, + prompt: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + extra_body: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract + raise _unsupported("", _VIDEO_INPUT_UNSUPPORTED.format(operation="Video remix")) + + def transform_video_remix_response( + self, + raw_response: httpx.Response, + logging_obj: "Logging", + custom_llm_provider: str | None = None, + ) -> VideoObject: + raise NotImplementedError("Video remix is not supported by Volcengine.") + + def transform_video_edit_request( + self, + prompt: str, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + video_file: object | None = None, + extra_body: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract + prefetched_source_data: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, Mapping[str, object], RequestFiles | None]: + raise _unsupported("", _VIDEO_INPUT_UNSUPPORTED.format(operation="Video edit")) + + def transform_video_extension_request( + self, + prompt: str, + video_id: str, + seconds: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + extra_body: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract + raise _unsupported("", _VIDEO_INPUT_UNSUPPORTED.format(operation="Video extension")) + + def transform_video_list_request( + self, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + after: str | None = None, + limit: int | None = None, + order: str | None = None, + extra_query: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract + """Ark pages by page_num/page_size with filters and has no cursor or order.""" + unsupported: Final = tuple(name for name, value in (("after", after), ("order", order)) if value is not None) + if unsupported and not _drops_params(litellm_params): + raise _droppable( + "", f"Ark's task list pages by page_num in extra_query and has no {' or '.join(unsupported)}" + ) + return f"{api_base}{VOLCENGINE_VIDEO_TASKS_PATH}", { + **({"page_size": str(limit)} if limit is not None else {}), + **(extra_query or {}), + } + + def transform_video_list_response( # pyright: ignore[reportIncompatibleMethodOverride] # base declares dict[str, str] but the payload is a heterogeneous list body + self, + raw_response: httpx.Response, + logging_obj: "Logging", + custom_llm_provider: str | None = None, + ) -> dict[str, object]: # mutable-ok: OpenAI list body served as JSON by the proxy + list_payload: Final = _parse(_TASK_LIST_ADAPTER, raw_response) + videos: Final = tuple(_video_object_from_task(task, raw_response, None) for task in list_payload.items) + data: Final = [ + video.model_dump() + | ( + {"id": encode_video_id_with_provider(video.id, custom_llm_provider, video.model)} + if custom_llm_provider + else {} + ) + for video in videos + ] + return { + "object": "list", + "data": data, + "total": list_payload.total if list_payload.total is not None else len(data), + "has_more": False, + **({"first_id": data[0]["id"], "last_id": data[-1]["id"]} if data else {}), + } + + def transform_video_delete_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract + raise NotImplementedError("Video delete is not supported by Volcengine.") + + def transform_video_delete_response( + self, + raw_response: httpx.Response, + logging_obj: "Logging", + ) -> VideoObject: + raise NotImplementedError("Video delete is not supported by Volcengine.") + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, str] | httpx.Headers, # mutable-ok: BaseVideoConfig contract + ) -> BaseLLMException: + return VolcEngineVideoError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/watsonx/common_utils.py b/litellm/llms/watsonx/common_utils.py index fb1fa8b9ffb..6d3686e9f85 100644 --- a/litellm/llms/watsonx/common_utils.py +++ b/litellm/llms/watsonx/common_utils.py @@ -141,7 +141,7 @@ async def _aconvert_watsonx_messages_core( model_prompt_dict: Final = custom_prompt_dict[model] return ptf.custom_prompt( messages=messages, - role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles")), + role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles", {})), initial_prompt_value=model_prompt_dict.get("initial_prompt_value", ""), final_prompt_value=model_prompt_dict.get("final_prompt_value", ""), bos_token=model_prompt_dict.get("bos_token", ""), @@ -173,7 +173,7 @@ def _convert_watsonx_messages_core( model_prompt_dict: Final = custom_prompt_dict[model] return ptf.custom_prompt( messages=messages, - role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles")), + role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles", {})), initial_prompt_value=model_prompt_dict.get("initial_prompt_value", ""), final_prompt_value=model_prompt_dict.get("final_prompt_value", ""), bos_token=model_prompt_dict.get("bos_token", ""), 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 d8e62ebcdad..23cd3c40250 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -52,6 +52,7 @@ from litellm.utils import ( get_litellm_params, get_optional_params, peek_reasoning_summary_aliases, + reasoning_effort_with_summary_alias, strip_reasoning_summary_aliases_from_optional_params, ) @@ -254,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 @@ -1211,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, ) @@ -5681,7 +5688,15 @@ def completion( "allowed_openai_params": allowed_openai_params, "base_model": base_model, } - optional_params = get_optional_params(**optional_param_args, **non_default_params) + provider_reasoning_effort: Final = reasoning_effort_with_summary_alias( + reasoning_effort, + non_default_params, + accepts_summary=provider_config is not None + and provider_config.accepts_reasoning_summary_in_reasoning_effort, + ) + optional_params = get_optional_params( + **{**optional_param_args, "reasoning_effort": provider_reasoning_effort}, **non_default_params + ) processed_non_default_params: Final = pre_process_non_default_params( model=model, passed_params=optional_param_args, @@ -5776,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 }, @@ -7851,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) @@ -7884,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 @@ -7911,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 @@ -7932,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), @@ -8165,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, @@ -8179,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, @@ -8501,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 a204c49967e..c3efaf168ab 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -342,26 +342,33 @@ }, "writer.palmyra-vision-7b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 4096, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "output_cost_per_token_flex": 3e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_vision": true }, "amazon.nova-lite-v1:0": { "cache_read_input_token_cost": 1.5e-08, "input_cost_per_token": 6e-08, + "input_cost_per_token_batches": 3e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 300000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 2.4e-07, + "output_cost_per_token_batches": 1.2e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -371,13 +378,19 @@ }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_priority": 1.3125e-07, "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_priority": 5.25e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_priority": 4.375e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -392,13 +405,17 @@ "input_cost_per_token": 1.25e-06, "input_cost_per_image_token": 1.25e-06, "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_token_batches": 6.25e-07, + "input_cost_per_token_priority": 2.1875e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1e-05, - "source": "https://aws.amazon.com/nova/pricing/", + "output_cost_per_token_batches": 5e-06, + "output_cost_per_token_priority": 1.75e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -489,15 +506,21 @@ "us.amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 8.25e-08, "cache_read_input_token_cost_flex": 4.125e-08, + "cache_read_input_token_cost_priority": 1.44375e-07, "input_cost_per_token": 3.3e-07, + "input_cost_per_token_batches": 1.65e-07, "input_cost_per_token_flex": 1.65e-07, + "input_cost_per_token_priority": 5.775e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.75e-06, + "output_cost_per_token_batches": 1.375e-06, "output_cost_per_token_flex": 1.375e-06, + "output_cost_per_token_priority": 4.8125e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -512,13 +535,17 @@ "input_cost_per_token": 1.375e-06, "input_cost_per_image_token": 1.375e-06, "input_cost_per_audio_token": 1.375e-06, + "input_cost_per_token_batches": 6.875e-07, + "input_cost_per_token_priority": 2.40625e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.1e-05, - "source": "https://aws.amazon.com/nova/pricing/", + "output_cost_per_token_batches": 5.5e-06, + "output_cost_per_token_priority": 1.925e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -548,12 +575,15 @@ "amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, "input_cost_per_token": 3.5e-08, + "input_cost_per_token_batches": 1.75e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.4e-07, + "output_cost_per_token_batches": 7e-08, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, @@ -561,13 +591,19 @@ }, "amazon.nova-premier-v1:0": { "cache_read_input_token_cost": 6.25e-07, + "cache_read_input_token_cost_priority": 1.09375e-06, "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, + "input_cost_per_token_priority": 4.375e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.25e-05, + "output_cost_per_token_batches": 6.25e-06, + "output_cost_per_token_priority": 2.1875e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -578,15 +614,21 @@ "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 8e-07, + "input_cost_per_token_batches": 4e-07, "input_cost_per_token_flex": 4e-07, + "input_cost_per_token_priority": 1.4e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 300000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 3.2e-06, + "output_cost_per_token_batches": 1.6e-06, "output_cost_per_token_flex": 1.6e-06, + "output_cost_per_token_priority": 5.6e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -2773,7 +2815,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2808,7 +2850,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2843,7 +2885,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2878,7 +2920,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2913,7 +2955,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2948,7 +2990,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -3743,7 +3785,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": 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, @@ -15494,7 +15536,7 @@ }, "supports_adaptive_thinking": true, "supports_legacy_thinking": true, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -20832,7 +20874,7 @@ "output_cost_per_token": 2.5000010000000002e-05, "output_dbu_cost_per_token": 0.000357143, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_legacy_thinking": true, "supports_prompt_caching": true, @@ -21047,7 +21089,7 @@ "output_cost_per_token": 1.5000020000000002e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_legacy_thinking": 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 @@ -29256,22 +29302,32 @@ "output_cost_per_token_priority": 1.8e-05 }, "gemini-3.1-pro-preview": { - "prompt_cache_min_tokens": 4096, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "output_cost_per_image": 0.00012, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096.0, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29289,33 +29345,18 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_url_context": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 3.6e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "output_cost_per_token_priority": 2.16e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "cache_read_input_token_cost_priority": 3.6e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_flex": 2e-07, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "web_search_billing_unit": "per_query" }, "gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -29428,15 +29469,25 @@ "vertex_ai/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_flex": 5e-08, - "input_cost_per_token": 5e-07, + "cache_read_input_token_cost_priority": 9e-08, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_batches": 5e-07, + "input_cost_per_audio_token_priority": 1.8e-06, + "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": 9e-07, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65535, "max_tokens": 65535, "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.4e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29454,6 +29505,7 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -29463,23 +29515,7 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 9e-07, - "input_cost_per_audio_token_priority": 1.8e-06, - "output_cost_per_token_priority": 5.4e-06, - "cache_read_input_token_cost_priority": 9e-08, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 5e-07, - "input_cost_per_token_batches": 2.5e-07, - "input_cost_per_token_flex": 2.5e-07, - "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -29807,22 +29843,32 @@ "output_cost_per_video_token": 1e-06 }, "vertex_ai/gemini-3.1-pro-preview": { - "prompt_cache_min_tokens": 4096, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "output_cost_per_image": 0.00012, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096.0, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29840,33 +29886,18 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_url_context": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 3.6e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "output_cost_per_token_priority": 2.16e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "cache_read_input_token_cost_priority": 3.6e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_flex": 2e-07, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -31666,9 +31697,15 @@ "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_flex": 5e-08, + "cache_read_input_token_cost_priority": 9e-08, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_batches": 5e-07, + "input_cost_per_audio_token_priority": 1.8e-06, "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": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65535, @@ -31676,6 +31713,10 @@ "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "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.4e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -31693,6 +31734,7 @@ ], "supports_audio_output": false, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -31703,23 +31745,7 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 9e-07, - "input_cost_per_audio_token_priority": 1.8e-06, - "output_cost_per_token_priority": 5.4e-06, - "cache_read_input_token_cost_priority": 9e-08, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 5e-07, - "input_cost_per_token_batches": 2.5e-07, - "input_cost_per_token_flex": 2.5e-07, - "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "web_search_billing_unit": "per_query" }, "gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, @@ -33610,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, @@ -33628,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, @@ -33646,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, @@ -33767,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, @@ -40354,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 }, @@ -40428,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 }, @@ -40839,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, @@ -40858,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", @@ -40884,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 @@ -41072,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, @@ -41108,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, @@ -41126,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, @@ -41144,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, @@ -41197,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, @@ -41241,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 @@ -41257,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 @@ -42076,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, @@ -42093,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 }, @@ -43197,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, @@ -43215,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, @@ -43232,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, @@ -43251,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, @@ -44751,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", @@ -44759,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, @@ -44774,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", @@ -44782,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, @@ -44792,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, @@ -44810,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, @@ -45674,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, @@ -45688,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, @@ -45816,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, @@ -45832,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 @@ -47862,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, @@ -47878,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, @@ -47892,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, @@ -48619,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, @@ -48668,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 }, @@ -48742,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 }, @@ -49196,7 +49365,7 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_prompt_caching": true, @@ -51026,7 +51195,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -51199,7 +51368,7 @@ "supported_regions": [ "global" ], - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -51220,7 +51389,7 @@ "supported_regions": [ "us-central1" ], - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -53450,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 @@ -53488,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 @@ -58842,7 +59023,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -58871,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", @@ -58879,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" @@ -58892,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", @@ -58900,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" @@ -58913,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", @@ -58921,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" ], @@ -58932,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", @@ -58940,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" ], @@ -59573,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, @@ -59612,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, @@ -59695,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, @@ -59734,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, @@ -59819,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, @@ -59854,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, @@ -59892,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, @@ -59927,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, @@ -60022,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", @@ -60030,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", @@ -60043,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", @@ -60051,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", @@ -60064,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", @@ -60072,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", @@ -60088,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", @@ -60097,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": [ @@ -60108,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, @@ -60487,6 +60702,61 @@ } ] }, + "volcengine/doubao-seedance-2-5-260628": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.216, + "output_cost_per_second_480p": 0.096075, + "output_cost_per_second_720p": 0.216, + "output_cost_per_second_1080p": 0.5346, + "supported_modalities": ["text", "image", "video", "audio"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "854x480", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "480x854", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s). List 70 CNY per million tokens at 480p and 720p and 77 at 1080p without a video input, matching Ark's per-resolution table and its 1080p 16:9 example of 3.74 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Accepts up to 50 reference assets (30 images, 10 videos, 10 audios) and clips of 4-30s; per-asset billing beyond the free allowance is not encoded. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, + "volcengine/doubao-seedance-2-0-260128": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.141942857, + "output_cost_per_second_480p": 0.066003429, + "output_cost_per_second_720p": 0.141942857, + "output_cost_per_second_1080p": 0.354085714, + "output_cost_per_second_4k": 0.722057143, + "supported_modalities": ["text", "image", "video", "audio"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 46 CNY per million tokens at 480p and 720p, 51 at 1080p, and 26 at 4k without a video input, matching Ark's per-resolution table and its 1080p and 4k 16:9 examples of 2.48 and 5.05 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s; 4k is exclusive to this tier. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, + "volcengine/doubao-seedance-2-0-fast-260128": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.114171429, + "output_cost_per_second_480p": 0.053089714, + "output_cost_per_second_720p": 0.114171429, + "supported_modalities": ["text", "image", "video", "audio"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 37 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, + "volcengine/doubao-seedance-2-0-mini-260615": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.070971429, + "output_cost_per_second_480p": 0.033001714, + "output_cost_per_second_720p": 0.070971429, + "supported_modalities": ["text", "image"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 23 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Only text and image inputs are declared; the mini tier's reference video and audio caps were not verified. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, "bedrock/us-east-1/zai.glm-5": { "input_cost_per_token": 1e-06, "output_cost_per_token": 3.2e-06, @@ -67489,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", @@ -67497,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", @@ -67515,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", @@ -67533,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", @@ -67551,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", @@ -67570,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", @@ -67588,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", @@ -67605,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", @@ -67623,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, @@ -72354,7 +72656,7 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -75183,5 +75485,45 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "together_ai/reflection/Beam-1.0-501B-A23B": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 4e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.ai/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "baseten/reflection/Beam-1.0": { + "cache_read_input_token_cost": 1.3e-07, + "input_cost_per_token": 1.3e-06, + "litellm_provider": "baseten", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.3e-06, + "source": "https://inference.baseten.co/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "nebius/reflection/beam": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "nebius", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://tokenfactory.nebius.com/models/catalog/text2text/reflection%2Fbeam", + "supports_function_calling": true, + "supports_reasoning": true } } diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 97188217f63..420f42c690c 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio import contextvars +import weakref from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Generator, Iterator from functools import partial from types import TracebackType @@ -26,6 +27,23 @@ from litellm.utils import client base_llm_http_handler = BaseLLMHTTPHandler() from .utils import BasePassthroughUtils +_provider_errors: Final[weakref.WeakSet[httpx.HTTPStatusError]] = weakref.WeakSet() + + +def raise_for_provider_status(response: httpx.Response) -> None: + """``response.raise_for_status()`` for the provider's response to a passthrough request, recording the raised + error as the provider's own.""" + try: + response.raise_for_status() + except httpx.HTTPStatusError as error: + _provider_errors.add(error) + raise + + +def is_provider_error(error: httpx.HTTPStatusError) -> bool: + """Whether ``error`` is the provider's own error response to a passthrough request.""" + return error in _provider_errors + async def _as_async_generator(iterable: AsyncIterator[bytes]) -> AsyncGenerator[bytes, bytes]: async for chunk in iterable: @@ -112,7 +130,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): self.status_code = self._response.status_code self._initialized = True try: - self._response.raise_for_status() + raise_for_provider_status(self._response) self._iterator = _as_async_generator(self._response.aiter_bytes()) except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic try: @@ -626,7 +644,7 @@ async def _async_passthrough_request( else: response: Final = await response_result await response.aread() - response.raise_for_status() + raise_for_provider_status(response) return response else: # Fallback for sync-like behavior (shouldn't happen in async path) 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..224f956bcfd --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py @@ -0,0 +1,216 @@ +"""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. + +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 +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_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, +) -> SSOIdentityAssertion: + """The assertion as captured when no server wants a refresh token, otherwise the assertion plus the + refresh token the first target's IdP client exchanged it for. An OIDC login's own SSO refresh token + does not count, it is bound to the SSO client and the ID-JAG client cannot redeem it.""" + if not targets or assertion.exchanged_refresh_token is not None: + return assertion + 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, + ) + + +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) + ) + 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..bd0813dd966 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,41 @@ 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 exchanged the login assertion for when the user signed in. Never the SSO + refresh token an OIDC login holds, that one is bound to the SSO client. The assertion's own expiry + does not apply here, 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..58494654082 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,85 @@ 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.""" + return renewed.model_copy( + update={ + "exchanged_refresh_token": 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 +313,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 +372,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 +417,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 +441,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 e4a0a5242b4..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", @@ -3814,6 +3815,16 @@ class ModelAccessDeniedProxyException(ProxyException): return self.internal_message.replace("\r", "").replace("\n", "") +class AwsErrorTypeProxyException(Exception): + """A ProxyException answered to an AWS SDK, which reads the error code from the ``x-amzn-errortype`` header and + the message from a top-level ``message`` field.""" + + def __init__(self, error: ProxyException, aws_error_type: str) -> None: + super().__init__(error.message) + self.error: Final = error + self.aws_error_type: Final = aws_error_type + + from litellm.types.proxy.common_proxy_errors import ( # noqa: E402 # public re-export CommonProxyErrors as CommonProxyErrors, # noqa: PLC0414 # public re-export ) @@ -4752,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. @@ -4932,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 1f17f0a42d2..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, ) @@ -29,6 +32,8 @@ from litellm.proxy.common_request_processing import ( create_response, 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 @@ -43,7 +48,8 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, with_litellm_call_id, ) -from litellm.types.utils import TokenCountResponse +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: @@ -340,51 +391,36 @@ async def count_tokens( messages: Final = data.get("messages", []) if not model_name: - raise HTTPException(status_code=400, detail={"error": "model parameter is required"}) + raise HTTPException(status_code=400, detail="model parameter is required") if not messages: - raise HTTPException(status_code=400, detail={"error": "messages parameter is required"}) + 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), ) - _token_response_dict: dict = {} - if isinstance(token_response, TokenCountResponse): - _token_response_dict = token_response.model_dump() - elif isinstance(token_response, dict): - _token_response_dict = token_response - - # Convert the internal response to Anthropic API format - return {"input_tokens": _token_response_dict.get("total_tokens", 0)} - - except HTTPException: - raise + except HTTPException as e: + return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) except ProxyException as e: - status_code: Final = int(e.code) if e.code and e.code.isdigit() else 500 - detail: Final = AnthropicExceptionMapping.transform_to_anthropic_error( - status_code=status_code, - raw_message=e.message, - ) - raise HTTPException( - status_code=status_code, - detail=detail, - ) + return _anthropic_error_json_response(e, request) except Exception as e: log_llm_api_exception(e, litellm_call_id) - raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"}) + return _anthropic_error_json_response(proxy_exception_from_route_error(e), request) + else: + return {"input_tokens": token_response.total_tokens} @router.post( 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 95d1123bfc9..45c98a6c0db 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, @@ -130,6 +131,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.router_utils.endpoint_eligibility import UnsupportedEndpointError from litellm.types.guardrails import GuardrailEventHooks @@ -506,6 +508,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: @@ -635,6 +641,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) @@ -1970,10 +2033,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/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 61330d12ac3..ccf2232e379 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -6,6 +6,7 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from enum import Enum +from itertools import chain from types import MappingProxyType from typing import Final, Generic, Literal, Protocol, TypeVar @@ -65,6 +66,7 @@ from litellm.repositories.team_repository import TeamRepository from litellm.repositories.unit_of_work import ( LinkedSpendResetWrites, budget_cascade_unit_of_work, + enduser_spend_reset_unit_of_work, spend_reset_unit_of_work, ) from litellm.repositories.verification_token_repository import ( @@ -222,8 +224,8 @@ def _budget_link_where( def _enduser_invalidation_where(budget_ids: Sequence[str]) -> dict[str, object]: """Customers whose cached spend a committed reset of these tiers invalidated. - Mirrors ``_queue_enduser_resets`` without its ``spend > 0`` filter, which - post-commit would match nobody. + ``_enduser_reset_where`` adds the ``spend > 0`` filter, which post-commit + would match nobody. """ linked: Final = _budget_link_where(budget_ids) default_budget_id: Final = litellm.max_end_user_budget_id @@ -251,26 +253,44 @@ def _queue_budget_linked_resets( writes.queue_spend_zero(where=_budget_link_where(plain_ids, extra)) -def _queue_enduser_resets(writes: LinkedSpendResetWrites, cascade: "_BudgetCascade") -> None: - """End users reset on the budget link like every other gated table, plus a - NULL-budget_id branch: rows created implicitly persist no link and ride the - default tier (litellm.max_end_user_budget_id). +def _enduser_reset_where(cascade: "_BudgetCascade") -> Mapping[str, object]: + """End users with spend to reset: linked to an expiring tier, or carrying + no link and riding the default tier (litellm.max_end_user_budget_id).""" + return {**_enduser_invalidation_where(cascade.budget_ids), **_SPENT_ROWS_WHERE} - Matching on the link rather than enumerating user ids keeps a statement's - bind count proportional to the expiring tiers instead of the customer - population, which past ~32,700 dependents exceeds PostgreSQL's per-statement - bind ceiling and wedges the cascade permanently (#40564). + +def _enduser_tier(row: _EndUserRow) -> str | None: + return row.budget_id if row.budget_id is not None else litellm.max_end_user_budget_id + + +def _enduser_ids_by_tier(rows: Sequence[_EndUserRow]) -> Mapping[str | None, tuple[str, ...]]: + tiers: Final = tuple(dict.fromkeys(_enduser_tier(row) for row in rows)) + return MappingProxyType({tier: tuple(row.user_id for row in rows if _enduser_tier(row) == tier) for tier in tiers}) + + +def _queue_enduser_page_resets( + writes: LinkedSpendResetWrites, rows: Sequence[_EndUserRow], cascade: "_BudgetCascade" +) -> None: + """The statements of ``_queue_budget_linked_resets`` for one page of end + users, keyed by their user ids so each statement touches at most one page. + + A statement matching the budget link reads the whole end-user table (no + index on budget_id) and updates every linked row at once, which past a few + million end users outlives the prisma client's 30s request timeout (#46404). + A page of ids stays under PostgreSQL's per-statement bind ceiling (#40564) + and the spend predicates still classify each row at write time. """ - _queue_budget_linked_resets(writes, cascade, extra=_SPENT_ROWS_WHERE) - default_budget_id: Final = litellm.max_end_user_budget_id - if default_budget_id is None or default_budget_id not in cascade.budget_ids: - return - cap: Final = cascade.rollover_caps.get(default_budget_id) + for tier, user_ids in _enduser_ids_by_tier(rows).items(): + _queue_enduser_tier_resets(writes, user_ids, cap=cascade.rollover_caps.get(tier) if tier is not None else None) + + +def _queue_enduser_tier_resets(writes: LinkedSpendResetWrites, user_ids: Sequence[str], cap: float | None) -> None: + page: Final = {"user_id": {"in": list(user_ids)}} # bounded-ok: one page of RESET_BUDGET_JOB_BATCH_SIZE ids if cap is None: - writes.queue_spend_zero(where={"budget_id": None, **_SPENT_ROWS_WHERE}) + writes.queue_spend_zero(where={**page, **_SPENT_ROWS_WHERE}) return - writes.queue_spend_zero(where={"budget_id": None, "spend": {"gt": 0, "lte": cap}}) - writes.queue_spend_decrement(where={"budget_id": None, "spend": {"gt": cap}}, amount=cap) + writes.queue_spend_zero(where={**page, "spend": {"gt": 0, "lte": cap}}) + writes.queue_spend_decrement(where={**page, "spend": {"gt": cap}}, amount=cap) @dataclass(frozen=True, slots=True) @@ -824,6 +844,49 @@ class ResetBudgetJob: ), ) + async def _reset_enduser_spend(self, cascade: _BudgetCascade) -> int: + """Reset the end users the expiring tiers gate, one committed page at a time. + + Runs before ``_commit_budget_cascade`` so a failure here leaves the + tiers due: the next tick resumes with the pages still carrying spend. + Each page invalidates its own caches right after its commit, so no + zeroed row sits behind a cached spend that still rejects requests. + Capless like ``_invalidate_enduser_caches``, for its reason. + """ + if not cascade.budget_ids: + return 0 + where: Final = _enduser_reset_where(cascade) + walk = _EndUserWalk() # rebind-ok: keyset cursor advances one committed page at a time + while walk.cursor is not None: + walk = await self._reset_enduser_page(cascade, where=where, cursor=walk.cursor, reached=walk.invalidated) + return walk.invalidated + + async def _reset_enduser_page( + self, cascade: _BudgetCascade, where: Mapping[str, object], cursor: str, reached: int + ) -> _EndUserWalk: + rows: Final = await self._fetch_enduser_page(where=where, cursor=cursor) + if not rows: + return _EndUserWalk(cursor=None, invalidated=reached) + await self._with_db_write_retry( + lambda: self._commit_enduser_page(rows, cascade), + reason="reset_budget_write_endusers_failure", + ) + await self._invalidate_caches( + counter_keys=tuple(_enduser_counter_key(row) for row in rows), + cache_keys=tuple(chain.from_iterable(_enduser_cache_keys(row) for row in rows)), + ) + walked: Final = reached + len(rows) + if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE: + return _EndUserWalk(cursor=None, invalidated=walked) + return _EndUserWalk(cursor=rows[-1].user_id, invalidated=walked) + + async def _commit_enduser_page(self, rows: Sequence[_EndUserRow], cascade: _BudgetCascade) -> None: + async with ( + db_span("reset_budget_endusers", "LiteLLM_EndUserTable"), + enduser_spend_reset_unit_of_work(self._new_batch) as writes, + ): + _queue_enduser_page_resets(writes, rows, cascade) + async def _commit_budget_cascade(self, cascade: _BudgetCascade) -> None: """Zero the gated spend and advance ``budget_reset_at`` in one transaction. @@ -851,7 +914,6 @@ class ResetBudgetJob: _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE) - _queue_enduser_resets(uow.endusers, cascade) for budget_id, budget_reset_at in cascade.budget_resets: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) @@ -878,6 +940,7 @@ class ResetBudgetJob: return _BudgetCascadeFailed(cascade=_EMPTY_CASCADE, error=e) try: + await self._reset_enduser_spend(cascade) await self._commit_budget_cascade(cascade) except Exception as e: return _BudgetCascadeFailed(cascade=cascade, error=e) @@ -896,9 +959,10 @@ class ResetBudgetJob: """ Resets the spend a budget tier gates (end users, team members, keys, orgs, tags, model access groups) and advances the tier's - budget_reset_at, atomically. + budget_reset_at. End users commit in pages first, everything else + lands with the window advance in one transaction. - Caches are invalidated only after the transaction commits, so a failed + Caches are invalidated only after the commit they follow, so a failed run cannot leave a zeroed counter in front of an un-reset DB row. """ await _run_phase_in_chunks(self._reset_budget_for_litellm_budget_table_chunk) @@ -929,8 +993,8 @@ class ResetBudgetJob: case _BudgetCascadeFailed(cascade=cascade, error=error): verbose_proxy_logger.exception( "Failed to reset the budget table cascade (team member, enduser, org, tag and model access " - "group spend, plus budget_reset_at); nothing was committed and the budgets stay due for the " - "next run: %s", + "group spend, plus budget_reset_at); the budgets stay due for the next run, which resumes " + "with the end-user pages not committed yet: %s", error, exc_info=error, ) 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/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 784e3d4719f..d382ff53ffb 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -1,16 +1,22 @@ from collections.abc import Mapping from itertools import chain +from types import MappingProxyType from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response -from fastapi.responses import ORJSONResponse +from fastapi.responses import JSONResponse, ORJSONResponse from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types from litellm.proxy._types import * -from litellm.proxy._types import TokenCountRequest +from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + log_llm_api_exception, + proxy_exception_from_http_exception, + resolve_litellm_call_id, +) 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 read_request_body, @@ -63,6 +69,37 @@ def _input_schema(declaration: _GoogleCountTokensFunctionDeclaration) -> object: return {"type": "object", "properties": {}} +_GOOGLE_RPC_STATUS: Final = MappingProxyType( + { + 400: "INVALID_ARGUMENT", + 401: "UNAUTHENTICATED", + 403: "PERMISSION_DENIED", + 404: "NOT_FOUND", + 409: "ABORTED", + 429: "RESOURCE_EXHAUSTED", + 499: "CANCELLED", + 501: "UNIMPLEMENTED", + 503: "UNAVAILABLE", + 504: "DEADLINE_EXCEEDED", + } +) + + +def _google_error_response(request: Request, error: ProxyException) -> JSONResponse: + from litellm.proxy.proxy_server import ( + _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does + ) + + status_code: Final = int(error.code) if error.code is not None and error.code.isdigit() else 500 + _close_dangling_otel_server_span(request, status_code, exc=error) + status: Final = _GOOGLE_RPC_STATUS.get(status_code, "INTERNAL" if status_code >= 500 else "FAILED_PRECONDITION") + return JSONResponse( + status_code=status_code, + content={"error": {"code": status_code, "message": error.message, "status": status}}, + headers=error.headers, + ) + + def _google_count_tokens_request(model_name: str, data: object) -> TokenCountRequest: from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter @@ -251,16 +288,26 @@ async def google_count_tokens(request: Request, model_name: str): data: Final = await read_request_body(request=request) try: token_request: Final = _google_count_tokens_request(model_name, data) - except ValidationError as error: - raise HTTPException(status_code=400, detail="Invalid Google countTokens request body") from error + except ValidationError: + return _google_error_response( + request, + ProxyException( + message="Invalid Google countTokens request body", type="invalid_request_error", param=None, code=400 + ), + ) - # Call the internal token counter function with direct request flag set to False - token_response: Final = await internal_token_counter( - request=token_request, - call_endpoint=True, - ) - if token_response is None: - return TokenCountDetailsResponse(totalTokens=0, promptTokensDetails=[]) + try: + token_response: Final = await internal_token_counter(request=token_request, call_endpoint=True) + except ProxyException as error: + return _google_error_response(request, error) + except HTTPException as error: + return _google_error_response(request, proxy_exception_from_http_exception(error, {})) + except Exception as error: + log_llm_api_exception(error, resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))) + return _google_error_response( + request, + ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500), + ) original_response: Final[dict] = token_response.original_response or {} return TokenCountDetailsResponse( totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0, diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index d713a18f470..1d5c530a616 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -1483,6 +1483,17 @@ async def get_guardrail_info(guardrail_id: str): raise HTTPException(status_code=500, detail=str(e)) +def _undeclared_logging_only( + guardrail_class: type[CustomGuardrail], hooks: Sequence[GuardrailEventHooks] +) -> tuple[GuardrailEventHooks, ...]: + """logging_only when the class accepts it at save time without declaring it in its hooks.""" + return ( + (GuardrailEventHooks.logging_only,) + if guardrail_class.accepts_logging_only_mode() and GuardrailEventHooks.logging_only not in hooks + else () + ) + + @router.get( "/guardrails/ui/add_guardrail_settings", tags=["Guardrails"], @@ -1515,7 +1526,7 @@ async def get_guardrail_ui_settings() -> GuardrailUIAddGuardrailSettings: ] supported_modes_by_provider: Final = { - provider: [hook.value for hook in hooks] + provider: [hook.value for hook in (*hooks, *_undeclared_logging_only(guardrail_class, hooks))] for provider, guardrail_class in guardrail_class_registry.items() if (hooks := guardrail_class.get_supported_event_hooks()) is not None } | MappingProxyType( 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/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index f959292d567..40e44151e65 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -6,6 +6,7 @@ to detect and block/mask sensitive content. """ import asyncio +import contextvars import itertools import json import os @@ -50,6 +51,7 @@ from litellm.types.utils import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.types.guardrails import ( BlockedWord, @@ -101,6 +103,13 @@ WORD_NUMBER_SEQUENCE_PATTERN: Final = re.compile( ) WORD_NUMBER_TOKEN_FINDER: Final = re.compile(rf"(?:{WORD_NUMBER_TOKEN_REGEX})", re.IGNORECASE) +_logging_only_scan_guardrail: Final[contextvars.ContextVar["ContentFilterGuardrail | None"]] = contextvars.ContextVar( + "litellm_content_filter_logging_only_scan", default=None +) +_logging_only_scan_would_reject: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar( + "litellm_content_filter_logging_only_scan_would_reject", default=False +) + class ConditionalCategoryConfig(TypedDict): identifier_words: Sequence[str] @@ -1209,6 +1218,18 @@ class ContentFilterGuardrail(CustomGuardrail): return (keyword, action, description) return None + def _is_logging_only_scan(self) -> bool: + return _logging_only_scan_guardrail.get() is self + + def _enforces_block(self, action: ContentFilterAction) -> bool: + return action == ContentFilterAction.BLOCK and not self._is_logging_only_scan() + + @staticmethod + def _asks_to_block(detection: ContentFilterDetection) -> bool: + if detection["type"] == "competitor_intent": + return detection["action_hint"] in ("refuse", "reframe") + return detection["action"] == ContentFilterAction.BLOCK.value + def _handle_conditional_match( self, matched_phrase: str, @@ -1228,7 +1249,7 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(category_detection) - if action == ContentFilterAction.BLOCK: + if self._enforces_block(action): error_msg: Final = ( f"Content blocked: {category_name} conditional match '{matched_phrase}' detected (severity: {severity})" ) @@ -1269,7 +1290,7 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(category_detection) - if action == ContentFilterAction.BLOCK: + if self._enforces_block(action): error_msg = f"Content blocked: {category_name} category keyword '{keyword}' detected (severity: {severity})" verbose_proxy_logger.warning(error_msg) raise HTTPException( @@ -1312,7 +1333,7 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(pattern_detection) - if action == ContentFilterAction.BLOCK: + if self._enforces_block(action): error_msg: Final = f"Content blocked: {pattern_name} pattern detected" verbose_proxy_logger.warning(error_msg) raise HTTPException( @@ -1346,7 +1367,7 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(blocked_word_detection) - if action == ContentFilterAction.BLOCK: + if self._enforces_block(action): error_msg = f"Content blocked: keyword '{keyword}' detected" if description: error_msg += f" ({description})" @@ -1619,7 +1640,7 @@ class ContentFilterGuardrail(CustomGuardrail): - Number of entities masked - Whether the guardrail blocked the request (max risk) """ - if status == "guardrail_intervened": + if status in ("guardrail_intervened", "guardrail_flagged"): return 10.0 total_masked: Final = sum(masked_entity_count.values()) if masked_entity_count else 0 @@ -1665,6 +1686,8 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(detection) + if self._is_logging_only_scan(): + return if action_hint_val == "refuse": msg = "Content blocked: competitor comparison or ranking intent detected." if self._competitor_intent_checker and getattr( @@ -1763,18 +1786,18 @@ class ContentFilterGuardrail(CustomGuardrail): return raw_name return None + def _reject(self, error: str) -> None: + """Raise the 400 an enforcing mode returns, or note it on the leg a logging_only scan is observing.""" + if self._is_logging_only_scan(): + _logging_only_scan_would_reject.set(True) + return + raise HTTPException(status_code=400, detail={"error": error}) + def _assert_argument_label_clean( self, text: str, detections: list[ContentFilterDetection], context_label: str ) -> None: if self._filter_single_text(text, detections=detections) != text: - raise HTTPException( - status_code=400, - detail={ - "error": ( - f"Content blocked: {context_label} argument matched a masking rule on a non-rewritable field" - ) - }, - ) + self._reject(f"Content blocked: {context_label} argument matched a masking rule on a non-rewritable field") def _filter_argument_value( self, @@ -1784,10 +1807,8 @@ class ContentFilterGuardrail(CustomGuardrail): depth: int = 0, ) -> object: if depth > DEFAULT_MAX_RECURSE_DEPTH: - raise HTTPException( - status_code=400, - detail={"error": f"Content blocked: {context_label} arguments exceed the maximum nesting depth"}, - ) + self._reject(f"Content blocked: {context_label} arguments exceed the maximum nesting depth") + return value if isinstance(value, str): return self._filter_single_text(value, detections=detections) if isinstance(value, (int, float)) and not isinstance(value, bool): @@ -1874,6 +1895,21 @@ class ContentFilterGuardrail(CustomGuardrail): if filtered_arguments != arguments: self._set_tool_call_arguments(tool_call, filtered_arguments) + async def _scan_logged_call( + self, + kwargs: dict[str, object], # mutable-ok: CustomLogger.async_logging_hook contract + response: object | None, + translation: "BaseTranslation", + output_translation: "BaseTranslation", + scratch_metadata: dict[str, object], # mutable-ok: apply_guardrail records its verdict into request metadata + ) -> None: + """Scan the logged call with every BLOCK match recorded instead of raised.""" + token: Final = _logging_only_scan_guardrail.set(self) + try: + await super()._scan_logged_call(kwargs, response, translation, output_translation, scratch_metadata) + finally: + _logging_only_scan_guardrail.reset(token) + async def apply_guardrail( self, inputs: "GenericGuardrailAPIInputs", @@ -1904,6 +1940,7 @@ class ContentFilterGuardrail(CustomGuardrail): masked_entity_count: Final[dict[str, int]] = {} status: GuardrailStatus = "success" exception_str: str = "" + would_reject_token: Final = _logging_only_scan_would_reject.set(False) try: texts: Final = inputs.get("texts", []) @@ -1947,11 +1984,13 @@ class ContentFilterGuardrail(CustomGuardrail): exception_str = str(e) raise e finally: + would_reject: Final = _logging_only_scan_would_reject.get() or any(map(self._asks_to_block, detections)) + _logging_only_scan_would_reject.reset(would_reject_token) # Log guardrail information self._log_guardrail_information( request_data=request_data, detections=detections, - status=status, + status="guardrail_flagged" if status == "success" and would_reject else status, start_time=start_time, masked_entity_count=masked_entity_count, exception_str=exception_str, 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/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index ef93809d47b..ef60a800cdd 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -57,13 +57,6 @@ class KeyManagementEventHooks: ) from litellm.proxy.proxy_server import litellm_proxy_admin_name - # Send email notification - non-blocking, independent operation - if data.send_invite_email is True: - try: - await KeyManagementEventHooks._send_key_created_email(response.model_dump(exclude_none=True)) - except Exception as e: - verbose_proxy_logger.warning("Failed to send key created email: %s", e) - if is_audit_logging_enabled(): _updated_values: Final = response.model_dump_json(exclude_none=True) asyncio.create_task( @@ -86,6 +79,13 @@ class KeyManagementEventHooks: ) ) + # Send email notification - non-blocking, independent operation + if data.send_invite_email is True: + try: + await KeyManagementEventHooks._send_key_created_email(response.model_dump(exclude_none=True)) + except Exception as e: + verbose_proxy_logger.warning("Failed to send key created email: %s", e) + # Store the generated key in the secret manager - non-blocking, independent operation try: await KeyManagementEventHooks._store_virtual_key_in_secret_manager( @@ -118,16 +118,6 @@ class KeyManagementEventHooks: ) from litellm.proxy.proxy_server import litellm_proxy_admin_name - if data.key_alias is not None and data.key_alias != existing_key_row.key_alias: - try: - await KeyManagementEventHooks._rename_virtual_key_in_secret_manager( - current_secret_name=existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}", - new_secret_name=data.key_alias, - team_id=existing_key_row.team_id, - ) - except Exception as e: - verbose_proxy_logger.warning("Failed to rename virtual key in secret manager: %s", e) - if is_audit_logging_enabled(): updated_fields: Final = { **data.model_dump(exclude_none=True), @@ -156,6 +146,16 @@ class KeyManagementEventHooks: ) asyncio.create_task(create_audit_log_for_update(request_data=request_data)) + if data.key_alias is not None and data.key_alias != existing_key_row.key_alias: + try: + await KeyManagementEventHooks._rename_virtual_key_in_secret_manager( + current_secret_name=existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}", + new_secret_name=data.key_alias, + team_id=existing_key_row.team_id, + ) + except Exception as e: + verbose_proxy_logger.warning("Failed to rename virtual key in secret manager: %s", e) + @staticmethod async def async_key_rotated_hook( data: RegenerateKeyRequest | None, @@ -171,6 +171,28 @@ class KeyManagementEventHooks: ) from litellm.proxy.proxy_server import litellm_proxy_admin_name + # store the audit log + if is_audit_logging_enabled() and existing_key_row.token is not None: + asyncio.create_task( + create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.token, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=existing_key_row.token, + action="rotated", + updated_values=response.model_dump_json(exclude_none=True), + before_value=existing_key_row.model_dump_json(exclude_none=True), + ) + ) + ) + # Store the generated key in the secret manager - non-blocking, independent operation if response.token_id is not None: try: @@ -204,28 +226,6 @@ class KeyManagementEventHooks: except Exception as e: verbose_proxy_logger.warning("Failed to send key rotated email: %s", e) - # store the audit log - if is_audit_logging_enabled() and existing_key_row.token is not None: - asyncio.create_task( - create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ), - changed_by_api_key=user_api_key_dict.token, - table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=existing_key_row.token, - action="rotated", - updated_values=response.model_dump_json(exclude_none=True), - before_value=existing_key_row.model_dump_json(exclude_none=True), - ) - ) - ) - @staticmethod async def async_key_deleted_hook( data: KeyRequest, diff --git a/litellm/proxy/hooks/user_management_event_hooks.py b/litellm/proxy/hooks/user_management_event_hooks.py index 6d978929c05..6618af70c0a 100644 --- a/litellm/proxy/hooks/user_management_event_hooks.py +++ b/litellm/proxy/hooks/user_management_event_hooks.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( from litellm.proxy.management_helpers.audit_logs import ( create_audit_log_for_update, is_audit_logging_enabled, + track_audit_task, ) from litellm.repositories.user_repository import UserRepository @@ -41,7 +42,18 @@ class UserManagementEventHooks: - Creating an audit log for the user creation - Sending a user invitation email to the user """ - from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + if is_audit_logging_enabled(): + track_audit_task( + asyncio.create_task( + UserManagementEventHooks._create_user_created_audit_log( + response=response, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) + ) + ) ######################################################### ########## Send User Invitation Email ################ @@ -52,9 +64,14 @@ class UserManagementEventHooks: user_api_key_dict=user_api_key_dict, ) - ######################################################### - ########## CREATE AUDIT LOG ################ - ######################################################### + @staticmethod + async def _create_user_created_audit_log( + response: NewUserResponse, + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str | None, + ) -> None: + from litellm.proxy.proxy_server import prisma_client + try: if prisma_client is None: raise Exception(CommonProxyErrors.db_not_connected_error.value) @@ -63,16 +80,14 @@ class UserManagementEventHooks: user_row: Final = await UserRepository(prisma_client).find_by_id(response.user_id) if user_row is None: raise Exception(f"no user row found for user_id={response.user_id}") - asyncio.create_task( - UserManagementEventHooks.create_internal_user_audit_log( - user_id=user_row.user_id, - action="created", - litellm_changed_by=user_api_key_dict.user_id, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - before_value=None, - after_value=user_row.model_dump_json(exclude_none=True), - ) + await UserManagementEventHooks.create_internal_user_audit_log( + user_id=user_row.user_id, + action="created", + litellm_changed_by=user_api_key_dict.user_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + before_value=None, + after_value=user_row.model_dump_json(exclude_none=True), ) except Exception as e: verbose_proxy_logger.warning("Unable to create audit log for user on `/user/new` - %s", e) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 6bd31a42e9e..3ad5caade1b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -27,6 +27,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm._uuid import uuid from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, @@ -187,6 +188,30 @@ def _caller_trace_field(data: Mapping[str, object], metadata_variable_name: str, return requester_value if isinstance(requester_value, str) and requester_value else None +def _caller_body_session_id(data: Mapping[str, object], metadata_variable_name: str) -> str | None: + top_level: Final = data.get("litellm_session_id") + if isinstance(top_level, str) and top_level: + return top_level + return _caller_trace_field(data, metadata_variable_name, "session_id") + + +def _caller_session_beside_trace_header( + normalized_headers: Mapping[str, str], data: Mapping[str, object], metadata_variable_name: str +) -> str | None: + session_headers: Final = {key: value for key, value in normalized_headers.items() if key != "x-litellm-trace-id"} + return get_chain_id_from_headers(session_headers) or _caller_body_session_id(data, metadata_variable_name) + + +def _drop_client_caller_session(request_body: Mapping[str, object]) -> None: + for metadata_key in ("metadata", "litellm_metadata"): + if isinstance(metadata := request_body.get(metadata_key), dict): + metadata.pop(CALLER_SESSION_ID_METADATA_KEY, None) + + +def _requester_metadata_snapshot(metadata: Mapping[str, object]) -> Mapping[str, object]: + return copy.deepcopy({key: value for key, value in metadata.items() if key != CALLER_SESSION_ID_METADATA_KEY}) + + def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None: """Only proxy-validated keys are stamped, proven by the unforgeable via_virtual_key marker AND a known non-secret shape: the sha256 hex digest @@ -920,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 @@ -930,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 @@ -1623,6 +1650,8 @@ class LiteLLMProxyRequestSetup: from litellm.proxy._types import LitellmMetadataFromRequestHeaders metadata_from_headers: Final = LitellmMetadataFromRequestHeaders() + request_body: Final = cast(Mapping[str, object], data) # cast-ok: request body is a str-keyed JSON object + _drop_client_caller_session(request_body) spend_logs_metadata: Final = LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers(headers) if spend_logs_metadata is not None: metadata_from_headers["spend_logs_metadata"] = spend_logs_metadata @@ -1631,6 +1660,9 @@ class LiteLLMProxyRequestSetup: # Finally update the requests metadata with the `metadata_from_headers` ######################################################################################### + normalized_headers: Final = MappingProxyType( + {k.lower(): v for k, v in headers.items() if isinstance(k, str) and isinstance(v, str)} + ) agent_id_from_header: Final = headers.get("x-litellm-agent-id") # Explicit litellm headers take precedence; fall back to any x-*-session-id header. chain_id: Final = get_chain_id_from_headers(dict(headers)) @@ -1639,6 +1671,13 @@ class LiteLLMProxyRequestSetup: metadata_from_headers["agent_id"] = agent_id_from_header verbose_proxy_logger.debug("Extracted agent_id from header: %s", agent_id_from_header) + spend_log_session_id: Final = ( + _caller_session_beside_trace_header(normalized_headers, request_body, _metadata_variable_name) + if litellm.request_correlation_in_logs + and chain_id is not None + and chain_id == normalized_headers.get("x-litellm-trace-id") + else None + ) if chain_id: metadata_from_headers["trace_id"] = chain_id metadata_from_headers["session_id"] = chain_id @@ -1663,7 +1702,6 @@ class LiteLLMProxyRequestSetup: # - but lets a caller's existing traceparent/baggage headers (from # real OTel instrumentation) correlate with litellm's own logs instead # of generating an unrelated trace_id. - normalized_headers: Final = MappingProxyType({k.lower(): v for k, v in headers.items() if isinstance(k, str)}) if "litellm_trace_id" not in data: traceparent: Final = normalized_headers.get("traceparent") if isinstance(traceparent, str): @@ -1689,11 +1727,7 @@ class LiteLLMProxyRequestSetup: baggage: Final = normalized_headers.get("baggage") if isinstance(baggage, str): session_id_from_baggage: Final = _session_id_from_baggage(baggage) - caller_session_id: Final = _caller_trace_field( - cast(Mapping[str, object], data), # cast-ok: request body is a str-keyed JSON object - _metadata_variable_name, - "session_id", - ) + caller_session_id: Final = _caller_trace_field(request_body, _metadata_variable_name, "session_id") if session_id_from_baggage and ( caller_session_id is None or caller_session_id == session_id_from_baggage ): @@ -1702,7 +1736,11 @@ class LiteLLMProxyRequestSetup: verbose_proxy_logger.debug("Extracted session_id from W3C baggage header") if isinstance(data[_metadata_variable_name], dict): - data[_metadata_variable_name].update(metadata_from_headers) + data[_metadata_variable_name].update( + {**metadata_from_headers, CALLER_SESSION_ID_METADATA_KEY: spend_log_session_id} + if spend_log_session_id + else metadata_from_headers + ) return data @staticmethod @@ -2417,7 +2455,7 @@ async def add_litellm_data_to_request( # them — from leaking into requester_metadata where guardrails and audit # paths may read from it. if "metadata" in data and isinstance(data["metadata"], dict): - data[_metadata_variable_name]["requester_metadata"] = copy.deepcopy(data["metadata"]) + data[_metadata_variable_name]["requester_metadata"] = _requester_metadata_snapshot(data["metadata"]) if _metadata_variable_name == "litellm_metadata": data[_metadata_variable_name].update( _promoted_trace_control_fields( @@ -2581,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/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7877d887409..592c732bfd3 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -7713,6 +7713,11 @@ async def block_key( code=status.HTTP_404_NOT_FOUND, ) + record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update( + where={"token": hashed_token}, + data=with_settings_updated_at({"blocked": True}), + ) + if is_audit_logging_enabled(): asyncio.create_task( create_audit_log_for_update( @@ -7734,11 +7739,6 @@ async def block_key( ) ) - record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update( - where={"token": hashed_token}, - data=with_settings_updated_at({"blocked": True}), - ) - ## UPDATE KEY CACHE - invalidate so next read re-fetches from DB await delete_cache_key_object( hashed_token=hashed_token, @@ -7827,6 +7827,11 @@ async def unblock_key( code=status.HTTP_404_NOT_FOUND, ) + record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update( + where={"token": hashed_token}, + data=with_settings_updated_at({"blocked": False}), + ) + if is_audit_logging_enabled(): asyncio.create_task( create_audit_log_for_update( @@ -7848,11 +7853,6 @@ async def unblock_key( ) ) - record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update( - where={"token": hashed_token}, - data=with_settings_updated_at({"blocked": False}), - ) - ## UPDATE KEY CACHE - invalidate so next read re-fetches from DB await delete_cache_key_object( hashed_token=hashed_token, 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/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index ecd6abea3c3..ef95e8e2c0f 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -5,7 +5,7 @@ Functions to create audit logs for LiteLLM Proxy import asyncio import json from datetime import datetime, timezone -from typing import Final +from typing import Final, TypeVar import litellm from litellm._logging import verbose_proxy_logger @@ -134,6 +134,56 @@ def _audit_log_task_done_callback(task: asyncio.Task) -> None: verbose_proxy_logger.error("Audit log callback task failed: %s", exc, exc_info=exc) +_T: Final = TypeVar("_T") + + +class AuditTaskRegistry: + __slots__ = ("_pending_tasks",) + + def __init__(self) -> None: + self._pending_tasks: Final[set[asyncio.Task[object]]] = set() # mutable-ok: background task tracking + + def track(self, task: asyncio.Task[_T]) -> asyncio.Task[_T]: + if task.done() or task in self._pending_tasks: + return task + self._pending_tasks.add(task) + task.add_done_callback(self._on_task_done) + return task + + def _on_task_done(self, task: asyncio.Task[object]) -> None: + self._pending_tasks.discard(task) + _audit_log_task_done_callback(task) + + async def drain(self, timeout: float = 5.0) -> None: + loop: Final = asyncio.get_running_loop() + deadline: Final = loop.time() + timeout + while self._pending_tasks: + _, pending = await asyncio.wait(tuple(self._pending_tasks), timeout=max(deadline - loop.time(), 0)) + if pending: + self._pending_tasks.difference_update(pending) + verbose_proxy_logger.warning( + "Timed out draining %d audit tasks on shutdown, dropping: %s", + len(pending), + ", ".join(sorted(task.get_name() for task in pending)), + ) + return + + +_audit_task_registry: Final = AuditTaskRegistry() + + +def track_audit_task(task: asyncio.Task[_T]) -> asyncio.Task[_T]: + return _audit_task_registry.track(task) + + +async def drain_audit_tasks(timeout: float = 5.0) -> None: + await _audit_task_registry.drain(timeout=timeout) + + +def _audit_task_name(request_data: LiteLLM_AuditLogs, destination: str) -> str: + return f"{request_data.action} {request_data.table_name.value} {request_data.object_id} to {destination}" + + async def _dispatch_audit_log_to_callbacks( request_data: LiteLLM_AuditLogs, ) -> None: @@ -153,8 +203,12 @@ async def _dispatch_audit_log_to_callbacks( continue if isinstance(resolved, CustomLogger): - task = asyncio.create_task(resolved.async_log_audit_log_event(payload)) - task.add_done_callback(_audit_log_task_done_callback) + track_audit_task( + asyncio.create_task( + resolved.async_log_audit_log_event(payload), + name=_audit_task_name(request_data, type(resolved).__name__), + ) + ) except Exception as e: verbose_proxy_logger.error("Failed dispatching audit log to callback: %s", e) @@ -204,6 +258,14 @@ async def create_object_audit_log( ) +async def _insert_audit_log_row(repository: AuditLogRepository, request_data: LiteLLM_AuditLogs) -> None: + try: + await repository.table.create(data={**request_data.model_dump(exclude_none=True)}) + except Exception as e: + # [Non-Blocking Exception. Do not allow blocking LLM API call] + verbose_proxy_logger.error("Failed Creating audit log %s", e) + + async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): """ Create an audit log for an object. @@ -231,14 +293,9 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): verbose_proxy_logger.error("prisma_client is None, cannot write audit log to DB") return - _request_data: Final = request_data.model_dump(exclude_none=True) - - try: - await AuditLogRepository(prisma_client).table.create( - data={ - **_request_data, - } + await track_audit_task( + asyncio.create_task( + _insert_audit_log_row(AuditLogRepository(prisma_client), request_data), + name=_audit_task_name(request_data, "the database"), ) - except Exception as e: - # [Non-Blocking Exception. Do not allow blocking LLM API call] - verbose_proxy_logger.error("Failed Creating audit log %s", e) + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 5d825099ceb..464cfee270a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -34,6 +34,7 @@ from litellm import get_llm_provider from litellm._logging import verbose_proxy_logger from litellm.constants import ( ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS, + AWS_ERROR_TYPE_HEADER, AZURE_SPEECH_BATCH_PATH_PREFIX, AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN, AZURE_SPEECH_CUSTOM_LLM_PROVIDER, @@ -45,7 +46,8 @@ from litellm.constants import ( BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES, TYPESAFE_DANGER_ZONE_HEADER_PREFIX, ) -from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming +from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException +from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers from litellm.llms.azure.passthrough.transformation import ( @@ -67,7 +69,7 @@ from litellm.llms.openai.common_utils import OpenAIError as LiteLLMOpenAIError from litellm.llms.openai.workload_identity import get_workload_identity_bearer_token_for_api_base from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, validate_oss_request from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -from litellm.passthrough.main import AsyncPassthroughStreamingResponse +from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import enforced_model_allowlists from litellm.proxy.auth.handle_jwt import JWTHandler @@ -80,7 +82,11 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul user_api_key_auth_websocket, user_api_key_auth_websocket_for_model, ) -from litellm.proxy.common_request_processing import model_deployments_support_stream_options, open_sse_before_first_byte +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + model_deployments_support_stream_options, + open_sse_before_first_byte, +) 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_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -1064,6 +1070,68 @@ def extract_model_from_bedrock_endpoint(endpoint: str) -> str: _extract_model_from_bedrock_endpoint: Final = extract_model_from_bedrock_endpoint +def _guardrail_refusal_status(error: Exception) -> int | None: + """The status a guardrail blocked the request with. None for a request check that answers with the same 400, 403 + or 422, for a guardrail that failed closed because it could not get a verdict, and for a guardrail block raised + with a plain-string detail. + + Guardrail hooks stamp ``guardrail_name`` onto the dict detail of an HTTPException they raise.""" + match error: + case HTTPException(): + detail: Final[object] = error.detail + stamped: Final = isinstance(detail, dict) and "guardrail_name" in detail + return error.status_code if stamped and is_guardrail_intervention(error) else None + case GuardrailRaisedException(blocked_content=True) | BlockedPiiEntityError(): + return error.status_code + case _: + return None + + +def _bedrock_error_response(error: Exception, handled: HTTPException) -> Response | None: + """Bedrock's own error response, for an error Bedrock answered the passthrough request with an AWS error + type.""" + if not isinstance(error, httpx.HTTPStatusError) or not is_provider_error(error): + return None + upstream: Final = error.response + if AWS_ERROR_TYPE_HEADER not in upstream.headers: + return None + passed_through: Final = { + name: upstream.headers[name] for name in (AWS_ERROR_TYPE_HEADER, "content-type") if name in upstream.headers + } + return Response( + content=upstream.content, + status_code=upstream.status_code, + headers={**(handled.headers or {}), **passed_through}, + ) + + +async def bedrock_passthrough_error_response( + processor: ProxyBaseLLMRequestProcessing, + error: Exception, + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLoggingType, +) -> Response: + """Run the shared error handling for a Bedrock passthrough request, then answer a guardrail block or an error + Bedrock returned the way Bedrock does, so an AWS SDK reads its code and message. Anything else is raised as the + shared handler raises it.""" + try: + await processor.handle_llm_api_exception( + e=error, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException as handled: + response: Final = _bedrock_error_response(error, handled) + if response is None: + raise + return response + except ProxyException as handled: + refusal_status: Final = _guardrail_refusal_status(error) + if refusal_status is not None and str(refusal_status) == handled.code: + raise AwsErrorTypeProxyException(handled, "ValidationException") from handled + raise + + async def handle_bedrock_passthrough_router_model( model: str, endpoint: str, @@ -1148,11 +1216,8 @@ async def handle_bedrock_passthrough_router_model( ) return result except Exception as e: - # Use common exception handling - raise await base_llm_response_processor.handle_llm_api_exception( - e=e, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, + return await bedrock_passthrough_error_response( + base_llm_response_processor, e, user_api_key_dict, proxy_logging_obj ) @@ -1351,10 +1416,8 @@ async def bedrock_llm_proxy_route( return result except Exception as e: - raise await base_llm_response_processor.handle_llm_api_exception( - e=e, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, + return await bedrock_passthrough_error_response( + base_llm_response_processor, e, user_api_key_dict, proxy_logging_obj ) 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 4b4c6a13b1d..3de3dccc532 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -67,6 +67,7 @@ from litellm.constants import ( AIOHTTP_NEEDS_CLEANUP_CLOSED, AIOHTTP_TTL_DNS_CACHE, AUDIO_SPEECH_CHUNK_SIZE, + AWS_ERROR_TYPE_HEADER, BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME, BASE_MCP_ROUTE, DAILY_TAG_SPEND_BATCH_MULTIPLIER, @@ -97,6 +98,7 @@ from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.proxy._types import ( UI_TEAM_ID, + AwsErrorTypeProxyException, CallbackDelete, CallInfo, CommonProxyErrors, @@ -169,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, @@ -1198,6 +1201,20 @@ async def _flush_spend_logs_queue_on_shutdown() -> None: verbose_proxy_logger.exception("Error flushing spend logs queue on shutdown: %s", e) +async def _drain_audit_tasks_on_shutdown() -> None: + try: + from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks + + await drain_audit_tasks() + except Exception as e: # noqa: BLE001 # shutdown must continue even if the drain fails + verbose_proxy_logger.exception("Error draining audit tasks on shutdown: %s", e) + + +async def _flush_request_accumulators(prisma: PrismaClient) -> None: + await flush_gateway_requests(prisma, gateway_request_accumulator) + await flush_request_errors(prisma, request_error_accumulator) + + async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = None) -> None: global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server") @@ -1208,10 +1225,11 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N # is lost, and a write attempted after disconnect raises # ClientNotConnectedError rather than persisting anything. Ordering this # inside the same guard is what keeps the two from drifting apart. - await flush_gateway_requests(prisma_client, gateway_request_accumulator) - await flush_request_errors(prisma_client, request_error_accumulator) + await asyncio.gather(_drain_audit_tasks_on_shutdown(), _flush_request_accumulators(prisma_client)) verbose_proxy_logger.debug("Disconnecting from Prisma") await prisma_client.disconnect() + else: + await _drain_audit_tasks_on_shutdown() await drain_post_call_redis_batches() if litellm.cache is not None: @@ -1507,7 +1525,6 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState await asyncio.sleep(5) asyncio.create_task(_run_agent_grant_id_migration()) - await install_background_interaction_settlement(prisma_client) ## A coordination_redis block saved from the admin UI lives in the database, ## which is only reachable once the prisma client exists. Apply it here, before @@ -1634,6 +1651,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState ## SYNC UI SETTINGS ## await ProxyStartupEvent.sync_ui_settings_to_general_settings() + await install_background_interaction_settlement( + prisma_client, router=lambda: llm_router, spend_logs_kept=lambda: disable_spend_logs is False + ) + # Start background health checks AFTER models are loaded and index is built if use_background_health_checks: asyncio.create_task(_run_background_health_check()) # start the background health check coroutine. @@ -2061,6 +2082,17 @@ async def openai_exception_handler(request: Request, exc: ProxyException): ) +@app.exception_handler(AwsErrorTypeProxyException) +async def aws_error_type_exception_handler(request: Request, exc: AwsErrorTypeProxyException) -> Response: + """The ProxyException response plus the ``x-amzn-errortype`` header and top-level ``message`` an AWS SDK reads.""" + response: Final = await openai_exception_handler(request, exc.error) + return JSONResponse( + status_code=response.status_code, + content={"message": exc.error.message, **JSON_OBJECT.validate_json(bytes(response.body))}, + headers={**exc.error.headers, AWS_ERROR_TYPE_HEADER: exc.aws_error_type}, + ) + + @app.exception_handler(StarletteHTTPException) async def otlp_http_exception_handler(connection: Request | WebSocket, exc: StarletteHTTPException) -> Response | None: if isinstance(connection, WebSocket): @@ -9981,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({}) @@ -10256,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 @@ -14425,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): @@ -14440,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) @@ -14460,6 +14502,49 @@ async def _try_provider_token_count( return result +_PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( + {litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value} +) + + +def _drops_params(deployment: Mapping[str, object] | None) -> bool: + litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) + 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 @@ -14490,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 @@ -14526,8 +14613,8 @@ async def count_request_tokens( litellm_model_name = litellm_model_name.split("/", 1)[1] model_to_use: str = ( - litellm_model_name or request.model - ) # use litellm model name, if it's not avalable then fallback to request.model + litellm_model_name if isinstance(litellm_model_name, str) and litellm_model_name else request.model + ) # Try provider-specific token counting first - only for non-direct requests (from provider endpoints) provider_counter: BaseTokenCounter | None = None @@ -14538,19 +14625,30 @@ async def count_request_tokens( if _model is not None: model_to_use = _model + conversation_empty: Final = not messages and not contents + if prompt and not conversation_empty and not _drops_params(deployment): + raise ProxyException( + message="prompt cannot be sent with messages or contents. Set drop_params to count them without the prompt", + type="invalid_request_error", + param="prompt", + code=400, + ) + prompt_only: Final = bool(prompt) and conversation_empty if provider_counter is not None: + prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES result: Final = await _try_provider_token_count( provider_counter=provider_counter, custom_llm_provider=custom_llm_provider, model_to_use=model_to_use, - messages=messages, - contents=contents, + messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages, + contents=None if prompt_as_message else contents, deployment=deployment, request_model=request.model, tools=tools, 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 @@ -14577,8 +14675,10 @@ async def count_request_tokens( tokenizer_used: Final = str(_tokenizer_used["type"]) system_message: Final = _system_message(system) + prompt_turn: Final = prompt_only and not (system_message is None and not tools) typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes - Sequence[AllMessageValues] | None, messages + Sequence[AllMessageValues] | None, + [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else messages), ) counted_messages: Final = ( typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages) @@ -14588,7 +14688,7 @@ async def count_request_tokens( ) total_tokens: Final = await offload_token_count(litellm.token_counter)( model=model_to_use, - text=prompt, + text=prompt if counted_messages is None else None, messages=counted_messages, tools=counted_tools, custom_tokenizer=_tokenizer_used, @@ -19645,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"], @@ -19680,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) @@ -19706,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, ) ) @@ -19718,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/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index ad82a1ed724..0a0e05b09c5 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.router_utils.endpoint_eligibility import UnsupportedEndpointError from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -70,7 +77,6 @@ from litellm.types.utils import TokenCountResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.utils import ProxyLogging - from litellm.router import Router router: Final = APIRouter() _RESPONSES_WS_CONFIG_VALUE_ADAPTER: Final[TypeAdapter[object | None]] = TypeAdapter(object | None) @@ -1377,6 +1383,15 @@ def _proxy_token_counter() -> _TokenCounter: _token_counter_dep: Final = Depends(_proxy_token_counter) +def _proxy_llm_router() -> Router | None: + from litellm.proxy.proxy_server import llm_router + + return llm_router + + +_llm_router_dep: Final = Depends(_proxy_llm_router) + + def _responses_invalid_request_response(message: str, param: str | None, code: str | None) -> JSONResponse: body: Final[_ResponsesApiErrorBody] = { "error": { @@ -1433,6 +1448,8 @@ def _responses_input_as_token_count_messages( async def responses_input_tokens( request: Request, token_counter: _TokenCounter = _token_counter_dep, + llm_router: Router | None = _llm_router_dep, + user_api_key_dict: UserAPIKeyAuth = _user_api_key_auth_dep, ): """ Count the input tokens of a Responses API request without calling the model. @@ -1465,6 +1482,13 @@ async def responses_input_tokens( code="missing_required_parameter", ) + counted_fields: Final = { + name: data[name] for name in ("model", "tools", "drop_params", "additional_drop_params") if name in data + } + token_fields: Final = resolve_counted_tool_fields(counted_fields, llm_router, user_api_key_dict.team_id) + if isinstance(token_fields, ProxyException): + return _responses_invalid_request_response(message=token_fields.message, param=token_fields.param, code=None) + try: payload: Final[_TokenCountPayload] = { "model": model_name, @@ -1472,9 +1496,13 @@ async def responses_input_tokens( input_value=input_value, instructions=data.get("instructions"), ), - "tools": data.get("tools"), + "tools": token_fields.get("tools"), } token_request: Final = TokenCountRequest.model_validate(payload) + except ValidationError as e: + return _responses_invalid_request_response( + message=first_validation_error(e), param=str(e.errors(include_url=False)[0]["loc"][0]), code=None + ) except Exception as e: return _responses_invalid_request_response( message=f"Invalid request for token counting: {e}", param=None, code=None @@ -1918,7 +1946,7 @@ async def _responses_websocket_session( reason="Pre-call error" if isinstance(e, UnsupportedEndpointError) else "Internal server error", ) finally: - await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict) + await proxy_logging_obj.arelease_responses_websocket_slot(user_api_key_dict) @router.websocket("/v1/responses") diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d40bac380aa..83a0e6f9e05 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1981,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement { outcome String? @@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at") + @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } model LiteLLM_Lens { diff --git a/litellm/proxy/spend_tracking/background_interaction_settlement.py b/litellm/proxy/spend_tracking/background_interaction_settlement.py index 165f813cd50..ffa8185619c 100644 --- a/litellm/proxy/spend_tracking/background_interaction_settlement.py +++ b/litellm/proxy/spend_tracking/background_interaction_settlement.py @@ -1,21 +1,25 @@ import asyncio import os import socket -from collections.abc import Awaitable, Mapping, Sequence -from dataclasses import dataclass -from datetime import datetime, timezone +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass, replace +from datetime import datetime, timedelta, timezone from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol, TypeVar +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeVar from pydantic import ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED +from litellm.constants import ( + BACKGROUND_INTERACTION_COST_POLLING_ENABLED, + BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS, +) from litellm.interactions.background_cost_polling import ( DEFAULT_POLL_SCHEDULE, BackgroundInteractionCreateContext, + BackgroundInteractionPollContext, FetchInteraction, PendingBackgroundInteraction, PollSchedule, @@ -24,10 +28,13 @@ from litellm.interactions.background_cost_polling import ( fetch_background_interaction, resume_unsettled_background_interactions, ) -from litellm.repositories.table_repositories import BackgroundInteractionSettlementRepository +from litellm.repositories.table_repositories import BackgroundInteractionSettlementRepository, SpendLogsRepository +from litellm.types.interactions import InteractionsAPIResponse +from litellm.types.router import CredentialLiteLLMParams if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient + from litellm.router import Router class _SettlementRow(Protocol): @@ -41,6 +48,8 @@ class _SettlementRow(Protocol): def created_at(self) -> datetime: ... @property def claimed_at(self) -> datetime | None: ... + @property + def settled_at(self) -> datetime | None: ... class _NewSettlementRow(TypedDict): @@ -59,8 +68,24 @@ class _UnclaimedRowKey(TypedDict): claimed_at: ReadOnly[None] -class _UnclaimedRows(TypedDict): - claimed_at: ReadOnly[None] +class _Before(TypedDict): + lt: ReadOnly[datetime] + + +class _ExpiredClaimRowKey(TypedDict): + interaction_id: ReadOnly[str] + settled_at: ReadOnly[None] + claimed_at: ReadOnly[_Before] + + +class _OwnClaimRowKey(TypedDict): + interaction_id: ReadOnly[str] + claimed_by: ReadOnly[str] + settled_at: ReadOnly[None] + + +class _UnsettledRows(TypedDict): + settled_at: ReadOnly[None] class _Claim(TypedDict): @@ -79,9 +104,19 @@ class _SettlementTableActions(Protocol): def find_unique(self, *, where: _RowKey) -> Awaitable[_SettlementRow | None]: ... - def find_many(self, *, where: _UnclaimedRows) -> Awaitable[Sequence[_SettlementRow]]: ... + def find_many(self, *, where: _UnsettledRows) -> Awaitable[Sequence[_SettlementRow]]: ... - def update_many(self, *, data: _Claim | _Outcome, where: _RowKey | _UnclaimedRowKey) -> Awaitable[int]: ... + def update_many( + self, *, data: _Claim | _Outcome, where: _RowKey | _UnclaimedRowKey | _ExpiredClaimRowKey | _OwnClaimRowKey + ) -> Awaitable[int]: ... + + +class _SpendLogKey(TypedDict): + request_id: ReadOnly[str] + + +class _SpendLogLookup(Protocol): + def find_unique(self, *, where: _SpendLogKey) -> Awaitable[object | None]: ... _CLEARED_CREATE_CONTEXT: Final[Mapping[str, object]] = MappingProxyType({}) @@ -103,11 +138,7 @@ def _json(data: Mapping[str, object]) -> object: return Json.keys(**data) -def _pending_rows(rows: Sequence[_SettlementRow]) -> tuple[PendingBackgroundInteraction, ...]: - return tuple(chain.from_iterable(_pending_row(row) for row in rows)) - - -def _pending_row(row: _SettlementRow) -> tuple[PendingBackgroundInteraction, ...]: +def _pending_row(row: _SettlementRow, claim_expires_at: datetime | None) -> tuple[PendingBackgroundInteraction, ...]: try: create_context: Final = BackgroundInteractionCreateContext.model_validate(row.create_context) except ValidationError: @@ -122,14 +153,31 @@ def _pending_row(row: _SettlementRow) -> tuple[PendingBackgroundInteraction, ... custom_llm_provider=row.custom_llm_provider, create_context=create_context, created_at=row.created_at, + claim_expires_at=claim_expires_at, ), ) @dataclass(frozen=True, slots=True) class PrismaBackgroundSettlementStore: + """ + A claim is a lease, not a settlement: the settler that claims a row bills + or releases it and then records the outcome, so a row claimed longer ago + than the lease with no outcome belongs to a settler that died in between, + and any settler may claim it again. A live settler renews its claim while + it settles, so it is never overtaken. The settler before may have billed + and then died or failed to record it, so a takeover bills only when the + interaction has no spend row yet, and otherwise records it as billed. The + spend row is read before the claim is taken, so a failed read takes + nothing and the claim is retried. With spend logs off there is no row to + tell, so a takeover bills nothing rather than risk billing twice. + """ + table: _SettlementTableActions + spend_logs: _SpendLogLookup + spend_logs_kept: Callable[[], bool] claimed_by: str + claim_lease_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS async def register(self, pending: PendingBackgroundInteraction) -> None: await self.table.create( @@ -143,24 +191,71 @@ class PrismaBackgroundSettlementStore: async def pending(self, interaction_id: str) -> PendingBackgroundInteraction | None: row: Final = await self._row(interaction_id) - if row is None or row.claimed_at is not None: + if row is None or self._is_owned(row): return None - return next(iter(_pending_row(row)), None) + return next(iter(_pending_row(row, claim_expires_at=None)), None) async def is_claimed(self, interaction_id: str) -> bool: row: Final = await self._row(interaction_id) - return row is not None and row.claimed_at is not None + return row is not None and self._is_owned(row) async def claim(self, interaction_id: str) -> bool: + if await self._claim_where(_UnclaimedRowKey(interaction_id=interaction_id, claimed_at=None)): + return True + settled_before: Final = await self._outcome_a_dead_settler_left(interaction_id) + took_over: Final = await self._claim_where( + _ExpiredClaimRowKey( + interaction_id=interaction_id, + settled_at=None, + claimed_at=_Before(lt=datetime.now(timezone.utc) - self._lease), + ) + ) + if not took_over or settled_before is None: + return took_over + if settled_before == "unsettled": + verbose_proxy_logger.warning( + "Background interaction %s lost its settler mid-settle, and with spend logs off nothing shows " + "whether it was billed, so it is left unbilled rather than risk billing it twice", + interaction_id, + ) + await self.record_outcome(interaction_id, settled_before) + return False + + async def renew_claim(self, interaction_id: str) -> None: + await self.table.update_many( + data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by), + where=_OwnClaimRowKey(interaction_id=interaction_id, claimed_by=self.claimed_by, settled_at=None), + ) + + async def _claim_where(self, where: _UnclaimedRowKey | _ExpiredClaimRowKey) -> bool: claimed_rows: Final = await _read_from_a_table_that_may_not_exist( self.table.update_many( - data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by), - where=_UnclaimedRowKey(interaction_id=interaction_id, claimed_at=None), + data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by), where=where ), when_missing=0, ) return claimed_rows == 1 + async def _outcome_a_dead_settler_left(self, interaction_id: str) -> Literal["billed", "unsettled"] | None: + if not self.spend_logs_kept(): + return "unsettled" + if await self.spend_logs.find_unique(where=_SpendLogKey(request_id=interaction_id)) is None: + return None + return "billed" + + @property + def _lease(self) -> timedelta: + return timedelta(seconds=self.claim_lease_seconds) + + def _claim_expires_at(self, row: _SettlementRow) -> datetime | None: + return None if row.claimed_at is None else row.claimed_at + self._lease + + def _is_owned(self, row: _SettlementRow) -> bool: + if row.settled_at is not None: + return True + claim_expires_at: Final = self._claim_expires_at(row) + return claim_expires_at is not None and claim_expires_at > datetime.now(timezone.utc) + async def _row(self, interaction_id: str) -> _SettlementRow | None: return await _read_from_a_table_that_may_not_exist( self.table.find_unique(where=_RowKey(interaction_id=interaction_id)), when_missing=None @@ -174,19 +269,78 @@ class PrismaBackgroundSettlementStore: where=_RowKey(interaction_id=interaction_id), ) - async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: - return _pending_rows(await self.table.find_many(where=_UnclaimedRows(claimed_at=None))) + async def unsettled(self) -> Sequence[PendingBackgroundInteraction]: + rows: Final = await self.table.find_many(where=_UnsettledRows(settled_at=None)) + return tuple(chain.from_iterable(_pending_row(row, self._claim_expires_at(row)) for row in rows)) + + +def _provider_credentials(router: "Router", deployment_id: str, custom_llm_provider: str) -> Mapping[str, object]: + credentials: Final = router.get_deployment_credentials_with_provider(model_id=deployment_id) + if credentials is None or credentials.get("custom_llm_provider") != custom_llm_provider: + return MappingProxyType({}) + return MappingProxyType(credentials) + + +def _group_deployment_credentials(router: "Router", context: BackgroundInteractionPollContext) -> Mapping[str, object]: + if context.model_group is None: + return MappingProxyType({}) + candidates: Final = tuple( + credentials + for deployment_id in router.get_model_ids(model_name=context.model_group) + if (credentials := _provider_credentials(router, deployment_id, context.custom_llm_provider)) + ) + return candidates[0] if len(candidates) == 1 else MappingProxyType({}) + + +def _deployment_credentials( + router: "Router | None", context: BackgroundInteractionPollContext +) -> CredentialLiteLLMParams | None: + if router is None: + return None + serving: Final = ( + _provider_credentials(router, context.deployment_id, context.custom_llm_provider) + if context.deployment_id is not None and router.get_deployment(model_id=context.deployment_id) is not None + else _group_deployment_credentials(router, context) + ) + return CredentialLiteLLMParams.model_validate(serving) if serving else None + + +def fetch_with_deployment_credentials( + router: Callable[[], "Router | None"], fetch_interaction: FetchInteraction = fetch_background_interaction +) -> FetchInteraction: + """ + A resumed poll has no request to take credentials from, and they are never + stored, so it fetches with the credentials of the deployment that served + the create, looked up by its id on the live router at each fetch so a + deployment loaded from the database after startup is found too. Editing a + deployment changes its id, so a missing id falls back to the create's model + group when exactly one of its deployments serves the same provider. Anything + else falls back to the provider's environment credentials, never to a + wildcard route, whose key belongs to whatever provider it routes to. + """ + + async def fetch(context: BackgroundInteractionPollContext) -> InteractionsAPIResponse: + credentials: Final = _deployment_credentials(router(), context) + if credentials is None: + return await fetch_interaction(context) + return await fetch_interaction(replace(context, credentials=credentials)) + + return fetch async def configure_background_interaction_settlement( table: _SettlementTableActions, + spend_logs: _SpendLogLookup, + spend_logs_kept: Callable[[], bool], claimed_by: str, fetch_interaction: FetchInteraction = fetch_background_interaction, schedule: PollSchedule = DEFAULT_POLL_SCHEDULE, ) -> tuple["asyncio.Task[SettlementOutcome | None]", ...]: if not BACKGROUND_INTERACTION_COST_POLLING_ENABLED: return () - store: Final = PrismaBackgroundSettlementStore(table=table, claimed_by=claimed_by) + store: Final = PrismaBackgroundSettlementStore( + table=table, spend_logs=spend_logs, spend_logs_kept=spend_logs_kept, claimed_by=claimed_by + ) configure_background_settlement_store(store) resumed: Final = await resume_unsettled_background_interactions(store, fetch_interaction, schedule) if resumed: @@ -194,11 +348,16 @@ async def configure_background_interaction_settlement( return resumed -async def install_background_interaction_settlement(prisma_client: "PrismaClient") -> None: +async def install_background_interaction_settlement( + prisma_client: "PrismaClient", router: Callable[[], "Router | None"], spend_logs_kept: Callable[[], bool] +) -> None: try: await configure_background_interaction_settlement( table=BackgroundInteractionSettlementRepository(prisma_client).table, + spend_logs=SpendLogsRepository(prisma_client).table, + spend_logs_kept=spend_logs_kept, claimed_by=f"{socket.gethostname()}:{os.getpid()}", + fetch_interaction=fetch_with_deployment_credentials(router), ) except Exception as e: # noqa: BLE001 # a boot step must survive any DB error; billing then settles in-process as before verbose_proxy_logger.warning( diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index efd0d42a70f..0c5c0a7fbf3 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -15,6 +15,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, CLI_SESSION_KEY_PREFIX, EMPTY_MAPPING, INTERNAL_CALL_ORIGIN_METADATA_KEY, @@ -33,6 +34,8 @@ from litellm.constants import ( from litellm.constants import ( MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB, ) +from litellm.integrations.otel.model.metadata import caller_session_id +from litellm.integrations.otel.model.trace_controls import TraceControls from litellm.litellm_core_utils.classifier_logging import ( CLASSIFIER_AUDIT_FIELDS, classifier_audit_fields, @@ -44,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, @@ -543,6 +547,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): @@ -574,11 +579,11 @@ def get_logging_payload( elif isinstance(_usage, dict): usage = _usage - # A request that failed mid-stream has no usable response_obj usage, but the - # streaming handler may have recovered the usage from the chunks already - # delivered. Honor that override so the partial usage lands in spend tracking. + # A request that failed mid-stream has no usable response_obj usage, and a stream + # continued on a fallback only has the fallback's. combined_usage_object holds the + # usage of every call the client's answer took, so it wins in both cases. _combined_usage: Final = kwargs.get("combined_usage_object") - if not usage and isinstance(_combined_usage, litellm.Usage): + if (not usage or "continued_partial_usage" in kwargs) and isinstance(_combined_usage, litellm.Usage): usage = _combined_usage.model_dump() id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs, metadata) @@ -813,7 +818,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 @@ -934,6 +939,16 @@ def _get_batch_trace_session_id(call_type: str | None, request_id: str | None) - return request_id.removesuffix(BATCH_COST_REQUEST_ID_SUFFIX) +def _caller_session_for_spend_log(kwargs: Mapping[str, object], metadata: Mapping[str, object] | None) -> str | None: + beside_trace_header: Final = metadata.get(CALLER_SESSION_ID_METADATA_KEY) if metadata else None + caller_session: Final = ( + beside_trace_header + if isinstance(beside_trace_header, str) and beside_trace_header + else caller_session_id(kwargs, TraceControls()) + ) + return caller_session if caller_session and "\x00" not in caller_session else None + + def _get_session_id_for_spend_log( kwargs: Mapping[str, object], metadata: Mapping[str, object] | None, @@ -941,9 +956,12 @@ def _get_session_id_for_spend_log( omit_when_missing: bool, batch_trace_session_id: str | None = None, ) -> str | None: - """Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session; `litellm_session_id` may - be a copied trace id. Batch call types carry a deterministic session derived from the batch id, which outranks - the per-request trace ids because those differ between the create call and the cost poller's row.""" + """Resolve the session id for the spend log row: `omit` honors only metadata.session_id, + batch sessions outrank everything, then the conversation id the caller sent (also one an x-litellm-trace-id header + displaced from the request's session), then trace ids. + + A session minted by `missing_session_id: generate` or copied from the request's trace id is not a conversation, so + it yields to the trace id, which the Responses bridge points at the previous turn's session to keep a chain whole.""" if omit_when_missing: session_id: Final = metadata.get("session_id") if metadata else None return str(session_id) if session_id else None @@ -952,11 +970,14 @@ def _get_session_id_for_spend_log( if batch_trace_session_id is not None: return batch_trace_session_id - if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None: - return str(standard_logging_payload.get("trace_id")) - if kwargs.get("litellm_trace_id") is not None: - return str(kwargs.get("litellm_trace_id")) - return str(uuid.uuid4()) + if litellm.request_correlation_in_logs and (caller_session := _caller_session_for_spend_log(kwargs, metadata)): + return caller_session + match standard_logging_payload: + case {"trace_id": trace_id} if trace_id is not None: + return str(trace_id) + case _: + litellm_trace_id: Final = kwargs.get("litellm_trace_id") + return str(litellm_trace_id) if litellm_trace_id is not None else str(uuid.uuid4()) def _get_request_duration_ms(start_time: datetime, end_time: datetime) -> int | None: 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 5d2ebe73cac..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 @@ -4405,12 +4430,20 @@ class ProxyLogging: double-decrement under the limiter's in-memory fallback. """ limiter: Final = self.get_proxy_hook("parallel_request_limiter") - if not isinstance(limiter, (PROXY_MaxParallelRequestsHandler_v3, _PROXY_MaxParallelRequestsHandler)): + if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): return await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) _arelease_max_parallel_requests_on_disconnect = arelease_max_parallel_requests_on_disconnect + async def arelease_responses_websocket_slot(self, user_api_key_dict: UserAPIKeyAuth) -> None: + """Release the slot a Responses WebSocket took at admission. Only call this once that admission succeeded: + the legacy limiter keeps no record of its acquisitions and decrements unconditionally.""" + limiter: Final = self.get_proxy_hook("parallel_request_limiter") + if not isinstance(limiter, (PROXY_MaxParallelRequestsHandler_v3, _PROXY_MaxParallelRequestsHandler)): + return + await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) + async def enforce_mcp_server_rate_limits( self, user_api_key_dict: UserAPIKeyAuth | None, 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/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 10f092c1cb2..95f4f5e81be 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -97,11 +97,13 @@ class SpendResetUnitOfWork: @dataclass(frozen=True, slots=True) class BudgetCascadeUnitOfWork: - """Every write a budget-tier reset performs, bound to one batch. + """Every write a budget-tier reset performs in its final batch. The dependent spend rows and the budget rows' ``budget_reset_at`` advance must land together: advancing the window without zeroing the spend it gates leaves the dependents pinned at their cap until the next window. + End users are the exception, paged and committed ahead of this batch by + ``enduser_spend_reset_unit_of_work`` because one tier can gate millions. """ team_memberships: LinkedSpendResetWrites @@ -110,7 +112,6 @@ class BudgetCascadeUnitOfWork: tags: LinkedSpendResetWrites model_access_groups: LinkedSpendResetWrites projects: LinkedSpendResetWrites - endusers: LinkedSpendResetWrites budgets: BudgetWindowWrites @@ -137,7 +138,17 @@ async def budget_cascade_unit_of_work( tags=LinkedSpendResetWrites(table=batch.litellm_tagtable), model_access_groups=LinkedSpendResetWrites(table=batch.litellm_modelaccessgroupbudgettable), projects=LinkedSpendResetWrites(table=batch.litellm_projecttable), - endusers=LinkedSpendResetWrites(table=batch.litellm_endusertable), budgets=BudgetWindowWrites(table=batch.litellm_budgettable), ) await batch.commit() + + +@asynccontextmanager +async def enduser_spend_reset_unit_of_work( + new_batch: Callable[[], PrismaBatch], +) -> AsyncGenerator[LinkedSpendResetWrites, None]: + """One page of end-user spend resets, committed on its own so a statement + never grows with the end-user population.""" + batch: Final = new_batch() + yield LinkedSpendResetWrites(table=batch.litellm_endusertable) + await batch.commit() diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index c2b4c105199..2552ee8c880 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -4,6 +4,7 @@ from urllib.parse import urlsplit import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.llms.azure_ai.common_utils import ( azure_ai_supports_native_responses, foundry_chat_rejects_function_tools_while_reasoning, @@ -11,7 +12,8 @@ from litellm.llms.azure_ai.common_utils import ( from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import OpenAIWebSearchOptions -from litellm.utils import get_model_info_helper +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager, get_model_info_helper _OPENAI_DEFAULT_API_BASE: Final = "https://api.openai.com/v1" _OPENAI_API_HOST: Final = "api.openai.com" @@ -36,6 +38,18 @@ def _resolve_openai_api_base(api_base: str | None) -> str: ) +def chat_target_accepts_reasoning_summary(model: str, custom_llm_provider: str | None) -> bool: + try: + local_model, provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) + chat_config: Final = ProviderConfigManager.get_provider_chat_config( + model=local_model, provider=LlmProviders(provider) + ) + except Exception as e: # noqa: BLE001 # a capability probe must never fail the request it probes for + verbose_logger.debug("chat target reasoning summary check failed: %s", e) + return False + return chat_config is not None and chat_config.accepts_reasoning_summary_in_reasoning_effort + + def responses_api_bridge_check( model: str, custom_llm_provider: str, 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 732c32ff2cb..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 @@ -281,7 +320,13 @@ class ResponsesSessionHandler: SQL query - SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id + SELECT session_id, api_key, "user", team_id FROM spend_logs WHERE request_id = previous_response_id, then + SELECT * FROM spend_logs with that session_id and the same owner + + Only rows with the referenced response's owner are returned: the same user and team, and for a key with + neither, the same key. The session id can be caller supplied, so matching on it alone would mix owners. A + user's or team's regenerated or second key keeps the history; a key with neither does not, and turns written + by another owner (an admin continuing the chain) are not replayed. A just-finished turn gets a short second chance: the worker that served it may still be writing its spend log when the follow-up arrives, and an empty result @@ -301,15 +346,27 @@ 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 + 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" - WHERE session_id IN (SELECT session_id FROM matching_session) + FROM "LiteLLM_SpendLogs" AS logs + WHERE EXISTS ( + SELECT 1 + FROM matching_session + WHERE matching_session.session_id = logs.session_id + AND matching_session."user" IS NOT DISTINCT FROM logs."user" + AND matching_session.team_id IS NOT DISTINCT FROM logs.team_id + AND ( + COALESCE(matching_session."user", '') <> '' + OR COALESCE(matching_session.team_id, '') <> '' + OR matching_session.api_key = logs.api_key + ) + ) ORDER BY "endTime" ASC; """ @@ -317,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 7f5d143e98e..35d4b8bf1ba 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -39,7 +39,11 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.caching import InMemoryCache -from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER +from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, + REDACTED_BY_LITELLM, + REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER, +) from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.get_supported_openai_params import ( @@ -56,6 +60,7 @@ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import bedrock_declared_cache_point_support from litellm.llms.openrouter.chat.transformation import openrouter_model_supports_cache_control from litellm.llms.snowflake.chat.transformation import is_snowflake_claude_model +from litellm.responses.bridge_check import chat_target_accepts_reasoning_summary from litellm.responses.litellm_completion_transformation.session_handler import ( ResponsesSessionHandler, ) @@ -562,6 +567,14 @@ class ChatCompletionSession(TypedDict, total=False): ########### End of Initialize Classes used for Responses API ########### +def _chain_session_metadata(chain_session: str, request_kwargs: Mapping[str, object]) -> Mapping[str, object]: + key: Final = "litellm_metadata" if request_kwargs.get("litellm_metadata") else "metadata" + metadata: Final = request_kwargs.get(key) + if not isinstance(metadata, dict): + return MappingProxyType({}) + return {key: {**_STR_KEY_DICT_ADAPTER.validate_python(metadata), CALLER_SESSION_ID_METADATA_KEY: chain_session}} + + class LiteLLMCompletionResponsesConfig: @staticmethod def get_supported_openai_params(model: str) -> list: @@ -785,8 +798,8 @@ class LiteLLMCompletionResponsesConfig: provider either rejects it or silently drops it, and dropping it turns reasoning off while still billing for the turn. ``summary`` has no chat equivalent, so it rides the ``reasoning_summary`` alias, which ``litellm.completion`` reassembles into ``{effort, - summary}`` when it bridges the model back onto the Responses API, and is sent to nothing - else. + summary}`` when it bridges the model back onto the Responses API or targets a chat provider + whose config maps a summary ask (Gemini's thought summaries), and is sent to nothing else. """ if not reasoning_param: return ResponsesReasoningChatForm(effort=None, summary=None) @@ -807,7 +820,10 @@ class LiteLLMCompletionResponsesConfig: reasoning_summary=summary, api_base=api_base, ) - return ResponsesReasoningChatForm(effort=effort, summary=summary if bridges_back else None) + forwards_summary: Final = bridges_back or ( + isinstance(effort, str) and chat_target_accepts_reasoning_summary(model, custom_llm_provider) + ) + return ResponsesReasoningChatForm(effort=effort, summary=summary if forwards_summary else None) @staticmethod def transform_responses_api_request_to_chat_completion_request( @@ -992,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 @@ -1006,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) @@ -1074,8 +1092,6 @@ class LiteLLMCompletionResponsesConfig: else: # Both are empty - this likely means function_call_output had empty/invalid call_id # Provide a helpful error message - import litellm - raise litellm.BadRequestError( message=( f"Unable to create messages for completion request. " @@ -1090,8 +1106,14 @@ class LiteLLMCompletionResponsesConfig: ) litellm_completion_request["messages"] = combined_messages - litellm_completion_request["litellm_trace_id"] = chat_completion_session.get("litellm_session_id") - return litellm_completion_request + chain_session: Final = chat_completion_session.get("litellm_session_id") + litellm_completion_request["litellm_trace_id"] = chain_session + chain_session_metadata: Final[Mapping[str, object]] = ( + _chain_session_metadata(chain_session, request_kwargs) + if chain_session and litellm.request_correlation_in_logs + else MappingProxyType({}) + ) + return {**litellm_completion_request, **chain_session_metadata} @staticmethod def _transform_response_input_param_to_chat_completion_message( 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/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7cf68332c07..89c32aed1fb 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base +from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( update_response_metadata, ) @@ -262,7 +263,9 @@ def stream_error_status_and_message(error_obj: object) -> tuple[int, str]: return _status_code_for_error_fields(error_type, error_code), message -def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_provider: str) -> Exception: +def _map_stream_error_to_exception( + error_obj: object, model: str, custom_llm_provider: str, headers: httpx.Headers | None = None +) -> Exception: from litellm.llms.base_llm.chat.transformation import BaseLLMException error_message, error_type, error_code = _error_event_fields(error_obj) @@ -272,6 +275,7 @@ def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_pro status_code=status_code, message=f"Error code: {status_code} - {{'error': {error_body}}}", body=error_body, + headers=headers, ) try: return litellm.exception_type( @@ -667,7 +671,12 @@ class BaseResponsesAPIStreamingIterator: ) def _map_error_event_exception(self, error_obj: object) -> Exception: - return _map_stream_error_to_exception(error_obj, self.model or "", self.custom_llm_provider or "") + return _map_stream_error_to_exception( + error_obj, + self.model or "", + self.custom_llm_provider or "", + headers=provider_request_id_headers(self._raw_response_headers), + ) def _maybe_raise_for_error_event(self, result: object) -> None: chunk_type: Final = getattr(result, "type", None) @@ -683,11 +692,17 @@ class BaseResponsesAPIStreamingIterator: mapped_exception: Final = self._map_error_event_exception(error_obj) if not _mid_stream_fallback_eligible(mapped_exception): raise mapped_exception + mapped_status: Final = getattr(mapped_exception, "status_code", None) raise MidStreamFallbackError( message=str(mapped_exception), model=self.model or "", llm_provider=self.custom_llm_provider or "", original_exception=mapped_exception, + response=httpx.Response( + status_code=int(mapped_status) if mapped_status is not None else 503, + headers=provider_request_id_headers(self._raw_response_headers), + request=httpx.Request(method="POST", url=f"https://{self.custom_llm_provider or ''}.com/v1/"), + ), generated_content=self._generated_content, is_pre_first_chunk=not self._yielded_first_chunk, ) diff --git a/litellm/router.py b/litellm/router.py index d73ec1b3180..64b5ec216ec 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -38,7 +38,7 @@ from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache, partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeAlias, TypeVar, Union, cast +from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast import anyio import httpx @@ -68,6 +68,7 @@ from litellm.constants import ( DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, DEFAULT_MAX_LRU_CACHE_SIZE, INTERNAL_CALL_ORIGIN_METADATA_KEY, + MID_STREAM_FAILURE_LOG_WAIT_SECONDS, OUTPUT_TOKEN_CEILING_PARAMS, PROMPT_CACHE_PIN_CACHE_MAX_SIZE, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, @@ -75,7 +76,7 @@ from litellm.constants import ( RUNTIME_UPDATABLE_ROUTER_SETTINGS, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) -from litellm.integrations.custom_guardrail import is_guardrail_intervention +from litellm.integrations.custom_guardrail import CustomGuardrail, is_guardrail_intervention from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.routing import routing_decision_attributes from litellm.integrations.otel.runtime import phase_attributes, phase_event, phase_span @@ -122,12 +123,14 @@ 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 ( RouterVectorStoreEmbeddingExecutor, vector_store_request_metadata, ) +from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.llms.openai_like.model_info import ( @@ -193,6 +196,7 @@ from litellm.router_utils.clientside_credential_handler import ( is_clientside_credential, ) from litellm.router_utils.common_utils import ( + MID_STREAM_CONTINUATION_KWARG, filter_team_based_models, filter_web_search_deployments, format_fallback_outcome_message, @@ -237,7 +241,10 @@ from litellm.router_utils.fallback_event_handlers import ( check_non_standard_fallback_format, clear_pre_routing_selection, committed_retry_budget_for_request, + continuation_fallback_targets, fallback_lookup_groups, + fallback_target_name, + fallback_target_overrides_request_params, fallbacks_disabled_for_request, get_fallback_model_group_for_lookup_groups, get_pre_routing_selection, @@ -262,6 +269,17 @@ from litellm.router_utils.handle_error import ( send_llm_exception_alert, ) from litellm.router_utils.health_state_cache import DeploymentHealthCache +from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_MARKER, + ContinuationOutput, + ContinuationPrefillDeploymentCheck, + ContinuedPartial, + constrains_continuation, + continuation_deployments, + continuation_output_ceilings, + continued_text, + deployment_can_continue, +) from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, warn_on_unknown_model_group_affinity_flags, @@ -307,6 +325,8 @@ from litellm.router_utils.routing_groups import ( ) from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingPrefetch, RoutingReadBatch from litellm.scheduler import FlowItem, Scheduler +from litellm.secret_managers.main import str_to_bool +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.litellm_params import RoutingStrategyName from litellm.types.llms.openai import ( AllMessageValues, @@ -333,6 +353,7 @@ from litellm.types.router import ( FallbackBudgetCheck, GuardrailTypedDict, LiteLLM_Params, + LiteLLMParamsTypedDict, MockRouterTestingParams, ModelGroupInfo, OptionalPreCallChecks, @@ -374,11 +395,13 @@ from litellm.utils import ( EmbeddingResponse, ModelResponse, Rules, + TextCompletionResponse, function_setup, get_llm_provider, get_non_default_completion_params, get_secret, get_utc_datetime, + is_generalized_model_info, is_region_allowed, provider_rejectable_params, set_live_deployment_replay, @@ -483,12 +506,15 @@ 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 CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding" +_PRE_CALL_CHECK_FILTERED_PARAMS: Final = frozenset({"response_format"}) _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType( { "prompt_caching": PromptCachingDeploymentCheck, @@ -519,6 +545,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]) @@ -528,6 +555,15 @@ _ROUTING_KWARGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeA _FALLBACK_HOP_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object) _DEPLOYMENT_SELECTED_EVENT: Final = "litellm.request.deployment_selected" +_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) + + +def _extend_trailing_text_block(blocks: Sequence[Mapping[str, object]], text: str) -> Sequence[Mapping[str, object]]: + """Continue the last text block in place so the prefill reads as the text the caller already received.""" + last: Final = blocks[-1] if blocks else None + if last is not None and last.get("type") == "text" and isinstance(last_text := last.get("text"), str): + return [*blocks[:-1], {**last, "text": f"{last_text}{text}"}] + return [*blocks, {"type": "text", "text": text}] def _is_fallback_hop(request_kwargs: Mapping[str, object]) -> bool: @@ -546,6 +582,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 {} @@ -996,6 +1047,49 @@ def as_output_cap(value: object) -> int | None: return cap if cap >= 0 else None +def _exact_base_model_max_input_tokens(base_model: str) -> int | None: + try: + model_info: Final = litellm.get_model_info(model=base_model) + except Exception: # noqa: BLE001 # get_model_info raises ModelNotMappedError or a bare Exception for an unmapped model + return None + if is_generalized_model_info(model_info): + return None + return model_info.get("max_input_tokens") + + +def _provider_model_id(model_name: str, custom_llm_provider: str) -> str: + model_id: Final = model_name.rsplit("/", 1)[-1] + return get_bedrock_base_model(model_id) if custom_llm_provider.startswith("bedrock") else model_id + + +def _names_route_model(base_model: str, route_key: str, custom_llm_provider: str) -> bool: + route_model_id: Final = _provider_model_id(route_key, custom_llm_provider) + base_model_id: Final = _provider_model_id(base_model, custom_llm_provider) + return route_model_id in (base_model_id, base_model_id.split("@", 1)[0]) + + +def _declared_base_model(user_model_info: Mapping[str, object]) -> str | None: + match user_model_info: + case {"base_model": str() as base_model}: + return base_model + case _: + return None + + +def _base_model_widened_window( + route_model_info: ModelMapInfo, base_model: str | None, custom_llm_provider: str +) -> Mapping[str, int]: + route_max_input_tokens: Final = route_model_info.get("max_input_tokens") + if base_model is None or route_max_input_tokens is None: + return MappingProxyType({}) + if not _names_route_model(base_model, route_model_info["key"], custom_llm_provider): + return MappingProxyType({}) + base_max_input_tokens: Final = _exact_base_model_max_input_tokens(base_model) + if base_max_input_tokens is None or base_max_input_tokens <= route_max_input_tokens: + return MappingProxyType({}) + return MappingProxyType({"max_input_tokens": base_max_input_tokens}) + + class Router: @property def _routing_groups(self) -> dict[str, RoutingGroup]: @@ -1077,6 +1171,7 @@ class Router: health_check_ignore_transient_errors: bool = False, background_health_check_model_groups: Sequence[str] | None = None, enable_weighted_failover: bool = False, + enable_mid_stream_fallback_continuation: bool = False, fallback_access_check: FallbackAccessCheck | None = None, fallback_budget_check: FallbackBudgetCheck | None = None, auto_router_capability_limit: AutoRouterCapabilityLimit | None = None, @@ -1118,6 +1213,7 @@ class Router: deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600. ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False. + enable_mid_stream_fallback_continuation (bool): When True, a chat-completions stream that breaks after plain assistant text has been delivered continues on a fallback deployment via assistant prefill instead of surfacing the error. Only deployments whose model supports assistant prefill are eligible, so the partial text is continued, not regenerated; if none is, the original error is surfaced. Streams that emitted tool calls, thinking blocks, audio/images, or constrained (JSON / forced tool_choice) output are never continued. Async-only. Defaults to False. fallback_access_check (Optional[FallbackAccessCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects is skipped. Defaults to None (every configured fallback is attempted). fallback_budget_check (Optional[FallbackBudgetCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects as over budget is skipped. Defaults to None (budget is not re-checked on fallback). Returns: @@ -1313,6 +1409,7 @@ class Router: self.fallback_on_unsupported_params = fallback_on_unsupported_params self.enable_health_check_routing = enable_health_check_routing self.enable_weighted_failover = enable_weighted_failover + self.enable_mid_stream_fallback_continuation = self._is_flag_enabled(enable_mid_stream_fallback_continuation) self.health_check_ignore_transient_errors = health_check_ignore_transient_errors self.background_health_check_model_groups: frozenset[str] | None = ( frozenset(background_health_check_model_groups) @@ -1507,6 +1604,14 @@ class Router: default_pre_call_checks: Final[OptionalPreCallChecks] = [] self.add_optional_pre_call_checks(default_pre_call_checks) + if self.enable_mid_stream_fallback_continuation: + self._register_continuation_prefill_check() + + @staticmethod + def _register_continuation_prefill_check() -> None: + # Process-global on purpose: Router.discard() must not drop a filter another router still needs. + litellm.logging_callback_manager.add_litellm_callback(ContinuationPrefillDeploymentCheck()) + def discard(self): """ Pseudo-destructor to be invoked to clean up global data structures when router is no longer used. @@ -3101,15 +3206,53 @@ class Router: usage: Final = cast(Usage | None, getattr(fallback_item, "usage", None)) usage_objects: Final = [usage] if usage is not None else [] - if ( - complete_response_object_usage is not None - and hasattr(complete_response_object_usage, "usage") - and complete_response_object_usage.usage is not None - ): + if complete_response_object_usage is not None: usage_objects.append(complete_response_object_usage) combined_usage: Final = BaseTokenUsageProcessor.combine_usage_objects(usage_objects=usage_objects) setattr(fallback_item, "usage", combined_usage) + @staticmethod + def _price_cut_partial( + cut_stream: CustomStreamWrapper, + partial_response: ModelResponse | TextCompletionResponse | None, + ) -> "ContinuedPartial | None": + """Price the cut stream's partial, plus any partial it inherited, before a fallback can reuse its logging.""" + from litellm.cost_calculator import BaseTokenUsageProcessor + + cut_logging: Final = cut_stream.logging_obj + partial_usage: Final = getattr(partial_response, "usage", None) + if partial_response is None or not isinstance(partial_usage, Usage): + return None + inherited_usage: Final = cut_logging.model_call_details.get("continued_partial_usage") + inherited_cost: Final = cut_logging.model_call_details.get("continued_partial_cost") + inherited_id: Final = cut_logging.model_call_details.get("continued_response_id") + priced_partial: Final = partial_response.model_copy( + update={"model": cut_stream.model} if cut_stream.model else {} + ) + priced_partial._hidden_params = { # pyright: ignore[reportPrivateUsage] # no public accessor + key: value for key, value in get_hidden_params_dict(partial_response).items() if key != "response_cost" + } + return ContinuedPartial( + usage=BaseTokenUsageProcessor.combine_usage_objects( + usage_objects=[partial_usage, inherited_usage] + if isinstance(inherited_usage, Usage) + else [partial_usage] + ), + response_cost=(cut_logging.response_cost_calculator(result=priced_partial) or 0.0) + + (inherited_cost if isinstance(inherited_cost, float) else 0.0), + response_id=inherited_id if isinstance(inherited_id, str) else partial_response.id, + ) + + @staticmethod + def _carry_partial_spend_to_fallback(fallback_response: object, cut_partial: "ContinuedPartial | None") -> None: + """Bill what the cut legs consumed on the continuation's success log.""" + fallback_logging: Final = getattr(fallback_response, "logging_obj", None) + if cut_partial is None or not isinstance(fallback_logging, LiteLLMLogging): + return + fallback_logging.record_continued_partial( + usage=cut_partial.usage, response_cost=cut_partial.response_cost, response_id=cut_partial.response_id + ) + @staticmethod def _prepare_fallback_hidden_params( fallback_response: object, @@ -3159,7 +3302,7 @@ class Router: self, model_response: CustomStreamWrapper, messages: list[dict[str, str]], - initial_kwargs: dict, + initial_kwargs: dict[str, Any], deployment_slot: contextlib.AsyncExitStack | None = None, first_chunk: _FirstChunk[AsyncIterator[object]] | None = None, ) -> CustomStreamWrapper: @@ -3222,16 +3365,22 @@ class Router: with anyio.CancelScope(shield=True): await close_model_response() await held_slot.aclose() - if not e.is_pre_first_chunk and ( - e.generated_content or _stream_chunks_have_generated_content(model_response.chunks) - ): - if e.original_exception is not None: - raise e.original_exception from e - raise + committed: Final = bool( + not e.is_pre_first_chunk + and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) + ) + cut_hidden_params: Final = get_hidden_params(model_response) or {} + continue_after_content: Final = committed and self._mid_stream_continuation_eligible( + e=e, request_kwargs=initial_kwargs, cut_deployment_id=cut_hidden_params.get("model_id") + ) + if committed and not continue_after_content: + self._raise_original_mid_stream_error(e) from litellm.main import stream_chunk_builder - complete_response_object: Final = stream_chunk_builder(chunks=model_response.chunks) + complete_response_object: Final = stream_chunk_builder( + chunks=model_response.chunks, messages=messages if continue_after_content else None + ) complete_response_object_usage: Final = cast( Usage | None, getattr(complete_response_object, "usage", None), @@ -3247,7 +3396,41 @@ class Router: "content_policy_fallbacks", self.content_policy_fallbacks ) initial_kwargs["original_function"] = self._acompletion - initial_kwargs["messages"] = messages + continuation_messages: Final = ( + self._build_completion_continuation_input(messages, e.generated_content) + if continue_after_content + else None + ) + continuation_output: Final = ( + ContinuationOutput( + e.generated_content, + continuation_messages[-1], + model_group, + complete_response_object, + checks_post_call_rules=initial_kwargs.get(MID_STREAM_CONTINUATION_KWARG) + is not MID_STREAM_CONTINUATION_MARKER, + ) + if continuation_messages is not None + else None + ) + cut_partial: Final = ( + Router._price_cut_partial(model_response, complete_response_object) + if continuation_output is not None + else None + ) + if continue_after_content: + emitted_tokens: Final = int( + getattr(complete_response_object_usage, "completion_tokens", 0) or 0 + ) + reduced_ceilings: Final = continuation_output_ceilings(initial_kwargs, emitted_tokens) + if reduced_ceilings is None or continuation_messages is None: + self._raise_original_mid_stream_error(e) + await model_response.wait_for_failure_logging(timeout=MID_STREAM_FAILURE_LOG_WAIT_SECONDS) + initial_kwargs.update(reduced_ceilings) + initial_kwargs["messages"] = continuation_messages + initial_kwargs[MID_STREAM_CONTINUATION_KWARG] = MID_STREAM_CONTINUATION_MARKER + else: + initial_kwargs["messages"] = messages self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = await self.async_function_with_fallbacks_common_utils( e=e, @@ -3263,11 +3446,17 @@ class Router: # If fallback returns a streaming response, iterate over it if hasattr(fallback_response, "__aiter__"): + Router._carry_partial_spend_to_fallback(fallback_response, cut_partial) prepared_fallback_hidden_params = Router._adopt_fallback_response_headers( wrapper_ref, fallback_response ) fallback_headers_are_settled = False - async for fallback_item in fallback_response: + fallback_items: Final = ( + continuation_output.splice(fallback_response) + if continuation_output is not None + else fallback_response + ) + async for fallback_item in fallback_items: if not fallback_headers_are_settled: fallback_headers_are_settled = True # a fallback that failed over again only repoints itself once it yields @@ -3280,7 +3469,10 @@ class Router: and isinstance(fallback_item, ModelResponseStream) and hasattr(fallback_item, "usage") ): - self._combine_fallback_usage(fallback_item, complete_response_object_usage) + self._combine_fallback_usage( + fallback_item, + complete_response_object_usage if continuation_output is not None else None, + ) yield fallback_item else: # If fallback returns a non-streaming response, yield None @@ -3437,6 +3629,97 @@ class Router: total_tokens=(partial_usage.total_tokens or 0) + (fb.total_tokens or 0), ) + @staticmethod + def _raise_original_mid_stream_error(e: "MidStreamFallbackError") -> NoReturn: + """Surface the provider error the stream wrapper carried instead of the internal MidStreamFallbackError.""" + if e.original_exception is not None: + raise e.original_exception from e + raise e + + def _params_sent_to_deployment( + self, request_kwargs: Mapping[str, object], deployment: Deployment | None + ) -> Mapping[str, object]: + """The request as the deployment received it, merged in the order _update_kwargs_with_deployment applies.""" + deployment_params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {} + defaults: Final = {key: value for key, value in self.default_litellm_params.items() if value is not None} + deployment_tool_choice: Final = ( + {"tool_choice": deployment_params["tool_choice"]} + if deployment_params.get("tool_choice") is not None + else {} + ) + return MappingProxyType({**deployment_params, **defaults, **deployment_tool_choice, **request_kwargs}) + + def _names_deployment_that_cannot_continue(self, target: object) -> bool: + name: Final = fallback_target_name(target) + if name is None or name in self.model_names or not self.has_model_id(name): + return False + deployment: Final = self.get_deployment(model_id=name) + return deployment is not None and not deployment_can_continue(deployment.model_dump(exclude_none=True)) + + def _mid_stream_continuation_eligible( + self, + e: "MidStreamFallbackError", + request_kwargs: Mapping[str, object], + cut_deployment_id: object = None, + ) -> bool: + """Whether a stream that broke after plain assistant text may continue via prefill.""" + if not self.enable_mid_stream_fallback_continuation: + return False + if not e.generated_content or e.emitted_disqualifying_content: + return False + if isinstance(e.original_exception, litellm.ContentPolicyViolationError): + return False + cut_deployment: Final = ( + self.get_deployment(model_id=cut_deployment_id) if isinstance(cut_deployment_id, str) else None + ) + if isinstance(cut_deployment_id, str) and cut_deployment is None: + return False + if constrains_continuation(self._params_sent_to_deployment(request_kwargs, cut_deployment)): + return False + requested_group: Final = request_kwargs.get("model") + lookup_groups: Final = fallback_lookup_groups( + request_kwargs, requested_group if isinstance(requested_group, str) else None + ) + targets: Final = continuation_fallback_targets(request_kwargs.get("fallbacks", self.fallbacks), lookup_groups) + if any(fallback_target_overrides_request_params(target) for target in targets): + return False + if any(self._names_deployment_that_cannot_continue(target) for target in targets): + return False + return not any( + isinstance(callback, CustomGuardrail) + and callback._event_hook_is_event_type( # pyright: ignore[reportPrivateUsage] # the only check of a guardrail's hooks + GuardrailEventHooks.post_call + ) + for callback in litellm.callbacks + ) + + @staticmethod + def _build_completion_continuation_input( + messages: Sequence[Mapping[str, object]], + generated_content: str, + ) -> Sequence[Mapping[str, object]] | None: + """Append the partial output as an assistant prefill, or extend a trailing `prefix: True` assistant turn in + place so the request never ends in two assistant messages. None when the request ends in a plain assistant + turn, in a prefill turn whose content has a shape this cannot extend, or when the cut came mid-echo of the + prefill.""" + last: Final = messages[-1] if messages else None + if last is None or last.get("role") != "assistant": + return [*messages, {"role": "assistant", "content": generated_content, "prefix": True}] + if last.get("prefix") is not True: + return None + new_text: Final = continued_text(last, generated_content) + if new_text is None: + return None + content: Final = last.get("content") + if content is None or isinstance(content, str): + return [*messages[:-1], {**last, "content": f"{content or ''}{new_text}", "prefix": True}] + try: + blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(content) + except ValidationError: + return None + extended: Final = _extend_trailing_text_block(blocks, new_text) + return [*messages[:-1], {**last, "content": extended, "prefix": True}] + @staticmethod def _build_responses_continuation_input( input_val: Union[str, "ResponseInputParam"] | None, @@ -3838,13 +4121,6 @@ class Router: raise e.original_exception from e raise - from litellm.main import stream_chunk_builder - - complete_response_object: Final = stream_chunk_builder(chunks=model_response.chunks) - complete_response_object_usage: Final = cast( - Usage | None, - getattr(complete_response_object, "usage", None), - ) try: model_group: Final = cast(str, initial_kwargs.get("model")) fallbacks: Final[list | None] = initial_kwargs.get("fallbacks", router_self.fallbacks) @@ -3890,7 +4166,7 @@ class Router: and isinstance(fallback_item, ModelResponseStream) and hasattr(fallback_item, "usage") ): - router_self._combine_fallback_usage(fallback_item, complete_response_object_usage) + router_self._combine_fallback_usage(fallback_item, None) yield fallback_item else: yield None @@ -4255,7 +4531,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). @@ -4266,11 +4544,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"] @@ -4291,7 +4570,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"] @@ -5783,6 +6062,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, ) @@ -5792,6 +6072,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 @@ -5878,11 +6166,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 @@ -5898,6 +6186,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 @@ -5917,6 +6206,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, @@ -10720,6 +11010,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 @@ -11356,6 +11647,7 @@ class Router: Augment litellm info with additional params set in `model_info`. For azure models, ignore the `model:`. Only set max tokens, cost values if base_model is set. + For every other provider, a base_model naming the same model as `model:` can only widen its context window. Returns - ModelInfo - If found -> typed dict with max tokens, input cost, etc. @@ -11452,20 +11744,23 @@ class Router: ## CHECK USER SET MODEL INFO raw_user_model_info: Final = deployment.get("model_info") - user_model_info: Final = ( + user_model_info: Final = _MODEL_INFO_ADAPTER.validate_python( raw_user_model_info.model_dump(exclude_none=True) if isinstance(raw_user_model_info, BaseModel) - else raw_user_model_info + else raw_user_model_info or {} ) # get_model_info() hands back an lru_cache'd dict, so merge into a copy; unset # values are skipped or Deployment's None pricing defaults would erase the map's merged_model_info: Final[ModelMapInfo] = { **copy.deepcopy(model_info), - **self.get_discovered_model_info((deployment.get("model_info") or {}).get("id")), - **MappingProxyType( - {key: value for key, value in (user_model_info or MappingProxyType({})).items() if value is not None} + **_base_model_widened_window( + model_info, + _declared_base_model(user_model_info) or _declared_base_model(litellm_params.model_extra or {}), + custom_llm_provider, ), + **self.get_discovered_model_info((deployment.get("model_info") or {}).get("id")), + **MappingProxyType({key: value for key, value in user_model_info.items() if value is not None}), } return merged_model_info @@ -12655,16 +12950,34 @@ class Router: str(deployment_params.get("model") or ""), deployment_params.get("custom_llm_provider") ): return True + return Router._deployment_supports_param(deployment, group, param) + + @staticmethod + def _capability_model_and_provider(deployment_params: LiteLLMParamsTypedDict, group: str) -> tuple[str, str]: + model_label: Final = str(deployment_params.get("model") or group) + authenticating_provider: Final = declared_authenticating_provider( + model_label, deployment_params.get("custom_llm_provider") + ) + if authenticating_provider is not None: + return model_label.removeprefix(f"{authenticating_provider}/"), authenticating_provider + model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=model_label, + custom_llm_provider=deployment_params.get("custom_llm_provider"), + litellm_params=LiteLLM_Params.model_validate(deployment_params), + ) + return model, custom_llm_provider + + @staticmethod + def _deployment_supports_param(deployment: DeploymentTypedDict, group: str, param: str) -> bool: + deployment_params: Final = deployment.get("litellm_params") + if not deployment_params: + return True deployment_model_info: Final = deployment.get("model_info") base_model: Final = ( deployment_model_info.get("base_model") if deployment_model_info else None ) or deployment_params.get("base_model") try: - model, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=deployment_params.get("model") or group, - custom_llm_provider=deployment_params.get("custom_llm_provider"), - litellm_params=LiteLLM_Params.model_validate(deployment_params), - ) + model, custom_llm_provider = Router._capability_model_and_provider(deployment_params, group) supported: Final = litellm.get_supported_openai_params( model=model, custom_llm_provider=custom_llm_provider, @@ -12672,7 +12985,7 @@ class Router: ) except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request verbose_router_logger.debug( - "litellm.router.py::_deployment_accepts_param: keeping %s for model=%s. Got - %s", param, group, e + "litellm.router.py::_deployment_supports_param: keeping %s for model=%s. Got - %s", param, group, e ) return True return supported is None or param in supported @@ -12948,6 +13261,7 @@ class Router: "retry_policy", "model_group_alias", "enable_weighted_failover", + "enable_mid_stream_fallback_continuation", "enable_tag_filtering", "tag_routing_prefix", "fallback_on_unsupported_params", @@ -13004,12 +13318,10 @@ class Router: rebuild_routing_groups = True elif var == "optional_pre_call_checks": self.set_optional_pre_call_checks(kwargs[var]) + elif var == "enable_mid_stream_fallback_continuation": + self._set_mid_stream_fallback_continuation(kwargs[var]) elif var == "retry_policy": - value = kwargs[var] - if isinstance(value, dict): - value = RetryPolicy(**value) - if value is None or isinstance(value, RetryPolicy): - setattr(self, var, value) + self._set_retry_policy(kwargs[var]) else: value = kwargs[var] # only run routing strategy init if it has changed @@ -13036,6 +13348,20 @@ class Router: self._routing_groups_input = routing_groups_input verbose_router_logger.debug("Updated Router settings: %s", self.get_settings()) + @staticmethod + def _is_flag_enabled(value: object) -> bool: + return value is True or (isinstance(value, str) and str_to_bool(value) is True) + + def _set_mid_stream_fallback_continuation(self, enabled: object) -> None: + self.enable_mid_stream_fallback_continuation = self._is_flag_enabled(enabled) + if self.enable_mid_stream_fallback_continuation: + self._register_continuation_prefill_check() + + def _set_retry_policy(self, value: object) -> None: + policy: Final = RetryPolicy.model_validate(value) if isinstance(value, dict) else value + if policy is None or isinstance(policy, RetryPolicy): + self.retry_policy = policy + def _get_client(self, deployment, kwargs, client_type=None): """ Returns the appropriate client based on the given deployment, kwargs, and client_type. @@ -13169,7 +13495,8 @@ class Router: Returns None when no deployment limits its context window, and when counting fails. The caller pairs this with ``skip_inline_token_count`` so neither case puts the count back on - the loop: a failed count leaves the deployments unfiltered, exactly as before. + the loop: a failed count only skips the context-window check - RPM, region and + invalid-params checks still run for every deployment. """ if messages is None and input is None: return None @@ -13187,6 +13514,66 @@ class Router: ) return None + def _resolve_pre_call_check_input_tokens( + self, + messages: Sequence[Mapping[str, str]] | None, + input: str | Sequence[object] | None, + request_kwargs: Mapping[str, object] | None, + input_token_count: int | None, + count_inline: bool, + ) -> int | None: + if input_token_count is not None: + return input_token_count + if not count_inline: + return None + try: + return self._count_pre_call_check_tokens( + messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter + input=cast(str | list[object] | None, input), # cast-ok: forwarded to the sync counter + request_kwargs=request_kwargs, + ) + except Exception as e: + verbose_router_logger.error( + "litellm.router.py::_pre_call_checks: failed to count tokens. " + "Skipping the context-window check only; RPM/region/param checks still apply. Got - %s", + e, + ) + return None + + @staticmethod + def _params_neither_dropped_nor_allowed(params: Mapping[str, object], candidates: frozenset[str]) -> frozenset[str]: + if params.get("drop_params") is True: + return frozenset() + dropped: Final = params.get("additional_drop_params") + dropped_names: Final = ( + frozenset(entry for entry in dropped if isinstance(entry, str)) + if isinstance(dropped, list) + else frozenset() + ) + return candidates - Router._declared_param_allowlist(params) - dropped_names + + @staticmethod + def _deployment_lacks_params(deployment: DeploymentTypedDict, group: str, params: frozenset[str]) -> bool: + return any(not Router._deployment_supports_param(deployment, group, param) for param in params) + + @staticmethod + def _deployment_rejects_params( + deployment: DeploymentTypedDict, group: str, params: frozenset[str], request_kwargs: Mapping[str, object] + ) -> bool: + call_settings: Final[Mapping[str, object]] = {**(deployment.get("litellm_params") or {}), **request_kwargs} + return Router._deployment_lacks_params( + deployment, group, Router._params_neither_dropped_nor_allowed(call_settings, params) + ) + + @staticmethod + def _pre_call_check_requested_chat_params(request_kwargs: Mapping[str, object]) -> frozenset[str]: + requested: Final = litellm.utils.get_non_default_params(passed_params=dict(request_kwargs)) + return frozenset( + param + for param, value in requested.items() + if param in _PRE_CALL_CHECK_FILTERED_PARAMS and not isinstance(value, str) + ) + def _pre_call_checks( self, model: str, @@ -13196,6 +13583,7 @@ class Router: request_kwargs: dict | None = None, input_token_count: int | None = None, skip_inline_token_count: bool = False, + request_deployments: Sequence[DeploymentTypedDict] | None = None, ): """ Filter out model in model group, if: @@ -13214,13 +13602,6 @@ class Router: invalid_model_indices: Final = set() # Use set for O(1) membership checks - # Token counting (tiktoken) is the dominant on-loop cost for large prompts. - # Only count when a deployment actually declares max_input_tokens, and count - # at most once; for model groups with no context-window limit it is skipped. - # Async callers pass the count in, already computed off the event loop, and set - # skip_inline_token_count so a failed off-loop count is not retried back on the loop. - input_tokens: int | None = input_token_count - _context_window_error = False _potential_error_str = "" _rate_limit_error = False @@ -13229,6 +13610,24 @@ class Router: has_countable_input: Final = (messages is not None or input is not None) and not compaction_pending( request_kwargs ) + deployment_max_input_tokens: Final = tuple( + self._deployment_max_input_tokens(model, deployment) for deployment in _returned_deployments + ) + input_tokens: Final = self._resolve_pre_call_check_input_tokens( + messages=messages, + input=cast(str | list[object] | None, input), # cast-ok: forwarded to the resolver + request_kwargs=cast(Mapping[str, object] | None, request_kwargs), # cast-ok: forwarded to the resolver + input_token_count=input_token_count, + count_inline=has_countable_input + and not skip_inline_token_count + and any(limit is not None for limit in deployment_max_input_tokens), + ) + request_settings: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else {} + requested_params: Final[frozenset[str]] = ( + self._pre_call_check_requested_chat_params(request_settings) + if request_kwargs is not None and litellm.drop_params is False + else frozenset() + ) ## get model group RPM ## dt: Final = get_utc_datetime() @@ -13242,39 +13641,22 @@ class Router: _litellm_params = deployment.get("litellm_params", {}) _model_info = deployment.get("model_info", {}) - # see if we have the info for this model - _deployment_model = None # per-deployment model name (avoids overwriting the outer `model` group name) - try: - base_model = _model_info.get("base_model", None) - if base_model is None: - base_model = _litellm_params.get("base_model", None) - _deployment_model = base_model or _litellm_params.get("model", None) - model_info = self.get_router_model_info(deployment=deployment, received_model_name=model) - - max_input_tokens = model_info.get("max_input_tokens") if isinstance(model_info, dict) else None - if isinstance(max_input_tokens, int) and has_countable_input: - if input_tokens is None: - if skip_inline_token_count: - return _returned_deployments - try: - input_tokens = self._count_pre_call_check_tokens( - messages=messages, input=input, request_kwargs=request_kwargs - ) - except Exception as e: - verbose_router_logger.error( - "litellm.router.py::_pre_call_checks: failed to count tokens. Returning initial list of deployments. Got - %s", - e, - ) - return _returned_deployments - if input_tokens > max_input_tokens: - invalid_model_indices.add(idx) - _context_window_error = True - _potential_error_str += ( - f"Model={_deployment_model}, Max Input Tokens={max_input_tokens}, Got={input_tokens}" - ) - continue - except Exception as e: - verbose_router_logger.exception("An error occurs - %s", e) + max_input_tokens = deployment_max_input_tokens[idx] + if ( + max_input_tokens is not None + and has_countable_input + and input_tokens is not None + and input_tokens > max_input_tokens + ): + invalid_model_indices.add(idx) + _context_window_error = True + deployment_model = ( + _model_info.get("base_model") or _litellm_params.get("base_model") or _litellm_params.get("model") + ) + _potential_error_str += ( + f"Model={deployment_model}, Max Input Tokens={max_input_tokens}, Got={input_tokens}" + ) + continue model_id = _model_info.get("id", "") ## RPM CHECK ## @@ -13306,53 +13688,57 @@ class Router: invalid_model_indices.add(idx) continue - ## INVALID PARAMS ## -> catch 'gpt-3.5-turbo-16k' not supporting 'response_format' param - if request_kwargs is not None and litellm.drop_params is False: - # get supported params — use per-deployment model to avoid overwriting the outer model group name - _dep_model_for_params: str = _deployment_model or model - try: - ( - _dep_model_for_params, - custom_llm_provider, - _, - _, - ) = litellm.get_llm_provider( - model=_dep_model_for_params, - litellm_params=LiteLLM_Params.model_validate(_litellm_params), - ) - except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not fail the request - verbose_router_logger.debug( - "litellm.router.py::_pre_call_checks: skipping supported-params check for model=%s. Got - %s", - _dep_model_for_params, - e, - ) - continue + typed_deployments: Final = tuple( + cast(DeploymentTypedDict, deployment) # cast-ok: router deployments are plain dicts + for deployment in _returned_deployments + ) + param_candidates: Final = tuple( + (idx, deployment) for idx, deployment in enumerate(typed_deployments) if idx not in invalid_model_indices + ) + candidates_lacking_params: Final = tuple( + self._deployment_lacks_params(deployment, model, requested_params) for _, deployment in param_candidates + ) + no_deployment_takes_params: Final = bool(candidates_lacking_params) and all(candidates_lacking_params) + candidates_failing_params: Final = ( + tuple( + self._deployment_rejects_params(deployment, model, requested_params, request_settings) + for _, deployment in param_candidates + ) + if no_deployment_takes_params + else candidates_lacking_params + ) + excluded_indices: Final = frozenset(invalid_model_indices) | frozenset( + idx for (idx, _), fails in zip(param_candidates, candidates_failing_params, strict=True) if fails + ) - supported_openai_params = litellm.get_supported_openai_params( - model=_dep_model_for_params, - custom_llm_provider=custom_llm_provider, - ) - - if supported_openai_params is None: - continue - else: - # check the non-default openai params in request kwargs - non_default_params = litellm.utils.get_non_default_params(passed_params=request_kwargs) - special_params = ["response_format"] - # check if all params are supported - for k, v in non_default_params.items(): - if k not in supported_openai_params and k in special_params: - # if not -> invalid model - verbose_router_logger.debug("INVALID MODEL INDEX @ REQUEST KWARG FILTERING, k=%s", k) - invalid_model_indices.add(idx) - - if len(invalid_model_indices) == len(_returned_deployments): + if len(excluded_indices) == len(_returned_deployments): """ - no healthy deployments available b/c context window checks or rate limit error - First check for rate limit errors (if this is true, it means the model passed the context window check but failed the rate limit check) """ + group_deployments: Final = ( + tuple(request_deployments) + if request_deployments is not None + else (tuple(self.get_all_deployments(model_name=model)) or typed_deployments) + ) + if ( + not _context_window_error + and group_deployments + and all( + self._deployment_rejects_params(deployment, model, requested_params, request_settings) + for deployment in group_deployments + ) + ): + raise litellm.UnsupportedParamsError( + message=( + f"no deployment of model={model} supports {sorted(requested_params)}. " + "To drop these, set `litellm.drop_params=True` or for proxy `litellm_settings: drop_params: true`, " + f"or send allowed_openai_params={sorted(requested_params)} to pass them through" + ), + model=model, + ) if _rate_limit_error is True: # allow generic fallback logic to take place raise RouterRateLimitErrorBasic( model=model, @@ -13364,9 +13750,9 @@ class Router: model=model, llm_provider="", ) - if len(invalid_model_indices) > 0: + if len(excluded_indices) > 0: # Single-pass filter using set for O(1) lookups (avoids O(n^2) from repeated pops) - _returned_deployments = [d for i, d in enumerate(_returned_deployments) if i not in invalid_model_indices] + _returned_deployments = [d for i, d in enumerate(_returned_deployments) if i not in excluded_indices] return _returned_deployments @@ -13388,6 +13774,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, @@ -13802,6 +14197,10 @@ class Router: ) return healthy_deployments + request_deployments: Final = cast( # cast-ok: router deployments are plain dicts + Sequence[DeploymentTypedDict], healthy_deployments + ) + # Health-check-based filtering (before cooldown) healthy_deployments = await self._async_filter_health_check_unhealthy_deployments( healthy_deployments=healthy_deployments, @@ -13861,6 +14260,7 @@ class Router: request_kwargs=request_kwargs, ), skip_inline_token_count=True, + request_deployments=request_deployments, ) # check if user wants to do tag based routing healthy_deployments = await get_deployments_for_tag( @@ -15011,9 +15411,8 @@ class Router: request_kwargs=request_kwargs, ) strategy, strategy_selector = self._get_routing_context(model, request_kwargs) - pick_attributes: Final = _deployment_pick_attributes( - model, _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs) - ) + routing_kwargs: Final = _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs) + pick_attributes: Final = _deployment_pick_attributes(model, routing_kwargs) if isinstance(healthy_deployments, dict): if (healthy_deployments.get("model_info") or {}).get("blocked") is True: @@ -15027,6 +15426,9 @@ class Router: return healthy_deployments parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(request_kwargs) + request_deployments: Final = cast( # cast-ok: router deployments are plain dicts + Sequence[DeploymentTypedDict], healthy_deployments + ) # Health-check-based filtering (before cooldown) healthy_deployments = self._filter_health_check_unhealthy_deployments( @@ -15048,7 +15450,9 @@ class Router: ) healthy_deployments = _pre_cooldown_deployments - healthy_deployments = self._filter_blocked_deployments(healthy_deployments) + healthy_deployments = continuation_deployments( + self._filter_blocked_deployments(healthy_deployments), routing_kwargs + ) # filter pre-call checks if self.enable_pre_call_checks and (messages is not None or input is not None): @@ -15058,6 +15462,7 @@ class Router: messages=messages, input=input, request_kwargs=request_kwargs, + request_deployments=request_deployments, ) ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) @@ -15221,6 +15626,9 @@ class Router: # 4. Apply health-check and cooldown filtering parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(request_kwargs) + request_deployments: Final = cast( # cast-ok: router deployments are plain dicts + Sequence[DeploymentTypedDict], pass_through_deployments + ) pass_through_deployments = self._filter_health_check_unhealthy_deployments( healthy_deployments=pass_through_deployments, parent_otel_span=parent_otel_span, @@ -15242,6 +15650,7 @@ class Router: messages=messages, input=input, request_kwargs=request_kwargs, + request_deployments=request_deployments, ) if self._is_priority_routing_group(model): 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 e12304c6ed1..9b9a3d185cc 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -27,7 +27,8 @@ from litellm.types.utils import LlmProviders, ProviderSpecificModelInfo _V = TypeVar("_V") -ROUTER_ONLY_CALL_KWARGS: Final = frozenset({"silent_model", "include_fallback_errors"}) +MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" +ROUTER_ONLY_CALL_KWARGS: Final = frozenset({"silent_model", "include_fallback_errors", MID_STREAM_CONTINUATION_KWARG}) _OBJECT_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapter(Mapping[object, object]) @@ -308,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. @@ -342,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/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 1f37c6c8cb1..06bd3329cc1 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -615,6 +615,32 @@ def get_fallback_model_group( return fallback_model_group, generic_fallback_idx +def _resolved_fallback_chain(entries: Sequence[object], lookup_groups: tuple[str, ...]) -> Sequence[object]: + group_entries: Final[Sequence[Mapping[str, list[str]]]] = tuple( # mutable-ok: mirrors the fallbacks config shape + entry for entry in entries if isinstance(entry, Mapping) and entry + ) + chain, _ = get_fallback_model_group_for_lookup_groups(fallbacks=group_entries, lookup_groups=lookup_groups) + return chain if isinstance(chain, list) else () + + +def continuation_fallback_targets(fallbacks: object, lookup_groups: tuple[str, ...]) -> Sequence[object]: + entries: Final[Sequence[object]] = fallbacks if isinstance(fallbacks, list) else () + return ( + entries + if check_non_standard_fallback_format(fallbacks=entries) + else _resolved_fallback_chain(entries, lookup_groups) + ) + + +def fallback_target_overrides_request_params(target: object) -> bool: + return isinstance(target, Mapping) and "model" in target and len(target) > 1 + + +def fallback_target_name(target: object) -> str | None: + name: Final = target.get("model") if isinstance(target, Mapping) else target + return name if isinstance(name, str) else None + + PROVIDER_SCOPED_RESOURCE_KEYS: Final = ("input_file_id", "training_file", "batch_id", "file_id", "fine_tuning_job_id") PROVIDER_SCOPED_RESOURCE_FUNCTION_NAMES: Final = frozenset( { diff --git a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py new file mode 100644 index 00000000000..8f211ef67db --- /dev/null +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -0,0 +1,264 @@ +"""Keeps a mid-stream continuation on deployments whose model supports assistant prefill and whose own params leave +the output unconstrained.""" + +from collections.abc import AsyncIterable, AsyncIterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import accumulate +from os.path import commonprefix +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +import litellm +from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.rules import Rules +from litellm.router_utils.common_utils import MID_STREAM_CONTINUATION_KWARG +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ModelResponse, ModelResponseStream, Usage +from litellm.utils import supports_assistant_prefill + + +class _ContinuationMarker: + """Only the router constructs one, so a JSON request body cannot forge the continuation flag.""" + + +MID_STREAM_CONTINUATION_MARKER: Final = _ContinuationMarker() + +_STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) +_UNCONSTRAINED_RESPONSE_FORMAT: Final[Mapping[str, str]] = MappingProxyType({"type": "text"}) +_OUTPUT_CEILING_KEYS: Final = ("max_tokens", "max_completion_tokens") + + +def _str_keyed(value: object) -> Mapping[str, object]: + try: + return _STR_KEYED_DICT_ADAPTER.validate_python(value) + except ValidationError: + return _NO_FIELDS + + +def deployment_supports_prefill(deployment: object) -> bool: + deployment_map: Final = _str_keyed(deployment) + model_info: Final = _str_keyed(deployment_map.get("model_info")) + declared: Final = model_info.get("supports_assistant_prefill") + if isinstance(declared, bool): + return declared + litellm_params: Final = _str_keyed(deployment_map.get("litellm_params")) + candidates: Final = (model_info.get("base_model"), litellm_params.get("base_model"), litellm_params.get("model")) + model: Final = next((candidate for candidate in candidates if isinstance(candidate, str) and candidate), None) + return model is not None and supports_assistant_prefill(model=model) + + +def constrains_continuation(params: Mapping[str, object]) -> bool: + sent: Final = {**params, **_str_keyed(params.get("extra_body"))} + choice_count: Final = sent.get("n") + response_format: Final = sent.get("response_format") + tool_choice: Final = sent.get("tool_choice") + return ( + (isinstance(choice_count, int) and choice_count > 1) + or (response_format is not None and response_format != _UNCONSTRAINED_RESPONSE_FORMAT) + or tool_choice == "required" + or isinstance(tool_choice, Mapping) + or sent.get("merge_reasoning_content_in_choices") is True + ) + + +def deployment_can_continue(deployment: object) -> bool: + litellm_params: Final = _str_keyed(_str_keyed(deployment).get("litellm_params")) + return deployment_supports_prefill(deployment) and not constrains_continuation(litellm_params) + + +def continuation_deployments( + healthy_deployments: list[dict[str, object]], # mutable-ok: router deployment-list contract + request_kwargs: Mapping[str, object] | None, +) -> list[dict[str, object]]: # mutable-ok: returns a mutable deployment list + marker: Final = (request_kwargs or {}).get(MID_STREAM_CONTINUATION_KWARG) + if not isinstance(marker, _ContinuationMarker): + return healthy_deployments + return [deployment for deployment in healthy_deployments if deployment_can_continue(deployment)] + + +class ContinuationPrefillDeploymentCheck(CustomLogger): + async def async_filter_deployments( + self, + model: str, + healthy_deployments: list[dict[str, object]], # mutable-ok: CustomLogger deployment-list contract + messages: Sequence[AllMessageValues] | None, + request_kwargs: Mapping[str, object] | None = None, + parent_otel_span: Span | None = None, + ) -> list[dict[str, object]]: # mutable-ok: returns a mutable deployment list + return continuation_deployments(healthy_deployments, request_kwargs) + + +def _delta_text(chunk: ModelResponseStream) -> str: + if not chunk.choices: + return "" + content: Final = chunk.choices[0].delta.content + return content if isinstance(content, str) else "" + + +_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) + + +def _prefill_text(prefill_message: Mapping[str, object]) -> str: + content: Final = prefill_message.get("content") + if isinstance(content, str): + return content + try: + blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(content) + except ValidationError: + return "" + return "".join( + text for block in blocks if block.get("type") == "text" and isinstance(text := block.get("text"), str) + ) + + +def _resendable_whitespace(prefill: str) -> str: + return prefill[len(prefill.rstrip()) :] if prefill.strip() else "" + + +def continued_text(prefill_message: Mapping[str, object], generated_content: str) -> str | None: + prefill: Final = _prefill_text(prefill_message) + if generated_content and len(generated_content) < len(prefill) and prefill.startswith(generated_content): + return None + echo: Final = len(prefill) if prefill and generated_content.startswith(prefill) else 0 + after_echo: Final = generated_content[echo:] + return after_echo[len(commonprefix((after_echo, _resendable_whitespace(prefill)))) :] + + +def _reduced_ceilings(params: Mapping[str, object], emitted_tokens: int) -> Mapping[str, int]: + return MappingProxyType( + {key: value - emitted_tokens for key in _OUTPUT_CEILING_KEYS if isinstance(value := params.get(key), int)} + ) + + +def continuation_output_ceilings( + request_kwargs: Mapping[str, object], emitted_tokens: int +) -> Mapping[str, object] | None: + """The caller's output ceilings, top-level and in extra_body, minus the tokens already emitted, or None once + nothing is left.""" + extra_body: Final = _str_keyed(request_kwargs.get("extra_body")) + top_level: Final = _reduced_ceilings(request_kwargs, emitted_tokens) + in_extra_body: Final = _reduced_ceilings(extra_body, emitted_tokens) + if any(ceiling <= 0 for ceiling in (*top_level.values(), *in_extra_body.values())): + return None + if not in_extra_body: + return top_level + return MappingProxyType({**top_level, "extra_body": {**extra_body, **in_extra_body}}) + + +def _with_first_delta(chunk: ModelResponseStream, update: Mapping[str, object]) -> ModelResponseStream: + choice: Final = chunk.choices[0] + delta: Final = choice.delta.model_copy(update=update) + return chunk.model_copy(update={"choices": [choice.model_copy(update={"delta": delta}), *chunk.choices[1:]]}) + + +def _without_leading_text(chunk: ModelResponseStream, count: int) -> ModelResponseStream: + if count <= 0: + return chunk + return _with_first_delta(chunk, {"content": _delta_text(chunk)[count:]}) + + +def _without_role(chunk: ModelResponseStream) -> ModelResponseStream: + if not chunk.choices or chunk.choices[0].delta.role is None: + return chunk + return _with_first_delta(chunk, {"role": None}) + + +def _stream_identity(primary_response: object) -> Mapping[str, object]: + if not isinstance(primary_response, ModelResponse): + return _NO_FIELDS + return MappingProxyType({"id": primary_response.id, "created": primary_response.created}) + + +def _drop_leading_text(chunks: tuple[ModelResponseStream, ...], count: int) -> tuple[ModelResponseStream, ...]: + lengths: Final = tuple(len(_delta_text(chunk)) for chunk in chunks) + offsets: Final = tuple(accumulate(lengths, initial=0))[:-1] + return tuple( + _without_leading_text(chunk, min(length, count - offset)) + for chunk, offset, length in zip(chunks, offsets, lengths, strict=True) + ) + + +@dataclass(frozen=True, slots=True) +class ContinuedPartial: + usage: Usage + response_cost: float + response_id: str + + +class ContinuationOutput: + """Splices a fallback's continuation onto the stream the client already received. Drops a leading echo of the + prefill, since some providers (Mistral) return the prefix inside the completion, and the prefill's trailing + whitespace sent again by a provider that strips it (Anthropic), keeps the primary's chunk id, created time and single role delta, and + runs litellm.post_call_rules over the joined text, since each stream wrapper only sees its own fragment. A + continuation nested in another leaves the rules to the outermost one, the only one holding the whole text.""" + + def __init__( + self, + emitted_text: str, + prefill_message: Mapping[str, object], + model: str, + primary_response: object, + checks_post_call_rules: bool = True, + ) -> None: + self._text = emitted_text + self._checks_post_call_rules: Final = checks_post_call_rules + self._prefill_text: Final = _prefill_text(prefill_message) + self._model: Final = model + self._identity: Final = _stream_identity(primary_response) + self._rules: Final = Rules() + self._held: tuple[ModelResponseStream, ...] = () + self._echo_settled = not self._prefill_text + self._seam = _resendable_whitespace(self._prefill_text) + + async def splice(self, stream: AsyncIterable[object]) -> AsyncIterator[object]: + async for item in stream: + if not isinstance(item, ModelResponseStream): + yield item + continue + for released in self._observed(self._past_seam(self._released(item))): + yield self._in_primary_stream(released) + for released in self._observed(self._past_seam(self._flushed())): + yield self._in_primary_stream(released) + + def _in_primary_stream(self, chunk: ModelResponseStream) -> ModelResponseStream: + return _without_role(chunk.model_copy(update=self._identity) if self._identity else chunk) + + def _flushed(self) -> tuple[ModelResponseStream, ...]: + held: Final = self._held + self._held = () + self._echo_settled = True + return held + + def _released(self, chunk: ModelResponseStream) -> tuple[ModelResponseStream, ...]: + if self._echo_settled: + return (chunk,) + held: Final = (*self._held, chunk) + text: Final = "".join(_delta_text(held_chunk) for held_chunk in held) + if len(text) < len(self._prefill_text) and self._prefill_text.startswith(text): + self._held = held + return () + self._held = () + self._echo_settled = True + if not text.startswith(self._prefill_text): + return held + return _drop_leading_text(held, len(self._prefill_text)) + + def _past_seam(self, chunks: tuple[ModelResponseStream, ...]) -> tuple[ModelResponseStream, ...]: + return tuple(self._without_seam(chunk) for chunk in chunks) + + def _without_seam(self, chunk: ModelResponseStream) -> ModelResponseStream: + text: Final = _delta_text(chunk) + matched: Final = len(commonprefix((text, self._seam))) + self._seam = self._seam[matched:] if matched == len(text) else "" + return _without_leading_text(chunk, matched) + + def _observed(self, chunks: tuple[ModelResponseStream, ...]) -> tuple[ModelResponseStream, ...]: + if not self._checks_post_call_rules or not litellm.post_call_rules: + return chunks + for chunk in chunks: + self._text += _delta_text(chunk) + self._rules.post_call_rules(input=self._text, model=self._model) + return chunks 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 096ad1c30dd..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 @@ -260,6 +261,7 @@ class PromptOptions: prompt_label: str | None = None litellm_system_prompt: str | None = None custom_prompt_dict: Mapping[str, object] | None = None + initial_prompt_value: str | None = None roles: Mapping[str, object] | None = None final_prompt_value: str | None = None bos_token: str | 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/llms/openai.py b/litellm/types/llms/openai.py index 8731bc2905a..8cc9e912e80 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1421,10 +1421,26 @@ class WebSearchToolUsage(LiteLLMBaseModel): num_requests: NonNegativeInt +class ImageGenTokenDetails(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + image_tokens: NonNegativeInt = 0 + text_tokens: NonNegativeInt = 0 + + +class ImageGenToolUsage(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + total_tokens: NonNegativeInt = 0 + input_tokens_details: ImageGenTokenDetails | None = None + output_tokens_details: ImageGenTokenDetails | None = None + + class ResponsesToolUsage(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) web_search: WebSearchToolUsage | None = None + image_gen: ImageGenToolUsage | None = None ResponsesAPIStatus = Literal["completed", "failed", "in_progress", "cancelled", "queued", "incomplete"] 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 6ae7a77d571..54bacfc84a3 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -497,6 +497,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 @@ -704,6 +705,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 @@ -1401,8 +1403,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 0bcb732a0f5..054a0a1c034 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 @@ -4045,7 +4046,9 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"}) -DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset({"supports_regex_lookaround"}) +DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset( + {"supports_regex_lookaround", "supports_assistant_prefill"} +) SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__) @@ -4475,6 +4478,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 c0c80f09e4a..3b43c570466 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -89,8 +89,10 @@ from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, ) +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 @@ -473,6 +475,7 @@ from .exceptions import ( BudgetExceededError, ContentPolicyViolationError, ContextWindowExceededError, + MalformedRequestFieldError, ModelNotMappedError, NotFoundError, OpenAIError, @@ -1373,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, @@ -1493,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. @@ -1504,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: @@ -3059,6 +3088,20 @@ def supports_prompt_cache_breakpoint(model: str, custom_llm_provider: str | None ) +def supports_assistant_prefill(model: str, custom_llm_provider: str | None = None) -> bool: + """ + Whether the model can continue a prefilled assistant message (Anthropic's + ``prefix: True`` trick and equivalents). Missing metadata reads as False, + so a mid-stream fallback continuation is only routed to a model known to + support it. + """ + return _supports_factory( + model=model, + custom_llm_provider=custom_llm_provider, + key="supports_assistant_prefill", + ) + + def supports_thinking_cache_preservation(model: str, custom_llm_provider: str | None = None) -> bool: return supports_factory( model=model, @@ -3402,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. @@ -3430,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 = {} @@ -3515,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 @@ -9236,7 +9238,7 @@ class ProviderConfigManager: return XAIAudioTranscriptionConfig() elif litellm.LlmProviders.OPENAI == provider: - if "gpt-4o" in model: + if "gpt-4o" in model or (model.startswith("gpt-") and "transcribe" in model): return litellm.OpenAIGPTAudioTranscriptionConfig() else: return litellm.OpenAIWhisperAudioTranscriptionConfig() @@ -9932,6 +9934,10 @@ class ProviderConfigManager: return MinimaxVideoConfig() elif LlmProviders.EDENAI == provider: return litellm.EdenAIVideoConfig() + elif LlmProviders.VOLCENGINE == provider: + from litellm.llms.volcengine.videos.transformation import VolcEngineVideoConfig + + return VolcEngineVideoConfig() return None @staticmethod @@ -9956,6 +9962,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 @@ -9964,7 +9971,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 @@ -10446,22 +10453,21 @@ def get_non_default_completion_params(kwargs: Mapping[str, object]) -> dict[str, return filter_out_litellm_params(kwargs, excluding=litellm.OPENAI_CHAT_COMPLETION_PARAMS) -def peek_reasoning_summary_aliases(optional_params: dict) -> object | None: - """Read AI-SDK-style reasoning summary from optional_params or nested extra_body. - - Uses key membership (not ``or`` chains) so falsy values like ``""`` are not skipped. - """ - if "reasoningSummary" in optional_params: - return optional_params["reasoningSummary"] - if "reasoning_summary" in optional_params: - return optional_params["reasoning_summary"] - extra_body: Final = optional_params.get("extra_body") - if isinstance(extra_body, dict): - if "reasoningSummary" in extra_body: - return extra_body["reasoningSummary"] - if "reasoning_summary" in extra_body: - return extra_body["reasoning_summary"] - return None +def reasoning_effort_with_summary_alias( + reasoning_effort: object, params: Mapping[str, object], accepts_summary: bool +) -> object: + summary: Final = peek_reasoning_summary_aliases(params) + if not accepts_summary or not isinstance(summary, str): + return reasoning_effort + match reasoning_effort: + case str(): + return {"effort": reasoning_effort, "summary": summary} + case {"summary": _}: + return reasoning_effort + case {"effort": str() as effort}: + return {"effort": effort, "summary": summary} + case _: + return reasoning_effort def strip_reasoning_summary_aliases_from_optional_params( 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 a204c49967e..c3efaf168ab 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -342,26 +342,33 @@ }, "writer.palmyra-vision-7b": { "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 7.5e-08, "input_cost_per_token_flex": 7.5e-08, + "input_cost_per_token_priority": 2.625e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 4096, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_batches": 3e-07, "output_cost_per_token_flex": 3e-07, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token_priority": 1.05e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_vision": true }, "amazon.nova-lite-v1:0": { "cache_read_input_token_cost": 1.5e-08, "input_cost_per_token": 6e-08, + "input_cost_per_token_batches": 3e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 300000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 2.4e-07, + "output_cost_per_token_batches": 1.2e-07, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -371,13 +378,19 @@ }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_priority": 1.3125e-07, "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_priority": 5.25e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_priority": 4.375e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -392,13 +405,17 @@ "input_cost_per_token": 1.25e-06, "input_cost_per_image_token": 1.25e-06, "input_cost_per_audio_token": 1.25e-06, + "input_cost_per_token_batches": 6.25e-07, + "input_cost_per_token_priority": 2.1875e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1e-05, - "source": "https://aws.amazon.com/nova/pricing/", + "output_cost_per_token_batches": 5e-06, + "output_cost_per_token_priority": 1.75e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -489,15 +506,21 @@ "us.amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 8.25e-08, "cache_read_input_token_cost_flex": 4.125e-08, + "cache_read_input_token_cost_priority": 1.44375e-07, "input_cost_per_token": 3.3e-07, + "input_cost_per_token_batches": 1.65e-07, "input_cost_per_token_flex": 1.65e-07, + "input_cost_per_token_priority": 5.775e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.75e-06, + "output_cost_per_token_batches": 1.375e-06, "output_cost_per_token_flex": 1.375e-06, + "output_cost_per_token_priority": 4.8125e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -512,13 +535,17 @@ "input_cost_per_token": 1.375e-06, "input_cost_per_image_token": 1.375e-06, "input_cost_per_audio_token": 1.375e-06, + "input_cost_per_token_batches": 6.875e-07, + "input_cost_per_token_priority": 2.40625e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.1e-05, - "source": "https://aws.amazon.com/nova/pricing/", + "output_cost_per_token_batches": 5.5e-06, + "output_cost_per_token_priority": 1.925e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -548,12 +575,15 @@ "amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, "input_cost_per_token": 3.5e-08, + "input_cost_per_token_batches": 1.75e-08, "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.4e-07, + "output_cost_per_token_batches": 7e-08, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, @@ -561,13 +591,19 @@ }, "amazon.nova-premier-v1:0": { "cache_read_input_token_cost": 6.25e-07, + "cache_read_input_token_cost_priority": 1.09375e-06, "input_cost_per_token": 2.5e-06, + "input_cost_per_token_batches": 1.25e-06, + "input_cost_per_token_priority": 4.375e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 1.25e-05, + "output_cost_per_token_batches": 6.25e-06, + "output_cost_per_token_priority": 2.1875e-05, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -578,15 +614,21 @@ "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 8e-07, + "input_cost_per_token_batches": 4e-07, "input_cost_per_token_flex": 4e-07, + "input_cost_per_token_priority": 1.4e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 300000, "max_output_tokens": 10000, "max_tokens": 10000, "mode": "chat", "output_cost_per_token": 3.2e-06, + "output_cost_per_token_batches": 1.6e-06, "output_cost_per_token_flex": 1.6e-06, + "output_cost_per_token_priority": 5.6e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", "supports_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -2773,7 +2815,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2808,7 +2850,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2843,7 +2885,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2878,7 +2920,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2913,7 +2955,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2948,7 +2990,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -3743,7 +3785,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": 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, @@ -15494,7 +15536,7 @@ }, "supports_adaptive_thinking": true, "supports_legacy_thinking": true, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -20832,7 +20874,7 @@ "output_cost_per_token": 2.5000010000000002e-05, "output_dbu_cost_per_token": 0.000357143, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_legacy_thinking": true, "supports_prompt_caching": true, @@ -21047,7 +21089,7 @@ "output_cost_per_token": 1.5000020000000002e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_legacy_thinking": 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 @@ -29256,22 +29302,32 @@ "output_cost_per_token_priority": 1.8e-05 }, "gemini-3.1-pro-preview": { - "prompt_cache_min_tokens": 4096, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "output_cost_per_image": 0.00012, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096.0, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29289,33 +29345,18 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_url_context": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 3.6e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "output_cost_per_token_priority": 2.16e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "cache_read_input_token_cost_priority": 3.6e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_flex": 2e-07, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "web_search_billing_unit": "per_query" }, "gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -29428,15 +29469,25 @@ "vertex_ai/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_flex": 5e-08, - "input_cost_per_token": 5e-07, + "cache_read_input_token_cost_priority": 9e-08, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_batches": 5e-07, + "input_cost_per_audio_token_priority": 1.8e-06, + "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": 9e-07, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65535, "max_tokens": 65535, "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.4e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29454,6 +29505,7 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -29463,23 +29515,7 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 9e-07, - "input_cost_per_audio_token_priority": 1.8e-06, - "output_cost_per_token_priority": 5.4e-06, - "cache_read_input_token_cost_priority": 9e-08, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 5e-07, - "input_cost_per_token_batches": 2.5e-07, - "input_cost_per_token_flex": 2.5e-07, - "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -29807,22 +29843,32 @@ "output_cost_per_video_token": 1e-06 }, "vertex_ai/gemini-3.1-pro-preview": { - "prompt_cache_min_tokens": 4096, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "output_cost_per_image": 0.00012, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096.0, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29840,33 +29886,18 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_url_context": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 3.6e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "output_cost_per_token_priority": 2.16e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "cache_read_input_token_cost_priority": 3.6e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_flex": 2e-07, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -31666,9 +31697,15 @@ "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_flex": 5e-08, + "cache_read_input_token_cost_priority": 9e-08, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_batches": 5e-07, + "input_cost_per_audio_token_priority": 1.8e-06, "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": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65535, @@ -31676,6 +31713,10 @@ "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "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.4e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -31693,6 +31734,7 @@ ], "supports_audio_output": false, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -31703,23 +31745,7 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 9e-07, - "input_cost_per_audio_token_priority": 1.8e-06, - "output_cost_per_token_priority": 5.4e-06, - "cache_read_input_token_cost_priority": 9e-08, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 5e-07, - "input_cost_per_token_batches": 2.5e-07, - "input_cost_per_token_flex": 2.5e-07, - "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "web_search_billing_unit": "per_query" }, "gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, @@ -33610,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, @@ -33628,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, @@ -33646,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, @@ -33767,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, @@ -40354,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 }, @@ -40428,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 }, @@ -40839,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, @@ -40858,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", @@ -40884,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 @@ -41072,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, @@ -41108,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, @@ -41126,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, @@ -41144,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, @@ -41197,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, @@ -41241,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 @@ -41257,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 @@ -42076,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, @@ -42093,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 }, @@ -43197,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, @@ -43215,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, @@ -43232,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, @@ -43251,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, @@ -44751,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", @@ -44759,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, @@ -44774,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", @@ -44782,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, @@ -44792,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, @@ -44810,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, @@ -45674,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, @@ -45688,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, @@ -45816,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, @@ -45832,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 @@ -47862,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, @@ -47878,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, @@ -47892,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, @@ -48619,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, @@ -48668,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 }, @@ -48742,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 }, @@ -49196,7 +49365,7 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_prompt_caching": true, @@ -51026,7 +51195,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -51199,7 +51368,7 @@ "supported_regions": [ "global" ], - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -51220,7 +51389,7 @@ "supported_regions": [ "us-central1" ], - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -53450,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 @@ -53488,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 @@ -58842,7 +59023,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -58871,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", @@ -58879,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" @@ -58892,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", @@ -58900,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" @@ -58913,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", @@ -58921,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" ], @@ -58932,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", @@ -58940,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" ], @@ -59573,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, @@ -59612,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, @@ -59695,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, @@ -59734,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, @@ -59819,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, @@ -59854,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, @@ -59892,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, @@ -59927,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, @@ -60022,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", @@ -60030,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", @@ -60043,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", @@ -60051,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", @@ -60064,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", @@ -60072,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", @@ -60088,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", @@ -60097,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": [ @@ -60108,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, @@ -60487,6 +60702,61 @@ } ] }, + "volcengine/doubao-seedance-2-5-260628": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.216, + "output_cost_per_second_480p": 0.096075, + "output_cost_per_second_720p": 0.216, + "output_cost_per_second_1080p": 0.5346, + "supported_modalities": ["text", "image", "video", "audio"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "854x480", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "480x854", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s). List 70 CNY per million tokens at 480p and 720p and 77 at 1080p without a video input, matching Ark's per-resolution table and its 1080p 16:9 example of 3.74 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Accepts up to 50 reference assets (30 images, 10 videos, 10 audios) and clips of 4-30s; per-asset billing beyond the free allowance is not encoded. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, + "volcengine/doubao-seedance-2-0-260128": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.141942857, + "output_cost_per_second_480p": 0.066003429, + "output_cost_per_second_720p": 0.141942857, + "output_cost_per_second_1080p": 0.354085714, + "output_cost_per_second_4k": 0.722057143, + "supported_modalities": ["text", "image", "video", "audio"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 46 CNY per million tokens at 480p and 720p, 51 at 1080p, and 26 at 4k without a video input, matching Ark's per-resolution table and its 1080p and 4k 16:9 examples of 2.48 and 5.05 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s; 4k is exclusive to this tier. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, + "volcengine/doubao-seedance-2-0-fast-260128": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.114171429, + "output_cost_per_second_480p": 0.053089714, + "output_cost_per_second_720p": 0.114171429, + "supported_modalities": ["text", "image", "video", "audio"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 37 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, + "volcengine/doubao-seedance-2-0-mini-260615": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.070971429, + "output_cost_per_second_480p": 0.033001714, + "output_cost_per_second_720p": 0.070971429, + "supported_modalities": ["text", "image"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 23 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Only text and image inputs are declared; the mini tier's reference video and audio caps were not verified. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, "bedrock/us-east-1/zai.glm-5": { "input_cost_per_token": 1e-06, "output_cost_per_token": 3.2e-06, @@ -67489,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", @@ -67497,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", @@ -67515,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", @@ -67533,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", @@ -67551,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", @@ -67570,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", @@ -67588,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", @@ -67605,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", @@ -67623,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, @@ -72354,7 +72656,7 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -75183,5 +75485,45 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "together_ai/reflection/Beam-1.0-501B-A23B": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 4e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.ai/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "baseten/reflection/Beam-1.0": { + "cache_read_input_token_cost": 1.3e-07, + "input_cost_per_token": 1.3e-06, + "litellm_provider": "baseten", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.3e-06, + "source": "https://inference.baseten.co/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "nebius/reflection/beam": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "nebius", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1e-06, + "source": "https://tokenfactory.nebius.com/models/catalog/text2text/reflection%2Fbeam", + "supports_function_calling": true, + "supports_reasoning": true } } diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 45104d40018..0325fe1f796 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2766,7 +2766,8 @@ "batches": false, "rerank": false, "a2a": true, - "interactions": true + "interactions": true, + "video_generations": true } }, "voyage": { 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/schema.prisma b/schema.prisma index d40bac380aa..83a0e6f9e05 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1981,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement { outcome String? @@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at") + @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } model LiteLLM_Lens { diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 4282485714a..597e0c6d4d8 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -51,6 +51,7 @@ longer signal it. ### Fixed +- **model**: `additional_drop_params` in `litellm_model.additional_litellm_params` is now sent to the proxy as the deployment's `additional_drop_params` list, the same list a `model_list` entry takes. Before, the provider never sent the list and instead deleted every key it named from `litellm_params`, so a `store` or `vertex_location` listed there was lost and requests still carried the parameters (#38349). The value must be a JSON list of parameter names such as `jsonencode(["store"])`, which `terraform plan` now checks, and `null` entries are skipped as before. When the configuration sets the list, refresh reads the stored list back from the proxy, so a list changed outside Terraform shows up as drift. On a model an older provider created, the proxy has no stored list, so the first plan after upgrading shows the list and every listed key the configuration sets as changes, since that apply writes them. A write-only secret such as `model_api_key_wo` is never shown in a plan - **user**: `litellm_user` updates no longer send `blocked`, which `/user/update` writes into the user table where no such column exists, so every update failed with a 400 (`Could not find field at upsertOneLiteLLM_UserTable.create.blocked`). Create still sends `blocked`, including `false`, so a `default_internal_user_params.blocked` on the proxy cannot block the auto-created key. Changing `blocked` after creation has no effect, since the proxy has no user-level block; use `litellm_key_block` instead - **model**: `litellm_model` refresh now reads the `{"data": [...]}` envelope `/model/info` returns, so `model_info` fields changed outside Terraform show up as drift instead of silently keeping the previous state - **key**: An update that changes `team_id` and fails because the key was already cascade-deleted along with its previous team now recovers by recreating the key under the new team, instead of aborting the apply. The key's absence is confirmed against the proxy first, so an unrelated failure still errors out, and a `team_id` change between two teams that both still exist stays a plain in-place update @@ -69,6 +70,7 @@ longer signal it. ### Changed +- **model** (breaking): a key named in `additional_litellm_params.additional_drop_params` is no longer deleted from the deployment's `litellm_params`. A configuration that also sets that key, such as `vertex_location`, stores it on the deployment at its next apply, and the proxy applies the drop list to requests instead. For `vertex_location` this changes routing. `/v1/messages` calls now go to the configured region, while chat completions drop a listed `vertex_location` and use the default region. To keep such a key off a deployment, remove its attribute and replace the model with `terraform apply -replace`, which gives it a new model id, since an update keeps a stored value the provider leaves out. `model`, `custom_llm_provider` and `merge_reasoning_content_in_choices` are always sent. A value that is not a JSON list of names, such as a bare `"store"` or an empty string, now fails the plan, and with it `terraform destroy`, instead of being stored as a string the proxy ignores. Rewrite it as `jsonencode([...])`, or remove the model without fixing it through `terraform destroy -refresh=false` - **key** (breaking): `model_max_budget` on `litellm_key` is now a JSON string of per-model budget objects (`jsonencode({"gpt-4o-mini" = {budget_limit = 50, time_period = "30d"}})`), matching `litellm_user`, `litellm_budget` and `litellm_tag`. The old `map(number)` form sent bare numbers to `/key/generate`, which the proxy rejects with a 500 (`'int' object is not iterable`), so every key with a non-empty `model_max_budget` failed to apply. Existing state upgrades automatically (schema version 1) and the attribute is refilled from the proxy on the next read; configurations still using the map form must be rewritten - **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying diff --git a/terraform/provider/docs/resources/model.md b/terraform/provider/docs/resources/model.md index b7c440b0e65..85c495bf1ca 100644 --- a/terraform/provider/docs/resources/model.md +++ b/terraform/provider/docs/resources/model.md @@ -183,16 +183,16 @@ The following arguments are supported: * Numeric strings are parsed first as integers; if integer parsing fails, parsed as floats (e.g., `"16384"` -> 16384, `"0.75"` -> 0.75) * JSON strings (starting with `[` or `{`) are parsed as JSON objects/arrays * Non-convertible strings remain strings - * Non-string map values (if supplied) are passed through unchanged. * The provider merges these keys into the `litellm_params` payload sent to the API. * Note: the remote API may not echo back all custom parameters; this provider preserves `additional_litellm_params` in state when present in configuration. - **Special parameter: `additional_drop_params`** - * When `additional_drop_params` is provided as a JSON array string, it specifies parameters to remove from the final `litellm_params` before sending to the API - * This allows you to override or remove built-in parameters if needed - * The `additional_drop_params` key itself is not included in the final parameters + **`additional_drop_params`** + * Give it a JSON array string, such as `jsonencode(["store"])`. The provider sends the list to the proxy as the deployment's `additional_drop_params`, and the proxy applies it the same as `additional_drop_params` in a `model_list` entry. `terraform plan` fails on any other value, an empty value, or an entry that is not a name, and `null` entries are skipped + * Other keys on the resource, such as `vertex_location`, are sent unchanged + * When the configuration sets the list, refresh reads the stored list back from the proxy, so a list changed outside Terraform shows up as drift. A model whose configuration has no list keeps a list set outside Terraform out of the plan, but the proxy clears it on the next update + * When the proxy has no stored list, as on a model an older provider created, the next plan shows the list and every listed key the configuration sets as changes, since the next apply writes them - Example showing booleans, integers, floats, strings, and parameter dropping: + Example showing booleans, integers, floats, strings, JSON, and dropped request parameters: ```hcl resource "litellm_model" "with_additional" { @@ -208,7 +208,7 @@ The following arguments are supported: "scale" = "0.75" # becomes float 0.75 "note" = "for testing" # stays string "complex_config" = "{\"nested\": {\"value\": 42}}" # parsed as JSON object - "additional_drop_params" = "[\"reasoningEffort\"]" # removes reasoningEffort parameter + "additional_drop_params" = "[\"reasoning_effort\"]" # proxy drops reasoning_effort from requests } } ``` diff --git a/terraform/provider/examples/model_additional_params.tf b/terraform/provider/examples/model_additional_params.tf index fb4981ec6fb..64c37993e04 100644 --- a/terraform/provider/examples/model_additional_params.tf +++ b/terraform/provider/examples/model_additional_params.tf @@ -9,8 +9,7 @@ provider "litellm" { # - Numeric strings will be parsed to integer (if possible) otherwise float. # - JSON strings (starting with [ or {) will be parsed as JSON objects/arrays. # - Non-convertible strings remain strings. -# - Non-string map values are passed through unchanged. -# - Use "additional_drop_params" as a JSON array to remove parameters from the final request. +# - "additional_drop_params" takes a JSON array of request parameters the proxy drops before calling the model provider. resource "litellm_model" "with_additional" { model_name = "custom-model" @@ -21,14 +20,12 @@ resource "litellm_model" "with_additional" { # Additional parameters not exposed as first-class arguments additional_litellm_params = { - "use_fine_tune" = "true" # becomes boolean true - "max_context" = "16384" # becomes integer 16384 - "temperature_scale" = "0.75" # becomes float 0.75 - "experimental_feature" = "enabled" # stays string "enabled" - "complex_config" = "{\"nested\": {\"value\": 42}}" # parsed as JSON object - "additional_drop_params" = "[\"reasoningEffort\"]" # removes reasoningEffort parameter - # You may also pass non-string values (they will be passed through unchanged) - # "raw_flag" = true + "use_fine_tune" = "true" # becomes boolean true + "max_context" = "16384" # becomes integer 16384 + "temperature_scale" = "0.75" # becomes float 0.75 + "experimental_feature" = "enabled" # stays string "enabled" + "complex_config" = "{\"nested\": {\"value\": 42}}" # parsed as JSON object + "additional_drop_params" = "[\"reasoning_effort\"]" # proxy drops reasoning_effort from requests } # Cost configuration (optional) @@ -47,9 +44,9 @@ resource "litellm_model" "azure_with_drop_params" { tier = "paid" mode = "completion" - # Drop the reasoningEffort parameter that might be automatically added + # Have the proxy drop reasoning_effort from requests before they reach Azure additional_litellm_params = { - "additional_drop_params" = "[\"reasoningEffort\"]" + "additional_drop_params" = "[\"reasoning_effort\"]" } input_cost_per_million_tokens = 0.25 diff --git a/terraform/provider/litellm/resource_model.go b/terraform/provider/litellm/resource_model.go index 1dc029b861c..e9582ad75fb 100644 --- a/terraform/provider/litellm/resource_model.go +++ b/terraform/provider/litellm/resource_model.go @@ -7,10 +7,11 @@ import ( func resourceLiteLLMModel() *schema.Resource { return &schema.Resource{ - Create: resourceLiteLLMModelCreate, - Read: resourceLiteLLMModelRead, - Update: resourceLiteLLMModelUpdate, - Delete: resourceLiteLLMModelDelete, + Create: resourceLiteLLMModelCreate, + Read: resourceLiteLLMModelRead, + Update: resourceLiteLLMModelUpdate, + Delete: resourceLiteLLMModelDelete, + CustomizeDiff: validateAdditionalDropParams, Importer: &schema.ResourceImporter{ StateContext: schema.ImportStatePassthroughContext, }, @@ -237,7 +238,8 @@ func resourceLiteLLMModel() *schema.Resource { Elem: &schema.Schema{ Type: schema.TypeString, }, - Description: "Additional parameters to pass to litellm_params beyond the standard ones", + DiffSuppressFunc: suppressEquivalentDropParams, + Description: "Additional parameters to pass to litellm_params beyond the standard ones", }, }, } diff --git a/terraform/provider/litellm/resource_model_crud.go b/terraform/provider/litellm/resource_model_crud.go index 3309dbe26dd..3db6f68760d 100644 --- a/terraform/provider/litellm/resource_model_crud.go +++ b/terraform/provider/litellm/resource_model_crud.go @@ -1,10 +1,13 @@ package litellm import ( + "context" "encoding/json" "fmt" "log" + "maps" "net/url" + "slices" "strconv" "strings" "time" @@ -68,6 +71,119 @@ func secretFromResource(d *schema.ResourceData, attr string) string { return d.Get(attr).(string) } +const additionalDropParamsKey = "additional_drop_params" + +func parseAdditionalDropParams(raw string) ([]string, error) { + var items []interface{} + if err := json.Unmarshal([]byte(raw), &items); err != nil { + return nil, fmt.Errorf("additional_litellm_params.%s must be a JSON list of parameter names, such as jsonencode([\"store\"]): %w", additionalDropParamsKey, err) + } + names := []string{} + for _, item := range items { + if item == nil { + continue + } + name, ok := item.(string) + if !ok { + return nil, fmt.Errorf("additional_litellm_params.%s entries must be parameter names, got %v", additionalDropParamsKey, item) + } + names = append(names, name) + } + return names, nil +} + +func validateAdditionalDropParams(_ context.Context, d *schema.ResourceDiff, _ interface{}) error { + params, _ := d.Get("additional_litellm_params").(map[string]interface{}) + raw, set := params[additionalDropParamsKey].(string) + if !set { + return nil + } + _, err := parseAdditionalDropParams(raw) + return err +} + +func suppressEquivalentDropParams(key, oldValue, newValue string, _ *schema.ResourceData) bool { + if key != "additional_litellm_params."+additionalDropParamsKey { + return false + } + oldNames, oldErr := parseAdditionalDropParams(oldValue) + newNames, newErr := parseAdditionalDropParams(newValue) + return oldErr == nil && newErr == nil && slices.Equal(oldNames, newNames) +} + +var attributeSentAsLiteLLMParam = map[string]string{ + "custom_llm_provider": "custom_llm_provider", + "tpm": "tpm", + "rpm": "rpm", + "input_cost_per_token": "input_cost_per_million_tokens", + "output_cost_per_token": "output_cost_per_million_tokens", + "api_key": "model_api_key", + "api_base": "model_api_base", + "api_version": "api_version", + "input_cost_per_pixel": "input_cost_per_pixel", + "output_cost_per_pixel": "output_cost_per_pixel", + "input_cost_per_second": "input_cost_per_second", + "output_cost_per_second": "output_cost_per_second", + "aws_access_key_id": "aws_access_key_id", + "aws_secret_access_key": "aws_secret_access_key", + "aws_region_name": "aws_region_name", + "aws_session_name": "aws_session_name", + "aws_role_name": "aws_role_name", + "vertex_project": "vertex_project", + "vertex_location": "vertex_location", + "vertex_credentials": "vertex_credentials", + "reasoning_effort": "reasoning_effort", + "thinking": "thinking_enabled", + "merge_reasoning_content_in_choices": "merge_reasoning_content_in_choices", +} + +func refreshAdditionalDropParams(d *schema.ResourceData, stored interface{}) { + params, _ := d.Get("additional_litellm_params").(map[string]interface{}) + raw, managed := params[additionalDropParamsKey].(string) + if !managed { + return + } + refreshed := maps.Clone(params) + if stored != nil { + encoded, _ := json.Marshal(stored) + refreshed[additionalDropParamsKey] = string(encoded) + d.Set("additional_litellm_params", refreshed) + return + } + delete(refreshed, additionalDropParamsKey) + listed, _ := parseAdditionalDropParams(raw) + for _, key := range listed { + delete(refreshed, key) + if attribute, ok := attributeSentAsLiteLLMParam[key]; ok { + d.Set(attribute, nil) + } + } + d.Set("additional_litellm_params", refreshed) +} + +func coerceAdditionalParamValue(strValue string) interface{} { + trimmedValue := strings.TrimSpace(strValue) + if strings.HasPrefix(trimmedValue, "[") || strings.HasPrefix(trimmedValue, "{") { + var parsedValue interface{} + if err := json.Unmarshal([]byte(strValue), &parsedValue); err == nil { + return parsedValue + } + } + switch strValue { + case "true": + return true + case "false": + return false + } + if intValue, err := strconv.Atoi(strValue); err == nil { + return intValue + } + if floatValue, err := strconv.ParseFloat(strValue, 64); err == nil { + return floatValue + } + return strValue +} + func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) error { client, ok := m.(*Client) if !ok { @@ -175,76 +291,17 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e litellmParams["thinking"] = thinking } - // Add additional parameters if provided if additionalParams, ok := d.GetOk("additional_litellm_params"); ok { - var dropParams []string - for key, value := range additionalParams.(map[string]interface{}) { - // Convert string values to appropriate types where possible - if strValue, ok := value.(string); ok { - // Check if it's JSON (starts with [ or {) - trimmedValue := strings.TrimSpace(strValue) - if strings.HasPrefix(trimmedValue, "[") || strings.HasPrefix(trimmedValue, "{") { - var parsedValue interface{} - if err := json.Unmarshal([]byte(strValue), &parsedValue); err == nil { - // Successfully parsed JSON - if key == "additional_drop_params" { - // Handle drop params specially - if dropList, ok := parsedValue.([]interface{}); ok { - for _, item := range dropList { - if paramStr, ok := item.(string); ok { - dropParams = append(dropParams, paramStr) - } - } - } - continue // Don't add to litellmParams - } else { - litellmParams[key] = parsedValue - } - } else { - // Not valid JSON, apply existing conversion logic - if strValue == "true" { - litellmParams[key] = true - } else if strValue == "false" { - litellmParams[key] = false - } else { - // Try to convert numeric strings - if intValue, err := strconv.Atoi(strValue); err == nil { - litellmParams[key] = intValue - } else if floatValue, err := strconv.ParseFloat(strValue, 64); err == nil { - litellmParams[key] = floatValue - } else { - // Keep as string - litellmParams[key] = strValue - } - } - } - } else { - // Apply existing conversion logic for non-JSON strings - if strValue == "true" { - litellmParams[key] = true - } else if strValue == "false" { - litellmParams[key] = false - } else { - // Try to convert numeric strings - if intValue, err := strconv.Atoi(strValue); err == nil { - litellmParams[key] = intValue - } else if floatValue, err := strconv.ParseFloat(strValue, 64); err == nil { - litellmParams[key] = floatValue - } else { - // Keep as string - litellmParams[key] = strValue - } - } - } - } else { - litellmParams[key] = value + if key != additionalDropParamsKey { + litellmParams[key] = coerceAdditionalParamValue(value.(string)) + continue } - } - - // Apply drop params at the end - for _, paramToDrop := range dropParams { - delete(litellmParams, paramToDrop) + dropParams, err := parseAdditionalDropParams(value.(string)) + if err != nil { + return err + } + litellmParams[key] = dropParams } } @@ -419,6 +476,7 @@ func resourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error { if _, ok := d.GetOk("additional_litellm_params"); ok { d.Set("additional_litellm_params", d.Get("additional_litellm_params")) } + refreshAdditionalDropParams(d, modelResp.LiteLLMParams.AdditionalDropParams) return nil } diff --git a/terraform/provider/litellm/resource_model_test.go b/terraform/provider/litellm/resource_model_test.go index 22c226adfd3..4c2d5d27442 100644 --- a/terraform/provider/litellm/resource_model_test.go +++ b/terraform/provider/litellm/resource_model_test.go @@ -3,8 +3,11 @@ package litellm import ( "context" "encoding/json" + "maps" "net/http" "net/http/httptest" + "reflect" + "slices" "testing" "github.com/hashicorp/go-cty/cty" @@ -116,6 +119,246 @@ func TestResourceLiteLLMModelCreateOmitsUnsetDisplayName(t *testing.T) { } } +func TestResourceLiteLLMModelCreateSendsAdditionalLiteLLMParams(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/model/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(modelInfoBody(""))) + case "/model/info": + w.Write([]byte(modelInfoBody(""))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMModel().Schema, map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "custom_llm_provider": "vertex_ai", + "base_model": "claude-haiku-4-5", + "vertex_project": "test-project", + "vertex_location": "eu", + "additional_litellm_params": map[string]interface{}{ + "additional_drop_params": `["store", "vertex_location"]`, + "tags": ` ["prod", "eu"]`, + "headers": `{"x-team": "search"}`, + "use_in_pass_through": "true", + "merge_reasoning_content_in_choices": "false", + "max_parallel_requests": "8", + "input_cost_per_token": "0.5", + "description": "[draft] haiku", + }, + }) + + if err := resourceLiteLLMModelCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + litellmParams, ok := createPayload["litellm_params"].(map[string]interface{}) + if !ok { + t.Fatalf("expected litellm_params object in create payload, got %v", createPayload["litellm_params"]) + } + want := map[string]interface{}{ + "additional_drop_params": []interface{}{"store", "vertex_location"}, + "vertex_project": "test-project", + "vertex_location": "eu", + "tags": []interface{}{"prod", "eu"}, + "headers": map[string]interface{}{"x-team": "search"}, + "use_in_pass_through": true, + "merge_reasoning_content_in_choices": false, + "max_parallel_requests": float64(8), + "input_cost_per_token": 0.5, + "description": "[draft] haiku", + } + for key, wantValue := range want { + if got := litellmParams[key]; !reflect.DeepEqual(got, wantValue) { + t.Errorf("litellm_params[%q] = %#v, want %#v", key, got, wantValue) + } + } +} + +const terraformUnknownValue = "74D93920-ED26-11E3-AC10-0800200C9A66" + +func TestResourceLiteLLMModelChecksAdditionalDropParams(t *testing.T) { + cases := map[string]struct { + dropParams string + want []interface{} + wantErr bool + }{ + "null entry from an unset variable is skipped": {dropParams: `["store", null]`, want: []interface{}{"store"}}, + "number entry is rejected": {dropParams: `["store", 1]`, wantErr: true}, + "object entry is rejected": {dropParams: `["store", {"name": "temperature"}]`, wantErr: true}, + "object is rejected": {dropParams: `{"store": true}`, wantErr: true}, + "bare name is rejected": {dropParams: "store", wantErr: true}, + "empty value is rejected": {dropParams: "", wantErr: true}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/model/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(modelInfoBody(""))) + case "/model/info": + w.Write([]byte(modelInfoBody(""))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + res := resourceLiteLLMModel() + config := map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "custom_llm_provider": "vertex_ai", + "base_model": "claude-haiku-4-5", + "additional_litellm_params": map[string]interface{}{"additional_drop_params": tc.dropParams}, + } + if _, err := res.Diff(context.Background(), nil, terraform.NewResourceConfigRaw(config), nil); (err != nil) != tc.wantErr { + t.Errorf("plan error = %v, want error %v", err, tc.wantErr) + } + + d := schema.TestResourceDataRaw(t, res.Schema, config) + err := resourceLiteLLMModelCreate(d, NewClient(srv.URL, "test-key", true)) + if tc.wantErr { + if err == nil { + t.Fatal("expected create to fail") + } + if createPayload != nil { + t.Errorf("expected no /model/new request, got %v", createPayload) + } + return + } + if err != nil { + t.Fatalf("create failed: %v", err) + } + litellmParams := createPayload["litellm_params"].(map[string]interface{}) + if got := litellmParams["additional_drop_params"]; !reflect.DeepEqual(got, tc.want) { + t.Errorf("additional_drop_params = %#v, want %#v", got, tc.want) + } + }) + } +} + +func TestResourceLiteLLMModelPlanAcceptsDropParamsUnknownUntilApply(t *testing.T) { + config := terraform.NewResourceConfigRaw(map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "custom_llm_provider": "vertex_ai", + "additional_litellm_params": map[string]interface{}{"additional_drop_params": terraformUnknownValue}, + }) + if _, err := resourceLiteLLMModel().Diff(context.Background(), nil, config, nil); err != nil { + t.Errorf("plan failed on a drop list not known until apply: %v", err) + } +} + +func TestResourceLiteLLMModelPlanComparesAdditionalDropParamsWithProxy(t *testing.T) { + const ( + listKey = "additional_litellm_params.additional_drop_params" + parallelKey = "additional_litellm_params.max_parallel_requests" + ) + cases := map[string]struct { + configured string + configLocation string + stored interface{} + storedLocation string + extraConfig map[string]interface{} + wantChanged []string + }{ + "list never sent by an older provider": {configured: `["store"]`, configLocation: "eu", storedLocation: "eu", wantChanged: []string{listKey}}, + "listed keys an older provider deleted": { + configured: `["store", "vertex_location", "tpm", "max_parallel_requests"]`, + configLocation: "eu", + wantChanged: []string{listKey, "vertex_location", "tpm", parallelKey}, + }, + "listed flag an older provider deleted": { + configured: `["merge_reasoning_content_in_choices"]`, + configLocation: "eu", + storedLocation: "eu", + extraConfig: map[string]interface{}{"merge_reasoning_content_in_choices": true}, + wantChanged: []string{listKey, "merge_reasoning_content_in_choices"}, + }, + "listed key an older provider still sent": { + configured: `["litellm_credential_name"]`, + configLocation: "eu", + storedLocation: "eu", + extraConfig: map[string]interface{}{"litellm_credential_name": "vertex-cred"}, + wantChanged: []string{listKey}, + }, + "list changed in the Admin UI": {configured: `["store"]`, configLocation: "eu", stored: []string{"temperature"}, storedLocation: "eu", wantChanged: []string{listKey}}, + "list set outside Terraform": {configLocation: "eu", stored: []string{"store"}, storedLocation: "eu"}, + "same list written with spaces": {configured: `["store", "vertex_location"]`, configLocation: "eu", stored: []string{"store", "vertex_location"}, storedLocation: "eu"}, + "null entry the proxy never received": {configured: `["store", null]`, configLocation: "eu", stored: []string{"store"}, storedLocation: "eu"}, + "listed key left out of the config": {configured: `["vertex_location"]`, stored: []string{"vertex_location"}, storedLocation: "us-east5"}, + "no list on the model or in the proxy": {configLocation: "eu", storedLocation: "eu"}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + litellmParams := map[string]interface{}{"model": "vertex_ai/claude-haiku-4-5", "custom_llm_provider": "vertex_ai"} + if tc.stored != nil { + litellmParams["additional_drop_params"] = tc.stored + } + if tc.storedLocation != "" { + litellmParams["vertex_location"] = tc.storedLocation + } + body, _ := json.Marshal(map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "litellm_params": litellmParams, + "model_info": map[string]interface{}{"id": "model-123", "db_model": true, "base_model": "claude-haiku-4-5"}, + }) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write(body) + })) + defer srv.Close() + + res := resourceLiteLLMModel() + additionalParams := map[string]interface{}{"max_parallel_requests": "8"} + if tc.configured != "" { + additionalParams["additional_drop_params"] = tc.configured + } + config := map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "custom_llm_provider": "vertex_ai", + "base_model": "claude-haiku-4-5", + "tpm": 1000, + "additional_litellm_params": additionalParams, + } + if tc.configLocation != "" { + config["vertex_location"] = tc.configLocation + } + maps.Copy(config, tc.extraConfig) + d := schema.TestResourceDataRaw(t, res.Schema, config) + d.SetId("model-123") + if err := resourceLiteLLMModelRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + diff, err := res.Diff(context.Background(), d.State(), terraform.NewResourceConfigRaw(config), nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + for _, key := range []string{listKey, "vertex_location", "tpm", parallelKey, "merge_reasoning_content_in_choices", "litellm_credential_name"} { + attr := (*terraform.ResourceAttrDiff)(nil) + if diff != nil { + attr = diff.Attributes[key] + } + got := attr != nil && attr.Old != attr.New + if want := slices.Contains(tc.wantChanged, key); got != want { + t.Errorf("plan shows %s change = %v, want %v", key, got, want) + } + } + }) + } +} + func TestResourceLiteLLMModelReadDisplayName(t *testing.T) { cases := map[string]struct { serverBody string diff --git a/terraform/provider/litellm/types.go b/terraform/provider/litellm/types.go index 3af4650de53..aff90cd64be 100644 --- a/terraform/provider/litellm/types.go +++ b/terraform/provider/litellm/types.go @@ -120,6 +120,7 @@ type LiteLLMParams struct { VertexProject string `json:"vertex_project,omitempty"` VertexLocation string `json:"vertex_location,omitempty"` VertexCredentials string `json:"vertex_credentials,omitempty"` + AdditionalDropParams interface{} `json:"additional_drop_params,omitempty"` } // ModelInfo represents information about a model. diff --git a/tests/benchmarks/test_prompt_template_benchmarks.py b/tests/benchmarks/test_prompt_template_benchmarks.py new file mode 100644 index 00000000000..c7858771d35 --- /dev/null +++ b/tests/benchmarks/test_prompt_template_benchmarks.py @@ -0,0 +1,21 @@ +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt + +LONG_RUN_OF_EMPTY_CACHED_USER_MESSAGES: Final[tuple[Mapping[str, object], ...]] = ( + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + *({"role": "user", "content": [], "cache_control": {"type": "ephemeral"}} for _ in range(1_000)), +) + + +@pytest.mark.benchmark +def test_anthropic_messages_pt_long_run_of_empty_cached_user_messages() -> None: + result: Final = anthropic_messages_pt( + messages=list(LONG_RUN_OF_EMPTY_CACHED_USER_MESSAGES), model="claude-opus-5", llm_provider="vertex_ai" + ) + + assert len(result[-1]["content"]) == 1_000 diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index ba2913bcfb7..e669107900d 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -112,8 +112,28 @@ 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 + "_set_mid_stream_fallback_continuation", # Tested through the update_settings continuation test in test_router.py + "_is_flag_enabled", # Tested through the string-flag continuation tests in test_router.py + "_params_sent_to_deployment", # Tested through the configured-constraint continuation tests in test_router.py + "_names_deployment_that_cannot_continue", # Tested through the deployment-id continuation tests in test_router.py + "_set_retry_policy", # Tested through update_settings in test_router_retry_policy_update.py + "_resolve_pre_call_check_input_tokens", # Tested through the token-count failure pre-call check tests in test_router.py + "_params_neither_dropped_nor_allowed", # Tested through the allowed-or-dropped param pre-call check tests in test_router.py + "_deployment_rejects_params", # Tested through the unsupported-param pre-call check tests in test_router.py + "_deployment_lacks_params", # Tested through the unsupported-param pre-call check tests in test_router.py + "_pre_call_check_requested_chat_params", # Tested through the unsupported-param pre-call check tests in test_router.py + "_deployment_supports_param", # Tested through the deployment-capability pre-call check tests in test_router.py + "_capability_model_and_provider", # Tested through the authenticating-provider pre-call check test in test_router.py + "_exact_base_model_max_input_tokens", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "_base_model_widened_window", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "_provider_model_id", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "_names_route_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "_declared_base_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py ] diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index e7a26c12db8..d267059cdc7 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -44,7 +44,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a fixed budget nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) -- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres transactions and Redis commands each request cost (`postgres_calls_per_request`, `redis_calls_per_request`, from `pg_stat_database` and Redis `INFO` read by `datastore.py` before and after the run). Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` +- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests diff --git a/tests/e2e/llm_translation/responses_helpers.py b/tests/e2e/llm_translation/responses_helpers.py index be0fac74bb6..29b923a7f34 100644 --- a/tests/e2e/llm_translation/responses_helpers.py +++ b/tests/e2e/llm_translation/responses_helpers.py @@ -8,10 +8,13 @@ AZURE_OPENAI_BACKEND: Final = "azure/gpt-5.4-nano" AZURE_OPENAI_API_VERSION: Final = "v1" -def azure_openai_params(api_version: str = AZURE_OPENAI_API_VERSION) -> LiteLLMParamsBody: +def azure_openai_params( + api_version: str = AZURE_OPENAI_API_VERSION, responses_api_version: str | None = None +) -> LiteLLMParamsBody: return LiteLLMParamsBody( model=AZURE_OPENAI_BACKEND, api_base="os.environ/AZURE_API_BASE", api_key="os.environ/AZURE_API_KEY", api_version=api_version, + responses_api_version=responses_api_version, ) diff --git a/tests/e2e/llm_translation/test_image_generation_e2e.py b/tests/e2e/llm_translation/test_image_generation_e2e.py index 15cf20aa65a..46561ccbc37 100644 --- a/tests/e2e/llm_translation/test_image_generation_e2e.py +++ b/tests/e2e/llm_translation/test_image_generation_e2e.py @@ -24,7 +24,10 @@ from sdk_clients import SdkClients pytestmark = pytest.mark.e2e OPENAI_IMAGE_BACKEND: Final = "openai/gpt-image-1-mini" -BEDROCK_IMAGE_BACKEND: Final = "bedrock/amazon.nova-canvas-v1:0" +BEDROCK_IMAGE_BACKEND: Final = "bedrock/stability.stable-image-core-v1:1" +# Nova Canvas reached end of life on 2026-09-30. Stable Image Core is active on Bedrock only in us-west-2: +# https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-stability-ai-stable-image-core-1-0.html (2026-10-16) +BEDROCK_IMAGE_REGION: Final = "us-west-2" class _OptionalImageBody(BaseModel): @@ -96,10 +99,10 @@ class TestImageGeneration: model=BEDROCK_IMAGE_BACKEND, aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", - aws_region_name="os.environ/AWS_REGION", + aws_region_name=BEDROCK_IMAGE_REGION, ), ) - images = sdk.openai(key).images.generate(model=model, prompt="Draw a cute cat", n=1, size="1024x1024") + images = sdk.openai(key).images.generate(model=model, prompt="Draw a cute cat") _assert_image_returned(images) @pytest.mark.skip(reason="stage red: product gap, /v1/images/generations 500s (aimage_generation TypeError) on missing prompt instead of 400") diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index ee4699703b2..bb75569ad99 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -40,7 +40,7 @@ from provider_edge import LiveEdge, start_provider_edge from provider_edge_bedrock import ConverseBodyCapture, bedrock_signer from proxy_client import ProxyClient from pydantic import BaseModel, TypeAdapter -from responses_helpers import AZURE_OPENAI_BACKEND, azure_openai_params +from responses_helpers import AZURE_OPENAI_API_VERSION, AZURE_OPENAI_BACKEND, azure_openai_params from sdk_clients import NO_PROXY_CACHE, SdkClients pytestmark = pytest.mark.e2e @@ -748,6 +748,70 @@ class TestResponses: f"/responses over azure openai preview returned no output text: {response.output!r}" ) + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.RESPONSES, + providers=(Provider.AZURE,), + models=(AZURE_OPENAI_BACKEND,), + mode=Mode.NONSTREAM, + ) + ) + def test_responses_azure_openai_dated_ga_chat_api_version_serves_responses( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model: Final = _register( + proxy, + resources, + azure_openai_params(api_version="2024-10-21"), + prefix="e2e-responses-azure-dated-chat", + ) + client: Final = sdk.openai(resources.key()) + + response: Final = client.responses.create( + model=model, + input="reply with one word", + instructions=INSTRUCTIONS, + extra_body=NO_PROXY_CACHE, + ) + assert response.output_text.strip(), ( + f"/responses over an azure deployment pinned to a dated GA chat api_version returned no output text: " + f"{response.output!r}" + ) + + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.RESPONSES, + providers=(Provider.AZURE,), + models=(AZURE_OPENAI_BACKEND,), + mode=Mode.NONSTREAM, + ) + ) + def test_responses_azure_openai_responses_api_version_overrides_a_pre_responses_api_version( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + # Azure rejects /openai/responses below api-version 2025-03-01-preview (probed live 2026-10-15), + # so this call only succeeds if the proxy sends responses_api_version instead of api_version. + model: Final = _register( + proxy, + resources, + azure_openai_params(api_version="2024-10-01-preview", responses_api_version=AZURE_OPENAI_API_VERSION), + prefix="e2e-responses-azure-override", + ) + client: Final = sdk.openai(resources.key()) + + response: Final = client.responses.create( + model=model, + input="reply with one word", + instructions=INSTRUCTIONS, + extra_body=NO_PROXY_CACHE, + ) + assert response.output_text.strip(), ( + f"/responses over an azure deployment with responses_api_version set returned no output text: " + f"{response.output!r}" + ) + @meta( Subject( domain=Domain.LLM_TRANSLATION, diff --git a/tests/e2e/mcp/test_mcp_tool_search_e2e.py b/tests/e2e/mcp/test_mcp_tool_search_e2e.py index 847239e8a66..c4a8e29eb82 100644 --- a/tests/e2e/mcp/test_mcp_tool_search_e2e.py +++ b/tests/e2e/mcp/test_mcp_tool_search_e2e.py @@ -3,6 +3,7 @@ from __future__ import annotations from typing import Final import pytest +from pydantic import BaseModel, ConfigDict, TypeAdapter from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp from e2e_config import provider_edge_base, unique_marker from e2e_http import unwrap @@ -17,6 +18,19 @@ MCP_TOOL_SEARCH_TOOL_NAME: Final = "mcp_tool_search" OVERSIZED_DESCRIPTION: Final = f"OVERSIZED {'weather ' * 12_000}" +class _SearchHit(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str + + +_SEARCH_HITS: Final = TypeAdapter(tuple[_SearchHit, ...]) + + +def _hits_tool(hits: tuple[_SearchHit, ...], tool_name: str) -> bool: + return any(hit.name.endswith(f"-{tool_name}") for hit in hits) + + class TestMcpToolSearch: @pytest.mark.covers("mcp.tool_search.api_key.oversized_description_omitted") def test_native_tool_search_omits_oversized_description( @@ -82,5 +96,6 @@ class TestMcpToolSearch: ) assert result.is_error is not True, result.all_text - assert normal_tool_name in result.all_text - assert oversized_tool_name not in result.all_text + hits: Final = _SEARCH_HITS.validate_json(result.all_text) + assert _hits_tool(hits, normal_tool_name), hits + assert not _hits_tool(hits, oversized_tool_name), hits diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 25b48201bef..700afa4641c 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 new file mode 100644 index 00000000000..89263977fb0 --- /dev/null +++ b/tests/e2e/perf/cases/mixed_size_requests_2rps.py @@ -0,0 +1,77 @@ +"""Mixed-size requests at 2 rps on the small gateway: three real model groups over the mock providers, +two teams plus users in no team, JWT and key callers, a pre and +a post guardrail, half the requests streaming, short chats, long contexts and tool calls. Budgets +are real money at cost-map prices (a 335k-token `claude-sonnet-5` request is about $0.70 and the +run spends roughly $600), so they are sized so no caller runs dry mid-run.""" + +from __future__ import annotations + +from typing import Final + +from setup.performance_test_case import ( + LONG, + SHORT, + Guardrail, + Limit, + MockLatency, + PerformanceTestCase, + Team, + Traffic, + User, + VirtualKey, +) + +CASE: Final = PerformanceTestCase( + name="mixed_size_requests_2rps", + teams=( + Team("team-a", models=("gpt-6-luna", "claude-sonnet-5"), max_budget=5000.0, rpm_limit=6000), + Team("team-b", models=("claude-sonnet-5-bedrock", "claude-sonnet-5"), max_budget=5000.0), + ), + users=( + User("alice", team="team-a"), + User("bob", team="team-a"), + User("carol", team="team-a"), + User("dave", team="team-b"), + User("erin", team="team-b"), + User("frank", team="team-b"), + User("grace"), + User("heidi"), + User("ivan"), + User("judy", max_budget=1000.0), + ), + virtual_keys=( + VirtualKey("key-a1", team="team-a", models=("gpt-6-luna",), rpm_limit=1200), + VirtualKey("key-a2", team="team-a", user="alice"), + VirtualKey("key-b1", team="team-b"), + VirtualKey("key-b2", team="team-b", user="dave"), + VirtualKey("key-plain"), + ), + model_groups=("gpt-6-luna", "claude-sonnet-5", "claude-sonnet-5-bedrock"), + mock_latency=MockLatency(seconds=1.0, jitter_seconds=0.25), + guardrails=(Guardrail("perf-pre", mode="pre_call"), Guardrail("perf-post", mode="post_call")), + traffic=Traffic( + seed=7, + rps=2.0, + total_requests=1200, + endpoints={"chat_completions": 0.5, "responses": 0.25, "messages": 0.25}, + shapes={SHORT: 0.75, LONG: 0.25}, + stream_share=0.5, + intervene_share=0.0, + block_share=0.0, + concurrency=100, + ), + limits=( + Limit("unexpected_status_count", max=3), + Limit("gateway_added_seconds", percentile=50, max=0.58), + Limit("gateway_added_seconds", percentile=99, max=6.1), + Limit("time_to_first_token_seconds", percentile=50, max=2.7), + Limit("time_to_first_token_seconds", percentile=99, max=8.8), + 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 52fe41623e3..73772465492 100644 --- a/tests/e2e/perf/cases/small_bodies_5rps.py +++ b/tests/e2e/perf/cases/small_bodies_5rps.py @@ -64,11 +64,13 @@ CASE: Final = PerformanceTestCase( ), limits=( Limit("unexpected_status_count", max=3), - Limit("gateway_added_seconds", percentile=50, max=0.38), - Limit("time_to_first_token_seconds", percentile=99, max=3.0), + Limit("gateway_added_seconds", percentile=50, max=0.29), + Limit("time_to_first_token_seconds", percentile=99, max=3.7), Limit("memory_growth_megabytes", max=100), - Limit("postgres_calls_per_request", max=13), - Limit("redis_calls_per_request", max=35), - Limit("cpu_seconds_per_request", max=0.19), + 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 e123d129cc5..54c6943a054 100644 --- a/tests/e2e/perf/setup/datastore.py +++ b/tests/e2e/perf/setup/datastore.py @@ -1,26 +1,95 @@ -"""Postgres and Redis work counters, read before and after the traffic phase so a case can put a ceiling -on how much database and cache work the gateway does per request. Both come from the servers' own -statistics (`pg_stat_database`, Redis `INFO`), so nothing here depends on how the gateway talks to them. -`DATABASE_URL` and `REDIS_HOST` are the runner's own env, as in the other suites; without them the -counters are None and any Limit on them reports as missing.""" +"""Postgres and Redis work the gateway did during the traffic phase, split into what a request waits on and +what runs after the response. Both come from the servers' own accounting, so nothing here depends on how the +gateway talks to them: Postgres statement counts from `pg_stat_statements` (read before and after the run) and +Redis commands from a `MONITOR` connection held open for the run. `DATABASE_URL` and `REDIS_HOST` are the +runner's own env, as in the other suites; without them the counters are None and any Limit on them reports as +missing. `E2E_DATABASE_ADMIN_URL`, when set, is used once to create the `pg_stat_statements` extension. Only the +runtime user's own statements are read: other users' rows come back with a NULL queryid and hidden text.""" from __future__ import annotations import os +import re +import threading +from collections import Counter +from collections.abc import Generator +from contextlib import contextmanager from dataclasses import dataclass -from typing import Final +from typing import TYPE_CHECKING, Final, Literal from pydantic import TypeAdapter -_DB_QUERY: Final = "SELECT xact_commit FROM pg_stat_database WHERE datname = current_database()" -_DB_ROW: Final = TypeAdapter(tuple[int]) +if TYPE_CHECKING: + import redis + +Bucket = Literal["awaited", "background", "plumbing", "unclassified"] + +_COMMITS_QUERY: Final = "SELECT xact_commit FROM pg_stat_database WHERE datname = current_database()" +_STATEMENTS_QUERY: Final = ( + "SELECT queryid, calls, query FROM pg_stat_statements " + "WHERE dbid = (SELECT oid FROM pg_database WHERE datname = current_database()) " + "AND userid = (SELECT usesysid FROM pg_user WHERE usename = current_user) AND queryid IS NOT NULL" +) +_CREATE_EXTENSION: Final = "CREATE EXTENSION IF NOT EXISTS pg_stat_statements" +_COMMITS_ROW: Final = TypeAdapter(tuple[int]) +_STATEMENT_ROWS: Final = TypeAdapter(list[tuple[int, int, str]]) _INFO: Final = TypeAdapter(dict[str, object]) _INT: Final = TypeAdapter(int) +_MONITOR_LINE: Final = TypeAdapter(dict[str, object]) + +_TABLE: Final = re.compile(r'"(LiteLLM_\w+)"') +_AWAITED_TABLES: Final = frozenset( + { + "LiteLLM_VerificationToken", + "LiteLLM_UserTable", + "LiteLLM_TeamTable", + "LiteLLM_TeamMembership", + "LiteLLM_OrganizationTable", + "LiteLLM_OrganizationMembership", + "LiteLLM_EndUserTable", + "LiteLLM_BudgetTable", + "LiteLLM_ObjectPermissionTable", + } +) +_PLUMBING_VERBS: Final = frozenset({"BEGIN", "COMMIT", "ROLLBACK", "DEALLOCATE", "SET", "SHOW"}) +_REDIS_READS: Final = frozenset({"GET", "MGET", "HGET", "HGETALL", "EXISTS", "TTL", "PTTL"}) +_REDIS_NAMESPACE: Final = "litellm.caching:" +_REDIS_WRITES: Final = frozenset( + {"SET", "SETEX", "DEL", "EXPIRE", "INCRBYFLOAT", "INCR", "RPUSH", "LPUSH", "LPOP", "RPOP", "LTRIM", "HSET", "SADD"} +) + + +@dataclass(frozen=True, slots=True) +class StatementCalls: + queryid: int + query: str + calls: int + + @property + def bucket(self) -> Bucket: + """A SELECT on the key, user, team, org, end-user or budget tables is what a request waits on at auth + and routing. Transaction control and session statements are plumbing. Everything else (spend and + audit writes, daily aggregates, config and model reloads) runs after the response or on a timer.""" + verb: Final = self.query.lstrip("(").split(maxsplit=1)[0].upper() if self.query.strip() else "" + tables: Final = _TABLE.findall(self.query) + if verb in _PLUMBING_VERBS or not tables: + 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] @dataclass(frozen=True, slots=True) class DbCounters: commits: int + statements: tuple[StatementCalls, ...] | None @dataclass(frozen=True, slots=True) @@ -29,36 +98,169 @@ class RedisCounters: connected_clients: int +@dataclass(frozen=True, slots=True) +class RedisCommand: + command: str + key: str + inside_script: bool + count: int + + @property + def bucket(self) -> Bucket: + """Budget reads, cooldown reads and the rate-limiter script (with every command it issues) happen + before the request is forwarded. Spend increments, token counters, config rewrites, cache + invalidations (PUBLISH, DEL), the spend buffers the pods flush and any other write happen after the + response or on a timer. Reads on keys these rules do not know stay unclassified and count in neither.""" + key: Final = self.key.removeprefix(_REDIS_NAMESPACE) + if key.startswith(("litellm_config:", "daily_", "cronjob_lock:", "litellm_spend", "litellm_daily")): + return "background" + if self.inside_script or self.command in {"EVALSHA", "EVAL"}: + return "awaited" + if self.command in {"PUBLISH", "DEL", "SCRIPT", "INFO", "MONITOR"}: + return "background" + if key.startswith("spend:"): + return "awaited" if self.command in _REDIS_READS else "background" + if key.startswith("deployment:") or key.endswith(("}:requests", "}:window")): + return "awaited" + if key.endswith("}:tokens") or self.command in _REDIS_WRITES: + return "background" + return "unclassified" + + @property + def summary(self) -> str: + key: Final = re.sub(r"[0-9a-f]{16,}|sk-[A-Za-z0-9_-]+|[0-9a-f-]{36}", "*", self.key)[:70] + return f"{self.command} {key}{' (in script)' if self.inside_script else ''}".rstrip() + + +def statement_deltas(before: DbCounters | None, after: DbCounters | None) -> tuple[StatementCalls, ...] | None: + """Calls per statement during the run: `after` minus `before` by queryid, new statements whole.""" + if before is None or after is None or before.statements is None or after.statements is None: + return None + earlier: Final = {statement.queryid: statement.calls for statement in before.statements} + return tuple( + StatementCalls( + queryid=statement.queryid, query=statement.query, calls=statement.calls - earlier.get(statement.queryid, 0) + ) + for statement in after.statements + if statement.calls > earlier.get(statement.queryid, 0) + ) + + +def calls_in(bucket: Bucket, items: tuple[StatementCalls, ...] | tuple[RedisCommand, ...] | None) -> int | None: + if items is None: + return None + return sum(_count(item) for item in items if item.bucket == bucket) + + +def _count(item: StatementCalls | RedisCommand) -> int: + return item.calls if isinstance(item, StatementCalls) else item.count + + def read_db_counters() -> DbCounters | None: url: Final = os.environ.get("DATABASE_URL") if url is None: return None import psycopg + _ensure_statements_extension(url) with psycopg.connect(url) as conn: - row: Final = conn.execute(_DB_QUERY).fetchone() - (commits,) = _DB_ROW.validate_python(row) - return DbCounters(commits=commits) + (commits,) = _COMMITS_ROW.validate_python(conn.execute(_COMMITS_QUERY).fetchone()) + try: + rows = _STATEMENT_ROWS.validate_python(conn.execute(_STATEMENTS_QUERY).fetchall()) + except psycopg.errors.UndefinedTable: + return DbCounters(commits=commits, statements=None) + return DbCounters( + commits=commits, + statements=tuple(StatementCalls(queryid=queryid, query=query, calls=calls) for queryid, calls, query in rows), + ) -def read_redis_counters() -> RedisCounters | None: - """Standalone Redis only, which is what the perf lane's chart-bundled Redis and a local run are.""" - host: Final = os.environ.get("REDIS_HOST") - if host is None: - return None +def _ensure_statements_extension(url: str) -> None: + """Creating the extension needs a superuser, so prefer the admin URL; the runtime user's attempt is + allowed to fail (a database without `shared_preload_libraries = pg_stat_statements` cannot have it).""" + import psycopg + + try: + with psycopg.connect(os.environ.get("E2E_DATABASE_ADMIN_URL", url), autocommit=True) as conn: + _ = conn.execute(_CREATE_EXTENSION) + except psycopg.Error: + return + + +def _redis_client() -> redis.Redis[str]: import redis - client: Final = redis.Redis( - host=host, + return redis.Redis( + host=os.environ["REDIS_HOST"], port=int(os.environ.get("REDIS_PORT", "6379")), ssl=os.environ.get("REDIS_SSL", "false").lower() == "true", password=os.environ.get("REDIS_PASSWORD") or None, decode_responses=True, socket_connect_timeout=2, ) + + +def read_redis_counters() -> RedisCounters | None: + """Standalone Redis only, which is what the perf lane's chart-bundled Redis and a local run are.""" + if os.environ.get("REDIS_HOST") is None: + return None + client: Final = _redis_client() stats: Final = _INFO.validate_python(client.info("stats")) clients: Final = _INFO.validate_python(client.info("clients")) return RedisCounters( commands=_INT.validate_python(stats["total_commands_processed"]), connected_clients=_INT.validate_python(clients["connected_clients"]), ) + + +class RedisMonitor: + """Every command Redis executed while the monitor was open, counted by command, key and whether a Lua + script issued it. `MONITOR` echoes each command to this connection, so it costs Redis one write per + command; fine at the rates these cases run.""" + + def __init__(self) -> None: + self._counts: Counter[tuple[str, str, bool]] = Counter() + self._thread: threading.Thread | None = None + self._pool: redis.ConnectionPool | None = None + self.commands: tuple[RedisCommand, ...] | None = None + + def _run(self) -> None: + import redis + + client: Final = _redis_client() + self._pool = client.connection_pool # pyright: ignore[reportAny] # redis.Redis leaves connection_pool untyped + with client.monitor() as monitor: + try: + while True: + line = _MONITOR_LINE.validate_python(monitor.next_command()) + words = str(line["command"]).split(" ", 2) + key = words[1].strip('"') if len(words) > 1 else "" + self._counts[(words[0].upper(), key, line["client_type"] == "lua")] += 1 + except (redis.ConnectionError, OSError, ValueError): + return + + def start(self) -> None: + self._thread = threading.Thread(target=self._run, name="redis-monitor", daemon=True) + self._thread.start() + + def stop(self) -> tuple[RedisCommand, ...]: + if self._pool is not None: + self._pool.disconnect() + if self._thread is not None: + self._thread.join(timeout=5) + return tuple( + RedisCommand(command=command, key=key, inside_script=inside_script, count=count) + for (command, key, inside_script), count in self._counts.items() + ) + + +@contextmanager +def redis_monitor() -> Generator[RedisMonitor, None, None]: + """Counts Redis commands for the duration of the block; `.commands` is set once the block ends.""" + monitor: Final = RedisMonitor() + if os.environ.get("REDIS_HOST") is not None: + monitor.start() + try: + yield monitor + finally: + monitor.commands = monitor.stop() if os.environ.get("REDIS_HOST") is not None else None diff --git a/tests/e2e/perf/setup/observe.py b/tests/e2e/perf/setup/observe.py index a322e5a7443..f634df7de84 100644 --- a/tests/e2e/perf/setup/observe.py +++ b/tests/e2e/perf/setup/observe.py @@ -15,7 +15,14 @@ from e2e_http import NoBody, get_external, post_external, unwrap from memory_readings import RssCapture, RssReading, WorkerKey from pydantic import BaseModel -from setup.datastore import DbCounters, RedisCounters +from setup.datastore import ( + DbCounters, + RedisCommand, + RedisCounters, + StatementCalls, + calls_in, + statement_deltas, +) from setup.mocks.reply_common import Provider from setup.performance_test_case import Limit, Metric, MockLatency, PerformanceTestCase from setup.wire import WireRecord @@ -106,6 +113,7 @@ class Observation: db_after: DbCounters | None = None redis_before: RedisCounters | None = None redis_after: RedisCounters | None = None + redis_commands: tuple[RedisCommand, ...] | None = None gateway_cpu_seconds: float | None = None def per_request(self, before: int | None, after: int | None) -> float | None: @@ -113,6 +121,28 @@ class Observation: return None return (after - before) / len(self.records) + def count_per_request(self, count: int | None) -> float | None: + return None if count is None or not self.records else count / len(self.records) + + @property + 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) @@ -180,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(), @@ -191,11 +223,23 @@ def observe(metric: Metric, percentile_rank: int | None, observation: Observatio if metric == "unexpected_status_count": return float(observation.unexpected_status_count()) if metric == "postgres_calls_per_request": - db_before, db_after = observation.db_before, observation.db_after - return observation.per_request(db_before and db_before.commits, db_after and db_after.commits) + 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": - redis_before, redis_after = observation.redis_before, observation.redis_after - return observation.per_request(redis_before and redis_before.commands, redis_after and redis_after.commands) + return observation.count_per_request(calls_in("awaited", observation.redis_commands)) + if metric == "redis_background_calls_per_request": + return observation.count_per_request(calls_in("background", observation.redis_commands)) if metric == "cpu_seconds_per_request": cpu: Final = observation.gateway_cpu_seconds return None if cpu is None or not observation.records else cpu / len(observation.records) @@ -208,6 +252,7 @@ class Row: percentile: int | None observed: float | None limit: float | None + detail: str = "" @property def label(self) -> str: @@ -234,7 +279,11 @@ REPORTED: Final[tuple[tuple[Metric, int | None], ...]] = ( ("time_to_first_token_seconds", 99), ("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), ) @@ -254,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 @@ -268,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 () @@ -298,6 +358,7 @@ def render(case: PerformanceTestCase, observation: Observation, rows: list[Row]) flag = " OVER" if row.over else " MISSING" if row.missing else "" lines.append(f" {row.label:32} {observed:>12} {limit:>12}{flag}") lines.extend(_diagnostics(observation)) + lines.extend(_datastore_lines(observation)) for record in failures: lines.append( f" wrong: seq {record.seq} {record.caller} {record.endpoint}{' stream' if record.stream else ''} status {record.status}: {record.error}" @@ -316,9 +377,49 @@ def _diagnostics(observation: Observation) -> list[str]: f" guardrail missed {observation.guardrail_missed() if observation.guardrail_calls else 'n/a'}", f" retried share {len(observation.retried_markers) / count:.4f}", f" spend logged share {'n/a' if share is None else f'{share:.3f}'}", + f" postgres transactions per request {_fmt(observation.per_request(_commits(observation.db_before), _commits(observation.db_after)))}", + f" redis commands per request {_fmt(observation.per_request(_commands(observation.redis_before), _commands(observation.redis_after)))}", f" redis connected clients {_clients(observation.redis_before)} -> {_clients(observation.redis_after)}", ] +def _datastore_lines(observation: Observation) -> list[str]: + """What the per-request Postgres and Redis counts are made of, so a red run names the statement or + command that grew. `plumbing` (BEGIN, COMMIT, session settings) and `unclassified` count in neither limit.""" + count: Final = max(1, len(observation.records)) + lines: list[str] = [] + for title, items in ( + ("postgres statements", observation.statements), + ("redis commands", observation.redis_commands), + ): + if items is None: + lines.append(f" {title}: n/a") + continue + for bucket in ("awaited", "background", "plumbing", "unclassified"): + in_bucket = sorted((item for item in items if item.bucket == bucket), key=_item_count, reverse=True) + total = sum(_item_count(item) for item in in_bucket) + if not in_bucket: + continue + lines.append(f" {title} {bucket}: {total / count:.3f} per request, {len(in_bucket)} distinct") + lines.extend(f" {_item_count(item) / count:8.3f} {item.summary}" for item in in_bucket[:8]) + return lines + + +def _item_count(item: StatementCalls | RedisCommand) -> int: + return item.calls if isinstance(item, StatementCalls) else item.count + + +def _commits(counters: DbCounters | None) -> int | None: + return None if counters is None else counters.commits + + +def _commands(counters: RedisCounters | None) -> int | None: + return None if counters is None else counters.commands + + +def _fmt(value: float | None) -> str: + return "n/a" if value is None else f"{value:.3f}" + + def _clients(counters: RedisCounters | None) -> str: return "n/a" if counters is None else str(counters.connected_clients) diff --git a/tests/e2e/perf/setup/performance_test_case.py b/tests/e2e/perf/setup/performance_test_case.py index d90b262329c..37ad640db62 100644 --- a/tests/e2e/perf/setup/performance_test_case.py +++ b/tests/e2e/perf/setup/performance_test_case.py @@ -22,8 +22,12 @@ Metric = Literal[ "time_to_first_token_seconds", "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", ] @@ -146,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/perf/test_perf_cases.py b/tests/e2e/perf/test_perf_cases.py index d077cb9f8c7..d89a8634ab3 100644 --- a/tests/e2e/perf/test_perf_cases.py +++ b/tests/e2e/perf/test_perf_cases.py @@ -9,11 +9,13 @@ import time from typing import Final import pytest +from cases import CASES from idp import Keycloak from lifecycle import ResourceManager from memory_readings import read_rss_everywhere from models import SpendLogsParams -from cases import CASES +from proxy_client import ProxyClient +from setup.datastore import read_db_counters, read_redis_counters, redis_monitor from setup.observe import ( Observation, evaluate, @@ -26,11 +28,9 @@ from setup.observe import ( ) from setup.performance_test_case import PerformanceTestCase from setup.plan import build_plan -from setup.datastore import read_db_counters, read_redis_counters from setup.podcpu import gateway_cpu_sampler from setup.provision import MOCK_PROVIDERS, build_setup, guardrail_url, mock_url from setup.run import run_plan -from proxy_client import ProxyClient pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] _SPEND_SAMPLE: Final = 40 @@ -68,9 +68,11 @@ def test_perf_case( db_before: Final = read_db_counters() redis_before: Final = read_redis_counters() run_started: Final = time.monotonic() - with gateway_cpu_sampler() as cpu: + with gateway_cpu_sampler() as cpu, redis_monitor() as monitor: outcome: Final = run_plan(case, plan, setup) run_finished: Final = time.monotonic() + db_after: Final = read_db_counters() + redis_after: Final = read_redis_counters() rss_after: Final = read_rss_everywhere(proxy) rss_read: Final = time.monotonic() assert outcome.records, f"locust wrote no records (exit {outcome.locust_exit}):\n{outcome.locust_output}" @@ -81,8 +83,6 @@ def test_perf_case( proxy, served_call_ids[:: max(1, len(served_call_ids) // _SPEND_SAMPLE)][:_SPEND_SAMPLE] ) spend_read: Final = time.monotonic() - db_after: Final = read_db_counters() - redis_after: Final = read_redis_counters() mock_requests: Final = tuple(seen for url in mock_urls.values() for seen in read_mock_requests(url)) guardrail_calls: Final = tuple(read_guardrail_calls(guardrail_url())) if case.guardrails else () mocks_read: Final = time.monotonic() @@ -98,6 +98,7 @@ def test_perf_case( db_after=db_after, redis_before=redis_before, redis_after=redis_after, + redis_commands=monitor.commands, gateway_cpu_seconds=cpu.seconds, ) rows, exceeded_limits = evaluate(case, observation) 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 998d7d9b87e..ad54c61aed0 100644 --- a/tests/e2e_harness/perf/setup/test_observe.py +++ b/tests/e2e_harness/perf/setup/test_observe.py @@ -5,7 +5,7 @@ from typing import Final from cases.small_bodies_5rps import CASE from memory_readings import RssCapture -from setup.datastore import DbCounters +from setup.datastore import DbCounters, StatementCalls from setup.observe import Observation, evaluate, junit_properties from setup.performance_test_case import Limit from setup.wire import WireRecord @@ -35,6 +35,10 @@ def served(seq: int, total_s: float) -> WireRecord: ) +def _key_lookup(calls: int) -> StatementCalls: + return StatementCalls(queryid=1, query='SELECT * FROM "LiteLLM_VerificationToken" WHERE token = $1', calls=calls) + + def test_junit_properties_carry_each_observed_metric_and_each_limit_including_unobserved_ones() -> None: case: Final = replace( CASE, @@ -50,8 +54,8 @@ def test_junit_properties_carry_each_observed_metric_and_each_limit_including_un rss_before=NO_RSS, rss_after=NO_RSS, wall_seconds=3.0, - db_before=DbCounters(commits=100), - db_after=DbCounters(commits=130), + db_before=DbCounters(commits=100, statements=(_key_lookup(calls=100),)), + db_after=DbCounters(commits=130, statements=(_key_lookup(calls=130),)), ) rows, _ = evaluate(case, observation) @@ -64,5 +68,102 @@ def test_junit_properties_carry_each_observed_metric_and_each_limit_including_un "perf_gateway_added_seconds_p50_limit": 0.625, "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 9a3885af259..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"} @@ -11192,6 +11365,87 @@ GPT_5_6_SLASH_ALIAS_STREAM: Final = GPT_5_6_STREAM.model_copy( update={"name": "gpt-5.6-slash_alias_stream", "alias": "openai/custom/gpt-5.6-$MARKER"} ) +GPT_5_6_RESPONSES_IMAGE_GENERATION_TOOL_TOKENS: Final = CostTrackingTestCase( + name="gpt-5.6-responses_image_generation_tool_tokens", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-5.6", + endpoint="/v1/responses", + request={ + "model": "$MODEL", + "input": "draw a duck", + "tools": [{"type": "image_generation", "model": "gpt-image-next", "quality": "low", "size": "1024x1024"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.6", + "tools": [ + { + "type": "image_generation", + "background": "auto", + "model": "gpt-image-next", + "moderation": "auto", + "n": 1, + "output_compression": 100, + "output_format": "png", + "quality": "low", + "size": "1024x1024", + } + ], + "output": [ + { + "type": "image_generation_call", + "id": "ig_$REQUEST_ID", + "status": "completed", + "action": "generate", + "background": "opaque", + "output_format": "png", + "quality": "low", + "revised_prompt": "a duck", + "size": "1024x1024", + "result": "AAAA", + }, + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted response", "annotations": []}], + }, + ], + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "total_tokens": 2252, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "tool_usage": { + "image_gen": { + "input_tokens": 32, + "input_tokens_details": {"image_tokens": 0, "text_tokens": 32}, + "output_tokens": 196, + "output_tokens_details": {"image_tokens": 196, "text_tokens": 0}, + "total_tokens": 228, + }, + "web_search": {"num_requests": 0}, + }, + }, + ), + expected=ExactExpected( + spend=0.01004232, + input_cost=0.00322, + output_cost=0.005768, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.00105432, + ), +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_SLASH_ALIAS, GPT_5_6_SLASH_ALIAS_STREAM, @@ -11336,6 +11590,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_RESPONSES_WEB_SEARCH_REPORTED_COUNT_MIXED_ACTIONS, GPT_5_6_RESPONSES_WEB_SEARCH_REPORTED_COUNT_OPEN_PAGE_ONLY, GPT_5_6_RESPONSES_WEB_SEARCH_REPORTED_COUNT_MIXED_ACTIONS_STREAM, + GPT_5_6_RESPONSES_IMAGE_GENERATION_TOOL_TOKENS, GPT_5_3_CODEX_RESPONSES_FILE_SEARCH, GPT_5_6_RESPONSES_SERVICE_TIER_FLEX, GPT_5_6_RESPONSES_SERVICE_TIER_PRIORITY, @@ -11403,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, ...]] = ( @@ -11431,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, ...]] = ( @@ -11463,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_bedrock_error_request_id.py b/tests/integration/observability/test_bedrock_error_request_id.py index 230619a2863..436dcc572d5 100644 --- a/tests/integration/observability/test_bedrock_error_request_id.py +++ b/tests/integration/observability/test_bedrock_error_request_id.py @@ -1,13 +1,208 @@ +import asyncio +import base64 import json import uuid -from typing import Final +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from types import MappingProxyType +from typing import Final, Literal, NamedTuple +from urllib.parse import unquote -from integration._support.client import Gateway, eventually, object_value +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually, object_value from integration._support.database import read_rows +from integration._support.responses_stream import ( + created, + delta, + error_event, + failed, + frame, + healthy_stream, + openai_error_event, + spend_row, +) +from integration._support.upstream import _aws_event_frame, aws_event_stream_frame from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id _MODEL: Final = "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0" +_BEDROCK_RESPONSES_MODEL: Final = "bedrock/us.openai.gpt-5.6-sol" _TOKEN: Final = "synthetic-bedrock-bearer" +_SHAPES: Final = ("nested-error", "top-level-error", "response.failed", "error-before-output") +_ERROR_CODES: Final = ("server_error", "rate_limit_exceeded", "invalid_request_error", "content_policy_violation") +_INVOKE_MODEL_ID: Final = "anthropic.claude-3-haiku-20240307-v1:0" +_BEDROCK_AWS: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + } +) +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_BEDROCK_STREAM_TARGETS: Final = frozenset( + ( + "/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse-stream", + f"/model/{_INVOKE_MODEL_ID}/invoke-with-response-stream", + ) +) + +BedrockFailure = Literal["exception-event", "exception-message", "no-known-event", "truncated-frame"] + + +class _FailureShape(NamedTuple): + bedrock_status: int + converse_class: str + after_output: tuple[bool, ...] + + +BedrockEndpoint = Literal["chat", "responses", "messages", "messages-invoke"] +_BEDROCK_ENDPOINTS: Final[tuple[BedrockEndpoint, ...]] = ("chat", "responses", "messages", "messages-invoke") +_BEDROCK_FAILURES: Final[Mapping[BedrockFailure, _FailureShape]] = MappingProxyType( + { + "exception-event": _FailureShape(500, "ServiceUnavailableError", (False, True)), + "exception-message": _FailureShape(429, "RateLimitError", (False, True)), + "no-known-event": _FailureShape(502, "BadGatewayError", (False,)), + "truncated-frame": _FailureShape(502, "BadGatewayError", (True,)), + } +) +_BEDROCK_CELLS: Final = tuple( + pytest.param( + endpoint, failure, after_output, id=f"{endpoint}-{failure}-{'after' if after_output else 'before'}-output" + ) + for endpoint in _BEDROCK_ENDPOINTS + for failure, shape in _BEDROCK_FAILURES.items() + for after_output in shape.after_output +) + + +def _stream(identity: str, shape: str, error_code: str, message: str) -> tuple[bytes, ...]: + opening: Final = (frame(created(identity)), frame(delta(identity, "partial output"))) + nested: Final = frame(error_event({"type": error_code, "code": error_code, "message": message})) + match shape: + case "nested-error": + return (*opening, nested) + case "top-level-error": + return (*opening, frame(openai_error_event(error_code, message))) + case "response.failed": + return (*opening, frame(failed(identity, error_code, message))) + case "error-before-output": + return (nested,) + raise AssertionError(shape) + + +def _bedrock_frame(event_type: str, payload: dict[str, JsonValue]) -> bytes: + return _aws_event_frame(event_type, payload, "sc", "u") + + +def _invoke_chunk(event: dict[str, JsonValue]) -> bytes: + return _bedrock_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}) + + +def _bedrock_output(invoke: bool, text: str = "partial output") -> tuple[bytes, ...]: + if not invoke: + return ( + _bedrock_frame("messageStart", {"role": "assistant"}), + _bedrock_frame("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}), + ) + started: Final[dict[str, JsonValue]] = { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "content": [], + "model": _INVOKE_MODEL_ID, + "stop_reason": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + return ( + _invoke_chunk({"type": "message_start", "message": started}), + _invoke_chunk({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + _invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}), + ) + + +def _bedrock_failure_frame(failure: BedrockFailure, message: str) -> bytes: + match failure: + case "exception-event": + return _bedrock_frame("internalServerException", {"message": message}) + case "exception-message": + headers: Final = { + ":exception-type": "throttlingException", + ":content-type": "application/json", + ":message-type": "exception", + } + return aws_event_stream_frame(headers, json.dumps({"message": message}).encode()) + case "no-known-event": + return _bedrock_frame("somethingBedrockAddedLater", {"message": message}) + case "truncated-frame": + return _bedrock_frame("contentBlockDelta", {"delta": {"text": message}, "contentBlockIndex": 0})[:-7] + + +def _bedrock_failure_text(failure: BedrockFailure, message: str) -> str: + return "undecoded bytes after" if failure == "truncated-frame" else message + + +def _bedrock_stream_request( + endpoint: BedrockEndpoint, model: str, prompt: str, *, stream: bool = True +) -> tuple[str, dict[str, JsonValue]]: + extra: Final[dict[str, JsonValue]] = {"stream": stream, "num_retries": 0, "cache": {"no-cache": True}} + turn: Final[list[JsonValue]] = [{"role": "user", "content": prompt}] + match endpoint: + case "chat": + return "/v1/chat/completions", {"model": model, "messages": turn, **extra} + case "responses": + return "/v1/responses", {"model": model, "input": prompt, **extra} + case "messages" | "messages-invoke": + return "/v1/messages", {"model": model, "messages": turn, "max_tokens": 16, **extra} + + +def _failure_information(call_id: str) -> dict[str, JsonValue]: + row: Final = spend_row(call_id) + assert row["status"] == "failure", row + metadata: Final = row["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + return object_value(parsed["error_information"]) + + +def _streamed_response(gateway: Gateway, model: str, prompt: str) -> httpx.Response: + response: Final = gateway.request( + "POST", + "/v1/responses", + {"model": model, "input": prompt, "stream": True, "num_retries": 0, "cache": {"no-cache": True}}, + ) + assert response.status_code == 200, response.text + return response + + +async def _consume_sdk_stream(asynchronous: bool, api_base: str, identity: str) -> None: + if asynchronous: + async_stream: Final = await litellm.aresponses( + model="openai/gpt-6", + input=identity, + stream=True, + api_base=api_base, + api_key="synthetic-provider-key", + num_retries=0, + ) + async for _ in async_stream: + pass + return + tuple( + litellm.responses( + model="openai/gpt-6", + input=identity, + stream=True, + api_base=api_base, + api_key="synthetic-provider-key", + num_retries=0, + ) + ) def test_bedrock_500_keeps_amzn_request_id_on_error_headers_and_failure_log(gateway: Gateway) -> None: @@ -54,3 +249,556 @@ def test_bedrock_500_keeps_amzn_request_id_on_error_headers_and_failure_log(gate parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) error_information: Final = object_value(parsed["error_information"]) assert error_information["error_provider_request_id"] == amzn_request_id, error_information + + +@pytest.mark.parametrize("shape", _SHAPES) +@pytest.mark.parametrize("error_code", _ERROR_CODES) +def test_responses_stream_failure_log_retains_provider_request_id( + gateway: Gateway, shape: str, error_code: str +) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + provider_request_id: Final = str(uuid.uuid4()) + message: Final = f"stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + assert json.loads(request.body)["stream"] is True + return Reply( + headers={"x-amzn-RequestId": provider_request_id}, + content_type="text/event-stream", + chunks=_stream(identity, shape, error_code, message), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + response: Final = _streamed_response(gateway, model, identity) + assert message in response.text, response.text + assert ("partial output" in response.text) is (shape != "error-before-output"), response.text + assert response.headers.get("llm_provider-x-amzn-requestid") == provider_request_id, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == provider_request_id, information + + +@pytest.mark.parametrize("shape", _SHAPES) +def test_responses_stream_failure_without_provider_request_id_logs_none(gateway: Gateway, shape: str) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + message: Final = f"stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + return Reply(content_type="text/event-stream", chunks=_stream(identity, shape, "server_error", message)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + response: Final = _streamed_response(gateway, model, identity) + assert message in response.text, response.text + assert "llm_provider-x-amzn-requestid" not in response.headers, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_class"] == "InternalServerError", information + assert information["error_provider_request_id"] is None, information + + +@pytest.mark.parametrize("shape", ("nested-error", "response.failed")) +def test_bedrock_responses_stream_server_error_logs_amzn_request_id(gateway: Gateway, shape: str) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/openai/v1/responses", request.target + assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers + assert json.loads(request.body)["stream"] is True + return Reply( + headers={"x-amzn-RequestId": amzn_request_id}, + content_type="text/event-stream", + chunks=_stream(identity, shape, "server_error", message), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=_BEDROCK_RESPONSES_MODEL, + api_key=_TOKEN, + api_base=wire.url, + aws_region_name="us-east-1", + num_retries=0, + ) + response: Final = _streamed_response(gateway, model, identity) + assert "partial output" in response.text and message in response.text, response.text + assert response.headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["llm_provider"] == "bedrock", information + assert information["error_provider_request_id"] == amzn_request_id, information + + +def test_responses_stream_failure_log_reads_openai_x_request_id(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + provider_request_id: Final = f"req_{uuid.uuid4().hex}" + message: Final = f"stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + headers={"x-request-id": provider_request_id}, + content_type="text/event-stream", + chunks=_stream(identity, "top-level-error", "server_error", message), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + response: Final = _streamed_response(gateway, model, identity) + assert message in response.text, response.text + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == provider_request_id, information + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("event_type", ("error", "response.failed")) +async def test_responses_sdk_stream_error_retains_provider_request_id(asynchronous: bool, event_type: str) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + provider_request_id: Final = str(uuid.uuid4()) + shape: Final = "nested-error" if event_type == "error" else "response.failed" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + headers={"x-amzn-RequestId": provider_request_id, "retry-after": "7"}, + content_type="text/event-stream", + chunks=_stream(identity, shape, "server_error", identity), + ) + + with wire_server(respond) as wire: + with pytest.raises(MidStreamFallbackError) as raised: + await _consume_sdk_stream(asynchronous, wire.url, identity) + assert raised.value.status_code == 500 + assert raised.value.response.status_code == 500 + assert str(raised.value.response.request.url) == "https://openai.com/v1/" + assert raised.value.generated_content == "partial output" + assert dict(raised.value.response.headers) == {"x-amzn-requestid": provider_request_id} + assert _get_provider_request_id(raised.value) == provider_request_id + assert raised.value.original_exception is not None + assert _get_provider_request_id(raised.value.original_exception) == provider_request_id + assert len(wire.drain()) == 1 + + +def _concurrent_reply(request: Request, failing: tuple[str, ...], recovery: str) -> Reply: + prompt: Final = json.loads(request.body)["input"] + assert prompt in (*failing, recovery) + chunks: Final = ( + healthy_stream(f"resp_{prompt}", "recovered") + if prompt == recovery + else _stream(f"resp_{prompt}", "nested-error", "server_error", prompt) + ) + return Reply(content_type="text/event-stream", headers={"x-amzn-requestid": prompt}, chunks=chunks) + + +def test_concurrent_responses_failures_keep_distinct_request_ids_and_recover(gateway: Gateway) -> None: + identity: Final = uuid.uuid4().hex + prompts: Final = tuple(f"{identity}-{index}" for index in range(4)) + recovery: Final = f"{identity}-recovery" + + with ( + wire_server(lambda request: _concurrent_reply(request, prompts, recovery)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + + def call(prompt: str) -> str: + response: Final = _streamed_response(gateway, model, prompt) + assert (prompt if prompt != recovery else "recovered") in response.text, response.text + return response.headers["x-litellm-call-id"] + + with ThreadPoolExecutor(max_workers=4) as workers: + call_ids: Final = tuple(workers.map(call, prompts)) + recovered_call_id: Final = call(recovery) + assert len(wire.drain()) == 5 + for prompt, call_id in zip(prompts, call_ids, strict=True): + assert _failure_information(call_id)["error_provider_request_id"] == prompt + assert spend_row(recovered_call_id)["status"] == "success" + + +def _bedrock_logged_error(endpoint: BedrockEndpoint, failure: BedrockFailure) -> tuple[str, str]: + shape: Final = _BEDROCK_FAILURES[failure] + if endpoint == "messages-invoke": + return "BedrockError", str(shape.bedrock_status) + return shape.converse_class, "503" if failure == "exception-event" else str(shape.bedrock_status) + + +def _bedrock_client_status(endpoint: BedrockEndpoint, failure: BedrockFailure, after_output: bool) -> int: + logged_code: Final = int(_bedrock_logged_error(endpoint, failure)[1]) + if after_output or endpoint in ("responses", "messages"): + return 200 + if endpoint == "messages-invoke" and failure == "no-known-event": + return 200 + return logged_code + + +@pytest.mark.parametrize(("endpoint", "failure", "after_output"), _BEDROCK_CELLS) +def test_bedrock_event_stream_failure_log_retains_amzn_request_id( + gateway: Gateway, endpoint: BedrockEndpoint, failure: BedrockFailure, after_output: bool +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {amzn_request_id}" + invoke: Final = endpoint == "messages-invoke" + frames: Final = (*(_bedrock_output(invoke) if after_output else ()), _bedrock_failure_frame(failure, message)) + + def respond(request: Request) -> Reply: + assert request.method == "POST" and unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply( + headers={"x-amzn-RequestId": amzn_request_id, "retry-after": "7"}, + content_type=_EVENT_STREAM, + body=b"".join(frames), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/invoke/{_INVOKE_MODEL_ID}" if invoke else _MODEL, + api_base=wire.url, + num_retries=0, + **_BEDROCK_AWS, + ) + path, body = _bedrock_stream_request(endpoint, model, message) + with gateway.client.stream( + "POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"} + ) as response: + text: Final = "\n".join(line for line in response.iter_lines() if line) + status: Final = response.status_code + response_headers: Final = dict(response.headers) + assert status == _bedrock_client_status(endpoint, failure, after_output), text + assert _bedrock_failure_text(failure, message) in text and ("partial output" in text) is after_output, text + assert len(wire.drain()) == 1 + information: Final = _failure_information(response_headers["x-litellm-call-id"]) + assert (information["error_class"], information["error_code"]) == _bedrock_logged_error(endpoint, failure) + assert information["error_provider_request_id"] == amzn_request_id, information + if status != 200 and not invoke: + assert response_headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, response_headers + assert "llm_provider-retry-after" not in response_headers, response_headers + + +def _bedrock_model(scenario: Scenario, endpoint: BedrockEndpoint, api_base: str) -> str: + return scenario.model( + model=f"bedrock/invoke/{_INVOKE_MODEL_ID}" if endpoint == "messages-invoke" else _MODEL, + api_base=api_base, + num_retries=0, + **_BEDROCK_AWS, + ) + + +def _bedrock_stream( + gateway: Gateway, endpoint: BedrockEndpoint, model: str, prompt: str +) -> tuple[int, str, dict[str, str]]: + path, body = _bedrock_stream_request(endpoint, model, prompt) + with gateway.client.stream("POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}) as response: + text: Final = "\n".join(line for line in response.iter_lines() if line) + return response.status_code, text, dict(response.headers) + + +@pytest.mark.parametrize("after_output", (False, True), ids=("before-output", "after-output")) +@pytest.mark.parametrize("endpoint", _BEDROCK_ENDPOINTS) +def test_bedrock_event_stream_failure_without_amzn_request_id_logs_none( + gateway: Gateway, endpoint: BedrockEndpoint, after_output: bool +) -> None: + message: Final = f"bedrock stream failure {uuid.uuid4()}" + frames: Final = ( + *(_bedrock_output(endpoint == "messages-invoke") if after_output else ()), + _bedrock_failure_frame("exception-event", message), + ) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply(headers={"retry-after": "7"}, content_type=_EVENT_STREAM, body=b"".join(frames)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + status, text, response_headers = _bedrock_stream( + gateway, endpoint, _bedrock_model(scenario, endpoint, wire.url), message + ) + assert status == _bedrock_client_status(endpoint, "exception-event", after_output), text + assert message in text, text + assert not any("amzn" in name for name in response_headers), response_headers + assert len(wire.drain()) == 1 + information: Final = _failure_information(response_headers["x-litellm-call-id"]) + assert (information["error_class"], information["error_code"]) == _bedrock_logged_error( + endpoint, "exception-event" + ) + assert information["error_provider_request_id"] is None, information + + +@pytest.mark.parametrize("endpoint", ("responses", "messages", "messages-invoke")) +def test_bedrock_non_stream_500_logs_amzn_request_id_on_every_endpoint( + gateway: Gateway, endpoint: BedrockEndpoint +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + targets: Final = frozenset( + ( + "/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse", + f"/model/{_INVOKE_MODEL_ID}/invoke", + ) + ) + + def respond(request: Request) -> Reply: + assert request.method == "POST" and unquote(request.target) in targets, request.target + return Reply( + status=500, + headers={"x-amzn-RequestId": amzn_request_id}, + body=b'{"message":"synthetic bedrock failure"}', + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + path, body = _bedrock_stream_request( + endpoint, _bedrock_model(scenario, endpoint, wire.url), amzn_request_id, stream=False + ) + response: Final = gateway.request("POST", path, body) + assert response.status_code == 503, response.text + if endpoint == "responses": + assert response.headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == amzn_request_id, information + + +@pytest.mark.parametrize("amzn_request_id", ("", "a" * 5000, "dup-1, dup-1"), ids=("empty", "5kb", "joined-duplicate")) +def test_bedrock_stream_failure_logs_odd_request_ids_like_the_non_stream_error( + gateway: Gateway, amzn_request_id: str +) -> None: + stream_body: Final = b"".join((*_bedrock_output(False), _bedrock_failure_frame("exception-event", "odd id"))) + + def respond(request: Request) -> Reply: + if unquote(request.target) in _BEDROCK_STREAM_TARGETS: + return Reply(headers={"x-amzn-RequestId": amzn_request_id}, content_type=_EVENT_STREAM, body=stream_body) + return Reply(status=500, headers={"x-amzn-RequestId": amzn_request_id}, body=b'{"message":"odd id"}') + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _bedrock_model(scenario, "chat", wire.url) + status, text, stream_headers = _bedrock_stream(gateway, "chat", model, "odd id") + path, body = _bedrock_stream_request("chat", model, "odd id", stream=False) + non_stream: Final = gateway.request("POST", path, body) + assert status == 200 and "odd id" in text, text + assert non_stream.status_code == 503, non_stream.text + assert len(wire.drain()) == 2 + streamed: Final = _failure_information(stream_headers["x-litellm-call-id"]) + unstreamed: Final = _failure_information(non_stream.headers["x-litellm-call-id"]) + assert streamed["error_provider_request_id"] == unstreamed["error_provider_request_id"], (streamed, unstreamed) + + +_BURST_CALLS: Final = tuple( + (endpoint, failure, after_output) + for endpoint in _BEDROCK_ENDPOINTS + for failure, shape in _BEDROCK_FAILURES.items() + for after_output in shape.after_output +) + + +def _healthy_bedrock_stream(invoke: bool) -> bytes: + closing: Final = ( + ( + _invoke_chunk({"type": "content_block_stop", "index": 0}), + _invoke_chunk( + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}} + ), + _invoke_chunk({"type": "message_stop"}), + ) + if invoke + else ( + _bedrock_frame("contentBlockStop", {"contentBlockIndex": 0}), + _bedrock_frame("messageStop", {"stopReason": "end_turn"}), + _bedrock_frame("metadata", {"usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}}), + ) + ) + return b"".join((*_bedrock_output(invoke, "recovered"), *closing)) + + +def test_concurrent_bedrock_stream_failures_on_every_endpoint_keep_distinct_request_ids(gateway: Gateway) -> None: + identity: Final = uuid.uuid4().hex + failing: Final = {f"{identity}-{index:02d}": call for index, call in enumerate(_BURST_CALLS)} + recoveries: Final = {f"{identity}-r{index}": endpoint for index, endpoint in enumerate(_BEDROCK_ENDPOINTS)} + + def respond(request: Request) -> Reply: + (token,) = (token for token in (*failing, *recoveries) if token in request.body.decode()) + if token in recoveries: + body = _healthy_bedrock_stream(recoveries[token] == "messages-invoke") + else: + endpoint, failure, after_output = failing[token] + opening = _bedrock_output(endpoint == "messages-invoke") if after_output else () + body = b"".join((*opening, _bedrock_failure_frame(failure, token))) + return Reply(headers={"x-amzn-RequestId": token, "retry-after": "7"}, content_type=_EVENT_STREAM, body=body) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + models: Final = {endpoint: _bedrock_model(scenario, endpoint, wire.url) for endpoint in _BEDROCK_ENDPOINTS} + + def call(token: str, endpoint: BedrockEndpoint) -> tuple[int, str, dict[str, str]]: + return _bedrock_stream(gateway, endpoint, models[endpoint], token) + + with ThreadPoolExecutor(max_workers=8) as workers: + outcomes: Final = dict( + zip(failing, workers.map(call, failing, (endpoint for endpoint, _, _ in failing.values())), strict=True) + ) + recovered: Final = {token: call(token, endpoint) for token, endpoint in recoveries.items()} + assert len(wire.drain()) == len(failing) + len(recoveries) + for token, (endpoint, failure, after_output) in failing.items(): + status, text, response_headers = outcomes[token] + assert status == _bedrock_client_status(endpoint, failure, after_output), (token, text) + assert _bedrock_failure_text(failure, token) in text, (token, text) + information = _failure_information(response_headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == token, (token, information) + for token, (status, text, response_headers) in recovered.items(): + assert status == 200 and "recovered" in text, (token, text) + assert spend_row(response_headers["x-litellm-call-id"])["status"] == "success" + + +def _sdk_stream_text( + client_kind: str, asynchronous: bool, gateway: Gateway, model: str, prompt: str +) -> tuple[str, str]: + extra_body: Final[dict[str, JsonValue]] = {"num_retries": 0, "cache": {"no-cache": True}} + turn: Final = [{"role": "user", "content": prompt}] + + def open_stream(client: openai.OpenAI | openai.AsyncOpenAI | anthropic.Anthropic | anthropic.AsyncAnthropic): + match client_kind: + case "openai-chat": + return client.chat.completions.with_raw_response.create( + model=model, messages=turn, stream=True, extra_body=extra_body + ) + case "openai-responses": + return client.responses.with_raw_response.create( + model=model, input=prompt, stream=True, extra_body=extra_body + ) + case "anthropic-messages": + return client.messages.with_raw_response.create( + model=model, messages=turn, max_tokens=16, stream=True, extra_body=extra_body + ) + raise AssertionError(client_kind) + + seen: Final[list[str]] = [] + base_url: Final = str(gateway.client.base_url).rstrip("/") + if not asynchronous: + with httpx.Client(timeout=15, trust_env=False) as transport: + client = ( + anthropic.Anthropic(api_key=gateway.key, base_url=base_url, max_retries=0, http_client=transport) + if client_kind == "anthropic-messages" + else openai.OpenAI(api_key=gateway.key, base_url=f"{base_url}/v1", max_retries=0, http_client=transport) + ) + raw = open_stream(client) + try: + seen.extend(repr(event) for event in raw.parse()) + except (openai.APIError, anthropic.APIError) as error: + seen.append(str(error)) + return raw.headers["x-litellm-call-id"], "\n".join(seen) + + async def run() -> str: + async with httpx.AsyncClient(timeout=15, trust_env=False) as transport: + client = ( + anthropic.AsyncAnthropic(api_key=gateway.key, base_url=base_url, max_retries=0, http_client=transport) + if client_kind == "anthropic-messages" + else openai.AsyncOpenAI( + api_key=gateway.key, base_url=f"{base_url}/v1", max_retries=0, http_client=transport + ) + ) + raw = await open_stream(client) + try: + async for event in raw.parse(): + seen.append(repr(event)) + except (openai.APIError, anthropic.APIError) as error: + seen.append(str(error)) + return raw.headers["x-litellm-call-id"] + + return asyncio.run(run()), "\n".join(seen) + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("client_kind", ("openai-chat", "openai-responses", "anthropic-messages")) +def test_sdk_clients_streaming_bedrock_failure_through_proxy_log_amzn_request_id( + gateway: Gateway, client_kind: str, asynchronous: bool +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {amzn_request_id}" + frames: Final = (*_bedrock_output(False), _bedrock_failure_frame("exception-event", message)) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply(headers={"x-amzn-RequestId": amzn_request_id}, content_type=_EVENT_STREAM, body=b"".join(frames)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + call_id, seen = _sdk_stream_text( + client_kind, asynchronous, gateway, _bedrock_model(scenario, "chat", wire.url), message + ) + assert "partial output" in seen and message in seen, seen + assert len(wire.drain()) == 1 + information: Final = _failure_information(call_id) + assert information["error_provider_request_id"] == amzn_request_id, information + + +async def _consume_bedrock_sdk_stream(api: Literal["chat", "responses"], asynchronous: bool, api_base: str) -> str: + common: Final[dict[str, JsonValue]] = { + "model": _MODEL, + "stream": True, + "api_base": api_base, + "num_retries": 0, + **_BEDROCK_AWS, + } + turn: Final[list[JsonValue]] = [{"role": "user", "content": "probe"}] + match api, asynchronous: + case "chat", True: + return "".join([repr(chunk) async for chunk in await litellm.acompletion(messages=turn, **common)]) + case "chat", False: + return "".join(repr(chunk) for chunk in litellm.completion(messages=turn, **common)) + case "responses", True: + return "".join([repr(event) async for event in await litellm.aresponses(input="probe", **common)]) + case "responses", False: + return "".join(repr(event) for event in litellm.responses(input="probe", **common)) + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("api", ("chat", "responses")) +async def test_bedrock_sdk_stream_exception_carries_amzn_request_id( + api: Literal["chat", "responses"], asynchronous: bool +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {uuid.uuid4()}" + frames: Final = (*_bedrock_output(False), _bedrock_failure_frame("exception-event", message)) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply( + headers={"x-amzn-RequestId": amzn_request_id, "retry-after": "7"}, + content_type=_EVENT_STREAM, + body=b"".join(frames), + ) + + with wire_server(respond) as wire: + with pytest.raises(MidStreamFallbackError) as raised: + await _consume_bedrock_sdk_stream(api, asynchronous, wire.url) + assert len(wire.drain()) == 1 + provider_error: Final = raised.value.original_exception + assert isinstance(provider_error, litellm.ServiceUnavailableError), provider_error + assert raised.value.generated_content == "partial output" + assert message in str(provider_error) + assert _get_provider_request_id(provider_error) == amzn_request_id + assert "retry-after" not in provider_error.response.headers, provider_error.response.headers + + +@pytest.mark.parametrize("api", ("chat", "responses")) +@pytest.mark.parametrize("healthy", (True, False), ids=("healthy", "failing")) +async def test_sync_bedrock_sdk_stream_with_non_ascii_request_id_header_keeps_streaming( + api: Literal["chat", "responses"], healthy: bool +) -> None: + body: Final = ( + _healthy_bedrock_stream(False) + if healthy + else b"".join((*_bedrock_output(False), _bedrock_failure_frame("exception-event", "non-ascii id failure"))) + ) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply(headers={"x-amzn-RequestId": "req-\u00fc-1"}, content_type=_EVENT_STREAM, body=body) + + with wire_server(respond) as wire: + if healthy: + assert "recovered" in await _consume_bedrock_sdk_stream(api, False, wire.url) + else: + with pytest.raises(MidStreamFallbackError) as raised: + await _consume_bedrock_sdk_stream(api, False, wire.url) + assert "non-ascii id failure" in str(raised.value.original_exception) + assert _get_provider_request_id(raised.value.original_exception) is None + assert len(wire.drain()) == 1 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_content_filter_logging_only.py b/tests/integration/observability/test_content_filter_logging_only.py new file mode 100644 index 00000000000..ab1779c497d --- /dev/null +++ b/tests/integration/observability/test_content_filter_logging_only.py @@ -0,0 +1,712 @@ +"""`litellm_content_filter` in `mode: logging_only`, driven through a real proxy. + +The provider is the only double, an owned wire peer. The proxy, its guardrail registry, Postgres and Redis run +for real with two workers. Every check reads the spend row, which is written after the logging-only scans finish +""" + +from __future__ import annotations + +import json +import uuid +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from datetime import datetime, timezone +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import anthropic +import openai +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value +from integration._support.database import read_rows +from integration._support.mcp import EntryPoint, McpCaller, echo_tool, register_mcp, scripted_peer, tool_calls +from integration._support.otlp_sink import AuditConfigWriter, SpanSinks, recorded_spans +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server + +AWS_KEY: Final = "AKIA" + "SYNTHETIC0000000" +GITHUB_TOKEN: Final = "ghp_" + "a1" * 18 +FORBIDDEN_WORD: Final = "synthetic-forbidden-word" +LEAK_MARK: Final = "SYNTHETIC-LEAK-TOKEN" +OUTAGE_MARK: Final = "SYNTHETIC-PROVIDER-OUTAGE" +OBSERVE_BLOCK: Final = "cf-observe-block" +OBSERVE_MASK: Final = "cf-observe-mask" +ENFORCE_BLOCK: Final = "cf-enforce-block" +ENFORCE_MASK: Final = "cf-enforce-mask" +OBSERVE_CATEGORY: Final = "cf-observe-category" +OBSERVE_COMPETITOR: Final = "cf-observe-competitor" +SECRETS: Final = f"AWS_ACCESS_KEY_ID={AWS_KEY}\nGITHUB_TOKEN={GITHUB_TOKEN}" + + +def _answer(marker: str) -> str: + return "synthetic answer " + marker + + +def _provider_answer(body: bytes) -> str: + text: Final = body.decode() + start: Final = text.find("mark-") + assert start >= 0, text + answer: Final = _answer(text[start : start + 37]) + return answer + " " + GITHUB_TOKEN if LEAK_MARK in text else answer + + +def _sse(events: tuple[dict[str, object], ...]) -> tuple[bytes, ...]: + return tuple(("data: " + json.dumps(event) + "\n\n").encode() for event in events) + + +def _chat_reply(answer: str, streaming: bool) -> Reply: + if not streaming: + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": 1, + "model": "gpt-5.4-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": answer}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + identity: Final = "chatcmpl-" + uuid.uuid4().hex + + def chunk(delta: dict[str, object], finish: str | None) -> dict[str, object]: + return { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-5.4-mini", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + } + + events: Final = (chunk({"role": "assistant", "content": answer[:9]}, None), chunk({"content": answer[9:]}, "stop")) + return Reply(chunks=(*_sse(events), b"data: [DONE]\n\n"), content_type="text/event-stream") + + +def _messages_reply(answer: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": "msg_" + uuid.uuid4().hex, + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": answer}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + ).encode() + ) + + +def _responses_reply(answer: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": "resp_" + uuid.uuid4().hex, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msgo_" + uuid.uuid4().hex, + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": answer, "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + +def _provider(request: Request) -> Reply: + if OUTAGE_MARK.encode() in request.body: + return Reply(status=503, body=b'{"error": {"message": "synthetic provider outage", "type": "server_error"}}') + answer: Final = _provider_answer(request.body) + if request.target.endswith("/v1/messages"): + return _messages_reply(answer) + if request.target.endswith("/v1/responses"): + return _responses_reply(answer) + assert request.target.endswith("/v1/chat/completions"), request.target + return _chat_reply(answer, bool(json.loads(request.body).get("stream"))) + + +def _content_filter(name: str, mode: str, action: str, *, default_on: bool) -> dict[str, object]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": mode, + "default_on": default_on, + "patterns": [ + {"pattern_type": "regex", "pattern": "AKIA[0-9A-Z]{16}", "name": "aws_access_key", "action": action}, + {"pattern_type": "regex", "pattern": "ghp_[A-Za-z0-9]{36}", "name": "github_pat", "action": action}, + ], + "blocked_words": [{"keyword": FORBIDDEN_WORD, "action": action}], + }, + } + + +def _observe_only(name: str, **rules: object) -> dict[str, object]: + return { + "guardrail_name": name, + "litellm_params": {"guardrail": "litellm_content_filter", "mode": "logging_only", "default_on": False, **rules}, + } + + +def _rig_config(root: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"]["cache"] = False + config["general_settings"]["store_prompts_in_spend_logs"] = True + config["guardrails"] = [ + _content_filter(OBSERVE_BLOCK, "logging_only", "BLOCK", default_on=True), + _content_filter(OBSERVE_MASK, "logging_only", "MASK", default_on=True), + _content_filter(ENFORCE_BLOCK, "pre_call", "BLOCK", default_on=False), + _content_filter(ENFORCE_MASK, "pre_call", "MASK", default_on=False), + _observe_only( + OBSERVE_CATEGORY, + categories=[ + {"category": "harmful_violence", "enabled": True, "action": "BLOCK", "severity_threshold": "medium"} + ], + blocked_words=[{"keyword": FORBIDDEN_WORD, "action": "BLOCK"}], + ), + _observe_only( + OBSERVE_COMPETITOR, + competitor_intent_config={ + "brand_self": ["emirates"], + "competitors": ["qatar airways"], + "domain_words": ["airline", "flight"], + "policy": {"competitor_comparison": "refuse"}, + "threshold_high": 0.5, + }, + ), + ] + path: Final = root / "content-filter-logging-only.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + provider: Wire + chat_model: str + claude_model: str + responses_model: str + scenario: Scenario + + def base(self) -> str: + return str(self.proxy.client.base_url).rstrip("/") + + def provider_bodies(self, marker: str) -> tuple[str, ...]: + return tuple(request.body.decode() for request in self.provider.drain() if marker.encode() in request.body) + + def entries(self, response_id: str, name: str) -> tuple[dict[str, object], ...]: + return _named_entries( + 'SELECT request_id, status, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (response_id,), name + ) + + +def _named_entries(query: str, parameters: tuple[str, ...], name: str) -> tuple[dict[str, object], ...]: + rows: Final = eventually(lambda: read_rows(query, parameters), lambda values: len(values) >= 1, seconds=70) + assert len(rows) == 1, rows + recorded: Final = object_value(rows[0]["metadata"]).get("guardrail_information") + assert isinstance(recorded, list), rows[0] + return tuple(entry for entry in (object_value(item) for item in recorded) if entry.get("guardrail_name") == name) + + +def _findings(entry: dict[str, object]) -> tuple[str, ...]: + response: Final = entry.get("guardrail_response") + assert isinstance(response, list), entry + return tuple( + str(detection.get("pattern_name") or detection.get("keyword") or detection.get("type")) + for detection in (object_value(item) for item in response) + ) + + +def _verdicts(entries: tuple[dict[str, object], ...]) -> list[tuple[object, object, tuple[str, ...]]]: + return [(entry["guardrail_mode"], entry["guardrail_status"], _findings(entry)) for entry in entries] + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + root: Final = tmp_path_factory.mktemp("content-filter-logging-only") + with ( + gateway_from_environment() as gateway, + wire_server(_provider) as provider, + owned_proxy_process(gateway, root, {}, config=_rig_config(root), workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + chat: Final = scenario.model( + model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key" + ) + claude: Final = scenario.model( + model="anthropic/claude-haiku-5-5", api_base=provider.url, api_key="synthetic-anthropic-key" + ) + responses: Final = scenario.model( + model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key" + ) + yield Rig(owned.gateway, provider, chat, claude, responses, scenario) + + +def _marker() -> str: + return "mark-" + uuid.uuid4().hex + + +def _sync_chat(rig: Rig, text: str) -> tuple[str, str]: + client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0) + completion: Final = client.chat.completions.create( + model=rig.chat_model, messages=[{"role": "user", "content": text}] + ) + return completion.id, completion.choices[0].message.content or "" + + +async def _async_chat(rig: Rig, text: str) -> tuple[str, str]: + client: Final = openai.AsyncOpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0) + completion: Final = await client.chat.completions.create( + model=rig.chat_model, messages=[{"role": "user", "content": text}] + ) + return completion.id, completion.choices[0].message.content or "" + + +def _streamed_chat(rig: Rig, text: str) -> tuple[str, str]: + client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0) + chunks: Final = tuple( + client.chat.completions.create(model=rig.chat_model, messages=[{"role": "user", "content": text}], stream=True) + ) + return chunks[0].id, "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) + + +def _messages(rig: Rig, text: str) -> tuple[str, str]: + client: Final = anthropic.Anthropic(base_url=rig.base(), api_key=rig.proxy.key, max_retries=0) + message: Final = client.messages.create( + model=rig.claude_model, max_tokens=64, messages=[{"role": "user", "content": text}] + ) + assert message.content[0].type == "text" + return message.id, message.content[0].text + + +def _responses(rig: Rig, text: str) -> tuple[str, str]: + client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0) + result: Final = client.responses.create(model=rig.responses_model, input=text) + return result.id, result.output_text + + +SYNC_CALLS: Final = {"chat": _sync_chat, "chat_stream": _streamed_chat, "messages": _messages, "responses": _responses} +EVERY_FINDING: Final = ("aws_access_key", "github_pat", FORBIDDEN_WORD) + + +@pytest.mark.parametrize("endpoint", list(SYNC_CALLS)) +def test_logging_only_block_records_every_match_without_touching_the_call(rig: Rig, endpoint: str) -> None: + marker: Final = _marker() + prompt: Final = f"{marker} what did the file contain?\n{SECRETS}\n{FORBIDDEN_WORD}" + response_id, answer = SYNC_CALLS[endpoint](rig, prompt) + assert answer == _answer(marker) + sent: Final = rig.provider_bodies(marker) + assert len(sent) == 1 and AWS_KEY in sent[0] and GITHUB_TOKEN in sent[0] and FORBIDDEN_WORD in sent[0], sent + + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "guardrail_flagged", EVERY_FINDING), + ("logging_only", "success", ()), + ] + + +@pytest.mark.asyncio +async def test_logging_only_block_records_every_match_for_an_async_client(rig: Rig) -> None: + marker: Final = _marker() + response_id, answer = await _async_chat(rig, f"{marker} {SECRETS}") + assert answer == _answer(marker) + + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "guardrail_flagged", ("aws_access_key", "github_pat")), + ("logging_only", "success", ()), + ] + + +def test_logging_only_block_still_scans_the_response_after_a_request_match(rig: Rig) -> None: + marker: Final = _marker() + response_id, answer = _sync_chat(rig, f"{marker} {LEAK_MARK} key {AWS_KEY}") + assert answer == _answer(marker) + " " + GITHUB_TOKEN + + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "guardrail_flagged", ("aws_access_key",)), + ("logging_only", "guardrail_flagged", ("github_pat",)), + ] + + +def _tool_call_history(text: str, arguments: str) -> list[dict[str, object]]: + return [ + {"role": "user", "content": text}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": arguments}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "done"}, + ] + + +@pytest.mark.parametrize( + ("arguments", "request_findings"), + ( + pytest.param(json.dumps({AWS_KEY: 1}), ("aws_access_key",), id="masked-argument-key"), + pytest.param("[" * 150 + "1" + "]" * 150, (), id="too-deep"), + ), +) +def test_logging_only_flags_tool_arguments_an_enforcing_mode_rejects_and_keeps_scanning( + rig: Rig, arguments: str, request_findings: tuple[str, ...] +) -> None: + marker: Final = _marker() + messages: Final = _tool_call_history(f"{marker} {LEAK_MARK}", arguments) + headers: Final = {"Authorization": f"Bearer {rig.proxy.key}"} + enforced: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={"model": rig.chat_model, "messages": messages, "guardrails": [ENFORCE_MASK]}, + headers=headers, + ) + assert enforced.status_code == 400, enforced.text + assert rig.provider_bodies(marker) == () + + observed: Final = rig.proxy.client.post( + "/v1/chat/completions", json={"model": rig.chat_model, "messages": messages}, headers=headers + ) + assert observed.status_code == 200, observed.text + completion: Final = object_value(observed.json()) + assert len(rig.provider_bodies(marker)) == 1 + + assert _verdicts(rig.entries(str(completion["id"]), OBSERVE_MASK)) == [ + ("logging_only", "guardrail_flagged", request_findings), + ("logging_only", "success", ("github_pat",)), + ] + + +@pytest.mark.parametrize("entry", ("mcp", "rest")) +def test_logging_only_block_records_every_match_on_an_mcp_tool_call(rig: Rig, entry: EntryPoint) -> None: + alias: Final = "cfmcp" + uuid.uuid4().hex[:8] + arguments: Final = {"note": f"{_marker()} {SECRETS}"} + with scripted_peer(echo_tool("echo")) as peer: + server_id: Final = register_mcp(rig.scenario, peer, alias) + key: Final = rig.scenario.key(object_permission={"mcp_servers": [server_id]}) + outcome: Final = McpCaller(rig.proxy, key, entry, alias).call(f"{alias}-echo", arguments, server_id=server_id) + assert (outcome.error, outcome.text) == (None, json.dumps(arguments, sort_keys=True)), outcome.raw + assert [call["body"]["params"]["arguments"] for call in tool_calls(peer.drain())] == [arguments] + + entries: Final = _named_entries( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND call_type = %s', + (sha256(key.encode()).hexdigest(), "call_mcp_tool"), + OBSERVE_BLOCK, + ) + assert _verdicts(entries) == [("logging_only", "guardrail_flagged", ("aws_access_key", "github_pat"))], entries + + +@pytest.mark.parametrize( + ("guardrail", "prompt", "findings"), + ( + pytest.param( + OBSERVE_CATEGORY, + f"how to kill someone {FORBIDDEN_WORD}", + ("how to kill someone", FORBIDDEN_WORD), + id="category-keyword", + ), + pytest.param( + OBSERVE_COMPETITOR, "Is Qatar Airways better than Emirates?", ("competitor_intent",), id="competitor-intent" + ), + ), +) +def test_logging_only_block_flags_category_and_competitor_matches_without_blocking( + rig: Rig, guardrail: str, prompt: str, findings: tuple[str, ...] +) -> None: + marker: Final = _marker() + response: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": f"{marker} {prompt}"}], + "guardrails": [guardrail], + }, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + assert response.status_code == 200, response.text + completion: Final = object_value(response.json()) + assert len(rig.provider_bodies(marker)) == 1 + + assert _verdicts(rig.entries(str(completion["id"]), guardrail)) == [ + ("logging_only", "guardrail_flagged", findings), + ("logging_only", "success", ()), + ] + + +def test_logging_only_mask_records_every_match(rig: Rig) -> None: + marker: Final = _marker() + response_id, answer = _sync_chat(rig, f"{marker} {SECRETS}\n{FORBIDDEN_WORD}") + assert answer == _answer(marker) + + entries: Final = rig.entries(response_id, OBSERVE_MASK) + assert _verdicts(entries) == [ + ("logging_only", "success", EVERY_FINDING), + ("logging_only", "success", ()), + ] + assert entries[0]["masked_entity_count"] == {"aws_access_key": 1, "github_pat": 1, "blocked_word": 1}, entries + + +def test_logging_only_records_nothing_for_a_clean_call(rig: Rig) -> None: + marker: Final = _marker() + response_id, answer = _sync_chat(rig, f"{marker} nothing secret here") + assert answer == _answer(marker) + + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "success", ()), + ("logging_only", "success", ()), + ] + + +def test_pre_call_block_still_rejects_the_request(rig: Rig) -> None: + marker: Final = _marker() + response: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": f"{marker} {SECRETS}"}], + "guardrails": [ENFORCE_BLOCK], + }, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + assert response.status_code == 400, response.text + assert "aws_access_key" in response.text + assert rig.provider_bodies(marker) == () + + +def test_apply_guardrail_endpoint_still_blocks_with_a_logging_only_content_filter(rig: Rig) -> None: + response: Final = rig.proxy.client.post( + "/guardrails/apply_guardrail", + json={"guardrail_name": OBSERVE_BLOCK, "text": f"key {AWS_KEY}"}, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + assert response.status_code == 400, response.text + assert "aws_access_key" in response.text + + +def _advertised_modes(rig: Rig) -> dict[str, object]: + response: Final = rig.proxy.client.get( + "/guardrails/ui/add_guardrail_settings", headers={"Authorization": f"Bearer {rig.proxy.key}"} + ) + assert response.status_code == 200, response.text + return object_value(response.json()["supported_modes_by_provider"]) + + +@pytest.mark.parametrize("provider", ["litellm_content_filter", "generic_guardrail_api", "bedrock", "akto"]) +def test_dashboard_offers_logging_only_where_the_config_accepts_it(rig: Rig, provider: str) -> None: + modes: Final = _advertised_modes(rig) + offered: Final = modes[provider] + assert isinstance(offered, list) and "logging_only" in offered, modes + assert offered.count("logging_only") == 1, offered + + +@pytest.mark.parametrize("provider", ["azure/prompt_shield", "lakera_v2"]) +def test_dashboard_does_not_offer_logging_only_where_the_config_rejects_it(rig: Rig, provider: str) -> None: + offered: Final = _advertised_modes(rig)[provider] + assert isinstance(offered, list) and "logging_only" not in offered, offered + + +def test_a_logging_only_content_filter_created_through_the_api_records_every_match(rig: Rig) -> None: + name: Final = "cf-observe-api-" + uuid.uuid4().hex[:8] + created: Final = rig.proxy.client.post( + "/guardrails", + json={"guardrail": {"guardrail_name": name, **_content_filter(name, "logging_only", "BLOCK", default_on=True)}}, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + assert created.status_code == 200, created.text + guardrail_id: Final = str(object_value(created.json())["guardrail_id"]) + try: + marker: Final = _marker() + response_id: Final = eventually( + lambda: _sync_chat(rig, f"{marker} {SECRETS}")[0], + lambda identity: bool(rig.entries(identity, name)), + seconds=70, + ) + assert _verdicts(rig.entries(response_id, name))[0] == ( + "logging_only", + "guardrail_flagged", + ("aws_access_key", "github_pat"), + ) + finally: + rig.proxy.client.delete(f"/guardrails/{guardrail_id}", headers={"Authorization": f"Bearer {rig.proxy.key}"}) + + +def test_guardrails_monitor_counts_observe_only_matches_as_flagged_not_blocked(rig: Rig) -> None: + name: Final = "cf-observe-usage-" + uuid.uuid4().hex[:8] + headers: Final = {"Authorization": f"Bearer {rig.proxy.key}"} + created: Final = rig.proxy.client.post( + "/guardrails", + json={"guardrail": {"guardrail_name": name, **_content_filter(name, "logging_only", "BLOCK", default_on=True)}}, + headers=headers, + ) + assert created.status_code == 200, created.text + guardrail_id: Final = str(object_value(created.json())["guardrail_id"]) + try: + eventually( + lambda: _sync_chat(rig, f"{_marker()} {SECRETS}")[0], + lambda identity: bool(rig.entries(identity, name)), + seconds=70, + ) + today: Final = datetime.now(timezone.utc).date().isoformat() + detail: Final = eventually( + lambda: object_value( + rig.proxy.client.get( + f"/guardrails/usage/detail/{guardrail_id}", + params={"start_date": today, "end_date": today}, + headers=headers, + ).json() + ), + lambda body: int(str(body.get("requestsEvaluated", 0))) >= 1, + seconds=60, + ) + assert (detail["failRate"], detail["status"]) == (0.0, "healthy"), detail + finally: + rig.proxy.client.delete(f"/guardrails/{guardrail_id}", headers=headers) + + +@pytest.mark.parametrize( + ("environment", "name_key", "status_key"), + ( + pytest.param({}, "guardrail_name", "guardrail_status", id="otel-v1"), + pytest.param({"LITELLM_OTEL_V2": "1"}, "litellm.guardrail.name", "litellm.guardrail.status", id="otel-v2"), + ), +) +def test_otel_guardrail_spans_carry_the_flagged_status_without_an_error( + gateway: Gateway, + audit_sinks: SpanSinks, + otel_audit_config: AuditConfigWriter, + tmp_path: Path, + environment: dict[str, str], + name_key: str, + status_key: str, +) -> None: + config_path: Final = otel_audit_config(tmp_path, {"cache": False}) + config: Final = yaml.safe_load(config_path.read_text()) + config["guardrails"] = [_content_filter(OBSERVE_BLOCK, "logging_only", "BLOCK", default_on=True)] + config_path.write_text(yaml.safe_dump(config)) + overrides: Final = {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": audit_sinks.operator, **environment} + with ( + wire_server(_provider) as provider, + owned_proxy_process(gateway, tmp_path, overrides, config=config_path, workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key" + ) + start, _ = recorded_spans(audit_sinks.operator) + marker: Final = _marker() + client: Final = openai.OpenAI( + base_url=str(owned.gateway.client.base_url).rstrip("/") + "/v1", api_key=owned.gateway.key, max_retries=0 + ) + completion: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": f"{marker} {SECRETS}"}] + ) + assert completion.choices[0].message.content == _answer(marker) + + spans: Final = eventually( + lambda: tuple( + span + for span in recorded_spans(audit_sinks.operator, start)[1] + if span["attributes"].get(name_key) == OBSERVE_BLOCK + ), + lambda found: len(found) >= 2, + seconds=60, + return_last_on_timeout=True, + ) + assert sorted( + (str(span["attributes"].get(status_key)), "error.type" in span["attributes"]) for span in spans + ) == [("guardrail_flagged", False), ("success", False)], spans + + +def _burst_call(rig: Rig, index: int) -> tuple[str, int, str, str]: + marker: Final = _marker() + if index % 4 == 0: + enforced: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": f"{marker} {SECRETS}"}], + "guardrails": [ENFORCE_BLOCK], + }, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + return "enforce", enforced.status_code, marker, "" + if index % 4 == 1: + failed: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": f"{marker} {OUTAGE_MARK} {SECRETS}"}], + }, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + return "outage", failed.status_code, marker, failed.headers.get("x-litellm-call-id", "") + endpoint: Final = list(SYNC_CALLS)[index % len(SYNC_CALLS)] + response_id, answer = SYNC_CALLS[endpoint](rig, f"{marker} {SECRETS}\n{FORBIDDEN_WORD}") + assert answer == _answer(marker), (endpoint, answer) + return "observe:" + endpoint, 200, marker, response_id + + +def _failure_row(call_id: str) -> tuple[object, tuple[dict[str, object], ...]]: + rows: Final = eventually( + lambda: read_rows('SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)), + lambda values: len(values) >= 1, + seconds=70, + ) + assert len(rows) == 1, rows + recorded: Final = object_value(rows[0]["metadata"]).get("guardrail_information") or [] + assert isinstance(recorded, list), rows[0] + return rows[0]["status"], tuple( + entry for entry in (object_value(item) for item in recorded) if entry.get("guardrail_name") == OBSERVE_BLOCK + ) + + +def test_concurrent_burst_with_a_provider_outage_records_every_observed_call_once(rig: Rig) -> None: + def call(index: int) -> tuple[str, int, str, str]: + return _burst_call(rig, index) + + with ThreadPoolExecutor(max_workers=12) as pool: + results: Final = tuple(pool.map(call, range(32))) + + assert sorted({(kind.split(":")[0], status) for kind, status, _, _ in results}) == [ + ("enforce", 400), + ("observe", 200), + ("outage", 503), + ], results + sent: Final = rig.provider_bodies("mark-") + for kind, _, marker, _ in results: + if kind == "enforce": + assert not any(marker in body for body in sent), marker + outages: Final = tuple(call_id for kind, _, _, call_id in results if kind == "outage") + assert len(outages) == 8 and all(outages), results + assert [_failure_row(call_id) for call_id in outages] == [("failure", ())] * 8 + observed: Final = tuple((kind, response_id) for kind, _, _, response_id in results if kind.startswith("observe")) + assert len(observed) == 16 + for kind, response_id in observed: + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "guardrail_flagged", EVERY_FINDING), + ("logging_only", "success", ()), + ], (kind, response_id) + + marker: Final = _marker() + response_id, answer = _sync_chat(rig, f"{marker} {SECRETS}") + assert answer == _answer(marker) + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK))[0] == ( + "logging_only", + "guardrail_flagged", + ("aws_access_key", "github_pat"), + ) 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_logging_only_scope_config.py b/tests/integration/observability/test_logging_only_scope_config.py index c8c765a68a2..354e3b464fe 100644 --- a/tests/integration/observability/test_logging_only_scope_config.py +++ b/tests/integration/observability/test_logging_only_scope_config.py @@ -179,16 +179,16 @@ def test_native_content_filter_scope_logs_without_blocking( ( identity, "logging_only", - "guardrail_intervened" if direction == blocked_side else "success", + "guardrail_flagged" if direction == blocked_side else "success", ) for direction in expected_directions ), entries - intervened_entries: Final = tuple( - entry for entry in entries if entry["guardrail_status"] == "guardrail_intervened" + flagged_entries: Final = tuple( + entry for entry in entries if entry["guardrail_status"] == "guardrail_flagged" ) - assert len(intervened_entries) == 1, entries - assert intervened_entries[0]["guardrail_response"] is not None, entries - assert intervened_entries[0]["guardrail_response"] == "REDACTED_BY_LITELM", entries + assert len(flagged_entries) == 1, entries + assert flagged_entries[0]["guardrail_response"] is not None, entries + assert flagged_entries[0]["guardrail_response"] == "REDACTED_BY_LITELM", entries finally: delete_scenario(upstream_handle) diff --git a/tests/integration/observability/test_otel_excluded_services_matrix.py b/tests/integration/observability/test_otel_excluded_services_matrix.py index 07646a0b189..8a7236f8626 100644 --- a/tests/integration/observability/test_otel_excluded_services_matrix.py +++ b/tests/integration/observability/test_otel_excluded_services_matrix.py @@ -358,6 +358,10 @@ def _has_root(spans: tuple[Span, ...]) -> bool: return any(span["kind"] == SERVER for span in spans) +def _has_model_span(spans: tuple[Span, ...]) -> bool: + return any("gen_ai.operation.name" in span["attributes"] for span in spans) + + def _trace_of_call(sink: str, call_id: str, since: int) -> tuple[Span, ...]: _, spans = recorded_spans(sink, since) traces: Final = {span["trace_id"] for span in spans if span["attributes"].get("litellm.call_id") == call_id} @@ -367,7 +371,7 @@ def _trace_of_call(sink: str, call_id: str, since: int) -> tuple[Span, ...]: def _operator_trace(rig: Rig, sent: Sent, cursors: Cursors) -> tuple[Span, ...]: trace: Final = eventually( lambda: _trace_of_call(rig.sinks.operator, sent.call_id, cursors.operator), - lambda spans: _has_root(spans) and "redis" in _db_systems(spans), + lambda spans: _has_root(spans) and "redis" in _db_systems(spans) and _has_model_span(spans), seconds=40, ) assert len({span["trace_id"] for span in trace}) == 1, _names(trace) @@ -405,7 +409,7 @@ def _assert_tenant_mirrors(rig: Rig, operator: tuple[Span, ...], cursors: Cursor def _assert_withheld(rig: Rig, sent: Sent, cursors: Cursors) -> tuple[Span, ...]: tenant: Final = _assert_tenant_mirrors(rig, _operator_trace(rig, sent, cursors), cursors) - assert any("gen_ai.operation.name" in span["attributes"] for span in tenant), _names(tenant) + assert _has_model_span(tenant), _names(tenant) return tenant @@ -719,7 +723,7 @@ def _assert_tenant_kept( return ( sum(1 for span in spans if span["kind"] == SERVER) == 1 and "redis" in _db_systems(spans) - and (not needs_model_span or any("gen_ai.operation.name" in span["attributes"] for span in spans)) + and (not needs_model_span or _has_model_span(spans)) ) tenant: Final = eventually( @@ -730,7 +734,7 @@ def _assert_tenant_kept( ) assert sum(1 for span in tenant if span["kind"] == SERVER) == 1, _names(tenant) assert "redis" in _db_systems(tenant), f"redis spans missing at the tenant: {_names(tenant)}" - assert not needs_model_span or any("gen_ai.operation.name" in span["attributes"] for span in tenant), _names(tenant) + assert not needs_model_span or _has_model_span(tenant), _names(tenant) return tenant 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_anthropic_empty_user_turn_wire.py b/tests/integration/providers/test_anthropic_empty_user_turn_wire.py new file mode 100644 index 00000000000..654a88c60bb --- /dev/null +++ b/tests/integration/providers/test_anthropic_empty_user_turn_wire.py @@ -0,0 +1,216 @@ +import json +from collections.abc import Iterator, Mapping, Sequence +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 Gateway, gateway_from_environment, object_value +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import ( + anthropic_deployment, + anthropic_peer, + final_text, + new_marker, +) +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MODEL: Final = "empty-user-turn" +_PLACEHOLDER: Final = "[System: Empty message content sanitised to satisfy protocol]" +_REPLY: Final = "Noted." +_EMPTY_ROWS: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "list": MappingProxyType({"role": "user", "content": []}), + "null": MappingProxyType({"role": "user", "content": None}), + "missing": MappingProxyType({"role": "user"}), + "empty-string": MappingProxyType({"role": "user", "content": ""}), + } +) +_CHAT_EMPTY_KINDS: Final = ("list", "null", "missing") + +Turn = tuple[str, tuple[str, ...]] + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + provider: Wire + + +def _empty_row(kind: str) -> dict[str, JsonValue]: + return dict(_EMPTY_ROWS[kind]) + + +def _guardrail_peer(request: Request) -> Reply: + kind: Final = request.target.split("/")[1] + rows: Final = _JSON_OBJECT.validate_json(request.body)["structured_messages"] + assert isinstance(rows, list), rows + verdict: Final = {"action": "GUARDRAIL_INTERVENED", "structured_messages": [*rows, _empty_row(kind)]} + return Reply(body=json.dumps(verdict).encode()) + + +def _guardrail(kind: str, api_base: str) -> dict[str, JsonValue]: + return { + "guardrail_name": f"append-{kind}", + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": f"{api_base}/{kind}", + "api_key": "synthetic-guardrail-key", + }, + } + + +def _config(directory: Path, provider_url: str, guardrail_url: str) -> Path: + stock: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final[dict[str, JsonValue]] = { + **stock, + "model_list": [anthropic_deployment(_MODEL, provider_url)], + "router_settings": {**object_value(stock["router_settings"]), "num_retries": 0}, + "guardrails": [_guardrail(kind, guardrail_url) for kind in _EMPTY_ROWS], + } + path: Final = directory / "empty-user-turn.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("empty-user-turn") + with ( + wire_server(anthropic_peer) as provider, + wire_server(_guardrail_peer) as guardrail, + gateway_from_environment() as gateway, + ): + config: Final = _config(directory, provider.url, guardrail.url) + with owned_proxy(gateway, directory, {}, config=config, workers=2) as owned: + yield _Rig(owned, provider) + + +def _block(block: JsonValue) -> str: + part: Final = object_value(block) + return str(part["text"]) if part.get("type") == "text" else str(part.get("type")) + + +def _turn(message: JsonValue) -> Turn: + row: Final = object_value(message) + content: Final = row["content"] + if isinstance(content, str): + return (str(row["role"]), (content,)) + assert isinstance(content, list), row + return (str(row["role"]), tuple(_block(block) for block in content)) + + +def _send(rig: _Rig, path: str, body: Mapping[str, JsonValue]) -> None: + response: Final = rig.gateway.client.post( + path, json=dict(body), headers={"Authorization": f"Bearer {rig.gateway.key}"} + ) + assert response.status_code == 200, response.text + + +def _provider_turns(rig: _Rig, marker: str) -> tuple[Turn, ...]: + needle: Final = marker.encode() + requests: Final = [request for request in rig.provider.drain() if needle in request.body] + assert len(requests) == 1, requests + messages: Final = _JSON_OBJECT.validate_json(requests[0].body)["messages"] + assert isinstance(messages, list), messages + return tuple(_turn(message) for message in messages) + + +def _ask_row(ask: str) -> dict[str, JsonValue]: + return {"role": "user", "content": ask} + + +def _chat(messages: Sequence[JsonValue]) -> dict[str, JsonValue]: + return {"model": _MODEL, "max_tokens": 64, "messages": list(messages)} + + +@pytest.mark.parametrize("kind", _CHAT_EMPTY_KINDS) +def test_chat_trailing_empty_user_message_keeps_a_user_turn(rig: _Rig, kind: str) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + _send( + rig, + "/v1/chat/completions", + _chat([_ask_row(ask), {"role": "assistant", "content": _REPLY}, _empty_row(kind)]), + ) + assert _provider_turns(rig, marker) == (("user", (ask,)), ("assistant", (_REPLY,)), ("user", (_PLACEHOLDER,))) + + +@pytest.mark.parametrize("kind", _CHAT_EMPTY_KINDS) +def test_chat_empty_user_message_beside_text_sends_only_the_text(rig: _Rig, kind: str) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + _send(rig, "/v1/chat/completions", _chat([_ask_row(ask), _empty_row(kind)])) + assert _provider_turns(rig, marker) == (("user", (ask,)),) + + +def _system_then_empty_turns(rig: _Rig, kind: str) -> tuple[Turn, ...]: + marker: Final = new_marker() + ask: Final = final_text(marker) + _send( + rig, + "/v1/chat/completions", + _chat( + [ + _ask_row(ask), + {"role": "assistant", "content": _REPLY}, + {"role": "system", "content": "Answer in French."}, + _empty_row(kind), + ] + ), + ) + turns: Final = _provider_turns(rig, marker) + assert turns[0] == ("user", (ask,)), turns + return turns[1:] + + +@pytest.mark.parametrize("kind", _CHAT_EMPTY_KINDS) +def test_chat_system_message_before_an_empty_user_message_is_placed_like_an_empty_string(rig: _Rig, kind: str) -> None: + assert _system_then_empty_turns(rig, kind) == _system_then_empty_turns(rig, "empty-string") + + +def test_responses_trailing_empty_input_message_keeps_a_user_turn(rig: _Rig) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + body: Final[dict[str, JsonValue]] = { + "model": _MODEL, + "max_output_tokens": 64, + "input": [ + _ask_row(ask), + {"role": "assistant", "content": _REPLY}, + {"role": "user", "content": []}, + ], + } + _send(rig, "/v1/responses", body) + assert _provider_turns(rig, marker) == (("user", (ask,)), ("assistant", (_REPLY,)), ("user", (_PLACEHOLDER,))) + + +@pytest.mark.parametrize("kind", tuple(_EMPTY_ROWS)) +def test_messages_guardrail_empty_user_row_is_sent_like_an_empty_string(rig: _Rig, kind: str) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + body: Final[dict[str, JsonValue]] = { + "model": _MODEL, + "max_tokens": 64, + "messages": [_ask_row(ask)], + "guardrails": [f"append-{kind}"], + } + _send(rig, "/v1/messages", body) + assert _provider_turns(rig, marker) == (("user", (ask,)), ("user", (_PLACEHOLDER,))) + + +def test_a_long_run_of_empty_user_messages_is_sent_as_one_text_turn(rig: _Rig) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + _send( + rig, "/v1/chat/completions", _chat([*({"role": "user", "content": []} for _ in range(10_000)), _ask_row(ask)]) + ) + assert _provider_turns(rig, marker) == (("user", (ask,)),) 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_passthrough_aws_errors.py b/tests/integration/providers/test_bedrock_passthrough_aws_errors.py new file mode 100644 index 00000000000..bb87b15b363 --- /dev/null +++ b/tests/integration/providers/test_bedrock_passthrough_aws_errors.py @@ -0,0 +1,443 @@ +"""Errors on the Bedrock passthrough reach the AWS SDK with their message (LIT-9617). + +boto3 reads an error's code from the ``x-amzn-errortype`` header and its message from the body's ``message``. A +guardrail block on ``/bedrock/model//...`` that happens before any frame was sent reaches the caller as a +``ValidationException`` carrying the guardrail's message, and an error Bedrock itself returned, such as a 429 +``ThrottlingException``, reaches the caller with Bedrock's status, ``x-amzn-errortype`` and body.""" + +from __future__ import annotations + +import json +import os +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import boto3 +import httpx +import pytest +import yaml +from botocore.config import Config +from botocore.exceptions import ClientError +from integration._support.budget import case_timeout_seconds +from integration._support.client import Gateway, JsonValue +from integration._support.process import owned_proxy_process +from integration._support.upstream import aws_event_stream_frame +from integration._support.wire import Reply, Request, Wire, wire_server + +MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" +ROUTER_MODEL: Final = "aws-errors" +EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +POST_CALL: Final = "aws-errors-post-call" +PRE_CALL: Final = "aws-errors-pre-call" +BEDROCK: Final = "aws-errors-bedrock" +VENDOR_DOWN: Final = "aws-errors-vendor-down" +CATO: Final = "aws-errors-cato" +UNGUARDED: Final = "unguarded" +OVER_BUDGET: Final = "over-budget" +ONE_RPM: Final = "one-rpm" +BLOCKED_WORD: Final = "FORBIDDEN" +BLOCK_REASON: Final = "synthetic policy says no" +BEDROCK_BLOCK_MESSAGE: Final = "Violated guardrail policy" +THROTTLE_MESSAGE: Final = "Too many requests, please wait before trying again." +THROTTLE_TYPE: Final = "ThrottlingException:http://internal.amazon.com/coral/com.amazon.bedrock/" +INVALID_MESSAGE: Final = "The provided model identifier is invalid." +INVALID_TYPE: Final = "ValidationException:http://internal.amazon.com/coral/com.amazon.bedrock/" +ANTHROPIC_THROTTLE: Final = b'{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}' + + +@dataclass(frozen=True, slots=True) +class UpstreamError: + status: int + errortype: str + message: str + + @property + def body(self) -> bytes: + return json.dumps({"message": self.message}).encode() + + @property + def code(self) -> str: + return self.errortype.split(":", 1)[0] + + +THROTTLED: Final = UpstreamError(429, THROTTLE_TYPE, THROTTLE_MESSAGE) +INVALID: Final = UpstreamError(400, INVALID_TYPE, INVALID_MESSAGE) +UPSTREAM_ERRORS: Final = MappingProxyType({"throttled": THROTTLED, "invalid": INVALID}) +ROUTES: Final = (pytest.param(ROUTER_MODEL, id="router"), pytest.param(MODEL_ID, id="direct")) + + +def upstream_marker(kind: str) -> str: + """A prompt the scripted Bedrock answers with the ``UPSTREAM_ERRORS`` entry named ``kind``.""" + return f"upstream-{kind}-{uuid.uuid4().hex}" + + +def upstream_error_of(marker: str) -> UpstreamError | None: + prefix, _, rest = marker.partition("-") + return UPSTREAM_ERRORS.get(rest.partition("-")[0]) if prefix == "upstream" else None + + +def event_frame(event_type: str, payload: Mapping[str, JsonValue]) -> bytes: + return aws_event_stream_frame( + {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, + json.dumps(payload, separators=(",", ":")).encode(), + ) + + +def converse_frames(text: str) -> tuple[bytes, ...]: + return ( + event_frame("messageStart", {"role": "assistant"}), + event_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": text}}), + event_frame("contentBlockStop", {"contentBlockIndex": 0}), + event_frame("messageStop", {"stopReason": "end_turn"}), + event_frame("metadata", {"usage": {"inputTokens": 9, "outputTokens": 4, "totalTokens": 13}, "metrics": {}}), + ) + + +def converse_reply(text: str) -> bytes: + return json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": text}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 9, "outputTokens": 4, "totalTokens": 13}, + "metrics": {"latencyMs": 1}, + } + ).encode() + + +def invoke_body(marker: str) -> str: + return json.dumps( + { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 16, + "messages": [{"role": "user", "content": [{"type": "text", "text": marker}]}], + } + ) + + +def marker_of(request: Request) -> str: + return str(json.loads(request.body)["messages"][0]["content"][0]["text"]) + + +def bedrock_texts(request: Request) -> tuple[str, ...]: + return tuple(str(item["text"]["text"]) for item in json.loads(request.body)["content"]) + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + provider: Wire + vendor: Wire + keys: Mapping[str, str] + + def provider_calls(self, marker: str) -> int: + return sum(1 for request in self.provider.drain() if marker.encode() in request.body) + + def client(self, guardrail: str, attempts: int = 1): # noqa: ANN201 # boto3 clients are untyped + client: Final = boto3.client( + "bedrock-runtime", + region_name="us-east-1", + endpoint_url=str(self.gateway.client.base_url).rstrip("/") + "/bedrock", + aws_access_key_id="AKIASYNTHETICCALLER", + aws_secret_access_key="synthetic-caller-secret", + config=Config( + read_timeout=case_timeout_seconds(), + retries={"total_max_attempts": attempts, "mode": "standard"}, + ), + ) + key: Final = self.keys[guardrail] + + def authenticate(request: object, **_: object) -> None: + getattr(request, "headers")["x-litellm-api-key"] = key + + client.meta.events.register("before-send.bedrock-runtime.*", authenticate) + return client + + def call( + self, guardrail: str, action: str, marker: str, attempts: int = 1, model: str = ROUTER_MODEL + ) -> ClientError | None: + client: Final = self.client(guardrail, attempts) + messages: Final = [{"role": "user", "content": [{"text": marker}]}] + try: + if action == "converse-stream": + for _ in client.converse_stream(modelId=model, messages=messages)["stream"]: + pass + elif action == "invoke": + client.invoke_model(modelId=model, body=invoke_body(marker)) + else: + client.converse(modelId=model, messages=messages) + except ClientError as error: + return error + return None + + def raw(self, guardrail: str, action: str, marker: str) -> httpx.Response: + return self.gateway.client.post( + f"/bedrock/model/{ROUTER_MODEL}/{action}", + json={"messages": [{"role": "user", "content": [{"text": marker}]}]}, + headers={"x-litellm-api-key": self.keys[guardrail]}, + timeout=case_timeout_seconds(), + ) + + +def _generic_guardrail(name: str, mode: str, vendor: Wire, **settings: JsonValue) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": False, + "api_base": f"{vendor.url}/{name}", + "api_key": "synthetic-guardrail-key", + **settings, + }, + } + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + def provider(request: Request) -> Reply: + if request.target == "/v1/messages": + return Reply(status=429, body=ANTHROPIC_THROTTLE) + assert request.target in tuple( + f"/model/{MODEL_ID}/{action}" for action in ("converse-stream", "converse", "invoke") + ), request.target + marker: Final = marker_of(request) + error: Final = upstream_error_of(marker) + if error is not None: + return Reply(status=error.status, body=error.body, headers={"x-amzn-errortype": error.errortype}) + text: Final = f"The answer to {marker} is {BLOCKED_WORD}" + if request.target.endswith("/converse-stream"): + return Reply(content_type=EVENT_STREAM, chunks=converse_frames(text)) + return Reply(body=converse_reply(text)) + + def vendor(request: Request) -> Reply: + if request.target.startswith(f"/{BEDROCK}/"): + if not any(BLOCKED_WORD in text for text in bedrock_texts(request)): + return Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()) + return Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "Sorry, the model cannot answer this."}], + "assessments": [ + {"wordPolicy": {"customWords": [{"match": BLOCKED_WORD, "action": "BLOCKED"}]}} + ], + } + ).encode() + ) + if request.target.startswith(f"/{VENDOR_DOWN}"): + return Reply(status=500, body=b'{"detail":"vendor is down"}') + if request.target.startswith(f"/{CATO}/"): + return Reply( + status=THROTTLED.status, body=THROTTLED.body, headers={"x-amzn-errortype": THROTTLED.errortype} + ) + texts: Final = [str(text) for text in json.loads(request.body)["texts"]] + if any(BLOCKED_WORD in text for text in texts): + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": BLOCK_REASON}).encode()) + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + with httpx.Client(base_url=os.environ["INTEGRATION_PROXY_URL"], timeout=30, trust_env=False) as root_client: + root: Final = Gateway( + root_client, + os.environ.get("INTEGRATION_MASTER_KEY", "sk-integration-master"), + os.environ["INTEGRATION_UPSTREAM_URL"], + ) + directory: Final = tmp_path_factory.mktemp("bedrock-aws-errors") + with wire_server(provider) as upstream, wire_server(vendor) as sink: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["guardrails"] = [ + _generic_guardrail(POST_CALL, "post_call", sink, streaming_buffer_until_moderated=True), + _generic_guardrail(PRE_CALL, "pre_call", sink), + _generic_guardrail(VENDOR_DOWN, "pre_call", sink), + { + "guardrail_name": CATO, + "litellm_params": { + "guardrail": "cato_networks", + "mode": "pre_call", + "default_on": False, + "api_base": f"{sink.url}/{CATO}", + "api_key": "synthetic-cato-key", + }, + }, + { + "guardrail_name": BEDROCK, + "litellm_params": { + "guardrail": "bedrock", + "mode": "post_call", + "default_on": False, + "guardrailIdentifier": "aws-errors", + "guardrailVersion": "1", + "aws_access_key_id": "AKIASYNTHETICGUARD", + "aws_secret_access_key": "synthetic-guardrail-secret", + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": f"{sink.url}/{BEDROCK}", + }, + }, + ] + config["model_list"] = [ + { + "model_name": name, + "litellm_params": { + "model": f"bedrock/{model}", + "api_base": upstream.url, + "aws_access_key_id": "AKIASYNTHETICPROVIDER", + "aws_secret_access_key": "synthetic-provider-secret", + "aws_region_name": "us-east-1", + }, + } + for name, model in ((ROUTER_MODEL, MODEL_ID), ("anthropic.*", "anthropic.*")) + ] + config["router_settings"] = {**config.get("router_settings", {}), "num_retries": 0} + config_path: Final = directory / "bedrock-aws-errors.yaml" + config_path.write_text(yaml.safe_dump(config)) + overrides: Final = {"ANTHROPIC_API_BASE": upstream.url, "ANTHROPIC_API_KEY": "sk-ant-synthetic"} + with owned_proxy_process(root, directory, overrides, config=config_path, workers=2) as owned: + with owned.gateway.scenario() as scenario: + yield Rig( + owned.gateway, + upstream, + sink, + { + UNGUARDED: scenario.key(), + OVER_BUDGET: scenario.key(max_budget=0), + ONE_RPM: scenario.key(rpm_limit=1), + **{ + name: scenario.key(guardrails=[name]) + for name in (POST_CALL, PRE_CALL, BEDROCK, VENDOR_DOWN, CATO) + }, + }, + ) + + +BLOCKS: Final = ( + pytest.param(POST_CALL, "converse-stream", BLOCK_REASON, 1, id="post-call-held-converse-stream"), + pytest.param(POST_CALL, "converse", BLOCK_REASON, 1, id="post-call-converse"), + pytest.param(PRE_CALL, "converse-stream", BLOCK_REASON, 0, id="pre-call-converse-stream"), + pytest.param(PRE_CALL, "converse", BLOCK_REASON, 0, id="pre-call-converse"), + pytest.param(BEDROCK, "converse-stream", BEDROCK_BLOCK_MESSAGE, 1, id="bedrock-guardrail-converse-stream"), + pytest.param(BEDROCK, "converse", BEDROCK_BLOCK_MESSAGE, 1, id="bedrock-guardrail-converse"), + pytest.param(PRE_CALL, "invoke", BLOCK_REASON, 0, id="pre-call-invoke"), +) + + +@pytest.mark.parametrize("model", ROUTES) +@pytest.mark.parametrize(("guardrail", "action", "message", "provider_calls"), BLOCKS) +def test_guardrail_block_reaches_boto3_as_a_validation_exception_with_the_guardrail_message( + rig: Rig, guardrail: str, action: str, message: str, provider_calls: int, model: str +) -> None: + marker: Final = f"{BLOCKED_WORD}-{uuid.uuid4().hex}" + error: Final = rig.call(guardrail, action, marker, model=model) + assert error is not None, f"{guardrail} {action} was not blocked" + assert error.response["ResponseMetadata"]["HTTPStatusCode"] == 400, error.response + assert (error.response["Error"]["Code"], error.response["Error"]["Message"]) == ( + "ValidationException", + message, + ), error.response + assert rig.provider_calls(marker) == provider_calls + + +@pytest.mark.parametrize("action", ("converse-stream", "converse")) +def test_guardrail_block_body_adds_the_aws_message_and_keeps_the_litellm_error(rig: Rig, action: str) -> None: + marker: Final = f"{BLOCKED_WORD}-{uuid.uuid4().hex}" + response: Final = rig.raw(POST_CALL, action, marker) + assert response.status_code == 400, response.text + assert response.headers.get("x-amzn-errortype") == "ValidationException", (dict(response.headers), response.text) + assert response.headers.get("content-type", "").startswith("application/json"), dict(response.headers) + body: Final = response.json() + assert body["message"] == BLOCK_REASON, body + assert body["error"]["message"] == BLOCK_REASON, body + assert response.headers.get("x-litellm-call-id"), dict(response.headers) + + +@pytest.mark.parametrize("kind", tuple(UPSTREAM_ERRORS)) +@pytest.mark.parametrize("action", ("converse-stream", "converse")) +@pytest.mark.parametrize("guardrail", (UNGUARDED, POST_CALL)) +def test_bedrock_error_reaches_the_caller_as_bedrock_sent_it(rig: Rig, kind: str, action: str, guardrail: str) -> None: + upstream: Final = UPSTREAM_ERRORS[kind] + marker: Final = upstream_marker(kind) + response: Final = rig.raw(guardrail, action, marker) + assert response.status_code == upstream.status, response.text + assert response.headers.get("x-amzn-errortype") == upstream.errortype, (dict(response.headers), response.text) + assert response.content == upstream.body, response.text + assert response.headers.get("content-type") == "application/json", dict(response.headers) + assert response.headers.get("x-litellm-call-id"), dict(response.headers) + assert rig.provider_calls(marker) == 1 + + +@pytest.mark.parametrize("model", ROUTES) +@pytest.mark.parametrize("action", ("converse-stream", "converse", "invoke")) +@pytest.mark.parametrize("guardrail", (UNGUARDED, POST_CALL)) +def test_boto3_sees_bedrock_throttling_and_retries_it_as_against_bedrock( + rig: Rig, action: str, guardrail: str, model: str +) -> None: + marker: Final = upstream_marker("throttled") + error: Final = rig.call(guardrail, action, marker, attempts=3, model=model) + assert error is not None, f"{action} was not throttled" + assert (error.response["Error"]["Code"], error.response["Error"]["Message"]) == ( + THROTTLED.code, + THROTTLE_MESSAGE, + ), error.response + assert error.response["ResponseMetadata"]["HTTPStatusCode"] == 429, error.response + assert error.response["ResponseMetadata"]["RetryAttempts"] == 2, error.response + assert rig.provider_calls(marker) == 3 + + +def converse_with_key(rig: Rig, model: str, key: str) -> httpx.Response: + return rig.gateway.client.post( + f"/bedrock/model/{model}/converse", + json={"messages": [{"role": "user", "content": [{"text": f"plain-{uuid.uuid4().hex}"}]}]}, + headers={"x-litellm-api-key": key}, + timeout=case_timeout_seconds(), + ) + + +def assert_litellm_error(response: httpx.Response, status: int) -> None: + assert response.status_code == status, response.text + assert "x-amzn-errortype" not in response.headers, dict(response.headers) + assert set(response.json()) == {"error"}, response.text + + +@pytest.mark.parametrize( + ("guardrail", "model", "status"), + ( + pytest.param(VENDOR_DOWN, ROUTER_MODEL, 500, id="guardrail-vendor-down"), + pytest.param(UNGUARDED, "not-a-configured-model", 400, id="unknown-model"), + pytest.param(OVER_BUDGET, ROUTER_MODEL, 422, id="over-budget"), + ), +) +def test_errors_that_are_not_guardrail_blocks_keep_the_litellm_error( + rig: Rig, guardrail: str, model: str, status: int +) -> None: + assert_litellm_error(converse_with_key(rig, model, rig.keys[guardrail]), status) + + +def test_guardrail_service_error_shaped_like_aws_keeps_the_litellm_error(rig: Rig) -> None: + response: Final = converse_with_key(rig, ROUTER_MODEL, rig.keys[CATO]) + assert response.status_code == THROTTLED.status, response.text + assert "x-amzn-errortype" not in response.headers, dict(response.headers) + assert response.content != THROTTLED.body, response.text + + +def test_unknown_key_keeps_the_litellm_error(rig: Rig) -> None: + assert_litellm_error(converse_with_key(rig, ROUTER_MODEL, f"sk-unknown-{uuid.uuid4().hex}"), 401) + + +def test_litellm_rate_limit_keeps_the_litellm_error(rig: Rig) -> None: + admitted: Final = converse_with_key(rig, ROUTER_MODEL, rig.keys[ONE_RPM]) + assert admitted.status_code == 200, admitted.text + assert_litellm_error(converse_with_key(rig, ROUTER_MODEL, rig.keys[ONE_RPM]), 429) + + +def test_other_provider_passthrough_errors_carry_no_aws_error_type(rig: Rig) -> None: + response: Final = rig.gateway.client.post( + "/anthropic/v1/messages", + json={"model": "claude-sonnet-5", "max_tokens": 8, "messages": [{"role": "user", "content": "hi"}]}, + headers={"x-litellm-api-key": rig.keys[UNGUARDED]}, + timeout=case_timeout_seconds(), + ) + assert response.status_code == 429, response.text + assert "x-amzn-errortype" not in response.headers, dict(response.headers) + assert response.content == ANTHROPIC_THROTTLE, response.text diff --git a/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py b/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py index 0860d79459e..54923fe0c97 100644 --- a/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py +++ b/tests/integration/providers/test_bedrock_runtime_chat_completions_sad_wire.py @@ -27,6 +27,7 @@ BAD_KEY: Final = "sk-synthetic-bad-key" NATIVE_TARGET: Final = "/openai/v1/chat/completions" CONVERSE_TARGET: Final = f"/model/{GPT}/converse" LONG_VERSION_GPT: Final = "openai.gpt-" + "1" * 30000 +LONG_UNVERSIONED_GPT: Final = "openai.gpt-x" + "1" * 30000 PNG_DATA_URL: Final = ( "data:image/png;base64," "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==" @@ -258,11 +259,12 @@ def test_thirty_thousand_digit_version_is_classified_quickly_and_served_by_conve marker: Final = uuid.uuid4().hex control_marker: Final = uuid.uuid4().hex with wire_server(respond) as wire, gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=1) as pool: - control_model: Final = _deployment(scenario, wire) + control_model: Final = _deployment(scenario, wire, model=f"bedrock/{LONG_UNVERSIONED_GPT}") model: Final = _deployment(scenario, wire, model=f"bedrock/{LONG_VERSION_GPT}") control: Final = timed(lambda: _chat(gateway, control_model, control_marker)) assert _content(control.value) == answer(control_marker), control.value.text - _only_request(wire, control_marker) + control_request: Final = _only_request(wire, control_marker) + assert target_of(control_request) == f"/model/{LONG_UNVERSIONED_GPT}/converse", control_request.target liveliness: Final = pool.submit(_timed_liveliness, gateway) served: Final = timed(lambda: _chat(gateway, model, marker)) response: Final = served.value diff --git a/tests/integration/providers/test_count_tokens_error_shapes_wire.py b/tests/integration/providers/test_count_tokens_error_shapes_wire.py new file mode 100644 index 00000000000..200f5e81673 --- /dev/null +++ b/tests/integration/providers/test_count_tokens_error_shapes_wire.py @@ -0,0 +1,205 @@ +import json +from collections.abc import Iterator +from typing import Final + +import anthropic +import openai +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this synthetic request."}] +_PROVIDER_KEY_VARIABLES: Final = ( + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "AZURE_AI_API_KEY", + "AZURE_API_KEY", + "AZURE_OPENAI_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_TENANT_ID", + "OPENAI_API_KEY", +) + + +@pytest.fixture(scope="module") +def keyless_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("count-tokens-error-shapes"), + {}, + remove_environment=_PROVIDER_KEY_VARIABLES, + workers=1, + ) as owned, + ): + yield owned + + +@pytest.fixture(scope="module") +def azure_environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("count-tokens-azure-environment-key"), + {"AZURE_API_KEY": "synthetic-azure-env-key"}, + remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name != "AZURE_API_KEY"), + workers=1, + ) as owned, + ): + yield owned + + +def _anthropic(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client + ) + + +def _never_called(request: Request) -> Reply: + raise AssertionError(f"count reached the provider without a key: {request.target}") + + +@pytest.mark.parametrize( + ("deployment", "missing"), + ( + ({"model": "anthropic/claude-sonnet-4-6"}, "No Anthropic API key is set for model "), + ({"model": "azure_ai/claude-sonnet-4-6"}, "No Azure AI API key or Entra ID credential is set for model "), + ), + ids=("anthropic", "azure_ai_anthropic"), +) +def test_anthropic_count_without_a_provider_credential_is_an_anthropic_401( + keyless_proxy: OwnedProxy, deployment: dict[str, JsonValue], missing: str +) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(**deployment, api_base=wire.url, api_key=None, num_retries=0) + with pytest.raises(anthropic.AuthenticationError) as raised: + _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert raised.value.body == { + "type": "error", + "error": {"type": "authentication_error", "message": raised.value.body["error"]["message"]}, + } + assert missing in raised.value.message, raised.value.message + assert wire.drain() == () + + +def test_azure_ai_anthropic_count_sends_the_azure_environment_key_as_x_api_key( + azure_environment_key_proxy: OwnedProxy, +) -> None: + gateway: Final = azure_environment_key_proxy.gateway + + def count(request: Request) -> Reply: + assert request.headers.get("x-api-key") == "synthetic-azure-env-key", request.headers + return Reply(body=json.dumps({"input_tokens": 58}).encode()) + + with wire_server(count) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/claude-sonnet-4-6", + api_base=wire.url, + api_key=None, + num_retries=0, + ) + result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert result.input_tokens == 58 + + +def test_azure_ai_anthropic_count_authenticates_with_an_entra_token(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + + def count(request: Request) -> Reply: + assert request.headers.get("authorization") == "Bearer synthetic-entra-token", request.headers + assert "x-api-key" not in request.headers, request.headers + return Reply(body=json.dumps({"input_tokens": 61}).encode()) + + with wire_server(count) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/claude-sonnet-4-6", + api_base=wire.url, + api_key=None, + azure_ad_token="synthetic-entra-token", + num_retries=0, + ) + result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert result.input_tokens == 61 + assert len(wire.drain()) == 1 + + +def test_anthropic_count_request_errors_use_the_anthropic_envelope(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + upstream_error: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted reject"}} + with ( + wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-6", api_base=wire.url, api_key="synthetic-key", num_retries=0 + ) + missing_messages: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model}) + assert missing_messages.status_code == 400, missing_messages.text + assert _JSON_OBJECT.validate_json(missing_messages.content) == { + "type": "error", + "error": {"type": "invalid_request_error", "message": "messages parameter is required"}, + } + + rejected: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _MESSAGES} + ) + assert rejected.status_code == 400, rejected.text + body: Final = _JSON_OBJECT.validate_json(rejected.content) + assert set(body) <= {"type", "error", "request_id"}, body + assert body["error"] == upstream_error["error"], body + assert len(wire.drain()) == 1 + + +def test_openai_input_tokens_without_a_provider_key_is_an_openai_401(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key=None, num_retries=0 + ) + client: Final = openai.OpenAI( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client + ) + with pytest.raises(openai.AuthenticationError) as raised: + client.responses.input_tokens.count(model=model, input="Count this synthetic request.") + assert f"No OpenAI API key is set for model {model}" in raised.value.message + assert wire.drain() == () + + +def test_google_count_tokens_errors_use_the_google_envelope(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + upstream_error: Final = { + "error": {"code": 400, "message": "API key not valid. Please pass a valid API key.", "status": "INVALID_ARGUMENT"} + } + with ( + wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="gemini/gemini-2.5-flash", api_base=wire.url, api_key="synthetic-key", num_retries=0 + ) + rejected: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + {"contents": [{"role": "user", "parts": [{"text": "Count this synthetic request."}]}]}, + ) + assert rejected.status_code == 400, rejected.text + error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"] + assert isinstance(error, dict) and error["code"] == 400 and error["status"] == "INVALID_ARGUMENT", error + assert "API key not valid" in str(error["message"]), error + assert len(wire.drain()) == 1 + + malformed: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", {"contents": "not-a-list"}) + assert malformed.status_code == 400, malformed.text + assert _JSON_OBJECT.validate_json(malformed.content) == { + "error": {"code": 400, "message": "Invalid Google countTokens request body", "status": "INVALID_ARGUMENT"} + } diff --git a/tests/integration/providers/test_custom_prompt_template_chaos.py b/tests/integration/providers/test_custom_prompt_template_chaos.py new file mode 100644 index 00000000000..c0da612af23 --- /dev/null +++ b/tests/integration/providers/test_custom_prompt_template_chaos.py @@ -0,0 +1,542 @@ +import asyncio +import json +import re +import threading +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "llama3-custom-template-chaos" +_API_KEY: Final = "synthetic-ollama-key" +_CONFIG_MODEL: Final = "ollama-custom-template-chaos" +_SYSTEM: Final = "Answer in one word." +_ROLES: Final[dict[str, JsonValue]] = { + "system": {"pre_message": "", "post_message": ""}, + "user": {"pre_message": "", "post_message": ""}, +} +_REJECTION: Final = "initial_prompt_value: Input should be a valid string" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_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 + +Endpoint = Literal["chat", "messages", "responses"] + + +@dataclass(frozen=True, slots=True) +class _Call: + endpoint: Endpoint + stream: bool + marker: str + rejected: bool = False + + +@dataclass(frozen=True, slots=True) +class _Served: + call: _Call + status: int + text: str + + +def _question(marker: str) -> str: + return f"What colour is the sky? marker-{marker}" + + +def _prompt(marker: str) -> str: + return f"{_SYSTEM}{_question(marker)}" + + +def _path(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "/v1/chat/completions" + case "messages": + return "/v1/messages" + case "responses": + return "/v1/responses" + + +def _body(model: str, call: _Call) -> dict[str, JsonValue]: + common: Final[dict[str, JsonValue]] = { + "model": model, + "stream": call.stream, + "num_retries": 0, + "cache": {"no-cache": True}, + **({"initial_prompt_value": 5} if call.rejected else {}), + } + match call.endpoint: + case "chat": + return { + **common, + "messages": [ + {"role": "system", "content": _SYSTEM}, + {"role": "user", "content": _question(call.marker)}, + ], + } + case "messages": + return { + **common, + "max_tokens": 64, + "system": _SYSTEM, + "messages": [{"role": "user", "content": _question(call.marker)}], + } + case "responses": + return {**common, "store": False, "instructions": _SYSTEM, "input": _question(call.marker)} + + +def _generate_reply(marker: str, stream: bool, drop_connection: bool = False) -> Reply: + done: Final = { + "model": _BACKEND, + "created_at": "2026-10-15T00:00:00Z", + "response": "", + "done": True, + "done_reason": "stop", + "prompt_eval_count": 30, + "eval_count": 12, + } + if not stream: + return Reply( + body=json.dumps({**done, "response": f"blue marker-{marker}"}).encode(), drop_connection=drop_connection + ) + frames: Final = ( + *( + {"model": _BACKEND, "created_at": "2026-10-15T00:00:00Z", "response": piece, "done": False} + for piece in ("blue ", f"marker-{marker}") + ), + done, + ) + return Reply( + content_type="application/x-ndjson", + chunks=tuple(json.dumps(frame).encode() + b"\n" for frame in frames), + drop_connection=drop_connection, + ) + + +def _is_generate(request: Request) -> bool: + return (request.method, request.target) == ("POST", "/api/generate") + + +def _marker_of(request: Request) -> str: + found: Final = _MARKER.search(request.body.decode()) + assert found is not None, request.body + return found.group(1) + + +def _echo(request: Request) -> Reply: + if not _is_generate(request): + return Reply(body=b"{}") + stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True + return _generate_reply(_marker_of(request), stream) + + +def _assert_each_prompt_rendered_once(received: tuple[Request, ...], markers: frozenset[str]) -> None: + generates: Final = tuple(request for request in received if _is_generate(request)) + assert sorted(_marker_of(request) for request in generates) == sorted(markers) + for request in generates: + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["prompt"] == _prompt(_marker_of(request)), body["prompt"] + assert "roles" not in body and "initial_prompt_value" not in body, sorted(body) + + +def _response_id(served: _Served) -> str | None: + if served.call.endpoint == "responses": + return None + if not served.call.stream: + identity: Final = _JSON_OBJECT.validate_json(served.text)["id"] + assert isinstance(identity, str) + return identity + for line in served.text.splitlines(): + if not line.startswith("data: ") or line == "data: [DONE]": + continue + payload: Final = _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + if served.call.endpoint == "chat": + first: Final = payload["id"] + assert isinstance(first, str) + return first + if payload.get("type") == "message_start": + message: Final = payload["message"] + assert isinstance(message, dict) and isinstance(message["id"], str) + return message["id"] + raise AssertionError(served.text) + + +def _spend_statuses(model: str, expected: int) -> MappingProxyType[str, str]: + rows: Final = eventually( + lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)), + lambda found: len(found) >= expected, + seconds=70, + ) + statuses: Final = MappingProxyType({str(row["request_id"]): str(row["status"]) for row in rows}) + assert len(statuses) == len(rows) == expected, rows + return statuses + + +def _successes(statuses: Mapping[str, str]) -> frozenset[str]: + return frozenset(identity for identity, status in statuses.items() if status == "success") + + +async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served: + async with client.stream( + "POST", + _path(call.endpoint), + json=_body(model, call), + headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"}, + ) as response: + raw: Final = await response.aread() + return _Served(call=call, status=response.status_code, text=raw.decode()) + + +async def _burst( + base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Served)) + + +def _calls(count: int, endpoints: tuple[Endpoint, ...], stream: Callable[[int], bool]) -> tuple[_Call, ...]: + return tuple( + _Call(endpoint=endpoints[index % len(endpoints)], stream=stream(index), marker=uuid.uuid4().hex) + for index in range(count) + ) + + +def _assert_answered_with_its_own_marker(served: _Served) -> None: + assert served.status == 200, served.text + assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text + + +def _assert_rejected_by_name(served: _Served) -> None: + assert served.status == 500, served.text + assert _REJECTION in served.text and "marker-" not in served.text, served.text + + +async def test_concurrent_burst_renders_each_template_and_rejects_only_the_bad_requests(gateway: Gateway) -> None: + calls: Final = tuple( + _Call(endpoint=call.endpoint, stream=call.stream, marker=call.marker, rejected=index % 5 == 4) + for index, call in enumerate(_calls(30, ("chat", "messages", "responses"), lambda index: index % 2 == 0)) + ) + accepted: Final = tuple(call for call in calls if not call.rejected) + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY, roles=_ROLES) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 30 + for item in served: + if item.call.rejected: + _assert_rejected_by_name(item) + else: + _assert_answered_with_its_own_marker(item) + _assert_each_prompt_rendered_once(wire.drain(), frozenset(call.marker for call in accepted)) + answered: Final = tuple(item for item in served if not item.call.rejected) + known: Final = frozenset(identity for identity in map(_response_id, answered) if identity is not None) + assert len(known) == 16, known + statuses: Final = _spend_statuses(model, 30) + assert len(_successes(statuses)) == len(accepted), statuses + assert known <= _successes(statuses) + + +async def test_dropped_ollama_connections_fail_their_callers_and_the_rest_keep_their_template( + gateway: Gateway, +) -> None: + calls: Final = _calls(12, ("chat",), lambda index: index % 2 == 1) + dropped: Final = frozenset(call.marker for index, call in enumerate(calls) if index % 3 == 0) + + def respond(request: Request) -> Reply: + if not _is_generate(request): + return Reply(body=b"{}") + marker: Final = _marker_of(request) + stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True + return _generate_reply(marker, stream, drop_connection=marker in dropped) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY, roles=_ROLES) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 12 + for item in served: + if item.call.marker in dropped: + assert item.status == 500, item.text + assert "APIConnectionError" in item.text and "marker-" not in item.text, item.text + else: + _assert_answered_with_its_own_marker(item) + recovery: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) + (recovered,) = await _burst(str(gateway.client.base_url), gateway.key, model, (recovery,)) + _assert_answered_with_its_own_marker(recovered) + _assert_each_prompt_rendered_once(wire.drain(), frozenset(call.marker for call in (*calls, recovery))) + answered: Final = tuple(item for item in served if item.call.marker not in dropped) + survivors: Final = frozenset( + identity for identity in map(_response_id, (*answered, recovered)) if identity is not None + ) + assert len(survivors) == 9, survivors + statuses: Final = _spend_statuses(model, 13) + assert _successes(statuses) == survivors, statuses + assert sum(status == "failure" for status in statuses.values()) == len(dropped), statuses + + +def _owned_config( + wires: tuple[Wire, ...], + tmp_path: Path, + template: Mapping[str, JsonValue], + litellm_settings: Mapping[str, JsonValue], + router_settings: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + deployments: Final = [ + { + "model_name": _CONFIG_MODEL, + "litellm_params": { + "model": f"ollama/{_BACKEND}", + "api_base": wire.url, + "api_key": _API_KEY, + "order": order, + **template, + }, + } + for order, wire in enumerate(wires, start=1) + ] + config: Final = { + **base, + "model_list": deployments, + "litellm_settings": {**base.get("litellm_settings", {}), **litellm_settings}, + "router_settings": {**base.get("router_settings", {}), **router_settings}, + } + path: Final = tmp_path / "ollama-custom-template-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_rendering_templates( + gateway: Gateway, tmp_path: Path +) -> None: + calls: Final = _calls(20, ("chat",), lambda _: False) + release: Final = threading.Event() + held_markers: Final[SimpleQueue[str]] = SimpleQueue() + + def held(request: Request) -> Reply: + if not _is_generate(request): + return Reply(body=b"{}") + held_markers.put(_marker_of(request)) + assert release.wait(timeout=60), "The burst was never released" + return _echo(request) + + with wire_server(held) as wire: + path: Final = _owned_config((wire,), tmp_path, {"roles": _ROLES}, {}) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + burst: Final = asyncio.create_task( + _burst( + str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, calls, tolerate_transport_errors=True + ) + ) + await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) + follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) + (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) + _assert_answered_with_its_own_marker(answered) + _assert_each_prompt_rendered_once(wire.drain(), frozenset(call.marker for call in (*calls, follow_up))) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +async def test_rejected_templates_never_cool_down_the_preferred_deployment(gateway: Gateway, tmp_path: Path) -> None: + rejected: Final = tuple( + _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex, rejected=True) for _ in range(4) + ) + valid: Final = _calls(8, ("chat",), lambda _: False) + with wire_server(_echo) as preferred, wire_server(_echo) as standby: + path: Final = _owned_config( + (preferred, standby), tmp_path, {"roles": _ROLES}, {}, {"disable_cooldowns": False, "allowed_fails": 0} + ) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + base_url: Final = str(owned.gateway.client.base_url) + for served in await _burst(base_url, owned.gateway.key, _CONFIG_MODEL, rejected): + _assert_rejected_by_name(served) + for served in await _burst(base_url, owned.gateway.key, _CONFIG_MODEL, valid): + _assert_answered_with_its_own_marker(served) + _assert_each_prompt_rendered_once(preferred.drain(), frozenset(call.marker for call in valid)) + assert not any(_is_generate(request) for request in standby.drain()) + + +def _cached_chat(candidate: Gateway, marker: str, field: str, value: str) -> httpx.Response: + return candidate.request( + "POST", + "/v1/chat/completions", + { + "model": _CONFIG_MODEL, + "messages": [{"role": "system", "content": _SYSTEM}, {"role": "user", "content": _question(marker)}], + field: value, + }, + ) + + +def _served_from_cache(candidate: Gateway, field: str) -> bool: + marker: Final = uuid.uuid4().hex + first: Final = _cached_chat(candidate, marker, field, "") + assert first.status_code == 200, first.text + assert "x-litellm-cache-key" not in first.headers, sorted(first.headers) + eventually( + lambda: _cached_chat(candidate, marker, field, ""), + lambda found: "x-litellm-cache-key" in found.headers, + seconds=20, + ) + changed: Final = _cached_chat(candidate, marker, field, "") + assert changed.status_code == 200, changed.text + return "x-litellm-cache-key" in changed.headers + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_initial_and_final_prompt_values_key_the_provider_specific_cache_alike( + gateway: Gateway, tmp_path: Path +) -> None: + template: Final[dict[str, JsonValue]] = { + "roles": _ROLES, + "initial_prompt_value": "", + "final_prompt_value": "", + } + with wire_server(_echo) as wire: + path: Final = _owned_config( + (wire,), tmp_path, template, {"enable_caching_on_provider_specific_optional_params": True} + ) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + initial: Final = _served_from_cache(owned.gateway, "initial_prompt_value") + final: Final = _served_from_cache(owned.gateway, "final_prompt_value") + assert initial == final, {"initial_prompt_value": initial, "final_prompt_value": final} + + +def _gemini_reply(_: Request) -> Reply: + candidate: Final[dict[str, JsonValue]] = { + "content": {"role": "model", "parts": [{"text": '{"colour": "blue"}'}]}, + "finishReason": "STOP", + } + usage: Final[dict[str, JsonValue]] = {"promptTokenCount": 30, "candidatesTokenCount": 12, "totalTokenCount": 42} + return Reply(body=json.dumps({"candidates": [candidate], "usageMetadata": usage}).encode()) + + +def _texts(contents: JsonValue) -> tuple[str, ...]: + assert isinstance(contents, list), contents + parts: Final = tuple(part for content in contents if isinstance(content, dict) for part in _parts(content)) + return tuple(text for part in parts if isinstance(text := part.get("text"), str)) + + +def _parts(content: dict[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + parts: Final = content.get("parts") + assert isinstance(parts, list), content + return tuple(part for part in parts if isinstance(part, dict)) + + +def _anthropic_text_reply(_: Request) -> Reply: + completion: Final[dict[str, JsonValue]] = { + "type": "completion", + "id": "compl_scripted", + "completion": "blue", + "stop_reason": "stop_sequence", + } + return Reply(body=json.dumps(completion).encode()) + + +def _boot_deployment(name: str, model: str, wire: Wire) -> dict[str, JsonValue]: + return {"model_name": name, "litellm_params": {"model": model, "api_base": wire.url, "api_key": _API_KEY}} + + +def _boot_chat(candidate: Gateway, model: str, extra: Mapping[str, JsonValue]) -> None: + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "system", "content": _SYSTEM}, {"role": "user", "content": _question(model)}], + "cache": {"no-cache": True}, + **extra, + }, + ) + assert response.status_code == 200, response.text + + +def _anthropic_text_prompts(wire: Wire) -> MappingProxyType[str, JsonValue]: + bodies: Final = tuple(_JSON_OBJECT.validate_json(request.body) for request in wire.drain()) + return MappingProxyType({str(body["model"]): body["prompt"] for body in bodies}) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_boot_time_custom_prompt_dict_renders_partial_templates(gateway: Gateway, tmp_path: Path) -> None: + templates: Final[dict[str, JsonValue]] = { + "response_schema_prompt": {"roles": {"user": {"pre_message": "", "post_message": ""}}}, + "claude-2.1": {"roles": _ROLES}, + "claude-instant-1.2": {"initial_prompt_value": "", "final_prompt_value": ""}, + } + schema: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"colour": {"type": "string"}}, + "required": ["colour"], + } + with wire_server(_gemini_reply) as gemini, wire_server(_anthropic_text_reply) as anthropic_text: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + deployments: Final = [ + _boot_deployment("gemini-schema", "gemini/scripted-template-model", gemini), + _boot_deployment("anthropic-text-roles", "anthropic_text/claude-2.1", anthropic_text), + _boot_deployment("anthropic-text-values", "anthropic_text/claude-instant-1.2", anthropic_text), + ] + settings: Final = {**base.get("litellm_settings", {}), "custom_prompt_dict": templates} + path: Final = tmp_path / "boot-time-custom-prompt-dict.yaml" + path.write_text(yaml.safe_dump({**base, "model_list": deployments, "litellm_settings": settings})) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + _boot_chat( + owned.gateway, + "gemini-schema", + {"response_format": {"type": "json_schema", "json_schema": {"name": "sky", "schema": schema}}}, + ) + posts: Final = tuple(request for request in gemini.drain() if request.method == "POST") + assert len(posts) == 1, [request.target for request in posts] + texts: Final = _texts(_JSON_OBJECT.validate_json(posts[0].body)["contents"]) + wrapped: Final = tuple(text for text in texts if text.startswith("")) + assert len(wrapped) == 1, texts + assert wrapped[0].endswith("") and "colour" in wrapped[0], wrapped[0] + _boot_chat(owned.gateway, "anthropic-text-roles", {}) + _boot_chat(owned.gateway, "anthropic-text-values", {}) + assert _anthropic_text_prompts(anthropic_text) == { + "claude-2.1": f"{_SYSTEM}{_question('anthropic-text-roles')}", + "claude-instant-1.2": f"{_SYSTEM}{_question('anthropic-text-values')}", + } diff --git a/tests/integration/providers/test_custom_prompt_template_wire.py b/tests/integration/providers/test_custom_prompt_template_wire.py new file mode 100644 index 00000000000..49a47a8c5e5 --- /dev/null +++ b/tests/integration/providers/test_custom_prompt_template_wire.py @@ -0,0 +1,701 @@ +import json +import uuid +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Final + +import anthropic +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "llama3-custom-template" +_API_KEY: Final = "synthetic-ollama-key" +_SYSTEM: Final = "Answer in one word." +_QUESTION: Final = "What colour is the sky?" +_ANSWER: Final = "blue" +_INITIAL: Final = "" +_FINAL: Final = "" +_ROLES: Final[dict[str, JsonValue]] = { + "system": {"pre_message": "", "post_message": ""}, + "user": {"pre_message": "", "post_message": ""}, + "assistant": {"pre_message": "", "post_message": ""}, +} +_ROLES_ONLY: Final[dict[str, JsonValue]] = {"roles": _ROLES} +_VALUES_ONLY: Final[dict[str, JsonValue]] = {"initial_prompt_value": _INITIAL, "final_prompt_value": _FINAL} +_FULL: Final[dict[str, JsonValue]] = {**_ROLES_ONLY, **_VALUES_ONLY} +_ROLES_PROMPT: Final = f"{_SYSTEM}{_QUESTION}" +_VALUES_PROMPT: Final = f"{_INITIAL}{_SYSTEM}{_QUESTION}{_FINAL}" +_FULL_PROMPT: Final = f"{_INITIAL}{_ROLES_PROMPT}{_FINAL}" +_TEMPLATES: Final[Mapping[str, tuple[dict[str, JsonValue], str]]] = { + "roles_only": (_ROLES_ONLY, _ROLES_PROMPT), + "values_only": (_VALUES_ONLY, _VALUES_PROMPT), +} +_TEMPLATE_FIELDS: Final = frozenset({"roles", "initial_prompt_value", "final_prompt_value", "bos_token", "eos_token"}) +_AWS: Final[dict[str, JsonValue]] = { + "aws_access_key_id": "AKIASYNTHETICKEY0001", + "aws_secret_access_key": "synthetic-secret-key-for-testing", + "aws_region_name": "us-east-1", +} +_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) + + +def _generate_reply(text: str) -> Reply: + return Reply( + body=json.dumps( + { + "model": _BACKEND, + "created_at": "2026-10-15T00:00:00Z", + "response": text, + "done": True, + "done_reason": "stop", + "prompt_eval_count": 30, + "eval_count": 12, + } + ).encode() + ) + + +def _streamed_reply(text: str) -> Reply: + frames: Final = ( + {"model": _BACKEND, "created_at": "2026-10-15T00:00:00Z", "response": text, "done": False}, + { + "model": _BACKEND, + "created_at": "2026-10-15T00:00:00Z", + "response": "", + "done": True, + "done_reason": "stop", + "prompt_eval_count": 30, + "eval_count": 12, + }, + ) + return Reply( + content_type="application/x-ndjson", chunks=tuple(json.dumps(frame).encode() + b"\n" for frame in frames) + ) + + +def _is_generate(request: Request) -> bool: + return (request.method, request.target) == ("POST", "/api/generate") + + +def _echo_answer(request: Request) -> Reply: + stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True + return _streamed_reply(_ANSWER) if stream else _generate_reply(_ANSWER) + + +@contextmanager +def _ollama_server(respond: Callable[[Request], Reply] = _echo_answer) -> Iterator[Wire]: + with wire_server(lambda request: respond(request) if _is_generate(request) else Reply(body=b"{}")) as wire: + yield wire + + +def _generate_bodies(wire: Wire) -> tuple[dict[str, JsonValue], ...]: + return tuple(_JSON_OBJECT.validate_json(request.body) for request in wire.drain() if _is_generate(request)) + + +def _assert_rendered(body: dict[str, JsonValue], prompt: str, *, stream: bool) -> None: + assert body["model"] == _BACKEND + assert body["prompt"] == prompt, body["prompt"] + assert body["stream"] is stream + options: Final = body.get("options", {}) + assert isinstance(options, dict) + assert not _TEMPLATE_FIELDS & (set(body) | set(options)), sorted(body) + + +def _assert_only_rendered(wire: Wire, prompt: str, *, stream: bool = False) -> None: + bodies: Final = _generate_bodies(wire) + assert len(bodies) == 1, bodies + _assert_rendered(bodies[0], prompt, stream=stream) + + +def _ollama_model(scenario: Scenario, wire: Wire, template: Mapping[str, JsonValue]) -> str: + return scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY, **template) + + +def _chat_messages() -> list[dict[str, JsonValue]]: + return [{"role": "system", "content": _SYSTEM}, {"role": "user", "content": _QUESTION}] + + +def _spend_row(identity: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT model_group, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + ), + lambda found: len(found) == 1, + seconds=70, + ) + return rows[0] + + +def _model_spend_rows(model: str, expected: int) -> list[dict[str, JsonValue]]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE model_group=%s', + (model,), + ), + lambda found: len(found) >= expected, + seconds=70, + ) + assert len({row["request_id"] for row in rows}) == len(rows) == expected, rows + return rows + + +def _billed(model: str) -> dict[str, JsonValue]: + return {"model_group": model, "status": "success", "prompt_tokens": 30, "completion_tokens": 12} + + +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 _async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) + + +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 _async_anthropic_client(gateway: Gateway) -> anthropic.AsyncAnthropic: + return anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def _post(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[int, str]: + response: Final = gateway.request("POST", path, {**body, "cache": _NO_CACHE}) + return response.status_code, response.text + + +def _chat(gateway: Gateway, model: str, **extra: JsonValue) -> tuple[int, str]: + return _post(gateway, "/v1/chat/completions", {"model": model, "messages": _chat_messages(), **extra}) + + +def _chat_content(text: str) -> str: + choices: Final = _JSON_LIST.validate_python(_JSON_OBJECT.validate_json(text)["choices"]) + message: Final = object_value(object_value(choices[0])["message"]) + return string_value(message["content"]) + + +def test_openai_sdk_chat_renders_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + completion: Final = _openai_client(gateway).chat.completions.create( + model=model, + messages=_chat_messages(), # pyright: ignore[reportArgumentType] # plain JSON messages + extra_body={"cache": _NO_CACHE}, + ) + assert completion.choices[0].message.content == _ANSWER + assert completion.choices[0].finish_reason == "stop" + assert completion.usage is not None + assert (completion.usage.prompt_tokens, completion.usage.completion_tokens) == (30, 12) + _assert_only_rendered(wire, _ROLES_PROMPT) + assert _spend_row(completion.id) == _billed(model) + + +async def test_async_openai_sdk_chat_stream_renders_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + stream: Final = await _async_openai_client(gateway).chat.completions.create( + model=model, + messages=_chat_messages(), # pyright: ignore[reportArgumentType] # plain JSON messages + stream=True, + stream_options={"include_usage": True}, + extra_body={"cache": _NO_CACHE}, + ) + chunks: Final = [chunk async for chunk in stream] + assert {chunk.id for chunk in chunks} == {chunks[0].id} + assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == _ANSWER + usages: Final = [chunk.usage for chunk in chunks if chunk.usage is not None] + assert [(usage.prompt_tokens, usage.completion_tokens) for usage in usages] == [(30, 12)] + _assert_only_rendered(wire, _ROLES_PROMPT, stream=True) + assert _spend_row(chunks[0].id) == _billed(model) + + +def test_openai_sdk_completions_render_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + completion: Final = _openai_client(gateway).completions.create( + model=model, prompt=_QUESTION, extra_body={"cache": _NO_CACHE} + ) + assert completion.choices[0].text == _ANSWER + _assert_only_rendered(wire, f"{_QUESTION}") + assert _spend_row(completion.id) == _billed(model) + + +def test_anthropic_sdk_messages_render_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + message: Final = _anthropic_client(gateway).messages.create( + model=model, + max_tokens=64, + system=_SYSTEM, + messages=[{"role": "user", "content": _QUESTION}], + extra_body={"cache": _NO_CACHE}, + ) + assert message.stop_reason == "end_turn" + assert [(block.type, getattr(block, "text", None)) for block in message.content] == [("text", _ANSWER)] + assert (message.usage.input_tokens, message.usage.output_tokens) == (30, 12) + _assert_only_rendered(wire, _ROLES_PROMPT) + assert _spend_row(message.id) == _billed(model) + + +async def test_async_anthropic_sdk_messages_stream_render_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + async with _async_anthropic_client(gateway).messages.stream( + model=model, + max_tokens=64, + system=_SYSTEM, + messages=[{"role": "user", "content": _QUESTION}], + extra_body={"cache": _NO_CACHE}, + ) as stream: + text: Final = "".join([piece async for piece in stream.text_stream]) + final: Final = await stream.get_final_message() + assert text == _ANSWER + assert final.stop_reason == "end_turn" + _assert_only_rendered(wire, _ROLES_PROMPT, stream=True) + assert _spend_row(final.id) == _billed(model) + + +def test_openai_sdk_responses_render_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + response: Final = _openai_client(gateway).responses.create( + model=model, instructions=_SYSTEM, input=_QUESTION, store=False, extra_body={"cache": _NO_CACHE} + ) + assert response.output_text == _ANSWER + assert response.usage is not None and (response.usage.input_tokens, response.usage.output_tokens) == (30, 12) + _assert_only_rendered(wire, _ROLES_PROMPT) + rows: Final = _model_spend_rows(model, 1) + assert (rows[0]["status"], rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == ("success", 30, 12) + + +async def test_async_openai_sdk_responses_stream_render_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + stream: Final = await _async_openai_client(gateway).responses.create( + model=model, + instructions=_SYSTEM, + input=_QUESTION, + store=False, + stream=True, + extra_body={"cache": _NO_CACHE}, + ) + events: Final = [event async for event in stream] + assert "".join(event.delta for event in events if event.type == "response.output_text.delta") == _ANSWER + completed: Final = [event for event in events if event.type == "response.completed"] + assert len(completed) == 1 and completed[0].response.output_text == _ANSWER + _assert_only_rendered(wire, _ROLES_PROMPT, stream=True) + assert _model_spend_rows(model, 1)[0]["status"] == "success" + + +def test_prompt_values_only_template_wraps_the_bare_messages(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _VALUES_ONLY) + code, text = _chat(gateway, model) + assert code == 200, text + assert _chat_content(text) == _ANSWER + _assert_only_rendered(wire, _VALUES_PROMPT) + + +def test_full_template_renders_every_part(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _FULL) + code, text = _chat(gateway, model) + assert code == 200, text + assert _chat_content(text) == _ANSWER + _assert_only_rendered(wire, _FULL_PROMPT) + + +@dataclass(frozen=True, slots=True) +class _Member: + deployment: Callable[[str], dict[str, JsonValue]] + reply: Callable[[str], Reply] + prompt_field: str + target: Callable[[str], bool] + + +def _codestral_reply(text: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": "cmpl-scripted", + "object": "chat.completion", + "model": "codestral-latest", + "created": 1, + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 30, "completion_tokens": 12, "total_tokens": 42}, + } + ).encode() + ) + + +def _predibase_reply(text: str) -> Reply: + return Reply(body=json.dumps({"generated_text": text, "details": {"finish_reason": "eos_token"}}).encode()) + + +def _sagemaker_reply(text: str) -> Reply: + return Reply(body=json.dumps([{"generated_text": text}]).encode()) + + +def _bedrock_meta_reply(text: str) -> Reply: + return Reply( + body=json.dumps( + {"generation": text, "prompt_token_count": 30, "generation_token_count": 12, "stop_reason": "stop"} + ).encode() + ) + + +_FAMILY: Final[Mapping[str, _Member]] = { + "codestral": _Member( + deployment=lambda url: { + "model": "text-completion-codestral/codestral-latest", + "api_base": f"{url}/v1/fim/completions", + "api_key": "scripted-key", + }, + reply=_codestral_reply, + prompt_field="prompt", + target=lambda target: target == "/v1/fim/completions", + ), + "predibase": _Member( + deployment=lambda url: { + "model": "predibase/llama-3-8b-instruct", + "api_base": url, + "api_key": "scripted-key", + "predibase_tenant_id": "scripted-tenant", + }, + reply=_predibase_reply, + prompt_field="inputs", + target=lambda target: target.startswith("/scripted-tenant/deployments/v2/llms/llama-3-8b-instruct/generate"), + ), + "sagemaker": _Member( + deployment=lambda url: { + "model": "sagemaker/scripted-endpoint", + "sagemaker_base_url": f"{url}/endpoints/scripted-endpoint/invocations", + **_AWS, + }, + reply=_sagemaker_reply, + prompt_field="inputs", + target=lambda target: target == "/endpoints/scripted-endpoint/invocations", + ), + "bedrock_invoke": _Member( + deployment=lambda url: {"model": "bedrock/invoke/meta.llama3-8b-instruct-v1:0", "api_base": url, **_AWS}, + reply=_bedrock_meta_reply, + prompt_field="prompt", + target=lambda target: target.endswith("/invoke"), + ), +} + + +@pytest.mark.parametrize("template", sorted(_TEMPLATES)) +@pytest.mark.parametrize("member", sorted(_FAMILY)) +def test_family_member_renders_a_partial_template(gateway: Gateway, member: str, template: str) -> None: + spec: Final = _FAMILY[member] + fields, prompt = _TEMPLATES[template] + with wire_server(lambda _: spec.reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(**spec.deployment(wire.url), **fields) + code, text = _chat(gateway, model) + received: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert code == 200, text + assert _chat_content(text) == _ANSWER + assert len(received) == 1, [request.target for request in received] + assert spec.target(received[0].target), received[0].target + body: Final = _JSON_OBJECT.validate_json(received[0].body) + assert body[spec.prompt_field] == prompt, body + assert not _TEMPLATE_FIELDS & set(body), sorted(body) + + +@pytest.mark.parametrize("provider", ("vllm", "petals")) +def test_in_process_family_member_fails_alone_and_the_rig_keeps_serving(gateway: Gateway, provider: str) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + broken: Final = scenario.model(model=f"{provider}/scripted-model", **_ROLES_ONLY) + healthy: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + failed, failure = _chat(gateway, broken) + assert failed >= 400, failure + code, text = _chat(gateway, healthy) + assert code == 200, text + _assert_only_rendered(wire, _ROLES_PROMPT) + + +def _watsonx_reply(text: str) -> Reply: + result: Final[dict[str, JsonValue]] = { + "generated_text": text, + "generated_token_count": 12, + "input_token_count": 30, + "stop_reason": "eos_token", + } + return Reply(body=json.dumps({"model_id": "ibm/granite", "results": [result]}).encode()) + + +def test_watsonx_text_deployment_with_a_template_still_answers(gateway: Gateway) -> None: + with wire_server(lambda _: _watsonx_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="watsonx_text/ibm/granite-chat-scripted", + api_base=wire.url, + token="scripted-token", + project_id="p", + **_ROLES_ONLY, + ) + code, text = _chat(gateway, model) + assert code == 200, text + assert _chat_content(text) == _ANSWER + posts: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert len(posts) == 1, [request.target for request in posts] + body: Final = _JSON_OBJECT.validate_json(posts[0].body) + assert _QUESTION in string_value(body["input"]), body + assert not _TEMPLATE_FIELDS & set(body), sorted(body) + + +def _openai_chat_reply(text: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": "chatcmpl-scripted", + "object": "chat.completion", + "created": 1, + "model": "scripted-compat", + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 30, "completion_tokens": 12, "total_tokens": 42}, + } + ).encode() + ) + + +def _converse_reply(text: str) -> Reply: + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": text}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 30, "outputTokens": 12, "totalTokens": 42}, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + +_FORWARDING: Final[Mapping[str, tuple[Callable[[str], dict[str, JsonValue]], Callable[[str], Reply]]]] = { + "openai_compatible": ( + lambda url: {"model": "openai/scripted-compat", "api_base": f"{url}/v1", "api_key": "scripted-key"}, + _openai_chat_reply, + ), + "bedrock_converse": ( + lambda url: {"model": "bedrock/converse/meta.llama3-8b-instruct-v1:0", "api_base": url, **_AWS}, + _converse_reply, + ), +} + + +def _keys_anywhere(value: JsonValue) -> frozenset[str]: + if isinstance(value, dict): + nested: Final = frozenset(key for child in value.values() for key in _keys_anywhere(child)) + return frozenset(value) | nested + if isinstance(value, list): + return frozenset(key for child in value for key in _keys_anywhere(child)) + return frozenset() + + +@pytest.mark.parametrize("source", ("deployment", "request")) +@pytest.mark.parametrize("member", sorted(_FORWARDING)) +def test_template_fields_never_reach_a_provider_that_does_not_render_them( + gateway: Gateway, member: str, source: str +) -> None: + deployment, reply = _FORWARDING[member] + with wire_server(lambda _: reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(**deployment(wire.url), **(_FULL if source == "deployment" else {})) + code, text = _chat(gateway, model, **(_FULL if source == "request" else {})) + assert code == 200, text + assert _chat_content(text) == _ANSWER + posts: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert len(posts) == 1, [request.target for request in posts] + body: Final = _JSON_OBJECT.validate_json(posts[0].body) + assert not _TEMPLATE_FIELDS & _keys_anywhere(body), body + assert _INITIAL not in posts[0].body.decode() and _FINAL not in posts[0].body.decode(), body + + +@pytest.mark.parametrize("value", (5, [_INITIAL]), ids=("int", "list")) +@pytest.mark.parametrize("source", ("deployment", "request")) +def test_non_string_initial_prompt_value_is_rejected_by_name(gateway: Gateway, source: str, value: JsonValue) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + template: Final[dict[str, JsonValue]] = {**_FULL, "initial_prompt_value": value} + broken: Final = _ollama_model(scenario, wire, template if source == "deployment" else _FULL) + healthy: Final = _ollama_model(scenario, wire, _FULL) + failed, failure = _chat(gateway, broken, **({"initial_prompt_value": value} if source == "request" else {})) + assert failed == 500, failure + assert "initial_prompt_value: Input should be a valid string" in failure, failure + assert _generate_bodies(wire) == () + code, text = _chat(gateway, healthy) + assert code == 200, text + _assert_only_rendered(wire, _FULL_PROMPT) + + +@pytest.mark.parametrize("value", ("", None), ids=("empty", "null")) +def test_blank_initial_prompt_value_in_the_request_drops_the_deployment_value( + gateway: Gateway, value: JsonValue +) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _FULL) + code, text = _chat(gateway, model, initial_prompt_value=value) + assert code == 200, text + _assert_only_rendered(wire, f"{_ROLES_PROMPT}{_FINAL}") + + +def test_request_initial_prompt_value_overrides_the_deployment_value(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _FULL) + code, text = _chat(gateway, model, initial_prompt_value="") + assert code == 200, text + _assert_only_rendered(wire, f"{_ROLES_PROMPT}{_FINAL}") + + +def test_five_kilobyte_initial_prompt_value_leads_the_prompt(gateway: Gateway) -> None: + large: Final = "x" * 5120 + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + code, text = _chat(gateway, model, initial_prompt_value=large) + assert code == 200, text + _assert_only_rendered(wire, f"{large}{_ROLES_PROMPT}") + + +def test_identical_requests_each_render_and_bill_once(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + first, first_text = _chat(gateway, model) + second, second_text = _chat(gateway, model) + assert (first, second) == (200, 200), (first_text, second_text) + bodies: Final = _generate_bodies(wire) + assert len(bodies) == 2, bodies + for body in bodies: + _assert_rendered(body, _ROLES_PROMPT, stream=False) + identities: Final = {string_value(_JSON_OBJECT.validate_json(text)["id"]) for text in (first_text, second_text)} + assert len(identities) == 2 + for identity in identities: + assert _spend_row(identity) == _billed(model) + + +def test_template_with_only_a_bos_token_names_the_fields_it_lacks(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, {"bos_token": ""}) + code, text = _chat(gateway, model) + assert code == 500, text + assert "it sets none of roles, initial_prompt_value, or final_prompt_value" in text, text + assert _generate_bodies(wire) == () + + +def test_roles_that_are_not_a_mapping_render_without_role_markers(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, {**_VALUES_ONLY, "roles": json.dumps(_ROLES)}) + code, text = _chat(gateway, model) + assert code == 200, text + _assert_only_rendered(wire, _VALUES_PROMPT) + + +def test_provider_error_reaches_the_caller_after_one_attempt(gateway: Gateway) -> None: + missing: Final = Reply(status=404, body=json.dumps({"error": f"model '{_BACKEND}' not found"}).encode()) + with _ollama_server(lambda _: missing) as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + code, text = _chat(gateway, model) + assert code == 404, text + assert f"model '{_BACKEND}' not found" in text, text + bodies: Final = _generate_bodies(wire) + assert len(bodies) == 1, bodies + _assert_rendered(bodies[0], _ROLES_PROMPT, stream=False) + + +def _template_keys(fields: Mapping[str, JsonValue]) -> tuple[bool, bool]: + return "initial_prompt_value" in fields, "final_prompt_value" in fields + + +def test_pass_through_body_treats_initial_and_final_prompt_values_alike(gateway: Gateway) -> None: + with wire_server(lambda _: _generate_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + path: Final = f"/integration-template-passthrough-{uuid.uuid4().hex}" + created: Final = gateway.post( + "/config/pass_through_endpoint", {"path": path, "target": f"{wire.url}/api/generate"} + ) + endpoint: Final = object_value(_JSON_LIST.validate_python(created["endpoints"])[0]) + scenario.cleanups.callback( + lambda: gateway.request( + "DELETE", "/config/pass_through_endpoint", params={"endpoint_id": str(endpoint["id"])} + ) + ) + response: Final = eventually( + lambda: gateway.request("POST", path, {"model": _BACKEND, "prompt": _QUESTION, "stream": False, **_FULL}), + lambda answered: answered.status_code != 404, + seconds=30, + ) + assert response.status_code == 200, response.text + posts: Final = tuple(request for request in wire.drain() if request.method == "POST") + forwarded: Final = _JSON_OBJECT.validate_json(posts[-1].body) + assert (forwarded["model"], forwarded["prompt"]) == (_BACKEND, _QUESTION), forwarded + initial, final = _template_keys(forwarded) + assert initial == final, sorted(forwarded) + + +def _a2a_reply(request: Request) -> Reply: + result: Final[dict[str, JsonValue]] = { + "kind": "message", + "role": "agent", + "messageId": f"{uuid.uuid4().hex}-out", + "parts": [{"kind": "text", "text": _ANSWER}], + } + outbound: Final = _JSON_OBJECT.validate_json(request.body) + return Reply(body=json.dumps({"jsonrpc": "2.0", "id": outbound["id"], "result": result}).encode()) + + +def test_a2a_params_treat_initial_and_final_prompt_values_alike(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + agent_name: Final = f"template-agent-{marker}" + + def peer(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps(_agent_card(wire.url, agent_name)).encode()) + return _a2a_reply(request) + + with wire_server(peer) as wire, gateway.scenario() as scenario: + created: Final = gateway.post( + "/v1/agents", {"agent_name": agent_name, "agent_card_params": _agent_card(wire.url, agent_name)} + ) + agent_id: Final = string_value(created["agent_id"]) + scenario.cleanups.callback(lambda: gateway.request("DELETE", f"/v1/agents/{agent_id}")) + message: Final[dict[str, JsonValue]] = { + "kind": "message", + "role": "user", + "messageId": f"{marker}-in", + "parts": [{"kind": "text", "text": _QUESTION}], + } + invoked: Final = gateway.request( + "POST", + f"/a2a/{agent_id}", + { + "jsonrpc": "2.0", + "id": marker, + "method": "message/send", + "params": {"message": message, "initial_prompt_value": _INITIAL, "final_prompt_value": _FINAL}, + }, + ) + assert invoked.status_code == 200, invoked.text + assert object_value(invoked.json()["result"])["parts"] == [{"kind": "text", "text": _ANSWER}], invoked.text + posts: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert len(posts) == 1, [request.target for request in posts] + params: Final = object_value(_JSON_OBJECT.validate_json(posts[0].body)["params"]) + assert _template_keys(params) == (False, False), sorted(params) + + +def _agent_card(url: str, name: str) -> dict[str, JsonValue]: + return { + "protocolVersion": "0.3", + "name": name, + "description": "Scripted A2A peer", + "version": "1.0.0", + "url": url + "/", + "capabilities": {"streaming": True}, + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "skills": [], + } diff --git a/tests/integration/providers/test_gemini_anyof_null_branches_wire.py b/tests/integration/providers/test_gemini_anyof_null_branches_wire.py new file mode 100644 index 00000000000..9fa99cb69f8 --- /dev/null +++ b/tests/integration/providers/test_gemini_anyof_null_branches_wire.py @@ -0,0 +1,681 @@ +import asyncio +import base64 +import json +import re +import threading +import uuid +from collections.abc import AsyncIterator, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from itertools import product +from types import MappingProxyType +from typing import Final, Literal + +import anthropic +import httpx +import openai +import pytest +import websockets +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.process import OwnedProxy, 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 integration._support.upstream import GEMINI_LIVE_PATH +from integration.providers._gemini_tts_voice_support import live_deployment, live_scenario +from integration.providers._vertex_batch_rows import encoded, gcs_peer, is_gcs_upload, jsonl, upload, vertex_model +from pydantic import JsonValue, TypeAdapter + +Provider = Literal["vertex_ai", "gemini"] +Endpoint = Literal["chat", "messages", "responses"] + +_BACKEND: Final = "gemini-3.8-flash" +_NO_JSON_SCHEMA_BACKEND: Final = "gemini-flash-latest" +_IMAGE_BACKEND: Final = "gemini-3.1-flash-image" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-central1" +_GEMINI_KEY: Final = "sk-scripted-gemini" +_COUNT: Final = 777 +_FAIL_MARKER: Final = "scripted-invalid-argument" +_STRING_ANYOF: Final = re.compile(rb'"anyOf":\s*"') +_SCALAR_BRANCH: Final = re.compile(rb'"anyOf":\s*\[\s*(?:true|false|-?\d)') +_PNG: Final = base64.b64encode(b"\x89PNG\r\n\x1a\nscripted").decode() +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) + +_NULL: Final[Mapping[str, JsonValue]] = {"type": "null"} +_STRING: Final[Mapping[str, JsonValue]] = {"type": "string"} +_INTEGER: Final[Mapping[str, JsonValue]] = {"type": "integer"} +_NULLABLE_STRING: Final[JsonValue] = {"anyOf": [{"type": "string", "nullable": True}]} +_NULLABLE_NULL: Final[JsonValue] = {"anyOf": [{"type": "null", "nullable": True}]} + +_UNIONS: Final[tuple[tuple[str, JsonValue, JsonValue], ...]] = ( + ("null-null-string", {"anyOf": [_NULL, _NULL, _STRING]}, _NULLABLE_STRING), + ("null-only", {"anyOf": [_NULL]}, _NULLABLE_NULL), + ("two-nulls-only", {"anyOf": [_NULL, _NULL]}, _NULLABLE_NULL), + ("string-null", {"anyOf": [_STRING, _NULL]}, _NULLABLE_STRING), + ( + "integer-null-null-string", + {"anyOf": [_INTEGER, _NULL, _NULL, _STRING]}, + {"anyOf": [{"type": "integer", "nullable": True}, {"type": "string", "nullable": True}]}, + ), + ("null-empty", {"anyOf": [_NULL, {}]}, {"anyOf": [{"type": "object", "nullable": True}]}), + ("empty-only", {"anyOf": [{}]}, {"anyOf": [{"type": "object"}]}), + ("empty-list", {"anyOf": []}, {"anyOf": []}), + ("object-anyof", {"anyOf": _STRING}, {"anyOf": _STRING}), + ("null-anyof", {"anyOf": None}, {"anyOf": None}), + ("five-hundred-nulls-then-string", {"anyOf": [*([_NULL] * 500), _STRING]}, _NULLABLE_STRING), + ("string-null-type-array", {"type": ["string", "null"]}, _NULLABLE_STRING), + ( + "null-inside-a-union-branch", + {"anyOf": [{"type": "object", "properties": {"x": {"anyOf": [_INTEGER, _NULL]}}}, _NULL]}, + { + "anyOf": [ + { + "type": "object", + "properties": {"x": {"anyOf": [{"type": "integer", "nullable": True}]}}, + "nullable": True, + } + ] + }, + ), +) +_FANOUT_UNIONS: Final = tuple(case for case in _UNIONS if case[0] in {"null-null-string", "null-only", "object-anyof"}) + + +def _text_reply() -> Mapping[str, JsonValue]: + return { + "candidates": [{"content": {"role": "model", "parts": [{"text": "OK"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 1, "totalTokenCount": 10}, + "modelVersion": _BACKEND, + } + + +def _image_reply() -> Mapping[str, JsonValue]: + return { + "candidates": [ + {"content": {"role": "model", "parts": [{"inlineData": {"mimeType": "image/png", "data": _PNG}}]}} + ], + "usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 1, "totalTokenCount": 10}, + } + + +def _rejection(message: str) -> Reply: + error: Final = {"code": 400, "message": message, "status": "INVALID_ARGUMENT"} + return Reply(status=400, body=json.dumps({"error": error}).encode()) + + +def _google(request: Request) -> Reply: + if is_gcs_upload(request): + return gcs_peer(request) + if _FAIL_MARKER.encode() in request.body: + return _rejection(f"{_FAIL_MARKER}: the scripted tool schema was refused") + if _STRING_ANYOF.search(request.body): + return _rejection("Invalid value at 'tools[0].function_declarations[0].parameters' (anyOf), not a list") + if _SCALAR_BRANCH.search(request.body): + return _rejection( + "Invalid value at 'tools[0].function_declarations[0].parameters.properties[0].value.any_of[0]'" + ) + path: Final = request.target.split("?")[0] + if path.endswith(":countTokens"): + return Reply(body=json.dumps({"totalTokens": _COUNT}).encode()) + if path.endswith(":streamGenerateContent"): + return Reply(content_type="text/event-stream", chunks=(f"data: {json.dumps(_text_reply())}\n\n".encode(),)) + if path.endswith(f"{_IMAGE_BACKEND}:generateContent"): + return Reply(body=json.dumps(_image_reply()).encode()) + if path.endswith(":generateContent"): + return Reply(body=json.dumps(_text_reply()).encode()) + return Reply(status=404, body=json.dumps({"error": {"code": 404, "message": request.target}}).encode()) + + +def _deployment(gateway: Gateway, scenario: Scenario, wire: Wire, provider: Provider, backend: str = _BACKEND) -> str: + if provider == "gemini": + return scenario.model(model=f"gemini/{backend}", api_base=wire.url, api_key=_GEMINI_KEY) + return scenario.model( + model=f"vertex_ai/{backend}", + api_base=wire.url, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + ) + + +def _parameters(value: JsonValue) -> dict[str, JsonValue]: + return {"type": "object", "properties": {"value": value}, "required": ["value"]} + + +def _declaration(description: str, value: JsonValue) -> dict[str, JsonValue]: + return {"name": "record", "description": description, "parameters": _parameters(value)} + + +def _tool(endpoint: Endpoint, description: str, parameters: JsonValue) -> dict[str, JsonValue]: + if endpoint == "messages": + return {"name": "record", "description": description, "input_schema": parameters} + if endpoint == "responses": + return {"type": "function", "name": "record", "description": description, "parameters": parameters} + return {"type": "function", "function": {"name": "record", "description": description, "parameters": parameters}} + + +def _request_body( + endpoint: Endpoint, model: str, description: str, parameters: JsonValue, stream: bool +) -> tuple[str, dict[str, JsonValue]]: + prompt: Final = f"Record the value hello by calling the record tool. {description}" + tools: Final[list[JsonValue]] = [_tool(endpoint, description, parameters)] + streaming: Final[dict[str, JsonValue]] = {"stream": True} if stream else {} + if endpoint == "messages": + messages_body: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": prompt}], + "tools": tools, + } + return "/v1/messages", {**messages_body, **streaming} + if endpoint == "responses": + return "/v1/responses", {"model": model, "input": prompt, "tools": tools, **streaming} + chat_body: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": prompt}], + "tools": tools, + } + return "/v1/chat/completions", {**chat_body, **streaming} + + +def _body(request: Request) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_json(request.body) + + +def _json(response: httpx.Response) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_json(response.content) + + +def _generate_requests(wire: Wire) -> tuple[Request, ...]: + return tuple( + request + for request in wire.drain() + if request.target.split("?")[0].endswith((":generateContent", ":streamGenerateContent")) + ) + + +def _sent_declarations(body: Mapping[str, JsonValue]) -> JsonValue: + tools: Final = _JSON_LIST.validate_python(body["tools"]) + assert len(tools) == 1, tools + tool: Final = object_value(tools[0]) + assert len(tool) == 1, tool + return next(iter(tool.values())) + + +def _description() -> str: + return f"Record a value {uuid.uuid4().hex[:12]}" + + +def _assert_caller_saw_ok(endpoint: Endpoint, response: httpx.Response, stream: bool) -> None: + assert response.status_code == 200, response.text + if stream: + assert "OK" in response.text, response.text + return + payload: Final = _json(response) + assert string_value(payload["id"]), payload + if endpoint == "chat": + choice: Final = object_value(_JSON_LIST.validate_python(payload["choices"])[0]) + assert object_value(choice["message"])["content"] == "OK", payload + return + assert "OK" in response.text, payload + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +@pytest.mark.parametrize(("union", "expected"), [case[1:] for case in _UNIONS], ids=[case[0] for case in _UNIONS]) +def test_chat_tool_union_reaches_google_with_null_branches_folded_into_nullable( + gateway: Gateway, provider: Provider, union: JsonValue, expected: JsonValue +) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, provider) + description: Final = _description() + path, body = _request_body("chat", model, description, _parameters(union), False) + response: Final = gateway.request("POST", path, body) + _assert_caller_saw_ok("chat", response, False) + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert requests[0].target.split("?")[0].endswith(f"{_BACKEND}:generateContent"), requests[0].target + assert _sent_declarations(_body(requests[0])) == [_declaration(description, expected)] + + +@pytest.mark.parametrize("stream", (False, True), ids=("whole", "stream")) +@pytest.mark.parametrize("endpoint", ("chat", "messages", "responses")) +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +@pytest.mark.parametrize( + ("union", "expected"), [case[1:] for case in _FANOUT_UNIONS], ids=[case[0] for case in _FANOUT_UNIONS] +) +def test_every_unified_endpoint_sends_the_same_folded_union( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, union: JsonValue, expected: JsonValue +) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, provider) + description: Final = _description() + path, body = _request_body(endpoint, model, description, _parameters(union), stream) + response: Final = gateway.request("POST", path, body) + _assert_caller_saw_ok(endpoint, response, stream) + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + method: Final = "streamGenerateContent" if stream else "generateContent" + assert requests[0].target.split("?")[0].endswith(f"{_BACKEND}:{method}"), requests[0].target + assert _sent_declarations(_body(requests[0])) == [_declaration(description, expected)] + + +def test_nested_item_and_property_unions_are_folded_at_every_depth(gateway: Gateway) -> None: + parameters: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "values": {"type": "array", "items": {"anyOf": [_NULL, _NULL, _INTEGER]}}, + "nested": {"type": "object", "properties": {"inner": {"anyOf": [_NULL, _NULL, _STRING]}}}, + }, + "required": ["values", "nested"], + } + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "vertex_ai") + description: Final = _description() + path, body = _request_body("chat", model, description, parameters, False) + response: Final = gateway.request("POST", path, body) + _assert_caller_saw_ok("chat", response, False) + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [ + { + "name": "record", + "description": description, + "parameters": { + "type": "object", + "properties": { + "values": {"type": "array", "items": {"anyOf": [{"type": "integer", "nullable": True}]}}, + "nested": {"type": "object", "properties": {"inner": _NULLABLE_STRING}}, + }, + "required": ["values", "nested"], + }, + } + ] + + +def test_a_string_anyof_reaches_google_and_its_400_reaches_the_caller(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "gemini") + description: Final = _description() + path, body = _request_body("chat", model, description, _parameters({"anyOf": "string"}), False) + response: Final = gateway.request("POST", path, body) + assert response.status_code == 400, response.text + assert "not a list" in response.text, response.text + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, {"anyOf": "string"})] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_a_google_rejection_of_a_folded_union_reaches_the_caller(gateway: Gateway, provider: Provider) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, provider) + description: Final = f"{_FAIL_MARKER} {uuid.uuid4().hex[:12]}" + path, body = _request_body("chat", model, description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}), False) + response: Final = gateway.request("POST", path, body) + assert response.status_code == 400, response.text + assert _FAIL_MARKER in response.text, response.text + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_STRING)] + + +@pytest.mark.parametrize( + ("union", "sent"), + (({"anyOf": [True, _NULL]}, {"anyOf": [True]}), ({"anyOf": [_NULL, 5]}, {"anyOf": [5]})), + ids=("true-null", "null-int"), +) +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_a_union_branch_that_is_not_an_object_reaches_google_and_its_400_reaches_the_caller( + gateway: Gateway, provider: Provider, union: JsonValue, sent: JsonValue +) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, provider) + rejected_description: Final = _description() + path, broken = _request_body("chat", model, rejected_description, _parameters(union), False) + failed: Final = gateway.request("POST", path, broken) + assert failed.status_code == 400, failed.text + assert "any_of[0]" in failed.text, failed.text + rejected: Final = _generate_requests(wire) + assert len(rejected) == 1, [request.target for request in rejected] + assert _sent_declarations(_body(rejected[0])) == [_declaration(rejected_description, sent)] + description: Final = _description() + _, valid = _request_body("chat", model, description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}), False) + recovered: Final = gateway.request("POST", path, valid) + _assert_caller_saw_ok("chat", recovered, False) + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_STRING)] + + +def test_a_repeated_identical_request_sends_the_same_folded_union_each_time(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "vertex_ai") + description: Final = _description() + path, body = _request_body("chat", model, description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}), False) + first: Final = gateway.request("POST", path, body) + second: Final = gateway.request("POST", path, body) + _assert_caller_saw_ok("chat", first, False) + _assert_caller_saw_ok("chat", second, False) + assert _json(first)["id"] != _json(second)["id"] + requests: Final = _generate_requests(wire) + assert len(requests) == 2, [request.target for request in requests] + assert [_sent_declarations(_body(request)) for request in requests] == [ + [_declaration(description, _NULLABLE_STRING)] + ] * 2 + + +def test_openai_sdk_chat_sends_the_folded_union(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "gemini") + description: Final = _description() + client: Final = openai.OpenAI( + base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", api_key=gateway.key, max_retries=0 + ) + completion: Final = client.chat.completions.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": "Record the value hello by calling the record tool."}], + tools=[_tool("chat", description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}))], # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping + ) + assert completion.choices[0].message.content == "OK" + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_STRING)] + + +async def _async_openai_stream_text(gateway: Gateway, model: str, description: str) -> str: + client: Final = openai.AsyncOpenAI( + base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", api_key=gateway.key, max_retries=0 + ) + stream: Final = await client.chat.completions.create( + model=model, + max_tokens=16, + stream=True, + messages=[{"role": "user", "content": "Record the value hello by calling the record tool."}], + tools=[_tool("chat", description, _parameters({"anyOf": [_NULL]}))], # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping + ) + return "".join([chunk.choices[0].delta.content or "" async for chunk in stream if chunk.choices]) + + +def test_async_openai_sdk_stream_sends_the_null_only_union_as_one_nullable_null(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "vertex_ai") + description: Final = _description() + assert asyncio.run(_async_openai_stream_text(gateway, model, description)) == "OK" + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert requests[0].target.split("?")[0].endswith(":streamGenerateContent"), requests[0].target + assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_NULL)] + + +def test_anthropic_sdk_messages_sends_the_folded_union(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "vertex_ai") + description: Final = _description() + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + message: Final = client.messages.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": "Record the value hello by calling the record tool."}], + tools=[_tool("messages", description, _parameters({"anyOf": [_NULL, _NULL]}))], # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping + ) + assert [block.type for block in message.content] == ["text"] + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_NULL)] + + +@pytest.mark.parametrize( + ("union", "expected"), [case[1:] for case in _FANOUT_UNIONS], ids=[case[0] for case in _FANOUT_UNIONS] +) +def test_vertex_count_tokens_sends_the_folded_union(gateway: Gateway, union: JsonValue, expected: JsonValue) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "vertex_ai") + description: Final = _description() + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + { + "model": model, + "messages": [{"role": "user", "content": "Record the value hello by calling the record tool."}], + "tools": [_tool("messages", description, _parameters(union))], + }, + ) + assert response.status_code == 200, response.text + assert _json(response)["input_tokens"] == _COUNT, response.text + requests: Final = tuple(request for request in wire.drain() if request.target.endswith(":countTokens")) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, expected)] + + +@pytest.mark.parametrize( + ("union", "expected"), [case[1:] for case in _FANOUT_UNIONS], ids=[case[0] for case in _FANOUT_UNIONS] +) +def test_gemini_token_counter_sends_the_folded_union(gateway: Gateway, union: JsonValue, expected: JsonValue) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "gemini") + description: Final = _description() + response: Final = gateway.request( + "POST", + "/utils/token_counter", + { + "model": model, + "contents": [ + {"role": "user", "parts": [{"text": "Record the value hello by calling the record tool."}]} + ], + "tools": [_tool("chat", description, _parameters(union))], + }, + params={"call_endpoint": "true"}, + ) + assert response.status_code == 200, response.text + assert _json(response)["total_tokens"] == _COUNT, response.text + requests: Final = tuple(request for request in wire.drain() if ":countTokens" in request.target) + assert len(requests) == 1, [request.target for request in requests] + sent: Final = object_value(_body(requests[0])["generateContentRequest"]) + assert _sent_declarations(sent) == [_declaration(description, expected)] + + +def test_response_format_on_a_model_without_json_schema_support_folds_the_union(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "gemini", _NO_JSON_SCHEMA_BACKEND) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "Return the value as JSON."}], + "response_format": { + "type": "json_schema", + "json_schema": {"name": "record", "schema": _parameters({"anyOf": [_NULL, _NULL, _STRING]})}, + }, + }, + ) + _assert_caller_saw_ok("chat", response, False) + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + generation_config: Final = object_value(_body(requests[0])["generationConfig"]) + assert generation_config["response_mime_type"] == "application/json", generation_config + assert generation_config["response_schema"] == {**_parameters(_NULLABLE_STRING), "propertyOrdering": ["value"]} + + +def test_native_generate_content_response_schema_folds_the_union(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "gemini", _NO_JSON_SCHEMA_BACKEND) + response: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:generateContent", + { + "contents": [{"role": "user", "parts": [{"text": "Return the value as JSON."}]}], + "generationConfig": { + "responseMimeType": "application/json", + "responseSchema": _parameters({"anyOf": [_NULL, _NULL, _STRING]}), + }, + }, + ) + assert response.status_code == 200, response.text + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + generation_config: Final = object_value(_body(requests[0])["generationConfig"]) + assert generation_config["responseSchema"] == {**_parameters(_NULLABLE_STRING), "propertyOrdering": ["value"]} + + +@pytest.mark.parametrize( + ("union", "expected"), [case[1:] for case in _FANOUT_UNIONS], ids=[case[0] for case in _FANOUT_UNIONS] +) +def test_image_generation_tools_send_the_folded_union(gateway: Gateway, union: JsonValue, expected: JsonValue) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "gemini", _IMAGE_BACKEND) + description: Final = _description() + response: Final = gateway.request( + "POST", + "/v1/images/generations", + { + "model": model, + "prompt": "A small red square on a white background", + "tools": [_tool("chat", description, _parameters(union))], + }, + ) + assert response.status_code == 200, response.text + assert len(_JSON_LIST.validate_python(_json(response)["data"])) == 1, response.text + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, expected)] + + +def test_vertex_batch_row_is_stored_with_the_folded_union(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url, _BACKEND) + description: Final = _description() + row: Final[dict[str, JsonValue]] = { + "custom_id": "anyof-row", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model, + "messages": [{"role": "user", "content": "Record the value hello by calling the record tool."}], + "tools": [_tool("chat", description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}))], + }, + } + response: Final = upload(gateway, "model", model, jsonl(encoded(row))) + assert response.status_code == 200, response.text + uploads: Final = tuple(request for request in wire.drain() if is_gcs_upload(request)) + assert len(uploads) == 1, [request.target for request in uploads] + stored: Final = tuple( + object_value(JSON_OBJECT.validate_json(line)["request"]) for line in uploads[0].body.splitlines() if line + ) + assert len(stored) == 1, stored + assert _sent_declarations(stored[0]) == [_declaration(description, _NULLABLE_STRING)] + + +async def _event_types_until_done(socket: websockets.ClientConnection) -> AsyncIterator[str]: + async for message in socket: + kind: Final = string_value(JSON_OBJECT.validate_json(message)["type"]) + yield kind + if kind in {"response.done", "error"}: + return + + +async def _realtime_session_with_tools(gateway: Gateway, model: str, tools: list[JsonValue]) -> tuple[str, ...]: + ws_base: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://") + headers: Final = {"Authorization": f"Bearer {gateway.key}"} + async with websockets.connect(f"{ws_base}/v1/realtime?model={model}", additional_headers=headers) as socket: + first: Final = JSON_OBJECT.validate_json(await socket.recv()) + await socket.send(json.dumps({"type": "session.update", "session": {"tools": tools}})) + item: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "say the line"}]} + await socket.send(json.dumps({"type": "conversation.item.create", "item": item})) + events: Final = [kind async for kind in _event_types_until_done(socket)] + return (string_value(first["type"]), *events) + + +def _realtime_types(gateway: Gateway, model: str, tools: list[JsonValue]) -> tuple[str, ...]: + async def bounded() -> tuple[str, ...]: + try: + return await asyncio.wait_for(_realtime_session_with_tools(gateway, model, tools), 30) + except TimeoutError: + return ("timeout",) + + return asyncio.run(bounded()) + + +def _live_frames(gateway: Gateway, scenario_id: str) -> tuple[dict[str, JsonValue], ...]: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + observed: Final = tuple( + map(object_value, _JSON_LIST.validate_python(_json(upstream.get("/__observations"))["requests"])) + ) + return tuple( + object_value(request["body"]) + for request in observed + if request["api_key"] == scenario_id + and request["method"] == "WEBSOCKET_FRAME" + and request["path"] == GEMINI_LIVE_PATH + ) + + +@pytest.fixture(scope="module") +def deferred_live(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + directory: Final = tmp_path_factory.mktemp("anyof-live") + with gateway_from_environment() as gateway: + with owned_proxy_process(gateway, directory, {"LITELLM_GEMINI_LIVE_DEFER_SETUP": "true"}, workers=2) as proxy: + yield proxy + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_deferred_realtime_setup_carries_session_tools_with_the_folded_union(deferred_live: OwnedProxy) -> None: + gateway: Final = deferred_live.gateway + with gateway.scenario() as scenario: + handle: Final = live_scenario(scenario) + model: Final = live_deployment(scenario, handle.scenario_id, gateway.upstream_url) + description: Final = _description() + tools: Final[list[JsonValue]] = [_tool("chat", description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}))] + types: Final = _realtime_types(gateway, model, tools) + assert types[0] == "session.created", types + assert types[-1] == "response.done", types + frames: Final = _live_frames(gateway, handle.scenario_id) + setups: Final = tuple(object_value(frame["setup"]) for frame in frames if "setup" in frame) + assert len(setups) == 1, [sorted(frame) for frame in frames] + assert _sent_declarations(setups[0]) == [_declaration(description, _NULLABLE_STRING)] + + +_PROVIDERS: Final[tuple[Provider, ...]] = ("vertex_ai", "gemini") +_ENDPOINTS: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses") +_BURST: Final[tuple[tuple[Provider, Endpoint, bool, bool], ...]] = tuple( + (provider, endpoint, stream, index % 4 == 0) + for index, (provider, endpoint, stream, _) in enumerate(product(_PROVIDERS, _ENDPOINTS, (False, True), range(2))) +) + + +def test_a_concurrent_mixed_burst_folds_every_union_once_and_keeps_each_status(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + models: Final = MappingProxyType( + {provider: _deployment(gateway, scenario, wire, provider) for provider in _PROVIDERS} + ) + shots: Final = tuple( + (provider, endpoint, stream, f"{_FAIL_MARKER if refused else 'burst'} {uuid.uuid4().hex}") + for provider, endpoint, stream, refused in _BURST + ) + + def send(shot: tuple[Provider, Endpoint, bool, str]) -> httpx.Response: + provider, endpoint, stream, description = shot + path, body = _request_body( + endpoint, models[provider], description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}), stream + ) + return gateway.request("POST", path, body) + + with ThreadPoolExecutor(max_workers=len(shots)) as pool: + responses: Final = tuple(pool.map(send, shots)) + expected_statuses: Final = [400 if _FAIL_MARKER in shot[3] else 200 for shot in shots] + assert [response.status_code for response in responses] == expected_statuses, [ + response.text[:120] for response in responses if response.status_code not in {200, 400} + ] + requests: Final = _generate_requests(wire) + assert len(requests) == len(shots), len(requests) + sent: Final = sorted((json.dumps(_sent_declarations(_body(request)), sort_keys=True) for request in requests)) + assert sent == sorted(json.dumps([_declaration(shot[3], _NULLABLE_STRING)], sort_keys=True) for shot in shots) 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_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index 3ab5826006f..2b23ccdf4f2 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -321,6 +321,48 @@ def test_messages_count_tokens_sends_the_input_v1_messages_inference_sends( assert count_body["model"] == inference_body["model"] +_UNCONVERTIBLE_BODY: Final[Mapping[str, JsonValue]] = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the notes."}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}}, + {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": [ + { + "type": "search_result", + "source": "https://example.com/hits", + "title": "Hits", + "content": [{"type": "text", "text": "hit"}], + } + ], + } + ], + }, + ], +} + + +def test_messages_count_tokens_rejects_blocks_it_cannot_convert(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, **_UNCONVERTIBLE_BODY}) + assert response.status_code == 400, response.text + assert "document with a text source, search_result, search_result inside a tool_result" in response.text + assert wire.drain() == () + + def test_messages_count_tokens_without_a_role_answers_400_before_calling_openai(counting_proxy: OwnedProxy) -> None: gateway: Final = counting_proxy.gateway with wire_server(_count_response) as wire, gateway.scenario() as scenario: 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 1a148b19f8d..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) @@ -65,7 +69,7 @@ class _GeminiPeer: return requests -def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str: +def _deployment(gateway: Gateway, scenario: Scenario, api_base: str, **parameters: JsonValue) -> str: return scenario.model( model=f"vertex_ai/{_MODEL}", api_base=api_base, @@ -73,6 +77,7 @@ def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str: vertex_project=_PROJECT, vertex_location=_LOCATION, vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + **parameters, ) @@ -167,6 +172,530 @@ def test_messages_and_count_tokens_send_matching_gemini_context( assert count_body == _count_tokens_request_from_generate_content(generated_body) +_LOOKUP_TOOL: Final[JsonValue] = { + "name": "lookup", + "description": "Look up a value", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, +} +_CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "base64-pdf": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the document."}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, + }, + ], + } + ], + }, + "base64-image": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + {"type": "text", "text": "What is this?"}, + ], + } + ], + }, + "tool-turns": { + "tools": [_LOOKUP_TOOL], + "messages": [ + {"role": "user", "content": "Look up the value."}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Looking it up."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "value"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}], + }, + ], + }, + } +) + + +def _ai_studio_deployment(scenario: Scenario, api_base: str, **parameters: JsonValue) -> str: + return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key", **parameters) + + +def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue: + if "contents" in count_body: + return count_body["contents"] + return _JSON_OBJECT.validate_python(count_body["generateContentRequest"])["contents"] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +@pytest.mark.parametrize("case", sorted(_CONTENT_BLOCK_BODIES)) +def test_count_tokens_sends_the_contents_messages_inference_sends(gateway: Gateway, provider: str, case: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = ( + _deployment(gateway, scenario, wire.url) + if provider == "vertex_ai" + else _ai_studio_deployment(scenario, wire.url) + ) + body: Final = {"model": model, **_CONTENT_BLOCK_BODIES[case]} + 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 == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text + requests: Final = peer.drain() + generated_body: Final = _json_body( + next(request for request in requests if request.target.endswith(":generateContent")) + ) + count_body: Final = _json_body(next(request for request in requests if request.target.endswith(":countTokens"))) + assert _counted_contents(count_body) == generated_body["contents"] + + +_UNCONVERTIBLE_MESSAGES: Final[JsonValue] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the notes."}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}}, + {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": [ + { + "type": "search_result", + "source": "https://example.com/hits", + "title": "Hits", + "content": [{"type": "text", "text": "hit"}], + } + ], + } + ], + }, +] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_count_tokens_rejects_blocks_it_cannot_convert(gateway: Gateway, provider: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = ( + _deployment(gateway, scenario, wire.url) + if provider == "vertex_ai" + else _ai_studio_deployment(scenario, wire.url) + ) + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _UNCONVERTIBLE_MESSAGES} + ) + assert response.status_code == 400, response.text + assert "document with a text source, search_result, search_result inside a tool_result" in response.text + 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: @@ -460,3 +989,110 @@ def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gatew tuple(_json_body(request) for request in requests) == ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9 ) + + +_CONTENTS: Final[JsonValue] = [{"role": "user", "parts": [{"text": _USER}]}] +_OPENAI_TOOLS: Final[JsonValue] = [ + { + "type": "function", + "function": {"name": "noop", "description": "Return no output", "parameters": {"type": "object"}}, + } +] + + +def _provider_deployment( + gateway: Gateway, scenario: Scenario, provider: str, api_base: str, **parameters: JsonValue +) -> str: + if provider == "vertex_ai": + return _deployment(gateway, scenario, api_base, **parameters) + return _ai_studio_deployment(scenario, api_base, **parameters) + + +def _count(gateway: Gateway, body: Mapping[str, JsonValue], call_endpoint: bool = True) -> httpx.Response: + return gateway.request( + "POST", "/utils/token_counter", body, params={"call_endpoint": "true" if call_endpoint else "false"} + ) + + +def _total_tokens(response: httpx.Response) -> JsonValue: + return _JSON_OBJECT.validate_json(response.content)["total_tokens"] + + +@pytest.mark.parametrize( + ("provider", "empty_conversation"), + (("vertex_ai", {}), ("gemini", {}), ("gemini", {"messages": []}), ("gemini", {"contents": []})), + ids=("vertex-prompt-only", "ai-studio-prompt-only", "ai-studio-empty-messages", "ai-studio-empty-contents"), +) +def test_utils_token_counter_sends_a_prompt_as_the_user_turn_of_its_messages_form( + gateway: Gateway, provider: str, empty_conversation: Mapping[str, JsonValue] +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _provider_deployment(gateway, scenario, provider, wire.url) + prompted: Final = _count(gateway, {"model": model, "prompt": _USER, **empty_conversation}) + prompted_requests: Final = _count_requests(wire) + as_message: Final = _count(gateway, {"model": model, "messages": _MESSAGES}) + message_requests: Final = _count_requests(wire) + assert prompted.status_code == 200, prompted.text + assert as_message.status_code == 200, as_message.text + assert _total_tokens(prompted) == _COUNT, prompted.text + assert len(prompted_requests) == 1 + assert len(message_requests) == 1 + assert _counted_contents(_json_body(prompted_requests[0])) == _counted_contents(_json_body(message_requests[0])) + + +@pytest.mark.parametrize( + ("conversation", "call_endpoint"), + (({"messages": _MESSAGES}, True), ({"contents": _CONTENTS}, True), ({"messages": _MESSAGES}, False)), + ids=("provider-count-messages", "provider-count-contents", "local-count-messages"), +) +def test_utils_token_counter_rejects_a_prompt_next_to_a_conversation_without_calling_the_provider( + gateway: Gateway, conversation: Mapping[str, JsonValue], call_endpoint: bool +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url) + response: Final = _count(gateway, {"model": model, "prompt": "hi", **conversation}, call_endpoint) + requests: Final = _count_requests(wire) + assert response.status_code == 400, response.text + assert _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_json(response.content)["error"])["param"] == "prompt" + assert requests == () + + +@pytest.mark.parametrize("call_endpoint", (True, False), ids=("provider-count", "local-count")) +def test_utils_token_counter_counts_only_the_conversation_when_drop_params_drops_the_prompt( + gateway: Gateway, call_endpoint: bool +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=True) + mixed: Final = _count(gateway, {"model": model, "prompt": "hi", "messages": _MESSAGES}, call_endpoint) + mixed_requests: Final = _count_requests(wire) + alone: Final = _count(gateway, {"model": model, "messages": _MESSAGES}, call_endpoint) + alone_requests: Final = _count_requests(wire) + assert mixed.status_code == 200, mixed.text + assert alone.status_code == 200, alone.text + assert _total_tokens(mixed) == _total_tokens(alone) + assert [_counted_contents(_json_body(request)) for request in mixed_requests] == [ + _counted_contents(_json_body(request)) for request in alone_requests + ] + + +@pytest.mark.parametrize( + "extra", + ({"system": "Follow the system instruction."}, {"tools": _OPENAI_TOOLS}), + ids=("system", "tools"), +) +def test_utils_token_counter_counts_a_local_prompt_with_system_or_tools_like_its_messages_form( + gateway: Gateway, extra: Mapping[str, JsonValue] +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url) + prompted: Final = _count(gateway, {"model": model, "prompt": _USER, **extra}, call_endpoint=False) + as_message: Final = _count(gateway, {"model": model, "messages": _MESSAGES, **extra}, call_endpoint=False) + requests: Final = _count_requests(wire) + assert prompted.status_code == 200, prompted.text + assert as_message.status_code == 200, as_message.text + assert _total_tokens(prompted) == _total_tokens(as_message) + assert requests == () diff --git a/tests/integration/routing/test_base_model_context_window_owned_proxy.py b/tests/integration/routing/test_base_model_context_window_owned_proxy.py new file mode 100644 index 00000000000..d09ec868c78 --- /dev/null +++ b/tests/integration/routing/test_base_model_context_window_owned_proxy.py @@ -0,0 +1,171 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import marker_of, owned_config +from pydantic import JsonValue + +_TOKEN: Final = "synthetic-bedrock-bearer" +_KIMI: Final = "moonshotai.kimi-k2.5" +_SONNET: Final = "global.anthropic.claude-sonnet-4-6" +_KIMI_BASE: Final = f"bedrock/us-east-1/{_KIMI}" +_KIMI_TARGET: Final = f"/model/{_KIMI}/converse" +_SONNET_TARGET: Final = f"/model/{_SONNET}/converse" +_PAST_ROUTE_ROW_LINES: Final = 16250 +_PAST_ESCALATION_BUFFER_LINES: Final = 15480 +_SMALL_LINES: Final = 1300 +_REPLY: Final = json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "8115"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 9, "outputTokens": 2, "totalTokens": 11}, + "metrics": {"latencyMs": 1}, + } +).encode() + + +@dataclass(frozen=True, slots=True) +class _Rig: + owned: OwnedProxy + wire: Wire + + +def _converse_peer(request: Request) -> Reply: + assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers + return Reply(body=_REPLY) + + +def _deployment( + name: str, + model: str, + api_base: str, + *, + model_info: Mapping[str, JsonValue] | None = None, + **params: JsonValue, +) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": f"bedrock/converse/{model}", + "api_key": _TOKEN, + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": api_base, + **params, + }, + "model_info": {"id": f"{name}-{uuid.uuid4().hex}", **(model_info or {})}, + } + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("base-model-window") + with gateway_from_environment() as gateway, wire_server(_converse_peer) as wire: + config: Final = owned_config( + directory, + [ + _deployment("base-in-model-info", _KIMI, wire.url, model_info={"base_model": _KIMI_BASE}), + _deployment("base-in-litellm-params", _KIMI, wire.url, base_model=_KIMI_BASE), + _deployment("no-base", _KIMI, wire.url), + _deployment("other-model-base", _KIMI, wire.url, model_info={"base_model": _SONNET}), + _deployment( + "declared-window", + _KIMI, + wire.url, + model_info={"base_model": _KIMI_BASE, "max_input_tokens": 257000}, + ), + _deployment( + "unmapped-base", _KIMI, wire.url, model_info={"base_model": f"{_KIMI_BASE}-not-in-cost-map"} + ), + _deployment( + "smaller-other-model-base", + _KIMI, + wire.url, + model_info={"base_model": "bedrock/us-east-1/meta.llama3-8b-instruct-v1:0"}, + ), + _deployment("one-million", _SONNET, wire.url), + { + "model_name": "complexity", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": { + "SIMPLE": "base-in-model-info", + "MEDIUM": "base-in-model-info", + "COMPLEX": "one-million", + "REASONING": "one-million", + }, + "enable_context_window_escalation": True, + }, + }, + }, + ], + litellm_settings={"cache": False}, + router_settings={"enable_pre_call_checks": True}, + ) + with owned_proxy_process(gateway, directory, {}, config=config, workers=2, client_timeout=120) as owned: + yield _Rig(owned, wire) + + +def _prompt(lines: int, marker: str) -> str: + listing: Final = "\n".join(f"Line {i}: the value of item {i} is {(i * 7) % 1000}." for i in range(1, lines + 1)) + return f"{listing}\n\nWhat is the value of item 12345? Answer with the number only. marker-{marker}" + + +def _request(surface: str, model: str, text: str) -> tuple[str, dict[str, JsonValue]]: + if surface == "messages": + return "/v1/messages", {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]} + if surface == "responses": + return "/v1/responses", {"model": model, "max_output_tokens": 16, "input": text} + return "/v1/chat/completions", {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]} + + +def _send(rig: _Rig, surface: str, model: str, lines: int) -> tuple[int, str, tuple[str, ...]]: + marker: Final = uuid.uuid4().hex + path, body = _request(surface, model, _prompt(lines, marker)) + response: Final = rig.owned.gateway.request("POST", path, body) + targets: Final = tuple(request.target for request in rig.wire.drain() if marker_of(request) == marker) + return response.status_code, response.text, targets + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +@pytest.mark.parametrize("surface", ["chat", "messages", "responses"]) +def test_a_same_model_base_model_row_widens_the_window_pre_call_checks_enforce(rig: _Rig, surface: str) -> None: + status, text, targets = _send(rig, surface, "base-in-model-info", _PAST_ROUTE_ROW_LINES) + assert status == 200, text + assert targets == (_KIMI_TARGET,) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_a_same_model_base_model_in_litellm_params_widens_the_window_too(rig: _Rig) -> None: + status, text, targets = _send(rig, "chat", "base-in-litellm-params", _PAST_ROUTE_ROW_LINES) + assert status == 200, text + assert targets == (_KIMI_TARGET,) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +@pytest.mark.parametrize("model", ["no-base", "other-model-base", "declared-window", "unmapped-base"]) +def test_a_base_model_that_cannot_widen_keeps_the_route_window_and_calls_no_upstream(rig: _Rig, model: str) -> None: + status, text, targets = _send(rig, "chat", model, _PAST_ROUTE_ROW_LINES) + assert status == 400, text + assert "ContextWindowExceededError" in text, text + assert targets == () + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_a_smaller_other_model_base_model_never_narrows_the_route_window(rig: _Rig) -> None: + status, text, targets = _send(rig, "chat", "smaller-other-model-base", _SMALL_LINES) + assert status == 200, text + assert targets == (_KIMI_TARGET,) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_complexity_router_keeps_a_prompt_on_the_tier_whose_base_model_row_holds_it(rig: _Rig) -> None: + status, text, targets = _send(rig, "chat", "complexity", _PAST_ESCALATION_BUFFER_LINES) + assert status == 200, text + assert targets == (_KIMI_TARGET,) diff --git a/tests/integration/routing/test_end_user_region_routing.py b/tests/integration/routing/test_end_user_region_routing.py index 16638d04bd8..d8ba2da042a 100644 --- a/tests/integration/routing/test_end_user_region_routing.py +++ b/tests/integration/routing/test_end_user_region_routing.py @@ -1,5 +1,5 @@ import uuid -from collections.abc import Iterator +from collections.abc import Callable, Iterator from pathlib import Path from typing import Final @@ -8,12 +8,17 @@ 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) MODEL: Final = "regional-model" UPSTREAM_BY_REGION: Final = {"eu": "regional-eu-upstream", "us": "regional-us-upstream"} CALLS: Final = 5 +UNCOUNTABLE_MODEL: Final = "regional-uncountable-model" +PINNED_REGION: Final = "eu" +WEIGHT_BY_REGION: Final = {"eu": 1, "us": 99} +SILENT_WAV: Final = "UklGRiQAAABXQVZFZm10IBAAAAABAAEAQB8AAIA+AAACABAAZGF0YQAAAAA=" @pytest.fixture(scope="module") @@ -22,16 +27,32 @@ def candidate(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: with gateway_from_environment() as base: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["model_list"] = [ - { - "model_name": MODEL, - "litellm_params": { - "model": f"openai/{upstream}", - "api_base": f"{base.upstream_url}/v1", - "api_key": "synthetic-region-key", - "region_name": region, - }, - } - for region, upstream in UPSTREAM_BY_REGION.items() + *( + { + "model_name": MODEL, + "litellm_params": { + "model": f"openai/{upstream}", + "api_base": f"{base.upstream_url}/v1", + "api_key": "synthetic-region-key", + "region_name": region, + }, + } + for region, upstream in UPSTREAM_BY_REGION.items() + ), + *( + { + "model_name": UNCOUNTABLE_MODEL, + "litellm_params": { + "model": f"openai/{upstream}", + "api_base": f"{base.upstream_url}/v1", + "api_key": "synthetic-region-key", + "region_name": region, + "weight": WEIGHT_BY_REGION[region], + }, + "model_info": {"id": f"{UNCOUNTABLE_MODEL}-{region}", "max_input_tokens": 128000}, + } + for region, upstream in UPSTREAM_BY_REGION.items() + ), ] config["router_settings"] = {**config["router_settings"], "enable_pre_call_checks": True} path: Final = directory / "region-routing.yaml" @@ -70,3 +91,55 @@ def test_an_end_users_allowed_region_pins_every_call_to_that_regions_deployment( assert [response.headers.get("x-litellm-model-region") for response in responses] == [region] * CALLS observed: Final = upstream.get("/__observations").json()["requests"] assert [request["body"]["model"] for request in observed] == [UPSTREAM_BY_REGION[region]] * CALLS + + +def _uncountable_chat(marker: str, end_user: str) -> tuple[str, dict[str, JsonValue]]: + return ( + "/v1/chat/completions", + { + "model": UNCOUNTABLE_MODEL, + "user": end_user, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": f"region {marker}"}, + {"type": "input_audio", "input_audio": {"data": SILENT_WAV, "format": "wav"}}, + ], + } + ], + }, + ) + + +def _uncountable_embedding(marker: str, end_user: str) -> tuple[str, dict[str, JsonValue]]: + return "/v1/embeddings", {"model": UNCOUNTABLE_MODEL, "user": end_user, "input": [f"region {marker}", marker]} + + +@pytest.mark.parametrize( + ("build", "recorded_chat_calls"), + [ + pytest.param(_uncountable_chat, CALLS, id="chat-input-audio"), + pytest.param(_uncountable_embedding, 0, id="embeddings-list-input"), + ], +) +def test_an_uncountable_prompt_still_pins_an_end_user_to_their_region( + candidate: Gateway, build: Callable[[str, str], tuple[str, dict[str, JsonValue]]], recorded_chat_calls: int +) -> None: + with ( + candidate.scenario() as scenario, + httpx.Client(base_url=candidate.upstream_url, timeout=5, trust_env=False) as upstream, + ): + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + candidate.post("/end_user/new", {"user_id": end_user, "allowed_model_region": PINNED_REGION}) + scenario.cleanups.callback(candidate.post, "/end_user/delete", {"user_ids": [end_user]}) + key: Final = scenario.key(models=[UNCOUNTABLE_MODEL]) + upstream.get("/__observations").raise_for_status() + requests: Final = tuple(build(uuid.uuid4().hex, end_user) for _ in range(CALLS)) + responses: Final = tuple(candidate.request("POST", path, body, key=key) for path, body in requests) + assert [response.status_code for response in responses] == [200] * CALLS, [r.text for r in responses] + pinned_deployment: Final = f"{UNCOUNTABLE_MODEL}-{PINNED_REGION}" + assert [response.headers.get("x-litellm-model-id") for response in responses] == [pinned_deployment] * CALLS + observed: Final = upstream.get("/__observations").json()["requests"] + pinned_upstream: Final = UPSTREAM_BY_REGION[PINNED_REGION] + assert [request["body"]["model"] for request in observed] == [pinned_upstream] * recorded_chat_calls 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_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py index 163f5983c3b..89a9d1c14f0 100644 --- a/tests/integration/routing/test_non_chat_model_endpoint_gate.py +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -1387,14 +1387,16 @@ def test_responses_websocket_uses_a_chat_fallback_without_calling_the_embedding_ owned_proxy_process(rig.gateway, tmp_path, {"SSL_CERT_FILE": str(certificate[0])}, workers=2) as owned, owned.gateway.scenario() as scenario, ): - fallback: Final = scenario.model(model=f"openai/{_CHAT_MODEL}", api_base=peer.url, api_key=_PROVIDER_KEY) - primary: Final = scenario.model( - model=f"openai/{_EMBEDDING_MODEL}", api_base=f"{rig.api_base}/{marker}", model_info=_EMBED_ONLY + fallback, fallback_ids = _deploy(scenario, peer.url, "ws-chat", _single({"mode": "chat"})) + primary, primary_ids = _deploy( + scenario, f"{rig.api_base}/{marker}", "ws-embed", _single(_EMBED_ONLY, f"openai/{_EMBEDDING_MODEL}") ) + _served_by_every_worker(owned.gateway, frozenset((*primary_ids, *fallback_ids))) key: Final = scenario.key(router_settings={"fallbacks": [{primary: [fallback]}]}) (received,) = asyncio.run( responses_ws_peer.session(str(owned.gateway.client.base_url), key, primary, (marker,)) ) + assert received[-1]["type"] == "response.completed", received assert responses_ws_peer.completed(received)["status"] == "completed" assert any(event.get("delta") == _SERVED_TEXT for event in received), received (forwarded,) = responses_ws_peer.drain(peer.frames) diff --git a/tests/integration/routing/test_pre_call_response_format_filter.py b/tests/integration/routing/test_pre_call_response_format_filter.py new file mode 100644 index 00000000000..24a4434fb5a --- /dev/null +++ b/tests/integration/routing/test_pre_call_response_format_filter.py @@ -0,0 +1,161 @@ +import uuid +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Final + +import httpx +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) + +LACKS_JSON_MODE: Final = "gpt-4" +TAKES_JSON_MODE: Final = "gpt-4o-mini" +UNSUPPORTED: Final = "json-mode-unsupported" +DROPPING: Final = "json-mode-dropped" +MIXED: Final = "json-mode-mixed" +FALLBACK_SOURCE: Final = "json-mode-fallback-source" +SMALL_WINDOW: Final = "json-mode-small-window" +FALLBACK_TARGET: Final = "json-mode-fallback-target" +JSON_MODE: Final[Mapping[str, JsonValue]] = {"type": "json_object"} +CALLS: Final = 4 +SMALL_WINDOW_TOKENS: Final = 50 + + +def _deployment(group: str, upstream_model: str, upstream_url: str, **params: JsonValue) -> dict[str, JsonValue]: + return { + "model_name": group, + "litellm_params": { + "model": f"openai/{upstream_model}", + "api_base": f"{upstream_url}/v1", + "api_key": "synthetic-json-mode-key", + **params, + }, + } + + +@pytest.fixture(scope="module") +def candidate(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("json-mode-filter") + with gateway_from_environment() as base: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + url: Final = base.upstream_url + config["model_list"] = [ + _deployment(UNSUPPORTED, LACKS_JSON_MODE, url), + _deployment(DROPPING, LACKS_JSON_MODE, url, drop_params=True), + _deployment(MIXED, LACKS_JSON_MODE, url), + _deployment(MIXED, TAKES_JSON_MODE, url), + _deployment(FALLBACK_SOURCE, LACKS_JSON_MODE, url), + { + **_deployment(SMALL_WINDOW, LACKS_JSON_MODE, url), + "model_info": {"max_input_tokens": SMALL_WINDOW_TOKENS}, + }, + _deployment(FALLBACK_TARGET, TAKES_JSON_MODE, url), + ] + config["router_settings"] = { + **config["router_settings"], + "enable_pre_call_checks": True, + "fallbacks": [{FALLBACK_SOURCE: [FALLBACK_TARGET]}], + "context_window_fallbacks": [{SMALL_WINDOW: [FALLBACK_TARGET]}], + } + path: Final = directory / "json-mode-filter.yaml" + path.write_text(yaml.safe_dump(config)) + with owned_proxy(base, directory, {}, config=path) as proxy: + yield proxy + + +def _json_mode_chat(group: str, text: str, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model": group, + "response_format": dict(JSON_MODE), + "messages": [{"role": "user", "content": f"Answer in JSON. {text}"}], + **extra, + } + + +def _drain(upstream: httpx.Client) -> list[dict[str, JsonValue]]: + response: Final = upstream.get("/__observations") + response.raise_for_status() + return [request["body"] for request in response.json()["requests"]] + + +@pytest.fixture +def upstream(candidate: Gateway) -> Iterator[httpx.Client]: + with httpx.Client(base_url=candidate.upstream_url, timeout=5, trust_env=False) as client: + _drain(client) + yield client + + +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +def test_json_mode_answers_400_without_an_upstream_call_when_no_deployment_takes_it( + candidate: Gateway, upstream: httpx.Client, stream: bool +) -> None: + response: Final = candidate.request( + "POST", "/v1/chat/completions", _json_mode_chat(UNSUPPORTED, uuid.uuid4().hex, stream=stream) + ) + + assert response.status_code == 400, response.text + assert _drain(upstream) == [] + + +@pytest.mark.parametrize( + ("group", "extra", "forwarded_format"), + [ + pytest.param(DROPPING, {}, None, id="deployment-drop-params"), + pytest.param(UNSUPPORTED, {"allowed_openai_params": ["response_format"]}, JSON_MODE, id="request-allowlist"), + pytest.param(UNSUPPORTED, {"drop_params": True}, None, id="request-drop-params"), + ], +) +def test_json_mode_reaches_a_deployment_that_lacks_it_once_drop_or_allow_is_set( + candidate: Gateway, + upstream: httpx.Client, + group: str, + extra: Mapping[str, JsonValue], + forwarded_format: Mapping[str, JsonValue] | None, +) -> None: + response: Final = candidate.request( + "POST", "/v1/chat/completions", _json_mode_chat(group, uuid.uuid4().hex, **extra) + ) + + assert response.status_code == 200, response.text + received: Final = _drain(upstream) + assert [body["model"] for body in received] == [LACKS_JSON_MODE] + assert received[0].get("response_format") == forwarded_format + + +def test_json_mode_in_a_mixed_group_always_reaches_the_deployment_that_takes_it( + candidate: Gateway, upstream: httpx.Client +) -> None: + responses: Final = tuple( + candidate.request("POST", "/v1/chat/completions", _json_mode_chat(MIXED, uuid.uuid4().hex)) + for _ in range(CALLS) + ) + + assert [response.status_code for response in responses] == [200] * CALLS, [r.text for r in responses] + received: Final = _drain(upstream) + assert [(body["model"], body.get("response_format")) for body in received] == [ + (TAKES_JSON_MODE, dict(JSON_MODE)) + ] * CALLS + + +@pytest.mark.parametrize( + ("group", "text"), + [ + pytest.param(FALLBACK_SOURCE, "short", id="generic-fallback"), + pytest.param(SMALL_WINDOW, "padding " * (4 * SMALL_WINDOW_TOKENS), id="context-window-fallback"), + ], +) +def test_json_mode_falls_back_to_a_group_that_takes_it( + candidate: Gateway, upstream: httpx.Client, group: str, text: str +) -> None: + response: Final = candidate.request( + "POST", "/v1/chat/completions", _json_mode_chat(group, f"{uuid.uuid4().hex} {text}") + ) + + assert response.status_code == 200, response.text + assert [(body["model"], body.get("response_format")) for body in _drain(upstream)] == [ + (TAKES_JSON_MODE, dict(JSON_MODE)) + ] 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_background_interaction_settlement.py b/tests/integration/spend/test_background_interaction_settlement.py index 16eb0a3434e..4408bd147f9 100644 --- a/tests/integration/spend/test_background_interaction_settlement.py +++ b/tests/integration/spend/test_background_interaction_settlement.py @@ -1,3 +1,4 @@ +import json import math import socket import time @@ -17,6 +18,8 @@ from integration._support.client import ( Gateway, eventually, gateway_from_environment, + held, + list_value, object_value, string_value, ) @@ -54,6 +57,9 @@ _CUSTOM_INPUT_RATE: Final = 2e-06 _CUSTOM_OUTPUT_RATE: Final = 4e-05 _ENV_KEY: Final = "integration-gemini-env-key" _DEPLOYMENT_KEY: Final = "integration-gemini-deployment-key" +_DATABASE_DEPLOYMENT_KEY: Final = "integration-gemini-database-deployment-key" +_ROTATED_DEPLOYMENT_KEY: Final = "integration-gemini-rotated-deployment-key" +_TAG: Final = "integration-settlement-team" _CREATOR_POLL: Final = {"BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS": "300"} _SETTLER_POLL: Final = { "BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS": "1", @@ -65,8 +71,10 @@ _RESUMER_POLL: Final = { "BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS": "1", "BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS": "120", } +_EAGER_RESUMER_POLL: Final = {**_RESUMER_POLL, "BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS": "0.01"} +_LATE_RESUMER_POLL: Final = {**_RESUMER_POLL, "BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS": "30"} _SPEND_QUERY: Final = ( - "SELECT request_id, spend, call_type, status, model, prompt_tokens, completion_tokens " + "SELECT request_id, spend, call_type, status, model, prompt_tokens, completion_tokens, request_tags " 'FROM "LiteLLM_SpendLogs" WHERE request_id = %s' ) _SETTLEMENT_QUERY: Final = ( @@ -74,6 +82,16 @@ _SETTLEMENT_QUERY: Final = ( 'settled_at IS NOT NULL AS settled, create_context FROM "LiteLLM_BackgroundInteractionSettlement" ' "WHERE interaction_id = %s" ) +_ABANDONED_CLAIM: Final = ( + 'UPDATE "LiteLLM_BackgroundInteractionSettlement" SET claimed_by = %s, ' + "claimed_at = (now() AT TIME ZONE 'UTC') - interval '10 minutes', settled_at = NULL, outcome = NULL, " + "create_context = %s::jsonb WHERE interaction_id = %s" +) +_DEAD_SETTLER: Final = "dead-settler:1" +_CREATED_AN_HOUR_AGO: Final = ( + 'UPDATE "LiteLLM_BackgroundInteractionSettlement" ' + "SET created_at = (now() AT TIME ZONE 'UTC') - interval '1 hour' WHERE interaction_id = %s" +) _SETTLEMENT_TABLE_PRESENT_QUERY: Final = "SELECT to_regclass(%s) IS NOT NULL AS present" _SETTLEMENT_TABLE: Final = '"LiteLLM_BackgroundInteractionSettlement"' _SETTLEMENT_BY_CALL_QUERY: Final = ( @@ -179,6 +197,40 @@ def _write_config(directory: Path, upstream_url: str, models: Deployments) -> Pa return path +def _variant_config(rig: Rig, name: str, document: Mapping[str, JsonValue]) -> Path: + path: Final = rig.directory / name + path.write_text(yaml.safe_dump(dict(document))) + return path + + +def _rig_config(rig: Rig) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_python(yaml.safe_load(rig.config.read_text())) + + +def _config_without_spend_logs(rig: Rig) -> Path: + config: Final = _rig_config(rig) + general_settings: Final = {**object_value(config["general_settings"]), "disable_spend_logs": True} + return _variant_config( + rig, "settlement_config_without_spend_logs.yaml", {**config, "general_settings": general_settings} + ) + + +def _config_with_a_rotated_key(rig: Rig) -> Path: + config: Final = _rig_config(rig) + model_list: Final = [_with_rotated_key(entry, rig.models.in_progress) for entry in list_value(config["model_list"])] + return _variant_config(rig, "settlement_config_rotated_key.yaml", {**config, "model_list": model_list}) + + +def _with_rotated_key(entry: JsonValue, model_name: str) -> JsonValue: + deployment: Final = object_value(entry) + if deployment["model_name"] != model_name: + return deployment + return { + **deployment, + "litellm_params": {**object_value(deployment["litellm_params"]), "api_key": _ROTATED_DEPLOYMENT_KEY}, + } + + def _interaction(status: str, usage: dict[str, JsonValue] | None, http_status: int = 200) -> JsonResponse: return JsonResponse( content_type="application/json", @@ -206,12 +258,14 @@ def _create( path: str = "/v1beta/interactions", background: bool = True, text: str | None = None, + headers: Mapping[str, str] | None = None, ) -> str: response: Final = replica.request( "POST", path, {"model": model, "input": text or f"settle {uuid.uuid4().hex}", "background": background}, key=key, + headers=headers, ) assert response.status_code == 200, response.text return string_value(JSON_OBJECT.validate_json(response.content)["id"]) @@ -252,6 +306,16 @@ def _settlement(interaction_id: str) -> dict[str, JsonValue] | None: return rows[0] if rows else None +def _registered(interaction_id: str) -> dict[str, JsonValue]: + row: Final = eventually(lambda: _settlement(interaction_id), lambda value: value is not None) + assert row is not None + return row + + +def _abandon_claim(interaction_id: str, create_context: JsonValue) -> None: + write_rows(_ABANDONED_CLAIM, (_DEAD_SETTLER, json.dumps(create_context), interaction_id)) + + def _settlement_table_present() -> bool: return read_rows(_SETTLEMENT_TABLE_PRESENT_QUERY, (_SETTLEMENT_TABLE,))[0]["present"] is True @@ -315,6 +379,11 @@ def _assert_billed(row: Mapping[str, JsonValue], rates: tuple[float, float]) -> return spend +def _assert_billed_with_tag(row: Mapping[str, JsonValue], rates: tuple[float, float]) -> None: + _assert_billed(row, rates) + assert _TAG in list_value(row["request_tags"]), row + + def _key_spend(replica: Gateway, key: str) -> float: spend: Final = object_value(replica.get("/key/info", {"key": key})["info"])["spend"] assert isinstance(spend, float | int), spend @@ -333,12 +402,12 @@ def _drain(rig: Rig) -> list[JsonValue]: return requests -def _calls(rig: Rig, interaction_id: str) -> tuple[tuple[str, str], ...]: - suffix: Final = f"/v1beta/interactions/{interaction_id}" +def _calls(rig: Rig, *interaction_ids: str) -> tuple[tuple[str, str], ...]: + suffixes: Final = tuple(f"/v1beta/interactions/{interaction_id}" for interaction_id in interaction_ids) return tuple( (string_value(object_value(entry)["method"]), string_value(object_value(entry)["api_key"])) for entry in _drain(rig) - if string_value(object_value(entry)["path"]).endswith(suffix) + if string_value(object_value(entry)["path"]).endswith(suffixes) ) @@ -447,7 +516,7 @@ def test_a_replica_booting_later_resumes_and_bills_unclaimed_interactions(rig: R model: Final = rig.models.in_progress key: Final = scenario.key() created_after: Final = time.time() - created: Final = tuple(_create(rig.creator, model, key) for _ in range(3)) + created: Final = tuple(_create(rig.creator, model, key, headers={"x-litellm-tags": _TAG}) for _ in range(3)) for item in created: _state(rig, item, _completed()) rates: Final = _rates(rig.creator, model) @@ -457,13 +526,113 @@ def test_a_replica_booting_later_resumes_and_bills_unclaimed_interactions(rig: R pids: Final = _worker_pids(resumer.process.pid) assert len(pids) == 2, pids for item in created: - _assert_billed(_await_spend_row(item, seconds=90), rates) + _assert_billed_with_tag(_await_spend_row(item, seconds=90), rates) claimer: Final = _claimer_pid(_await_outcome(item, "billed")) assert claimer in pids or _booted_after(claimer, created_after), (claimer, pids) + assert {key for _, key in _calls(rig, *created)} == {_DEPLOYMENT_KEY} for item in created: assert len(_spend_rows(item)) == 1 +def test_a_replica_booting_later_resumes_a_database_model_with_its_own_key(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = scenario.model( + model=_MODEL, api_key=_DATABASE_DEPLOYMENT_KEY, api_base=f"{rig.upstream.url}/{rig.models.in_progress}" + ) + created: Final = _create(rig.creator, model, scenario.key()) + _state(rig, created, _completed()) + with owned_proxy_process( + rig.gateway, rig.directory, rig.environment(**_EAGER_RESUMER_POLL), config=rig.config, workers=2 + ): + _await_spend_row(created, seconds=90) + _await_outcome(created, "billed") + assert {key for _, key in _calls(rig, created)} == {_DATABASE_DEPLOYMENT_KEY} + + +def test_a_replica_booting_later_takes_over_lapsed_claims_and_bills_only_the_unbilled(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = rig.models.in_progress + key: Final = scenario.key() + unbilled: Final = _create(rig.creator, model, key) + billed: Final = _create(rig.settler, model, key) + contexts: Final = {item: _registered(item)["create_context"] for item in (unbilled, billed)} + for item in (unbilled, billed): + _state(rig, item, _completed()) + rates: Final = _rates(rig.creator, model) + spend: Final = _assert_billed(_await_spend_row(billed), rates) + _await_outcome(billed, "billed") + for item, create_context in contexts.items(): + _abandon_claim(item, create_context) + with owned_proxy_process( + rig.gateway, rig.directory, rig.environment(**_EAGER_RESUMER_POLL), config=rig.config, workers=2 + ): + _assert_billed(_await_spend_row(unbilled, seconds=90), rates) + assert _await_outcome(unbilled, "billed")["claimed_by"] != _DEAD_SETTLER + assert _await_outcome(billed, "billed", seconds=90)["claimed_by"] != _DEAD_SETTLER + held( + lambda: _key_spend(rig.creator, key), + lambda observed: math.isclose(observed, 2 * spend, rel_tol=1e-9), + holding=5, + seconds=30, + ) + for item in (unbilled, billed): + assert len(_spend_rows(item)) == 1 + + +def test_a_lapsed_claim_with_spend_logs_off_is_recorded_unsettled_and_never_billed(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = rig.models.in_progress + key: Final = scenario.key() + created: Final = _create(rig.creator, model, key) + _abandon_claim(created, _registered(created)["create_context"]) + _state(rig, created, _completed()) + with owned_proxy_process( + rig.gateway, + rig.directory, + rig.environment(**_EAGER_RESUMER_POLL), + config=_config_without_spend_logs(rig), + workers=2, + ): + assert _await_outcome(created, "unsettled", seconds=90)["claimed_by"] != _DEAD_SETTLER + assert _spend_rows(created) == [] + assert _key_spend(rig.creator, key) == 0 + + +def test_a_resumed_interaction_whose_deployment_key_rotated_fetches_with_the_new_key(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = rig.models.in_progress + created: Final = _create(rig.creator, model, scenario.key()) + _state(rig, created, _completed()) + with owned_proxy_process( + rig.gateway, + rig.directory, + rig.environment(**_EAGER_RESUMER_POLL), + config=_config_with_a_rotated_key(rig), + workers=2, + ) as resumer: + assert _model_info(resumer.gateway, model)["id"] != _model_info(rig.creator, model)["id"] + _assert_billed(_await_spend_row(created, seconds=90), _rates(rig.creator, model)) + _await_outcome(created, "billed") + assert {key for _, key in _calls(rig, created)} == {_ROTATED_DEPLOYMENT_KEY} + + +def test_a_late_restart_retries_a_failed_fetch_within_a_full_poll_window(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = rig.models.in_progress + created: Final = _create(rig.creator, model, scenario.key()) + _registered(created) + write_rows(_CREATED_AN_HOUR_AGO, (created,)) + _state(rig, created, InteractionState(status="completed", usage=_USAGE, get_status=503)) + _drain(rig) + with owned_proxy_process( + rig.gateway, rig.directory, rig.environment(**_LATE_RESUMER_POLL), config=rig.config, workers=2 + ): + eventually(lambda: _calls(rig, created), lambda calls: calls != (), seconds=90) + _state(rig, created, _completed()) + _assert_billed(_await_spend_row(created), _rates(rig.creator, model)) + _await_outcome(created, "billed") + + def test_deletes_on_the_creating_proxy_bill_each_interaction_once(rig: Rig) -> None: with rig.settler.scenario() as scenario: model: Final = rig.models.in_progress 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/spend/test_spend_log_caller_session.py b/tests/integration/spend/test_spend_log_caller_session.py new file mode 100644 index 00000000000..213583d21a1 --- /dev/null +++ b/tests/integration/spend/test_spend_log_caller_session.py @@ -0,0 +1,783 @@ +import itertools +import json +import re +import threading +import uuid +from urllib.parse import urlsplit +from collections import deque +from collections.abc import Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from dataclasses import dataclass, replace +from datetime import datetime, timedelta, timezone +from hashlib import sha256 +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import openai +import psutil +import pytest +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + list_value, + object_value, + string_value, +) +from integration._support.budget import readiness_seconds +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * readiness_seconds(2) + 2 * graceful_stop_seconds() + 240) + +MARKER: Final = re.compile(rb"callersess-[0-9a-f]{32}") +CALLER_SESSION_KEY: Final = "litellm_caller_session_id" +PATHS: Final = MappingProxyType( + {"chat": "/v1/chat/completions", "messages": "/v1/messages", "responses": "/v1/responses"} +) +ENDPOINT_STREAMS: Final = tuple(itertools.product(PATHS, (False, True))) +Headers = Mapping[str, str] | Sequence[tuple[str, str]] + + +def _marker() -> str: + return "callersess-" + uuid.uuid4().hex + + +def _token(label: str) -> str: + return f"{label}-{uuid.uuid4().hex}" + + +def _frame(payload: Mapping[str, JsonValue]) -> bytes: + return b"data: " + json.dumps(payload).encode() + b"\n\n" + + +def _chat_reply(text: str, stream: bool) -> Reply: + identity: Final = "chatcmpl-" + uuid.uuid4().hex + usage: Final[dict[str, JsonValue]] = {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9} + if not stream: + return Reply( + body=json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"} + ], + "usage": usage, + } + ).encode() + ) + chunk: Final[dict[str, JsonValue]] = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + } + return Reply( + content_type="text/event-stream", + chunks=( + _frame({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": text}}]}), + _frame({**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": usage}), + b"data: [DONE]\n\n", + ), + ) + + +def _upstream(request: Request) -> Reply: + markers: Final = tuple(found.group(0).decode() for found in MARKER.finditer(request.body)) + if request.method != "POST" or not request.target.endswith("/chat/completions") or not markers: + return Reply(status=404, body=b'{"error":"unexpected request"}') + return _chat_reply(f"answer {markers[-1]}", JSON_OBJECT.validate_json(request.body).get("stream") is True) + + +def _user_markers(body: Mapping[str, JsonValue]) -> tuple[str, ...]: + user_contents: Final = tuple( + json.dumps(object_value(message)["content"]) + for message in list_value(body.get("messages", [])) + if object_value(message)["role"] == "user" + ) + return tuple(found.group(0).decode() for content in user_contents if (found := MARKER.search(content.encode()))) + + +def _spans(batch: Request) -> Iterator[dict[str, JsonValue]]: + for resource in list_value(JSON_OBJECT.validate_json(batch.body)["resourceSpans"]): + for scope in list_value(object_value(resource)["scopeSpans"]): + yield from (object_value(span) for span in list_value(object_value(scope)["spans"])) + + +def _attributes(span: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + string_value(object_value(attribute)["key"]): object_value(attribute)["value"] + for attribute in list_value(span.get("attributes", [])) + } + + +def _batch_attributes(batches: Sequence[Request]) -> Iterator[dict[str, JsonValue]]: + return (_attributes(span) for span in itertools.chain.from_iterable(_spans(batch) for batch in batches)) + + +@dataclass(frozen=True, slots=True) +class Collector: + wire: Wire + accepted: deque[Request] + guard: threading.Lock + + def span_attributes(self, response_id: str) -> tuple[dict[str, JsonValue], ...]: + with self.guard: + batches: Final = tuple(self.accepted) + attributes: Final = tuple(_batch_attributes(batches)) + return tuple( + found + for found in attributes + if object_value(found.get("gen_ai.response.id", {})).get("stringValue") == response_id + ) + + def conversation_ids(self, conversations: frozenset[str]) -> tuple[str, ...]: + with self.guard: + batches: Final = tuple(self.accepted) + responded: Final = ( + object_value(found.get("gen_ai.conversation.id", {})).get("stringValue") + for found in _batch_attributes(batches) + if "gen_ai.response.id" in found + ) + return tuple( + conversation + for conversation in responded + if isinstance(conversation, str) and conversation in conversations + ) + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + scenario: Scenario + model: str + upstream: Wire + process: OwnedProxy + + def body(self, endpoint: str, marker: str, extra: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + match endpoint: + case "chat": + return {"model": self.model, "messages": [{"role": "user", "content": marker}], **extra} + case "messages": + return { + "model": self.model, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + **extra, + } + case _: + return {"model": self.model, "input": marker, **extra} + + def post( + self, + endpoint: str, + marker: str, + fields: Mapping[str, JsonValue] = MappingProxyType({}), + *, + stream: bool = False, + headers: Headers = (), + key: str | None = None, + ) -> httpx.Response: + payload: Final = self.body(endpoint, marker, fields) | ({"stream": True} if stream else {}) + request_headers: Final = ( + ("Authorization", f"Bearer {self.proxy.key if key is None else key}"), + *(headers.items() if isinstance(headers, Mapping) else headers), + ) + with self.proxy.client.stream("POST", PATHS[endpoint], json=payload, headers=request_headers) as response: + response.read() + return response + + def send( + self, + endpoint: str, + marker: str, + fields: Mapping[str, JsonValue] = MappingProxyType({}), + *, + stream: bool = False, + headers: Headers = (), + key: str | None = None, + ) -> httpx.Response: + response: Final = self.post(endpoint, marker, fields, stream=stream, headers=headers, key=key) + assert response.status_code == 200, response.text + return response + + def upstream_bodies(self, marker: str) -> tuple[dict[str, JsonValue], ...]: + bodies: Final = tuple(JSON_OBJECT.validate_json(request.body) for request in self.upstream.drain()) + return tuple(body for body in bodies if _user_markers(body)[-1:] == (marker,)) + + +@dataclass(frozen=True, slots=True) +class Rigs: + flag_off: Rig + flag_on: Rig + sink: Collector + + +def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT session_id, request_id, metadata::text AS metadata_text FROM "LiteLLM_SpendLogs" ' + "WHERE litellm_call_id=%s", + (call_id,), + ) + + +def _spend_row(response: httpx.Response) -> dict[str, JsonValue]: + call_id: Final = response.headers["x-litellm-call-id"] + rows: Final = eventually(lambda: _spend_rows(call_id), lambda values: len(values) >= 1, seconds=70) + assert len(rows) == 1, rows + return rows[0] + + +def _session(response: httpx.Response) -> JsonValue: + return _spend_row(response)["session_id"] + + +def _model(scenario: Scenario, upstream: Wire) -> str: + return scenario.model( + model="hosted_vllm/gpt-4o-mini", + api_base=f"{upstream.url}/v1", + use_chat_completions_api=True, + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + + +def _rig_config( + directory: Path, name: str, litellm_settings: Mapping[str, JsonValue], extra: Mapping[str, JsonValue] +) -> Path: + shared: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "general_settings": {**object_value(shared["general_settings"]), "store_prompts_in_spend_logs": True}, + "litellm_settings": {**object_value(shared["litellm_settings"]), **litellm_settings}, + **extra, + } + path: Final = directory / f"{name}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def rigs(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rigs]: + accepted: Final[deque[Request]] = deque() # mutable-ok: sink thread appends each accepted OTLP batch + guard: Final = threading.Lock() + + def sink(request: Request) -> Reply: + with guard: + accepted.append(request) + return Reply() + + directory: Final = tmp_path_factory.mktemp("caller-session") + with ExitStack() as stack: + upstream: Final = stack.enter_context(wire_server(_upstream)) + collector: Final = Collector(stack.enter_context(wire_server(sink)), accepted, guard) + shared: Final = stack.enter_context(gateway_from_environment()) + flag_off: Final = stack.enter_context( + owned_proxy_process(shared, directory, {}, config=_rig_config(directory, "flag-off", {}, {}), workers=2) + ) + flag_on: Final = stack.enter_context( + owned_proxy_process( + shared, + directory, + {"LITELLM_OTEL_V2": "1", "OTEL_BSP_SCHEDULE_DELAY": "300"}, + config=_rig_config( + directory, + "flag-on", + {"request_correlation_in_logs": True, "callbacks": ["otel"]}, + { + "callback_settings": { + "otel": {"exporter": "http/json", "endpoint": collector.wire.url, "mapper_names": ["genai"]} + } + }, + ), + workers=2, + ) + ) + off_scenario: Final = stack.enter_context(flag_off.gateway.scenario()) + on_scenario: Final = stack.enter_context(flag_on.gateway.scenario()) + yield Rigs( + Rig(flag_off.gateway, off_scenario, _model(off_scenario, upstream), upstream, flag_off), + Rig(flag_on.gateway, on_scenario, _model(on_scenario, upstream), upstream, flag_on), + collector, + ) + + +@pytest.mark.parametrize(("endpoint", "stream"), ENDPOINT_STREAMS) +def test_spend_log_session_is_the_body_session_beside_a_trace_header(rigs: Rigs, endpoint: str, stream: bool) -> None: + marker: Final = _marker() + session: Final = _token("convo") + response: Final = rigs.flag_on.send( + endpoint, + marker, + stream=stream, + headers={"x-litellm-trace-id": _token("trace")}, + fields={"litellm_session_id": session}, + ) + + row: Final = _spend_row(response) + upstream: Final = rigs.flag_on.upstream_bodies(marker) + + assert row["session_id"] == session, row + assert CALLER_SESSION_KEY not in string_value(row["metadata_text"]), row + assert len(upstream) == 1, upstream + assert CALLER_SESSION_KEY not in json.dumps(upstream[0]), upstream + + +def test_sdk_clients_record_the_body_session_beside_a_trace_header(rigs: Rigs) -> None: + rig: Final = rigs.flag_on + base_url: Final = str(rig.proxy.client.base_url) + traces: Final = tuple(_token("trace") for _ in range(3)) + sessions: Final = tuple(_token("convo") for _ in range(3)) + markers: Final = tuple(_marker() for _ in range(3)) + chat_client: Final = openai.OpenAI(base_url=base_url + "/v1", api_key=rig.proxy.key, max_retries=0) + messages_client: Final = anthropic.Anthropic(base_url=base_url, api_key=rig.proxy.key, max_retries=0) + + with chat_client.chat.completions.with_streaming_response.create( + model=rig.model, + messages=[{"role": "user", "content": markers[0]}], + stream=True, + extra_headers={"x-litellm-trace-id": traces[0]}, + extra_body={"litellm_session_id": sessions[0]}, + ) as chat_stream: + chat_call: Final = chat_stream.headers["x-litellm-call-id"] + chat_text: Final = "".join( + chunk.choices[0].delta.content or "" for chunk in chat_stream.parse() if chunk.choices + ) + message: Final = messages_client.messages.with_raw_response.create( + model=rig.model, + max_tokens=16, + messages=[{"role": "user", "content": markers[1]}], + extra_headers={"x-litellm-trace-id": traces[1]}, + extra_body={"litellm_session_id": sessions[1]}, + ) + with chat_client.responses.with_streaming_response.create( + model=rig.model, + input=markers[2], + stream=True, + extra_headers={"x-litellm-trace-id": traces[2]}, + extra_body={"litellm_session_id": sessions[2]}, + ) as responses_stream: + responses_call: Final = responses_stream.headers["x-litellm-call-id"] + completed: Final = tuple(event for event in responses_stream.parse() if event.type == "response.completed") + + first_block: Final = message.parse().content[0] + assert chat_text == f"answer {markers[0]}", chat_text + assert isinstance(first_block, anthropic.types.TextBlock) and first_block.text == f"answer {markers[1]}", message + assert len(completed) == 1, completed + calls: Final = (chat_call, message.headers["x-litellm-call-id"], responses_call) + landed: Final = tuple( + eventually(lambda call=call: _spend_rows(call), lambda values: len(values) == 1, seconds=70)[0]["session_id"] + for call in calls + ) + assert landed == sessions, (landed, calls) + + +def test_conversation_span_keeps_the_trace_header_and_no_internal_key(rigs: Rigs) -> None: + trace: Final = _token("trace") + response: Final = rigs.flag_on.send( + "chat", _marker(), headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": _token("convo")} + ) + response_id: Final = string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + spans: Final = eventually(lambda: rigs.sink.span_attributes(response_id), lambda found: len(found) == 1, seconds=30) + + assert spans[0]["gen_ai.conversation.id"] == {"stringValue": trace}, spans + assert CALLER_SESSION_KEY not in json.dumps(spans[0]), spans + + +def test_a_continued_chains_conversation_spans_keep_each_turns_trace_header(rigs: Rigs) -> None: + traces: Final = (_token("trace"), _token("trace")) + first: Final = _turn(rigs.flag_on, _marker(), headers={"x-litellm-trace-id": traces[0]}) + _turn( + rigs.flag_on, + _marker(), + first, + headers={"x-litellm-trace-id": traces[1]}, + fields={"litellm_session_id": _token("convo")}, + ) + + conversations: Final = eventually( + lambda: rigs.sink.conversation_ids(frozenset(traces)), lambda found: len(found) == 2, seconds=30 + ) + + assert sorted(conversations) == sorted(traces), conversations + + +@pytest.mark.parametrize("endpoint", list(PATHS)) +def test_flag_off_keeps_the_trace_id_beside_a_body_session(rigs: Rigs, endpoint: str) -> None: + trace: Final = _token("trace") + + response: Final = rigs.flag_off.send( + endpoint, _marker(), headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": _token("convo")} + ) + + assert _session(response) == trace + + +@pytest.mark.parametrize("endpoint", list(PATHS)) +def test_response_cache_hit_records_the_body_session_beside_a_trace_header(rigs: Rigs, endpoint: str) -> None: + marker: Final = _marker() + sessions: Final = (_token("convo"), _token("convo")) + responses: Final = tuple( + rigs.flag_on.send( + endpoint, marker, headers={"x-litellm-trace-id": _token("trace")}, fields={"litellm_session_id": session} + ) + for session in sessions + ) + + assert tuple(_session(response) for response in responses) == sessions + assert len(rigs.flag_on.upstream_bodies(marker)) == 1 + + +def test_metadata_session_beside_a_new_trace_per_turn_groups_both_turns(rigs: Rigs) -> None: + session: Final = _token("convo") + responses: Final = tuple( + rigs.flag_on.send( + "chat", + _marker(), + headers={"x-litellm-trace-id": _token("trace")}, + fields={"metadata": {"session_id": session}}, + ) + for _ in range(2) + ) + + assert tuple(_session(response) for response in responses) == (session, session) + now: Final = datetime.now(timezone.utc) + page: Final = rigs.flag_on.proxy.get( + "/spend/logs/ui", + params={ + "session_id": session, + "start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + "end_date": (now + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + }, + ) + listed: Final = sorted(string_value(object_value(row)["request_id"]) for row in list_value(page["data"])) + assert listed == sorted( + string_value(JSON_OBJECT.validate_json(response.content)["id"]) for response in responses + ), page + + +def _shape(shape: str, trace: str, session: str, header: str) -> tuple[Headers, dict[str, JsonValue], str]: + traced: Final = {"x-litellm-trace-id": trace} + match shape: + case "session-header-wins": + return {"x-litellm-session-id": header}, {"litellm_session_id": session}, header + case "trace-header-only": + return traced, {}, trace + case "body-session-only": + return {}, {"litellm_session_id": session}, session + case "int-session": + return traced, {"litellm_session_id": 7}, trace + case "list-session": + return traced, {"litellm_session_id": [session]}, trace + case "empty-session": + return traced, {"litellm_session_id": ""}, trace + case "null-session": + return traced, {"litellm_session_id": None}, trace + case "nul-session": + return traced, {"litellm_session_id": session + "\x00x"}, trace + case "two-trace-headers": + return ( + (("x-litellm-trace-id", trace), ("x-litellm-trace-id", header)), + {"litellm_session_id": session}, + session, + ) + case _: + return traced, {"litellm_session_id": session + "s" * 5120}, session + "s" * 5120 + + +@pytest.mark.parametrize( + "shape", + [ + "session-header-wins", + "trace-header-only", + "body-session-only", + "int-session", + "list-session", + "empty-session", + "null-session", + "nul-session", + "two-trace-headers", + "5kb-session", + ], +) +def test_spend_log_session_for_each_session_shape(rigs: Rigs, shape: str) -> None: + headers, fields, expected = _shape(shape, _token("trace"), _token("convo"), _token("header")) + neighbor_trace: Final = _token("neighbor") + + response: Final = rigs.flag_on.send("chat", _marker(), fields, headers=headers) + neighbor: Final = rigs.flag_on.send("chat", _marker(), headers={"x-litellm-trace-id": neighbor_trace}) + + assert _session(response) == expected + assert _session(neighbor) == neighbor_trace + + +@pytest.mark.parametrize("endpoint", ["chat", "messages"]) +def test_client_sent_internal_session_key_is_ignored(rigs: Rigs, endpoint: str) -> None: + forged: Final = _token("forged") + trace: Final = _token("trace") + forged_metadata: Final[dict[str, JsonValue]] = { + "metadata" if endpoint == "chat" else "litellm_metadata": {CALLER_SESSION_KEY: forged} + } + + untraced: Final = rigs.flag_on.send(endpoint, _marker(), forged_metadata) + traced: Final = rigs.flag_on.send(endpoint, _marker(), forged_metadata, headers={"x-litellm-trace-id": trace}) + + assert _session(untraced) not in (forged, None) + assert _session(traced) == trace + + +def test_identical_requests_each_log_once_with_the_body_session(rigs: Rigs) -> None: + session: Final = _token("convo") + trace: Final = _token("trace") + marker: Final = _marker() + + responses: Final = tuple( + rigs.flag_on.send("chat", marker, headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": session}) + for _ in range(3) + ) + + assert len({response.headers["x-litellm-call-id"] for response in responses}) == 3 + assert tuple(_session(response) for response in responses) == (session, session, session) + + +def _pass_through_route(rig: Rig, target: str) -> str: + path: Final = f"/caller-session-{uuid.uuid4().hex}" + created: Final = rig.proxy.post("/config/pass_through_endpoint", {"path": path, "target": target, "auth": True}) + endpoint_id: Final = string_value(object_value(list_value(created["endpoints"])[0])["id"]) + rig.scenario.cleanups.callback( + lambda: rig.proxy.request("DELETE", "/config/pass_through_endpoint", params={"endpoint_id": endpoint_id}) + ) + return path + + +@pytest.mark.parametrize("field", ["litellm_session_id", "metadata"]) +def test_pass_through_spend_log_session_is_the_body_session_beside_a_trace_header(rigs: Rigs, field: str) -> None: + marker: Final = _marker() + session: Final = _token("convo") + rig: Final = rigs.flag_on + path: Final = _pass_through_route(rig, f"{rig.upstream.url}/v1/chat/completions") + key: Final = rig.scenario.key(allowed_passthrough_routes=[path]) + fields: Final[dict[str, JsonValue]] = ( + {"litellm_session_id": session} if field == "litellm_session_id" else {"metadata": {"session_id": session}} + ) + + response: Final = rig.proxy.request( + "POST", + path, + rig.body("chat", marker, fields), + key=key, + headers={"x-litellm-trace-id": _token("trace")}, + ) + + rows: Final = eventually( + lambda: read_rows( + 'SELECT session_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),) + ), + lambda found: len(found) >= 1, + seconds=70, + ) + assert response.status_code == 200, response.text + assert f"answer {marker}" in response.text + assert len(rig.upstream_bodies(marker)) == 1 + assert rows == [{"session_id": session}] + + +def _turn( + rig: Rig, + marker: str, + previous: str | None = None, + *, + key: str | None = None, + headers: Headers = (), + fields: Mapping[str, JsonValue] = MappingProxyType({}), +) -> str: + chained: Final[dict[str, JsonValue]] = {"previous_response_id": previous} if previous else {} + response: Final = rig.send("responses", marker, {**chained, **fields}, headers=headers, key=key) + _spend_row(response) + return string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + +def _turn_fields(shape: str, session: str) -> tuple[dict[str, JsonValue], dict[str, JsonValue]]: + match shape: + case "session-every-turn": + return {"litellm_session_id": session}, {"litellm_session_id": session} + case "session-from-turn-2": + return {}, {"litellm_session_id": session} + case _: + return {"litellm_session_id": _token("convo")}, {"litellm_session_id": _token("convo")} + + +@pytest.mark.parametrize( + ("correlated", "shape", "replays_turn_1"), + [ + (True, "session-every-turn", True), + (True, "session-from-turn-2", True), + (True, "new-session-per-turn", True), + (False, "session-from-turn-2", False), + ], +) +def test_bridged_chain_history(rigs: Rigs, correlated: bool, shape: str, replays_turn_1: bool) -> None: + rig: Final = rigs.flag_on if correlated else rigs.flag_off + first_fields, second_fields = _turn_fields(shape, _token("convo")) + markers: Final = (_marker(), _marker(), _marker()) + first: Final = _turn(rig, markers[0], headers={"x-litellm-trace-id": _token("trace")}, fields=first_fields) + second: Final = _turn(rig, markers[1], first, headers={"x-litellm-trace-id": _token("trace")}, fields=second_fields) + + _turn(rig, markers[2], second) + + replayed: Final = rig.upstream_bodies(markers[2]) + assert len(replayed) == 1, replayed + assert _user_markers(replayed[0]) == (markers if replays_turn_1 else markers[1:]) + + +def test_bridged_chains_sharing_a_metadata_session_replay_each_others_turns(rigs: Rigs) -> None: + shared: Final[dict[str, JsonValue]] = {"metadata": {"session_id": _token("convo")}} + markers: Final = (_marker(), _marker(), _marker()) + first: Final = _turn(rigs.flag_on, markers[0], fields=shared) + _turn(rigs.flag_on, markers[1], fields=shared) + + _turn(rigs.flag_on, markers[2], first) + + replayed: Final = rigs.flag_on.upstream_bodies(markers[2]) + assert len(replayed) == 1, replayed + assert _user_markers(replayed[0]) == markers + + +@dataclass(frozen=True, slots=True) +class Owners: + first: str + second: str + + +def _owners(scenario: Scenario, model: str, shape: str) -> Owners: + user: Final = _token("user") + match shape: + case "other-user": + return Owners( + scenario.key(user_id=user, models=[model]), scenario.key(user_id=_token("user"), models=[model]) + ) + case "same-user": + return Owners(scenario.key(user_id=user, models=[model]), scenario.key(user_id=user, models=[model])) + case "same-team": + team: Final = scenario.team() + return Owners(scenario.key(team_id=team, models=[model]), scenario.key(team_id=team, models=[model])) + case _: + return Owners(scenario.key(models=[model]), scenario.key(models=[model])) + + +@pytest.mark.parametrize( + ("shape", "replays_first_owner"), + [("other-user", False), ("same-user", True), ("same-team", True), ("ownerless-keys", False)], +) +def test_bridged_history_follows_the_previous_responses_owner( + rigs: Rigs, shape: str, replays_first_owner: bool +) -> None: + rig: Final = rigs.flag_off + owners: Final = _owners(rig.scenario, rig.model, shape) + session: Final = _token("convo") + markers: Final = (_marker(), _marker(), _marker()) + _turn(rig, markers[0], key=owners.first, fields={"litellm_session_id": session}) + second: Final = _turn(rig, markers[1], key=owners.second, fields={"litellm_session_id": session}) + + _turn(rig, markers[2], second, key=owners.second) + + replayed: Final = rig.upstream_bodies(markers[2]) + assert len(replayed) == 1, replayed + assert _user_markers(replayed[0]) == (markers if replays_first_owner else markers[1:]) + + +BURST: Final = ENDPOINT_STREAMS * 4 + + +def test_concurrent_burst_across_endpoints_logs_each_session_once(rigs: Rigs) -> None: + sessions: Final = tuple(_token("convo") for _ in BURST) + + def one(index: int) -> httpx.Response: + endpoint, stream = BURST[index] + return rigs.flag_on.send( + endpoint, + _marker(), + stream=stream, + headers={"x-litellm-trace-id": _token("trace")}, + fields={"litellm_session_id": sessions[index]}, + ) + + with ThreadPoolExecutor(max_workers=12) as pool: + responses: Final = tuple(pool.map(one, range(len(BURST)))) + + assert tuple(_session(response) for response in responses) == sessions + + +def _open_upstream_connections(pid: int, upstream: Wire) -> int: + port: Final = urlsplit(upstream.url).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def _post_or_none(rig: Rig, session: str) -> httpx.Response | None: + try: + return rig.post( + "chat", _marker(), headers={"x-litellm-trace-id": _token("trace")}, fields={"litellm_session_id": session} + ) + except httpx.HTTPError: + return None + + +def test_killing_one_worker_mid_burst_keeps_logging_the_body_session(rigs: Rigs) -> None: + release: Final = threading.Event() + arrived: Final[SimpleQueue[bytes]] = SimpleQueue() + + def held(request: Request) -> Reply: + arrived.put(request.body) + assert release.wait(timeout=60), "the burst was never released" + return _upstream(request) + + process: Final = rigs.flag_on.process + sessions: Final = tuple(_token("convo") for _ in range(20)) + after_session: Final = _token("convo") + with wire_server(held) as wire: + rig: Final = replace(rigs.flag_on, model=_model(rigs.flag_on.scenario, wire), upstream=wire) + workers: Final = eventually(process.worker_pids, lambda found: len(found) == 2, seconds=30) + with ThreadPoolExecutor(max_workers=len(sessions)) as pool: + pending: Final = tuple(pool.submit(_post_or_none, rig, session) for session in sessions) + eventually(arrived.qsize, lambda size: size == len(sessions), seconds=60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire) for pid in workers}) + victim, survivor = sorted(workers, key=held_by.__getitem__) + with process.replacing_worker(victim) as kill: + kill.freeze() + kill() + release.set() + results: Final = tuple(future.result() for future in pending) + after: Final = rig.send( + "chat", + _marker(), + headers={"x-litellm-trace-id": _token("trace")}, + fields={"litellm_session_id": after_session}, + ) + + served: Final = tuple( + (response, session) + for response, session in zip(results, sessions, strict=True) + if response is not None and response.status_code == 200 + ) + assert sum(held_by.values()) == len(sessions), held_by + assert len(served) == held_by[survivor] >= len(sessions) // 2, (held_by, len(served)) + assert _session(after) == after_session + assert tuple(_session(response) for response, _ in served) == tuple(session for _, session in served) 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_mid_stream_fallback_continuation.py b/tests/integration/streaming/test_mid_stream_fallback_continuation.py new file mode 100644 index 00000000000..8ec52f26dde --- /dev/null +++ b/tests/integration/streaming/test_mid_stream_fallback_continuation.py @@ -0,0 +1,749 @@ +from __future__ import annotations + +import functools +import hashlib +import itertools +import json +import math +import operator +import threading +import uuid +from collections.abc import Callable, Generator, Iterable, Iterator, Mapping, Sequence +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import openai +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, gateway_from_environment, object_value +from integration._support.database import read_rows, scratch_database +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from openai.types.chat import ChatCompletionChunk +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(4 * graceful_stop_seconds() + 120) + +SETTING: Final = "enable_mid_stream_fallback_continuation" +PARTIAL: Final = ("Alpha beta", " gamma") +JOINED_PARTIAL: Final = "".join(PARTIAL) +RELAY_TEXT: Final = " delta" +BACKUP_TEXT: Final = " omega." +PLAIN_TEXT: Final = "plain answer" +TOOL_PREAMBLE: Final = "Let me check" +BACKUP_USAGE: Final = (40, 6) +MAX_TOKENS: Final = 200 +PRIMARY_PRICES: Final = (0.001, 0.002) +BACKUP_PRICES: Final = (0.003, 0.005) +SHARED_BACKEND: Final = "openai/gpt-4o-mini" +SEPARATE_BACKEND: Final = "openai/gpt-4o" +WEATHER_TOOL: Final[Mapping[str, JsonValue]] = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} +TOOL_CALL_DELTA: Final[Mapping[str, JsonValue]] = { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": ""}, +} +NOTHING_SEEN: Final[tuple[ChatCompletionChunk, ...]] = () +ALLOW_EVERYTHING: Final = "def apply_guardrail(inputs, request_data, input_type):\n return allow()\n" +EVENT_BATCH: Final = TypeAdapter(list[dict[str, JsonValue]]) + + +@functools.cache +def cut_gate(marker: str, hop: str, attempt: int) -> threading.Event: + return threading.Event() + + +@functools.cache +def attempts(marker: str, hop: str) -> Iterator[int]: + return itertools.count() + + +def frame(identity: str, delta: Mapping[str, JsonValue], finish: str | None = None) -> bytes: + body: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish}], + } + return b"data: " + json.dumps(body).encode() + b"\n\n" + + +def complete_stream(identity: str, text: str, usage: tuple[int, int]) -> Reply: + usage_frame: Final[dict[str, JsonValue]] = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [], + "usage": {"prompt_tokens": usage[0], "completion_tokens": usage[1], "total_tokens": sum(usage)}, + } + return Reply( + content_type="text/event-stream", + chunks=( + frame(identity, {"role": "assistant", "content": text}), + frame(identity, {}, finish="stop"), + b"data: " + json.dumps(usage_frame).encode() + b"\n\n", + b"data: [DONE]\n\n", + ), + ) + + +def cut_after_caller_sees(marker: str, hop: str, *frames: bytes) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=(b"".join(frames), b"data: never sent\n\n"), + abort_after=1, + gate_after_first=cut_gate(marker, hop, next(attempts(marker, hop))), + ) + + +def first_user_text(request: Request) -> str: + messages: Final = JSON_OBJECT.validate_json(request.body)["messages"] + assert isinstance(messages, list), messages + first: Final = object_value(messages[0])["content"] + assert isinstance(first, str), first + return first + + +def marker_of(request: Request) -> str: + return first_user_text(request).split()[-1] + + +def is_chat(request: Request) -> bool: + return request.target.endswith("/chat/completions") + + +def model_list() -> Reply: + return Reply(body=b'{"object": "list", "data": []}') + + +def primary(request: Request) -> Reply: + if not is_chat(request): + return model_list() + scenario: Final = first_user_text(request).split()[0] + marker: Final = marker_of(request) + identity: Final = f"chatcmpl-{marker}-primary" + if scenario == "reject": + return Reply(status=500, body=b'{"error": {"message": "primary is down", "type": "server_error"}}') + if scenario == "tool-cut": + return cut_after_caller_sees( + marker, + "primary", + frame(identity, {"role": "assistant", "content": TOOL_PREAMBLE}), + frame(identity, {"tool_calls": [dict(TOOL_CALL_DELTA)]}), + ) + return cut_after_caller_sees( + marker, + "primary", + frame(identity, {"role": "assistant", "content": PARTIAL[0]}), + frame(identity, {"content": PARTIAL[1]}), + ) + + +def relay(request: Request) -> Reply: + if not is_chat(request): + return model_list() + marker: Final = marker_of(request) + return cut_after_caller_sees( + marker, "relay", frame(f"chatcmpl-{marker}-relay", {"role": "assistant", "content": RELAY_TEXT}) + ) + + +def backup(request: Request) -> Reply: + if not is_chat(request): + return model_list() + return complete_stream(f"chatcmpl-{marker_of(request)}-backup", BACKUP_TEXT, BACKUP_USAGE) + + +def no_prefill(request: Request) -> Reply: + if not is_chat(request): + return model_list() + return complete_stream(f"chatcmpl-{marker_of(request)}-plain", PLAIN_TEXT, (9, 2)) + + +def accept_events(request: Request) -> Reply: + return Reply() + + +def registered_guardrail(name: str, mode: str) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "custom_code", + "mode": mode, + "default_on": False, + "custom_code": ALLOW_EVERYTHING, + }, + } + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + database_url: str + primary: Wire + relay: Wire + backup: Wire + no_prefill: Wire + events: Wire + + +def joined_text(chunks: Iterable[ChatCompletionChunk]) -> str: + return "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) + + +def has_tool_call(chunks: Iterable[ChatCompletionChunk]) -> bool: + return any(choice.delta.tool_calls for chunk in chunks for choice in chunk.choices) + + +@dataclass(frozen=True, slots=True) +class Release: + gate: threading.Event + ready: Callable[[tuple[ChatCompletionChunk, ...]], bool] + + +def after_text(marker: str, hop: str, text: str, attempt: int = 0) -> Release: + return Release(cut_gate(marker, hop, attempt), lambda seen: joined_text(seen) == text) + + +@dataclass(frozen=True, slots=True) +class Streamed: + chunks: tuple[ChatCompletionChunk, ...] + error: openai.APIError | None + + @property + def text(self) -> str: + return joined_text(self.chunks) + + @property + def ids(self) -> frozenset[str]: + return frozenset(chunk.id for chunk in self.chunks) + + @property + def role_deltas(self) -> int: + return sum(1 for chunk in self.chunks for choice in chunk.choices if choice.delta.role is not None) + + @property + def finish_reasons(self) -> tuple[str, ...]: + return tuple(choice.finish_reason for chunk in self.chunks for choice in chunk.choices if choice.finish_reason) + + @property + def usage(self) -> tuple[int, int]: + reported: Final = tuple(chunk.usage for chunk in self.chunks if chunk.usage is not None) + assert reported, self.chunks + return (reported[-1].prompt_tokens, reported[-1].completion_tokens) + + +def deployment( + name: str, + wire: Wire, + prices: tuple[float, float], + *, + backend: str = SHARED_BACKEND, + prefill: bool | None = None, +) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": backend, + "api_base": wire.url + "/v1", + "api_key": "synthetic-continuation-key", + "input_cost_per_token": prices[0], + "output_cost_per_token": prices[1], + }, + **({"model_info": {"supports_assistant_prefill": prefill}} if prefill is not None else {}), + } + + +def write_config( + directory: Path, + model_entries: Sequence[Mapping[str, JsonValue]], + fallbacks: Sequence[Mapping[str, JsonValue]], + router_settings: Mapping[str, JsonValue], + guardrails: Sequence[Mapping[str, JsonValue]], +) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "model_list": [dict(entry) for entry in model_entries], + "litellm_settings": { + **object_value(base["litellm_settings"]), + "callbacks": ["generic_api"], + "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, + }, + "guardrails": [dict(entry) for entry in guardrails], + "router_settings": { + **object_value(base["router_settings"]), + "num_retries": 0, + "fallbacks": [dict(entry) for entry in fallbacks], + **router_settings, + }, + } + path: Final = directory / f"continuation-{uuid.uuid4().hex[:8]}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@contextmanager +def booted_rig( + directory: Path, + fallbacks: Sequence[Mapping[str, JsonValue]], + router_settings: Mapping[str, JsonValue], + workers: int, + guardrails: Sequence[Mapping[str, JsonValue]] = (), +) -> Generator[Rig]: + with ExitStack() as stack: + shared: Final = stack.enter_context(gateway_from_environment()) + database_url: Final = stack.enter_context(scratch_database()) + wires: Final = tuple( + stack.enter_context(wire_server(respond)) for respond in (primary, relay, backup, no_prefill, accept_events) + ) + primary_wire, relay_wire, backup_wire, no_prefill_wire, events_wire = wires + config: Final = write_config( + directory, + ( + *( + deployment(name, primary_wire, PRIMARY_PRICES) + for name in ("primary", "primary-skip", "primary-none", "primary-relay", "primary-shared") + ), + deployment("relay", relay_wire, PRIMARY_PRICES, prefill=True), + deployment("backup", backup_wire, BACKUP_PRICES, prefill=True), + deployment("no-prefill", no_prefill_wire, PRIMARY_PRICES, backend=SEPARATE_BACKEND), + deployment("no-prefill-shared", no_prefill_wire, PRIMARY_PRICES), + ), + fallbacks, + router_settings, + guardrails, + ) + gateway: Final = stack.enter_context( + owned_proxy( + shared, + directory, + {"DATABASE_URL": database_url, "GENERIC_LOGGER_ENDPOINT": events_wire.url}, + config=config, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=workers, + ) + ) + yield Rig(gateway, database_url, primary_wire, relay_wire, backup_wire, no_prefill_wire, events_wire) + + +@pytest.fixture(scope="module") +def configured(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + fallbacks: Final[tuple[Mapping[str, JsonValue], ...]] = ( + {"primary": ["backup"]}, + {"primary-skip": ["no-prefill", "backup"]}, + {"primary-none": ["no-prefill"]}, + {"primary-relay": ["relay"]}, + {"relay": ["backup"]}, + {"primary-shared": ["no-prefill-shared"]}, + ) + with booted_rig( + tmp_path_factory.mktemp("continuation-configured"), + fallbacks, + {SETTING: True}, + 2, + (registered_guardrail("input-check", "pre_call"),), + ) as rig: + yield rig + + +@pytest.fixture(scope="module") +def output_guarded(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with booted_rig( + tmp_path_factory.mktemp("continuation-output-guarded"), + ({"primary": ["backup"]},), + {SETTING: True}, + 2, + (registered_guardrail("output-check", "post_call"),), + ) as rig: + yield rig + + +@pytest.fixture(scope="module") +def unset(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with booted_rig( + tmp_path_factory.mktemp("continuation-unset"), ({"primary": ["no-prefill", "backup"]},), {}, 1 + ) as rig: + yield rig + + +def new_key(rig: Rig) -> str: + key: Final = rig.gateway.post("/key/generate", {})["key"] + assert isinstance(key, str), key + return key + + +def appended(seen: tuple[ChatCompletionChunk, ...], chunk: ChatCompletionChunk) -> tuple[ChatCompletionChunk, ...]: + return (*seen, chunk) + + +def observed( + response: Iterable[ChatCompletionChunk], releases: Sequence[Release] +) -> Iterator[ChatCompletionChunk | openai.APIError]: + try: + for seen in itertools.accumulate(response, appended, initial=NOTHING_SEEN): + for release in releases: + if release.ready(seen): + release.gate.set() + yield from seen[-1:] + except openai.APIError as error: + yield error + + +def stream( + rig: Rig, + key: str, + model: str, + prompt: str, + releases: Sequence[Release] = (), + extra: Mapping[str, JsonValue] | None = None, +) -> Streamed: + base_url: Final = f"{rig.gateway.client.base_url}".rstrip("/") + "/v1" + try: + with openai.OpenAI(base_url=base_url, api_key=key, max_retries=0, timeout=60) as client: + response: Final = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + stream=True, + stream_options={"include_usage": True}, + max_tokens=MAX_TOKENS, + extra_body=dict(extra or {}), + ) + outcome: Final = tuple(observed(response, releases)) + finally: + for release in releases: + release.gate.set() + errors: Final = tuple(item for item in outcome if isinstance(item, openai.APIError)) + return Streamed( + tuple(item for item in outcome if isinstance(item, ChatCompletionChunk)), errors[0] if errors else None + ) + + +def chat_bodies(wire: Wire, marker: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + JSON_OBJECT.validate_json(request.body) + for request in wire.drain() + if is_chat(request) and marker_of(request) == marker + ) + + +def spend_rows(rig: Rig, key: str) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + 'SELECT request_id, status, spend, prompt_tokens, completion_tokens, metadata FROM "LiteLLM_SpendLogs" ' + "WHERE api_key=%s", + (hashlib.sha256(key.encode()).hexdigest(),), + database_url=rig.database_url, + ), + lambda rows: len(rows) >= 1, + seconds=70, + ) + + +def amount(value: JsonValue) -> float: + assert isinstance(value, (int, float)), value + return float(value) + + +def key_spend(rig: Rig, key: str) -> float: + return amount(object_value(rig.gateway.get("/key/info", {"key": key})["info"])["spend"]) + + +def close(observed_value: float, expected: float) -> bool: + return math.isclose(observed_value, expected, rel_tol=1e-9, abs_tol=1e-12) + + +def cost(tokens: tuple[int, int], prices: tuple[float, float]) -> float: + return tokens[0] * prices[0] + tokens[1] * prices[1] + + +def prefill_turn(text: str) -> dict[str, JsonValue]: + return {"role": "assistant", "content": text, "prefix": True} + + +@dataclass(frozen=True, slots=True) +class Billed: + rows: list[dict[str, JsonValue]] + partial: tuple[int, int] + + +def billed_once(rig: Rig, key: str, received: Streamed) -> Billed: + rows: Final = spend_rows(rig, key) + assert [row["status"] for row in rows] == ["success"], rows + row_tokens: Final = (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) + assert isinstance(row_tokens[0], int) and isinstance(row_tokens[1], int), rows + assert received.usage == row_tokens, (received.usage, row_tokens) + partial: Final = (row_tokens[0] - BACKUP_USAGE[0], row_tokens[1] - BACKUP_USAGE[1]) + assert partial[0] > 0 and partial[1] > 0, (row_tokens, BACKUP_USAGE) + return Billed(rows, partial) + + +def assert_spend_prices_every_token(rig: Rig, key: str, billed: Billed) -> None: + partial_cost: Final = cost(billed.partial, PRIMARY_PRICES) + breakdown: Final = object_value(object_value(billed.rows[0]["metadata"])["cost_breakdown"]) + spend: Final = amount(billed.rows[0]["spend"]) + reported_partial_cost: Final = amount(object_value(breakdown["additional_costs"])["mid_stream_partial_cost"]) + assert close(spend, cost(BACKUP_USAGE, BACKUP_PRICES) + partial_cost) and close( + reported_partial_cost, partial_cost + ), f"partial={billed.partial} spend={spend} breakdown={json.dumps(breakdown)}" + eventually(lambda: key_spend(rig, key), lambda value: close(value, spend), seconds=70) + assert spend_rows(rig, key) == billed.rows + + +def test_text_cut_continues_on_the_prefill_fallback_and_bills_both_calls_once(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + key: Final = new_key(configured) + prompt: Final = f"text-cut {marker}" + received: Final = stream(configured, key, "primary", prompt, (after_text(marker, "primary", JOINED_PARTIAL),)) + assert received.error is None, received.error + assert received.text == JOINED_PARTIAL + BACKUP_TEXT, received.chunks + assert received.ids == {f"chatcmpl-{marker}-primary"}, received.ids + assert received.role_deltas == 1, received.chunks + assert received.finish_reasons == ("stop",), received.chunks + assert len(chat_bodies(configured.primary, marker)) == 1 + continued: Final = chat_bodies(configured.backup, marker) + assert len(continued) == 1, continued + assert continued[0]["messages"] == [{"role": "user", "content": prompt}, prefill_turn(JOINED_PARTIAL)] + assert "_mid_stream_continuation" not in continued[0], continued[0] + billed: Final = billed_once(configured, key, received) + assert continued[0]["max_tokens"] == MAX_TOKENS - billed.partial[1], (continued[0], billed.partial) + assert_spend_prices_every_token(configured, key, billed) + + +def test_spend_row_of_a_continued_stream_carries_the_id_the_caller_received(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + key: Final = new_key(configured) + received: Final = stream( + configured, key, "primary", f"text-cut {marker}", (after_text(marker, "primary", JOINED_PARTIAL),) + ) + assert received.error is None, received.error + assert [row["request_id"] for row in spend_rows(configured, key)] == [received.chunks[0].id] + + +def test_tool_call_before_the_cut_keeps_the_original_error(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + configured, + new_key(configured), + "primary", + f"tool-cut {marker}", + (Release(cut_gate(marker, "primary", 0), has_tool_call),), + {"tools": [dict(WEATHER_TOOL)]}, + ) + assert received.error is not None, received.chunks + assert (received.text, has_tool_call(received.chunks)) == (TOOL_PREAMBLE, True), received.chunks + assert len(chat_bodies(configured.primary, marker)) == 1 + assert chat_bodies(configured.backup, marker) == () + + +def test_continuation_skips_a_fallback_without_prefill_support(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + prompt: Final = f"text-cut {marker}" + received: Final = stream( + configured, new_key(configured), "primary-skip", prompt, (after_text(marker, "primary", JOINED_PARTIAL),) + ) + assert received.error is None, received.error + assert received.text == JOINED_PARTIAL + BACKUP_TEXT, received.chunks + assert chat_bodies(configured.no_prefill, marker) == () + assert [body["messages"] for body in chat_bodies(configured.backup, marker)] == [ + [{"role": "user", "content": prompt}, prefill_turn(JOINED_PARTIAL)] + ] + + +def test_no_prefill_capable_fallback_keeps_the_original_error(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + configured, + new_key(configured), + "primary-none", + f"text-cut {marker}", + (after_text(marker, "primary", JOINED_PARTIAL),), + ) + assert received.error is not None, received.chunks + assert "No deployments available" not in str(received.error), received.error + assert received.text == JOINED_PARTIAL, received.chunks + assert len(chat_bodies(configured.primary, marker)) == 1 + assert chat_bodies(configured.no_prefill, marker) == () + + +def test_undeclared_fallback_sharing_a_backend_with_a_prefill_deployment_keeps_the_original_error( + configured: Rig, +) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + configured, + new_key(configured), + "primary-shared", + f"text-cut {marker}", + (after_text(marker, "primary", JOINED_PARTIAL),), + ) + assert (received.error is not None, received.text) == (True, JOINED_PARTIAL), received.chunks + assert chat_bodies(configured.no_prefill, marker) == () + + +def test_failure_before_content_still_falls_back_to_a_fallback_without_prefill_support(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + prompt: Final = f"reject {marker}" + received: Final = stream(configured, new_key(configured), "primary-none", prompt) + assert received.error is None, received.error + assert received.text == PLAIN_TEXT, received.chunks + assert [body["messages"] for body in chat_bodies(configured.no_prefill, marker)] == [ + [{"role": "user", "content": prompt}] + ] + + +def stream_with_two_cuts(rig: Rig, key: str, marker: str) -> Streamed: + return stream( + rig, + key, + "primary-relay", + f"text-cut {marker}", + (after_text(marker, "primary", JOINED_PARTIAL), after_text(marker, "relay", JOINED_PARTIAL + RELAY_TEXT)), + ) + + +def test_second_cut_continues_with_both_partials_in_one_prefill(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + key: Final = new_key(configured) + received: Final = stream_with_two_cuts(configured, key, marker) + assert received.error is None, received.error + assert received.text == JOINED_PARTIAL + RELAY_TEXT + BACKUP_TEXT, received.chunks + assert received.role_deltas == 1, received.chunks + user_turn: Final = {"role": "user", "content": f"text-cut {marker}"} + relayed: Final = chat_bodies(configured.relay, marker) + continued: Final = chat_bodies(configured.backup, marker) + assert [body["messages"] for body in relayed] == [[user_turn, prefill_turn(JOINED_PARTIAL)]] + assert [body["messages"] for body in continued] == [[user_turn, prefill_turn(JOINED_PARTIAL + RELAY_TEXT)]] + partial: Final = billed_once(configured, key, received).partial + relay_ceiling: Final = relayed[0]["max_tokens"] + assert isinstance(relay_ceiling, int), relayed[0] + assert 0 < MAX_TOKENS - relay_ceiling < partial[1], (relayed[0], partial) + assert continued[0]["max_tokens"] == MAX_TOKENS - partial[1], (continued[0], partial) + + +def test_second_cut_bills_every_call_on_one_row(configured: Rig) -> None: + key: Final = new_key(configured) + received: Final = stream_with_two_cuts(configured, key, uuid.uuid4().hex) + assert received.error is None, received.error + assert_spend_prices_every_token(configured, key, billed_once(configured, key, received)) + + +def test_repeating_a_continued_request_never_returns_only_the_continuation(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + key: Final = new_key(configured) + prompt: Final = f"text-cut {marker}" + answers: Final = tuple( + stream(configured, key, "primary", prompt, (after_text(marker, "primary", JOINED_PARTIAL, attempt),)) + for attempt in range(2) + ) + assert tuple((answer.error, answer.text) for answer in answers) == ((None, JOINED_PARTIAL + BACKUP_TEXT),) * 2 + + +def test_without_the_setting_a_text_cut_keeps_the_original_error(unset: Rig) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + unset, new_key(unset), "primary", f"text-cut {marker}", (after_text(marker, "primary", JOINED_PARTIAL),) + ) + assert received.error is not None, received.chunks + assert received.text == JOINED_PARTIAL, received.chunks + assert len(chat_bodies(unset.primary, marker)) == 1 + assert chat_bodies(unset.no_prefill, marker) == () + assert chat_bodies(unset.backup, marker) == () + + +def router_setting(rig: Rig) -> JsonValue: + return object_value(rig.gateway.get("/router/settings")["current_values"]).get(SETTING) + + +def test_turning_the_setting_on_at_runtime_continues_on_a_prefill_fallback(unset: Rig) -> None: + rig: Final = unset + rig.gateway.post("/config/update", {"router_settings": {SETTING: True}}) + try: + assert router_setting(rig) is True + marker: Final = uuid.uuid4().hex + prompt: Final = f"text-cut {marker}" + received: Final = stream(rig, new_key(rig), "primary", prompt, (after_text(marker, "primary", JOINED_PARTIAL),)) + assert received.error is None, received.error + assert received.text == JOINED_PARTIAL + BACKUP_TEXT, received.chunks + assert chat_bodies(rig.no_prefill, marker) == () + assert [body["messages"] for body in chat_bodies(rig.backup, marker)] == [ + [{"role": "user", "content": prompt}, prefill_turn(JOINED_PARTIAL)] + ] + finally: + rig.gateway.post("/config/update", {"router_settings": {SETTING: False}}) + + +def test_a_registered_output_guardrail_keeps_the_original_error(output_guarded: Rig) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + output_guarded, + new_key(output_guarded), + "primary", + f"text-cut {marker}", + (after_text(marker, "primary", JOINED_PARTIAL),), + ) + assert (received.error is not None, received.text) == (True, JOINED_PARTIAL), received.chunks + assert len(chat_bodies(output_guarded.primary, marker)) == 1 + assert chat_bodies(output_guarded.backup, marker) == () + + +def test_a_request_fallback_bringing_its_own_messages_keeps_the_original_error(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + configured, + new_key(configured), + "primary", + f"text-cut {marker}", + (after_text(marker, "primary", JOINED_PARTIAL),), + {"fallbacks": [{"model": "backup", "messages": [{"role": "user", "content": f"start over {marker}"}]}]}, + ) + assert (received.error is not None, received.text) == (True, JOINED_PARTIAL), received.chunks + assert len(chat_bodies(configured.primary, marker)) == 1 + assert chat_bodies(configured.backup, marker) == () + + +def carries_tag(event: Mapping[str, JsonValue], marker: str) -> bool: + tags: Final = event.get("request_tags") + return isinstance(tags, list) and marker in tags + + +def tagged_events(batches: Iterable[Request], marker: str) -> tuple[dict[str, JsonValue], ...]: + events: Final = itertools.chain.from_iterable(EVENT_BATCH.validate_json(batch.body) for batch in batches) + return tuple(event for event in events if carries_tag(event, marker)) + + +def logged_events(sink: Wire, marker: str) -> Iterator[tuple[dict[str, JsonValue], ...]]: + return itertools.accumulate((tagged_events(sink.drain(), marker) for _ in itertools.count()), operator.add) + + +def test_the_cut_call_logs_its_failure_under_its_own_deployment_and_price(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + configured, + new_key(configured), + "primary", + f"text-cut {marker}", + (after_text(marker, "primary", JOINED_PARTIAL),), + {"metadata": {"tags": [marker]}}, + ) + assert received.error is None, received.error + polls: Final = logged_events(configured.events, marker) + events: Final = eventually(lambda: next(polls), lambda logged: len(logged) >= 2, seconds=30) + assert sorted(str(event["status"]) for event in events) == ["failure", "success"], events + failure: Final = next(event for event in events if event["status"] == "failure") + assert failure["model_group"] == "primary", failure + assert str(failure["api_base"]).startswith(configured.primary.url), failure + tokens: Final = (failure["prompt_tokens"], failure["completion_tokens"]) + assert isinstance(tokens[0], int) and isinstance(tokens[1], int) and tokens[0] > 0 and tokens[1] > 0, failure + assert close(amount(failure["response_cost"]), cost((tokens[0], tokens[1]), PRIMARY_PRICES)), failure diff --git a/tests/unit/embeddings/__init__.py b/tests/integration/translation/chat_completions/reasoning/__init__.py similarity index 100% rename from tests/unit/embeddings/__init__.py rename to tests/integration/translation/chat_completions/reasoning/__init__.py diff --git a/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py b/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py new file mode 100644 index 00000000000..194e6972453 --- /dev/null +++ b/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py @@ -0,0 +1,1041 @@ +import json +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace +from pathlib import Path +from typing import Final +from unittest.mock import ANY +from uuid import uuid4 + +import pytest +from integration._support.client import Gateway, eventually, list_value, object_value, string_value +from integration._support.process import graceful_stop_seconds +from integration._support.provider import SharedProvider +from integration._support.wire import Reply +from integration.caching.response_cache_case import CACHE_KEY, redis_store +from integration.translation.case import SSEFrame, TranslationTestCase +from integration.translation.chat_completions.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE +from integration.translation.chat_completions.bases.vertex_ai import ( + GEMINI_3_8_FLASH_TEST_CASE as VERTEX_GEMINI_3_8_FLASH_TEST_CASE, +) +from integration.translation.gemini_thoughts import ( + LONG_SUMMARY, + gemini_thoughts_proxy, + sent_thinking_config, + thoughts_off, + with_thinking_config, +) +from integration.translation.messages.bases.gemini import ( + GEMINI_3_8_FLASH_TEST_CASE as MESSAGES_GEMINI_3_8_FLASH_TEST_CASE, +) +from integration.translation.responses.bases.gemini import ( + GEMINI_3_8_FLASH_TEST_CASE as RESPONSES_GEMINI_3_8_FLASH_TEST_CASE, +) +from integration.translation.runner import ( + assert_provider_request, + assert_translation, + encode_sse_frame, + parse_sse_frames, +) +from pydantic import JsonValue + +__all__ = ("thoughts_off",) + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT: Final = "**My Inspection Strategy: Validation First**\n\nOkay, here's how I'm approaching this project review. The instruction specifically mentions considering zero, negative, and non-integer quantities. To efficiently address that, I'm going to start with `shop/validation.py`. My reasoning is simple: the validation logic is the first line of defense. Understanding how the system *expects* to handle these edge cases will inform everything else. It's the logical place to begin my investigation to understand the code. I'll use the `read_file` tool to examine that file. I am **not** changing any files. I will call the `chat_completions` client for this project review. I am now proceeding with reading `shop/validation.py`.\n\n\n" + +GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE: Final = ( + "EowMCokMAWkUfRPuOEs7DW4IXhoH/SNh9qYezwNSqhXT2tKoI1mcXDG+yx7R+12sMmzvjHh2OSghAlsvmR6welOzIM1LMw6c" + "QrSX2IPRo/xbtUvDLlo6aqufzRtini5EF0yxQdmyn9r9xznzDyuEf4eqrbK0KQAArEVM13rm3CTDIjyDRQPawLALfUh/e61/" + "XVcWrr1PZLR5Jj46ImyrobGGLmJuoytm3ezTEihTb2b0q3/vHOW9cO/gJzMcbCI8koexQLuBB1KLSHUCaZZQYPj4jLx/g0de" + "2l3H/DYR7Cdroh0P/R3mewkhu+3D2ELnHKt85InrdQOqdXp/XhICsivPgsIEWHMRr6kQ3ZLIUJvJbJ7zo7/4bYnHOm+JqmoX" + "E+VMobI2YZ+ff9GOYg4vsJRTeF7ff9kZHYBWRM9XvTsZMbeCZHRbCOLNmSSJPfn2ToDOgMxU1u/GqHYmVSV97ybT3Y0mxwtf" + "2WcVlPUCkFX0GW3GP5IHVkEISo0TznJSI3/yYb+5UmQFMuuIedDdhKa74eU1/K59VuTnd3DzVX/lfhpFbGLDkOSurRsDsViN" + "69zq+6VbLi+5ARi6SJqUU+1U02QaSIQVp8kxtrhEIsPedjCIYHQQNxR8hFPMhnEssR51THNELWH6sVccxpJIbFfGjl0TZS9d" + "rfYL2c+IBX56lrH8IX2gz1iMOsESKwTyAECmqJIL7d/HcVdTk404Hd9SYycB+2MP2GN2XA7ivoO/zfkkFcf+PGmXPed8dHlW" + "bAxrQcOAz89p8kPzb1bou5QUT3feeRG06aJL31gX7BSYj19mvLeo0AmMzx+Wb91NFlwUvlJYlhJo4NfZWg/ymcGclgM8r7iE" + "MmWZa8wul8g4hbJOg3GE9NV70zvKqEpQrJLTa9jHca0DBvj8N3C+XvpeEg+Dys0cxE7Ak0VVYbvu9r0OHKSWtC7NDYrJsF2n" + "VFmdavSxpyPEqPaAdWdiewDH0nhP5GCR/ciikokWFQMF5iClWRfL72qEbFx3cQkctqEf+kK2CevIWKPf8bvkrKVDdPqbjk0G" + "9NWuud+YgNxl4dSCkAWoMrAFlxXozEqMmanMAtZk/7nPeOVFbYEL/xOpSCfxaEKS1zEmMwTb/UNVAogsNY9YdBSFYVNSrS8F" + "N/IrHBvsP66+YEFuR46Di6cWuuZQdSOfB2uQe1eSC8IgrQ6XUhQHWeXKUXfbcb/1abPOi5PUMYyLBQ7r52UKCEfOBc1j0oNL" + "wYzI/vWAz74gq8NblQhHB2VlrInJ9rRVY1MehInqUH0kcG6ysQDQLiLJiS30TV0bZh05I86RfcjaN1twbRir82YchRjXG8/1" + "6iFRJ54Sk7oVzbdnbFOugDea8qIblFmgUTq0+SRLGBFcDMvAqYLNcR2ZO7o0xinMG7IuNRwoSMahnmreWYr9saVJAScd3All" + "i6RQuxjGIdxG0s0UNuoWiSNIf1Z2ngoJrvZHTc0iZK+A2zv3CBqWAaS8WGwSfcItnlc6hMB0bTWasoCelmRjfv+1jQnKzKMj" + "DzZlO6skMkv5x9GzLrAvpYoulLYa/wp9+exw/Ll7MxSB22aeAGfai9C7Vj310y8nApsG3QxaoUUO4jg1eCjrrluBLgoVY8cE" + "/O95IAUa2rFzVrOWbHuob7TEI/8+BwieGeyeHdPGn4fNeKYMkRjFeVJ1GA/2+iCHCjCbSNm4sdzh9188/Es3q4acRfQUiqwR" + "iQzW2E1fy9mXRDY9y5fdXuzjHtuCkH8MqR0cioN1p4OGZzPi6xXFaqV576NLjySBpfpQEYeXRA3j5NVwTwbeYKnGOot4bEtJ" + "JQfelwkXLXeAfxzCTD0/vGtdhilzEK8wSYM9a+8llr5D4XwtEsu/68YbXJ2jwzsS3rE/Frulrf6paz9fS45ZXj98+OT5EpFO" + "o5XHfTXKb5Brzr0tKfPSB2Mx6fXbYP3lZRwfJIaUq8Alwdo7g2jnHQUhCNBB8hkb/IkQrnhsAGS6f0ZEjwf6ErVOa8DkfizN" + "y4WTdS1D3gOUJxOcszm8J5b085tsfFxEJVgA0VjI60DBY0owWn2I" +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "To understand how zero, negative, and non-integer quantities (such as floats or invalid types) are currently handled and where the checkout rejection is enforced, we need to inspect `shop/validation.py` next.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "path": "shop/validation.py", + "end_line": 50, + "start_line": 1, + }, + "id": "call_330625", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + "finishMessage": "Model generated function call(s).", + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 77, + "totalTokenCount": 793, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 358, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "xy3RapaNEbT4-sAPhMrt6AU", +} + +GEMINI_3_8_FLASH_THINKING_OMITTED_THOUGHT_SIGNATURE: Final = ( + "EscECsQEAWkUfROlWYT6hX0DWOKSCg4MSuLzOV7u43EgJDn0dfS2R8v/MVGGcNolGA9rteb/cV3dyCSQ7D2EO2CB5whDkblj" + "fGAAw6NBhqnidYNQiXbsg7oqynxS116ZQnk302S4KzvAbIvo4gpRbahPVClXKkP8PCQWA5eRQ/e9oQeHr97gbBB7+PLsl0NK" + "m7YPV6lP3fAXx/r3EADbdNL1diHylKd6tCWUKRea/x2QIOBQZdY5jHSU73z/pnPOotVwggEiasDLSLA+uUR1KllgGc8WKZ4i" + "uRM3L1gDr1QZ4Rht3quiF8AwfopWE5m2ON+FD78aHUDG/lgDs0vXnJluzri9LpPaexfbpvwMPtAx8Bxrq1/dCeekR2VnVx1c" + "qZ7TSoWMqk4x98mPpyUFz1hSBRnXaOGvQl5Wk7Iart49tH2LoIBYoV+PkCgNZ2MOi2uzc7OlN/8gvOeBM05xbGAplsLiuKd4" + "jkgv7Nz8Jby3v/hxtQT9+Ow3mxwSaiIfEkuaxkPiK1cX1nwT92gFzpmhVsrGqpiCUdAfos39Wd2/rcrFbtl5iMAVcHG5JGjw" + "8AjubFqiWDMqWPC58FfsX9TLC2BMUbPf2xi3hrw+1QWDtSuNl9s+dofkCogXFdLA9lAFc0oRj3MVhzYUhn5eixApJdggDsJF" + "iJaQtRihs4dstbXX9QXYFnx22Ql3x9TWJibiH8jz4fJVOf0wreMWZ2xVXBiw2xKdy5dB7xRDnrQD/sdcpTpCWURYXVCXglJ9" + "DD7e7/OCNMLCPw==" +) + +GEMINI_3_8_FLASH_THINKING_OMITTED_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "To understand how the checkout validates quantities\u2014including zero, negative, and non-integer quantities\u2014I need to inspect `shop/validation.py` next.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "start_line": 1, + "path": "shop/validation.py", + "end_line": 40, + }, + "id": "call_745953", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THINKING_OMITTED_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + "finishMessage": "Model generated function call(s).", + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 65, + "totalTokenCount": 531, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 108, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "zS3RasS-Lp2-jMcPoay98As", +} + +GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT_SIGNATURE: Final = ( + "EqUECqIEAWkUfRNZEW6qD3+V/1LAc1riV0fvqkg4Bldd7YYDfln+ZHNg6Lbw2T31UZHTum0qYvrAOor848+zSkFG7qBDfdZ9" + "69FGl/dQGq2llr1yrv23blnbwC7tO0uevr8225EZ/4Pv+hbfP0teFcSZsasE4LQangtxYkewghc6cH1Acz293SIz1cqF/3CM" + "Ta+FtJ5x106oauWy79+z1VbV9TNyHpNtsrUy3ADR2PFNBjXKf2TvUWygzR4w2RJg581+dfx0sI97O3UlZrhSKRe+p/6siUP7" + "qnysURAia7k+ew5f05Nuo1aDwD2G20G2s9v7L0PxnCCAay5L4YTub5LXxvUGBXMhSHVb9MmO8+2rNPqDtPWh0bOQdKbkWAV/" + "dMckMLTKANDbnzBjtCK1HbN7Rw0sIBSnBVy59z7h+0ierHFgA8Vntfp/nvAy0wDpuZu09agdB3T6FO2HCvWCFG2Y+0B4lGJV" + "ErOknvn50gtn9Tpn1QSGv4VDRrrcv6nCOyohoL9CsBeO1Bs+RItyZSAoTA3bVB5IjcF2kKK27+/+VRJjN0cO4Cm5RD/ETb4+" + "kI101AFWJ3CYwKEu91SAVTtXLewQYd5LBiJ3Btt68m8d5uUp0a1kAn2cH6g2C7TOfoJ7BgU5+vU81mS8hk9QOFGb9RYUXs8Y" + "cJl21l6WfMMZri5bMI9qCAFn+YZdBBeIbGq+JJcIgjr/R2tULmU584+5beoEydkS" +) + +GEMINI_3_8_FLASH_THINKING_SHOWN_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT, + "thought": True, + }, + { + "text": "I need to inspect `shop/validation.py` next to examine how zero, negative, and non-integer quantities are currently handled in `validate_quantity`.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "end_line": 40, + "start_line": 1, + "path": "shop/validation.py", + }, + "id": "call_386884", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + "finishMessage": "Model generated function call(s).", + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 66, + "totalTokenCount": 529, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 105, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "0C3RavXPIJWX9MoPtorosAs", +} + +GEMINI_2_5_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "I will read the `shop/validation.py` file to understand the quantity validation logic.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "start_line": 1, + "path": "shop/validation.py", + "end_line": 20, + }, + }, + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + "finishMessage": "Model generated function call(s).", + }, + ], + "usageMetadata": { + "promptTokenCount": 334, + "candidatesTokenCount": 53, + "totalTokenCount": 387, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 334, + }, + ], + "serviceTier": "standard", + }, + "modelVersion": "gemini-2.5-flash", + "responseId": "0y3Rar76LO_K-8YPh6Kd0Ak", +} + +VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE: Final = ( + "AY89a18W+yVUlI0b2XsXQujlNKztnSpKXMf0+1JFiprxfp7kc0bG2g9WoPITuNkCqkuBaO5SkZaKDrKQTlJFJKj+H6SnETLk" + "JrR3SAzp2GTjwM2QRsJjI5ky54oylVT93Jh51PWCMp6vl8tZstMJrois8b1i4/eTOyKaVprypLW8DOIImqN3Cbkea8nEDe9r" + "XsdpdvsR7i+NvD0eiCW+Z4ZIflx5NVENbyEezjyJCL3lg3bV2Mfb0ytsACHXuqgGQEcH9HBKQKGkF5n7mA0GNLw6wAqlk8FE" + "JKttP6NgS+HzS/U+xaqxci6B50l0hIMw3pTOHu8Gg5fRKxf/dzvks+s709/ODCGBDrhMzpf/ilnPAMrB67cQrHvQOhnI2pGk" + "cvkZv+ZGwAQxjJHs/r0aXDGU2mwN9RY+/rSEv8y4MWd5c1KIs7xa1GytDBaaG8Zwh3kZzaSOYjVKIk7cy31mPJ01jCNdW9Pg" + "1p2r2iwbzJrBHo5Ku3s1yK4rADeY1WOU63sM6evVkWJEtu6KPxAqE9JkQKQoAPOFtoP89RtiTSl0nZZaO0stbqG7G+AkE2iO" + "5WkKCB1o55R7EEZDalDgFPput8khce5hPrhIGikN3yo11sHsAA2PAuDcHo4L3eWaYPExFiZtB+gVTbBwx42ST3UuBB9lrHtk" + "GCrp7tlOBHJJY+RlXj1G0TJsSxUZDSEswCB+gjG+1tdEchKGRRKOBUjWGARCozBsIEuBOR0F1ugSCLYn7CumS2bLwmMhKY7+" + "mfKXEgF2U84MkmwDEOBSTubNbRhy0qNhVPcLT8pJvGvtTvLny1C90TBTHErVzOKb3Z095fFpUyW7MmFQj5J9rX3A160W0SNQ" + "RKXf2Yr6ikyVhLBSFNvQ00n81iXlyg7bp++ChbbCJBFCUcrGF5WnPLk3+isixodP+G0b+pNtCNo8xMkrvqGqL5AMlmGJ9Qm4" + "a9S6QgDRaw5AzcV/W90a7/BcaIeqkIly3Z9M0+d3Eep4igYj6AK0JzX1BhchMogq9g/dquTTZZk4edVKflBmQcwg4qGGvioI" + "X/ae4tKJVqB2U0Ed+A6t3JZ4Ggdpopdnuy29XvdEyTmUg+fF2sTCgAdK07hueKYO2pkhqj86mSy88iHV0/secWQs0ZuHtcSW" + "D9lAG57yvpMf6RIwccLTyVmM4jcO0A2OB0sTnNIVKYYtswPtUqbrPMTR5saBfEfnULipI521kpb3BseODBKar+G1FISLrb0Q" + "OjMhNnCWCoCSVDk1FbDmtEvSjfiUzeJtFYoqM8PIXyGiBZnsmjGaOw3QMNvLfkoy4gqE6nvFVgwIk3+J6Z1GdFecCPSSCatz" + "LRwwQ8B5ZdlJ6+xjTDWIuyaOByD5BsB98e0YHyuLSYP17w6+0pSGXf3M6AdBJXK7hcL/vu3sGzftYARFHgRY1PQsYCqW8RVy" + "uJL/GwEAT4EgzkDDWufizfUHJEAUVBpQKkxZW7eTbyHyPS/Uiu5a74f5L8lYyep1pAdqt3iK4lvmQ3T9h3DAQIQ15T8PLV78" + "FjzJvZDQ2uMryvoRS1MGEcK6dxc4VvcWalQTjyczE4jQIxEEiN3U5Yo5CaM6FFGDnjIJBiy0fUD16P8e4ZaEaRFQCY4QZ0sA" + "bhwSHWhnU7QWvA/CVnqOVni5TpRsY1ShtEaQXKTslUuz96lB6XwEeaDkPqsPdw+mWY9QmmLqXnXERnTq13rQY1s=" +) + +VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "text": "To understand how zero, negative, and non-integer quantities are handled during checkout validation, we need to inspect `shop/validation.py` next.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "start_line": 1, + "path": "shop/validation.py", + "end_line": 50, + }, + "id": "call_1272842", + }, + "thoughtSignature": VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + }, + ], + }, + "finishReason": "STOP", + }, + ], + "usageMetadata": { + "promptTokenCount": 246, + "candidatesTokenCount": 64, + "totalTokenCount": 583, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 246, + }, + ], + "candidatesTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 64, + }, + ], + "thoughtsTokenCount": 273, + }, + "modelVersion": "gemini-3.8-flash", + "createTime": "2026-10-15T20:02:49.484984Z", + "responseId": "aTHRavjMHYvi6tkPt6OroAo", +} + +GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE: Final = ( + "ErkLCrYLAWkUfRO3TXF/xF4wcqKCnAFmIW9MEkIMZeaFFG/BmstoEFBw+0v0jari9LOu/TkcHKNFvewXanY1T+6rm512Fsvm" + "3cbxU5C7oP8kGPeEHIBOs2BIeocBFx3ZXSmIn25pOLd2YOqDaTzu/De/QX2JUkno3r2Ozbvuy0wvS5EsD3816rYYnaYc7+4E" + "iZRQhf5r6notr70qcHJ3DuKIP8ghDEz0gNSN1bAUa8N5Hvjh7ITbHLoSkeKO0pOjXaSWfQ86xv9F2zhpMZfW+muxcm1+g9tT" + "5QkyJxMtn+cnHi/w4MbQy+9tbJWdkkLo2H6z81pBsmjvIZAQnjjF0zP8aL91/+S9Po8HNV0LsrGs4P4kUgRR33I6UPMCVnHf" + "B7IxmVSEIGeFQ8gg1mJRDj7yZJFBgf61ya97bVWAhiAm6t/EHdnMi9xIPFEYSrY6ed8ZGUnf05x/GRMVOfrmeCePZIHXpftk" + "UfI2W4+qIM1YGR8+GD7KOtQEkcsrjRGHvhmcv1GW3NlLVpZ4LWIqX6ZtmZrOjDFvgnCg/jYdzDHL5j2vQiiSlcpsHJ4UZWRj" + "NmqUyUY0OZC77erTqensEeWJggUfyizdp3InLZtWM/C5B6rhLGku3LgGwAxdSf5aY25i0f+Ayl/xE/nnCkXS81D64z/+Xqdn" + "W4jsZn2tg8TGb11T6GncBK82aE/T9zVUHT7mbUtEljdinxNc/NS+9mWEkhtLSOAQqQmGF24oxMFlCcejVhGZz5QnFpjnu2Uk" + "COv2AZ7S9iYJQa/xT1LyQ18IibvAUWQPRBpjwRnOf0wi0beuE9Q/h0Y2itFjCt+mz5JlFRyQRJcpEOupOsfF47ULZCLeG0Uz" + "Af8JQzHMUWL3ESlOs30cLQQBJwtHW3QD3ZA1uwMCZnJcpBvIr7GhRdF/VhKZ7C8aYZLMOk32M1i+iGe5lYOM2WuoEzbprTGz" + "aVQeQBWEIdyT+RhXjpvLrrsWOIWI/OAw1Y8VEtouLt246yDjDrVuKUdJfXPz9VfnJ/7dKmsYWAaOCBnQ6edZ9HY8vzEKyuma" + "mcaJWSrwuq1mIouQ1Pw4fvazi6JE2xx17hEuOpS5sLD8E89V9uUbkcDne4XFoHNF9mm3nVgTplNFqo5J4kSiEZJ6qJveiMyF" + "sw3F7SbuU0f95AE/HVwMfDhBbnT3j0R5s2zejxLW6mWuASodcn+J9WGbRfMDLZV+aozCPy+rrXrIfeVQw8nLfjQbwY/iQLFM" + "rQCM4m9F5MLrIM0jGKePUeLRmkhJ92LiqEmDMkOZct09jVaZpSW6dRO+ANGfxqJ0y8KppBA7jgWy/PDKMHfwQFWh3Goxc86m" + "otxTUbF2i15mhhPGN2Z5xsRe1lsYRsFV7LNIkqLXcUVQ7uQGmlzi7wCGiMasvxDgzTLdTGAJ6FsaWzr+GAHH/RFSyRA9a6Qr" + "/Z/h9tm5sx5l6hTQrYspK3qK+cYI7BJCuiYYiIMWEAV6qb5SNhMZ88UlTayryuoolQ7pBco59MBV+P6oAoWAXFh69D4yDAe8" + "L2iNa6UTLt2BSMpAgqE5HdrsifGPfwOwNrwtfv/t4eXIBpBVgqC4csZSrN4zUsP6fH3mPZ4oPgwkDFaLu/x1S/PdxKvF0mYg" + "e3ORBogDHyO00QzHPP/9sl2PEgzIQmQv+0tkmpmAeTUfU8FMlRouFsgbrO2wUx1P8KoaFO4ZToQHeNaoqhI2yZ+GB9EquuBp" + "Xpu46qaU1Wrf3w7vtI0aLUPDKnudfZ3b/me3LVm1cgutomL5o33J86U0v62TwJz8xnLB0881oHQVXX4VfaPV4BVk74QCMcqK" + "/SRt3GpK9V1OGQh47ubs5MUFBnqauOnOKmttTxerI5QNi/t5vSwIvdy2MBcXnCXPJSjt3MpZjSZo34PQHK+1iuR1BCbmmHQk" + "v2UCtcYUZfnJ4FuKNXaZ/AFDT7VxbdnOaS8HVA==" +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES: Final = ( + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "To understand how zero, negative, and non-integer quantities are handled during checkout, I need to inspect `shop/validation.py", + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 27, + "totalTokenCount": 673, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 288, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "yi3RaufbH6bXjMcPt5mF2Q0", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "` next.", + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 30, + "totalTokenCount": 676, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 288, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "yi3RaufbH6bXjMcPt5mF2Q0", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "functionCall": { + "name": "read_file", + "args": { + "end_line": 50, + "path": "shop/validation.py", + "start_line": 1, + }, + "id": "call_317677", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 63, + "totalTokenCount": 709, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 288, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "yi3RaufbH6bXjMcPt5mF2Q0", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "", + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 63, + "totalTokenCount": 709, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 288, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "yi3RaufbH6bXjMcPt5mF2Q0", + } + ), +) + +GEMINI_2_5_FLASH_REQUEST: Final = with_thinking_config( + { + **GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, + "contents": [ + list_value(GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request["contents"])[0], + { + "role": "model", + "parts": [ + {"text": "I will locate the quantity validation and its tests."}, + { + "function_call": { + "name": "search_files", + "args": {"query": "quantity", "paths": ["shop/", "tests/"]}, + } + }, + ], + }, + { + "role": "user", + "parts": [ + { + "function_response": { + "name": "search_files", + "response": { + "matches": [ + { + "path": "shop/validation.py", + "line": 18, + "text": "def validate_quantity(quantity):", + }, + { + "path": "tests/test_validation.py", + "line": 7, + "text": "def test_zero_quantity():", + }, + ], + }, + }, + }, + ], + }, + list_value(GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request["contents"])[3], + ], + }, + {"thinkingBudget": 4096, "includeThoughts": False}, +) + + +def gemini_chat_response( + *, + content: str, + tool_call: Mapping[str, JsonValue], + message_fields: Mapping[str, JsonValue], + usage: Mapping[str, JsonValue], +) -> Mapping[str, JsonValue]: + return { + **GEMINI_3_8_FLASH_TEST_CASE.expected_litellm_response, + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "content": content, + "role": "assistant", + "tool_calls": [dict(tool_call)], + "images": [], + **message_fields, + }, + "provider_specific_fields": {"native_finish_reason": "STOP"}, + }, + ], + "usage": dict(usage), + } + + +def signed_tool_call(call_id: str, arguments: str, signature: str) -> Mapping[str, JsonValue]: + return { + "index": 0, + "provider_specific_fields": {"thought_signature": signature}, + "function": {"arguments": arguments, "name": "read_file"}, + "id": f"{call_id}__thought__{signature}", + "type": "function", + } + + +THINKING_REQUEST: Final[Mapping[str, JsonValue]] = { + key: value for key, value in GEMINI_3_8_FLASH_TEST_CASE.litellm_request.items() if key != "reasoning_effort" +} + +GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="thoughts_off", + expected_provider_request=with_thinking_config( + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"thinkingLevel": "high", "includeThoughts": False} + ), + mock_provider_response=GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE, + expected_litellm_response=gemini_chat_response( + content=( + "To understand how zero, negative, and non-integer quantities (such as floats or invalid types) are " + "currently handled and where the checkout rejection is enforced, we need to inspect " + "`shop/validation.py` next." + ), + tool_call=signed_tool_call( + "call_330625", + '{"path": "shop/validation.py", "end_line": 50, "start_line": 1}', + GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + ), + message_fields={ + "thinking_blocks": [], + "provider_specific_fields": {"thought_signatures": [GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE]}, + }, + usage={ + "completion_tokens": 435, + "prompt_tokens": 358, + "total_tokens": 793, + "completion_tokens_details": {"reasoning_tokens": 358, "text_tokens": 77}, + "prompt_tokens_details": {"text_tokens": 358}, + }, + ), +) + +GEMINI_2_5_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="thoughts_off", + litellm_request={**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "model": "gemini/gemini-2.5-flash"}, + expected_provider_endpoint="/models/gemini-2.5-flash:generateContent", + expected_provider_request=GEMINI_2_5_FLASH_REQUEST, + mock_provider_response=GEMINI_2_5_FLASH_THOUGHTS_OFF_RESPONSE, + expected_litellm_response={ + **gemini_chat_response( + content="I will read the `shop/validation.py` file to understand the quantity validation logic.", + tool_call={ + "index": 0, + "function": { + "arguments": '{"start_line": 1, "path": "shop/validation.py", "end_line": 20}', + "name": "read_file", + }, + "id": ANY, + "type": "function", + }, + message_fields={"thinking_blocks": []}, + usage={ + "completion_tokens": 53, + "prompt_tokens": 334, + "total_tokens": 387, + "completion_tokens_details": {"text_tokens": 53}, + "prompt_tokens_details": {"text_tokens": 334}, + }, + ), + "model": "gemini/gemini-2.5-flash", + }, +) + +VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace( + VERTEX_GEMINI_3_8_FLASH_TEST_CASE, + scenario="thoughts_off", + expected_provider_request=with_thinking_config( + VERTEX_GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"thinkingLevel": "high", "includeThoughts": False} + ), + mock_provider_response=VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE, + expected_litellm_response={ + **gemini_chat_response( + content=( + "To understand how zero, negative, and non-integer quantities are handled during checkout validation, " + "we need to inspect `shop/validation.py` next." + ), + tool_call=signed_tool_call( + "call_1272842", + '{"start_line": 1, "path": "shop/validation.py", "end_line": 50}', + VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + ), + message_fields={ + "thinking_blocks": [], + "provider_specific_fields": { + "thought_signatures": [VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE] + }, + }, + usage={ + "completion_tokens": 337, + "prompt_tokens": 246, + "total_tokens": 583, + "completion_tokens_details": {"reasoning_tokens": 273, "text_tokens": 64}, + "prompt_tokens_details": {"text_tokens": 246}, + }, + ), + "model": "vertex_ai/gemini-3.8-flash", + }, +) + +GEMINI_3_8_FLASH_THINKING_OMITTED_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="thinking_omitted", + litellm_request={ + **THINKING_REQUEST, + "thinking": {"type": "enabled", "budget_tokens": 1024, "display": "omitted"}, + }, + expected_provider_request=with_thinking_config( + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"includeThoughts": False} + ), + mock_provider_response=GEMINI_3_8_FLASH_THINKING_OMITTED_RESPONSE, + expected_litellm_response=gemini_chat_response( + content=( + "To understand how the checkout validates quantities—including zero, negative, and non-integer " + "quantities—I need to inspect `shop/validation.py` next." + ), + tool_call=signed_tool_call( + "call_745953", + '{"start_line": 1, "path": "shop/validation.py", "end_line": 40}', + GEMINI_3_8_FLASH_THINKING_OMITTED_THOUGHT_SIGNATURE, + ), + message_fields={ + "thinking_blocks": [], + "provider_specific_fields": {"thought_signatures": [GEMINI_3_8_FLASH_THINKING_OMITTED_THOUGHT_SIGNATURE]}, + }, + usage={ + "completion_tokens": 173, + "prompt_tokens": 358, + "total_tokens": 531, + "completion_tokens_details": {"reasoning_tokens": 108, "text_tokens": 65}, + "prompt_tokens_details": {"text_tokens": 358}, + }, + ), +) + +GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="thinking_shown", + litellm_request={ + **THINKING_REQUEST, + "thinking": {"type": "enabled", "budget_tokens": 1024, "display": "summarized"}, + }, + expected_provider_request=with_thinking_config( + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"includeThoughts": True} + ), + mock_provider_response=GEMINI_3_8_FLASH_THINKING_SHOWN_RESPONSE, + expected_litellm_response=gemini_chat_response( + content=( + "I need to inspect `shop/validation.py` next to examine how zero, negative, and non-integer quantities " + "are currently handled in `validate_quantity`." + ), + tool_call=signed_tool_call( + "call_386884", + '{"end_line": 40, "start_line": 1, "path": "shop/validation.py"}', + GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT_SIGNATURE, + ), + message_fields={ + "reasoning_content": GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT, + "thinking_blocks": [{"type": "thinking", "thinking": GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT}], + "provider_specific_fields": {"thought_signatures": [GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT_SIGNATURE]}, + }, + usage={ + "completion_tokens": 171, + "prompt_tokens": 358, + "total_tokens": 529, + "completion_tokens_details": {"reasoning_tokens": 105, "text_tokens": 66}, + "prompt_tokens_details": {"text_tokens": 358}, + }, + ), +) + + +def with_request(case: TranslationTestCase, scenario: str, **fields: JsonValue) -> TranslationTestCase: + return replace(case, scenario=scenario, litellm_request={**case.litellm_request, **fields}) + + +def with_thinking(case: TranslationTestCase, scenario: str, display: JsonValue) -> TranslationTestCase: + return with_request(case, scenario, thinking={"type": "enabled", "budget_tokens": 1024, "display": display}) + + +def with_summary(case: TranslationTestCase, scenario: str, summary: JsonValue) -> TranslationTestCase: + return with_request(case, scenario, reasoning_effort={"effort": "high", "summary": summary}) + + +@pytest.mark.parametrize( + "case", + ( + GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, + GEMINI_2_5_FLASH_THOUGHTS_OFF_TEST_CASE, + VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, + ), + ids=( + GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.id, + GEMINI_2_5_FLASH_THOUGHTS_OFF_TEST_CASE.id, + VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.id, + ), +) +def test_gemini_keeps_thinking_and_billing_reasoning_without_thought_text_when_the_gateway_turns_it_off( + case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, thoughts_off, provider) + + +SUMMARY_ASKS: Final = ( + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_auto", "auto"), + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_empty", ""), + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_long", LONG_SUMMARY), + with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_int", 7), + with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_list", ["auto"]), + with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_null", None), + with_request(GEMINI_3_8_FLASH_TEST_CASE, "summary_alias", reasoning_summary="auto"), + with_request(GEMINI_3_8_FLASH_TEST_CASE, "summary_alias_camel", reasoningSummary="detailed"), + with_request( + GEMINI_3_8_FLASH_TEST_CASE, + "summary_alias_effort_dict", + reasoning_effort={"effort": "high"}, + reasoning_summary="auto", + ), + with_request(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_alias_null", reasoning_summary=None), + with_request( + GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, + "thinking_summary", + thinking={"type": "enabled", "budget_tokens": 1024, "summary": "auto"}, + ), +) + + +@pytest.mark.parametrize("case", SUMMARY_ASKS, ids=tuple(case.id for case in SUMMARY_ASKS)) +def test_only_a_string_reasoning_summary_brings_thought_text_back_over_the_gateway_default( + case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, thoughts_off, provider) + + +def without_closed_schema(tool: JsonValue) -> JsonValue: + function: Final = object_value(object_value(tool)["function"]) + parameters: Final = object_value(function["parameters"]) + open_parameters: Final = {name: value for name, value in parameters.items() if name != "additionalProperties"} + return {**object_value(tool), "function": {**function, "parameters": open_parameters}} + + +def cacheable_run(case: TranslationTestCase, marker: str) -> TranslationTestCase: + messages: Final = list_value(case.litellm_request["messages"]) + system: Final = object_value(messages[0]) + instruction: Final = object_value(case.expected_provider_request["system_instruction"]) + (part,) = list_value(instruction["parts"]) + text: Final = f"{string_value(system['content'])} Review run {marker}." + return replace( + case, + litellm_request={ + **{name: value for name, value in case.litellm_request.items() if name != "cache"}, + "messages": [{**system, "content": text}, messages[1]], + "tools": [without_closed_schema(tool) for tool in list_value(case.litellm_request["tools"])], + }, + expected_provider_request={ + **case.expected_provider_request, + "contents": [list_value(case.expected_provider_request["contents"])[0]], + "system_instruction": {**instruction, "parts": [{**object_value(part), "text": text}]}, + }, + ) + + +def response_cache_keys() -> frozenset[str]: + keys: Final = (raw.decode() for raw in redis_store().scan_iter()) + return frozenset(key for key in keys if CACHE_KEY.match(key)) + + +def stored_reply_key(before: frozenset[str]) -> str: + (key,) = eventually(lambda: response_cache_keys() - before, lambda added: len(added) == 1) + return key + + +GEMINI_3_8_FLASH_SUMMARY_ALIAS_TEST_CASE: Final = with_request( + GEMINI_3_8_FLASH_TEST_CASE, "summary_alias", reasoning_summary="auto" +) + + +@pytest.mark.parametrize( + ("first", "second"), + ( + (GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, GEMINI_3_8_FLASH_SUMMARY_ALIAS_TEST_CASE), + (GEMINI_3_8_FLASH_SUMMARY_ALIAS_TEST_CASE, GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE), + ), + ids=("plain_then_summary_alias", "summary_alias_then_plain"), +) +def test_a_summary_alias_ask_and_a_plain_request_never_share_a_cached_reply( + first: TranslationTestCase, second: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider +) -> None: + marker: Final = uuid4().hex + before: Final = response_cache_keys() + assert_translation(cacheable_run(first, marker), thoughts_off, provider) + first_key: Final = stored_reply_key(before) + assert_translation(cacheable_run(second, marker), thoughts_off, provider) + second_key: Final = stored_reply_key(before | {first_key}) + repeat: Final = thoughts_off.request("POST", second.litellm_endpoint, cacheable_run(second, marker).litellm_request) + assert repeat.status_code == 200, repeat.text + assert repeat.headers.get("x-litellm-cache-key") == second_key, sorted(repeat.headers) + assert provider.received() == () + + +DISPLAY_ASKS: Final = ( + GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, + with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_updates", "updates"), + with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_int", 7), + with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_empty", ""), + with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_list", ["omitted"]), + with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_long", LONG_SUMMARY), + GEMINI_3_8_FLASH_THINKING_OMITTED_TEST_CASE, +) + + +@pytest.mark.parametrize("case", DISPLAY_ASKS, ids=tuple(case.id for case in DISPLAY_ASKS)) +def test_any_thinking_display_but_omitted_brings_thought_text_back_over_the_gateway_default( + case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, thoughts_off, provider) + + +DEFAULT_GATEWAY_DISPLAYS: Final = ( + with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_omitted_default", "omitted"), + GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, +) + + +@pytest.mark.parametrize("case", DEFAULT_GATEWAY_DISPLAYS, ids=tuple(case.id for case in DEFAULT_GATEWAY_DISPLAYS)) +def test_thinking_display_never_turns_thought_text_off_on_the_default_gateway( + case: TranslationTestCase, gateway: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, gateway, provider) + + +def test_a_streamed_turn_keeps_thinking_and_billing_reasoning_without_thought_text( + thoughts_off: Gateway, provider: SharedProvider +) -> None: + provider.expect( + Reply( + content_type="text/event-stream", + chunks=tuple(encode_sse_frame(frame) for frame in GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES), + ) + ) + response: Final = thoughts_off.request( + "POST", + "/v1/chat/completions", + {**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "stream": True, "stream_options": {"include_usage": True}}, + ) + assert_provider_request( + provider.received(), + "/models/gemini-3.8-flash:streamGenerateContent?alt=sse", + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_headers, + GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.expected_provider_request, + ) + assert response.status_code == 200, response.text + frames: Final = tuple(frame.data for frame in parse_sse_frames(response.text) if frame.data != "[DONE]") + deltas: Final = tuple( + object_value(object_value(choice)["delta"]) + for frame in frames + for choice in list_value(frame.get("choices", [])) + ) + assert "".join(str(delta.get("content") or "") for delta in deltas) == ( + "To understand how zero, negative, and non-integer quantities are handled during checkout, I need to " + "inspect `shop/validation.py` next." + ) + assert [delta for delta in deltas if delta.get("reasoning_content") or delta.get("thinking_blocks")] == [] + usage: Final = [object_value(frame["usage"]) for frame in frames if frame.get("usage")][-1] + assert {field: value for field, value in usage.items() if field != "cost"} == { + "completion_tokens": 351, + "prompt_tokens": 358, + "total_tokens": 709, + "completion_tokens_details": {"reasoning_tokens": 288, "text_tokens": 63}, + "prompt_tokens_details": {"text_tokens": 358}, + } + assert isinstance(usage["cost"], float) and usage["cost"] > 0 + + +@pytest.mark.parametrize( + ("configured", "case"), + (("false", GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE), (None, GEMINI_3_8_FLASH_TEST_CASE)), + ids=("string-false", "null"), +) +def test_a_flag_written_as_text_or_left_null_still_sends_gemini_a_boolean( + configured: JsonValue, case: TranslationTestCase, gateway: Gateway, provider: SharedProvider, tmp_path: Path +) -> None: + with gemini_thoughts_proxy(gateway, tmp_path, configured) as owned: + assert_translation(case, owned.gateway, provider) + + +BURST_ASKS: Final[tuple[tuple[str, Mapping[str, JsonValue], Mapping[str, JsonValue]], ...]] = ( + ( + "/v1/chat/completions", + {**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "reasoning_effort": {"effort": "low", "summary": "auto"}}, + {"thinkingLevel": "low", "includeThoughts": True}, + ), + ( + "/v1/chat/completions", + GEMINI_3_8_FLASH_TEST_CASE.litellm_request, + {"thinkingLevel": "high", "includeThoughts": False}, + ), + ( + "/v1/messages", + MESSAGES_GEMINI_3_8_FLASH_TEST_CASE.litellm_request, + {"thinkingLevel": "high", "includeThoughts": True}, + ), + ( + "/v1/messages", + {**MESSAGES_GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "thinking": {"type": "adaptive"}}, + {"thinkingLevel": "high", "includeThoughts": False}, + ), + ( + "/v1/responses", + {**RESPONSES_GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "reasoning": {"effort": "high", "summary": "auto"}}, + {"thinkingLevel": "high", "includeThoughts": True}, + ), + ( + "/v1/responses", + RESPONSES_GEMINI_3_8_FLASH_TEST_CASE.litellm_request, + {"thinkingLevel": "high", "includeThoughts": False}, + ), +) + + +def _assert_each_request_follows_its_own_ask(gateway: Gateway, provider: SharedProvider) -> None: + burst: Final = BURST_ASKS * 4 + for _ in burst: + provider.expect(Reply(body=json.dumps(GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE).encode())) + with ThreadPoolExecutor(max_workers=12) as pool: + responses: Final = tuple(pool.map(lambda ask: gateway.request("POST", ask[0], ask[1]), burst)) + assert [response.status_code for response in responses] == [200] * len(burst), [ + response.text for response in responses if response.status_code != 200 + ] + sent: Final = sorted(json.dumps(sent_thinking_config(request), sort_keys=True) for request in provider.received()) + assert sent == sorted(json.dumps(expected, sort_keys=True) for _, _, expected in burst) + + +def test_every_worker_honors_the_flag_on_every_endpoint_under_a_burst_and_after_a_worker_respawns( + gateway: Gateway, provider: SharedProvider, tmp_path: Path +) -> None: + with gemini_thoughts_proxy(gateway, tmp_path, False, workers=2) as owned: + _assert_each_request_follows_its_own_ask(owned.gateway, provider) + owned.kill_worker(min(owned.worker_pids())) + _assert_each_request_follows_its_own_ask(owned.gateway, provider) diff --git a/tests/integration/translation/gemini_thoughts.py b/tests/integration/translation/gemini_thoughts.py new file mode 100644 index 00000000000..e0d89c3d43e --- /dev/null +++ b/tests/integration/translation/gemini_thoughts.py @@ -0,0 +1,63 @@ +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Request +from pydantic import JsonValue + +GEMINI_2_5_FLASH_DEPLOYMENT: Final[Mapping[str, JsonValue]] = { + "model_name": "gemini/gemini-2.5-flash", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_base": "http://127.0.0.1:8191", + "api_key": "synthetic-gemini-key", + }, +} +LONG_SUMMARY: Final = "detailed " * 600 + + +def with_thinking_config( + request: Mapping[str, JsonValue], thinking_config: Mapping[str, JsonValue] +) -> Mapping[str, JsonValue]: + return { + **request, + "generationConfig": {**object_value(request["generationConfig"]), "thinkingConfig": dict(thinking_config)}, + } + + +def sent_thinking_config(request: Request) -> Mapping[str, JsonValue]: + body: Final = JSON_OBJECT.validate_json(request.body) + return object_value(object_value(body["generationConfig"])["thinkingConfig"]) + + +@contextmanager +def gemini_thoughts_proxy( + gateway: Gateway, directory: Path, default: JsonValue, *, workers: int = 1 +) -> Iterator[OwnedProxy]: + stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = directory / "gemini-thoughts.yaml" + config.write_text( + yaml.safe_dump( + { + **stock, + "model_list": [*list_value(stock["model_list"]), dict(GEMINI_2_5_FLASH_DEPLOYMENT)], + "litellm_settings": { + **object_value(stock["litellm_settings"]), + "gemini_include_thoughts_default": default, + }, + } + ) + ) + with owned_proxy_process(gateway, directory, {}, config=config, workers=workers) as owned: + yield owned + + +@pytest.fixture(scope="module") +def thoughts_off(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + with gemini_thoughts_proxy(session_gateway, tmp_path_factory.mktemp("gemini-thoughts-off"), False) as owned: + yield owned.gateway diff --git a/tests/integration/translation/messages/reasoning/test_messages_reasoning_gemini.py b/tests/integration/translation/messages/reasoning/test_messages_reasoning_gemini.py new file mode 100644 index 00000000000..c350d6e2a10 --- /dev/null +++ b/tests/integration/translation/messages/reasoning/test_messages_reasoning_gemini.py @@ -0,0 +1,373 @@ +from collections.abc import Mapping +from dataclasses import replace +from typing import Final + +import pytest +from integration._support.client import Gateway, object_value +from integration._support.process import graceful_stop_seconds +from integration._support.provider import SharedProvider +from integration._support.wire import Reply +from integration.translation.case import SSEFrame, TranslationTestCase +from integration.translation.gemini_thoughts import LONG_SUMMARY, thoughts_off, with_thinking_config +from integration.translation.messages.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE, GEMINI_MESSAGES_RESPONSE +from integration.translation.runner import ( + assert_provider_request, + assert_translation, + encode_sse_frame, + parse_sse_frames, +) +from pydantic import JsonValue + +__all__ = ("thoughts_off",) + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE: Final = ( + "EscLCsQLAWkUfRP4jtXWjNAugx5FeAwxasL/kKzH9gXUHHB2cTK13h4C8C6tKVYMKP5rXoPq4yhoXSrsDiIjUEpuAjdQYYCz" + "Wvx21tGAN4iBA4VELBqPOi92vo9YVGpczJ/XVREq9KNRFIkyMzWyEzR1lBbJFA4detKHzSVJ7VQ1iRBn9jk6xvknTCxQhbnl" + "Pct8TfMo4XluFsM2PY3qAqoX1DjkrjQNe4nrN9Yoglq8Bwr2ayyc3DUNkNJDu0DQN20W6HU7fbeqSizEOsvGIvJVEUZhMOFF" + "0r8NyzUGkNiiAZ+TEupIV+G0eTr6sVlMPLS1N74ryQ55sevdj6irenHlP19IwkiDNhy6ApRIoblY5I1uJVt4lj4tQyACGMAE" + "VOTg3Gb78hPgJ+8cF3ukmRF0H1bpYH3n/mpnkqiggfK3vm7fVnsCy03W7Keyl7xd7Gfto/HL7vPCta/sO6ZmZm28JvdCr/Ce" + "X/W82nHuEqcjiR8fNPk413rL/geE4Lppru1H9ESWBICXzwl6pLxaBBNr8BD/jGMvlYA/ZEcBfKWtYyzow3+yp+4M0fsY1Ah7" + "pxVCcCbivMRv2X5J5m9btU1IIm2REJicjMSWCxh4lQMEuuk2FaPzge+wXROC/Zx0HwNw2iPw6VpziTOByACQxXOE7RdWULq1" + "3rkbyXCMkuAfG281SiylOnJ3I3tkbRcFoTaY5DcMBKYCCmoJJZPO1HKJ2DMw5Spz1HDftQ7lg9df5FDoZbmehUZKfY+j/2LN" + "//Wr5OEUmuBpojC5ht3YYgXaxclo/pa70KnLHlQWHPZ3nqdB2ru0tqgOaXSv9KZNhtvXgagLsmffQ/rm5QiRoEeG4oZ5E6Ga" + "VLQB9OvWCDNdzgNVzss0tveNcG/cttJBnkSd5/azHSvzkCDQYh4J/GT3zDVvOyOayWTqzGSLgIoFGpTES7Vsbpph62apefvO" + "+ZAlxPNNo/3tHncHzGc8Pku07fF1yhiKdGRV1OjYndfQv6UuDc3G7n2oU90iCKWRoCqOYAotzw1yCh4lbsSpEWO2iB9kWW93" + "mJsdqQBdk6e9UiuHdJIOfQ+cOVNOrltXXtZZ+T2mKFTFaWaN1F7MMH+xIX10dcnsWzq5o0Qz/tpwEfFOaQ5j43ez5BJdE5yD" + "wj4sDIdY57jAQ4141jU2oHpJCjY5Obhd4692k/FRtN7IMWMn5rQS+uTctWOh/27qatdy+QeAMPZseT4Yoi3W6W//Um9+sJY4" + "RjJ0Q27Fjb/wEqBkGsqx1sIidQc19sZWa2jQxVgG3T8xHYCMbgkqT/OHJaHGa4gxo78ajhdQnGYNWtvuRlGDpdUsxcNfF+6z" + "kbqWgcxKleRnQ1jKKZKS9GeiLoJnBRk8wIgmnrleblMB0GRj8i124HKLTfAm3WHlqy3xpACI3AIxPbk1GrATdkrb496eecVG" + "ko9YTrsdgOn5RV6Ggy/Rxr5wAAvPKBrqx/u1Whyr05pge9AZcGAtFRIZMWRCBnNg+oHE8M1EVipx3wQQXFDFO27Jpr3IandO" + "GgpbjbpwcW0fnDKLw7Ndt8SCyyVUEZwZ4Vp4tXJCOv9rzIMHJi+68Upwp5ILirrBFT5vKs//ub6l5hHzEfZty66+DvXyrpqT" + "VLYRnYMEO10Te24QyYVCO4qSnbJzfwTv5Kp4ODP2ONdSuJgkyYy78MzZ926H1/8rOlTEbWpAYb+QJwrt1nBJBkWRNDi2QNrl" + "UlFKTXZNLNXTP/rHHqfw1us8bhwe4Up2BbECPO+8F0tgSifXusnShtsRwpW7o+z/iXpnwuWePH92YiORo43mdoG4MJ54Emwf" + "B7x9ftguHpq1ZDOBoa3frE8l8a5RIEHemkzQsLlPyZSB+eQgnXrToXwUeEfR3EHG0vSCfokMLk0HinhWGa6hgIyJZQAdUIg4" + "GALvjSsrGEHwJaeLNNn/9EscJOEGyFQFMEy+sVIkL6Pj2+iZuG0OKahk" +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "I need to inspect `shop/validation.py` next to examine how `validate_quantity` handles zero, negative, and non-integer quantities.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "end_line": 40, + "path": "shop/validation.py", + "start_line": 1, + }, + "id": "call_412671", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + "finishMessage": "Model generated function call(s).", + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 64, + "totalTokenCount": 739, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 321, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "1C3RatW-Jtuc_uMPhaLroA0", +} + +GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE: Final = ( + "Eo0LCooLAWkUfRMuNip0QePQVa4NCyMrc0dwJR2qI/PSRcyScXfwpZhjBJVTorCs5ei6D4xzVf9attfTHx8bOBCX9wXynQ6V" + "JxU6uN562z5VAyrgoGv2nVGJk/mT7SnzIFiOc4RN+xebNGO0BIDf95xyOMeYP+ZHXquqtx1SnQXw7V6p81wTA4cGMnqjNE7F" + "Vs7OqlcpvzeFUmE8gl/Cy1a+u5SgLHwfnvOR3uzIX3u8kdk//dzvCKqXYaRStDf0dsrJICex/9WoOBfZKCujLBMK7CCjKEWF" + "nqERtFxXayTUKz1dvvBkcZ+MbhtrG3RFkBegqokQn9lus3icwBAfkZ9b3ps42j/mYw7DfjR+3Dg/Pyz8mgJVRC312QRxt+70" + "yOCU8POm/GTHJAZvvxdm9U4eiYBs3789F5oV34YzZl1+pcE4WVRwiFkB0vi17xy5XJFYTtzuk9d9ySr5Fz1de6nTKWk6TpuX" + "PM4q8bhX51kJWPCr1W4oVn1GhKpCtRyrXzKhuMsWcHXu/yKHUVjxyEFY+cTsVtHfrMHsIhMHyNnBqz2JclGXtFZScaq6AZ8g" + "0JZPqqAXgfgNyx2RaxoaCtITznOlBwMC1hTAZCrtjhHQNvGsAEdaRBAnZRFgeDHkuL7r7D9tr38UdxaUf/R5vfl5eO0ksj4X" + "bdeUG/LAnWP2CLUVjIoF50o4ZHKRy5UT/Mhr45986ePmtWnIN9t6dXOCB5I4qoEaljwkDn+llwZMbPq1orHdSrUAR6flcl3h" + "qVutnVOFgN3NeHvNkrkjUCvckiQL/ZBiG5YeDKfMP+XECDWIoTS2/DCJ7SX6d44B2yJ/PO2PH3nm7hr81U7r7xYKly4/2YFX" + "8bAsFJt4lGihfDHXFOr/VbQQd/CobZtLV9UZP7m1wUGXfL0knP31bdFWDKmSp/2giddbpIsGQf6j7h+YiE3M6QEfvRi7kaDX" + "CQ16O/4XJzG85/BxPw8ssR1cN+xLM95pZkDQO2MN1rxLVogF+T+Xc0M2SSjwbahM4nwYmisOqXv4Wuy0V7hQYJUySCvHabdH" + "nmTi+3mIfs9/zzHT3QUGZvAumBCZYfz1kqyz6mytXSwQbhhktmDLgIaj1abGm2gSn10fxl8pDm/mTq8/owS9aNPyDYChNBMr" + "NGYEBMXA4DtZy9TjybzMXKGPi4d/M7EZM5sC33WnCl2JovveyeZUNPWkZTV0+TjSSqk8C488huCOzv+PgJMvtBq6XiQW9ZV1" + "HQWk8XRMK7VMMhuv9sRqy1Lh/h8kIlSWtaCCNwpqIgUlZ+xybbY9dRGdsWq5FoxCU/OLmlSv7BcqXAamxxfIK9inXRKd/xkI" + "bR7OCeIaFq97Z5M08LfUpu9WjcXGgvtsppOUB4Rik2uCBVSzA/OqQdY1/N3q+BPYCvhOVVEBOJEEAvH9iAQOtCw4peaCtSLL" + "7vLh7s/lIZGBYj5I4zjoOfoSuT9g1x9uvv9h2sJUeP1HsN9Mp95weKBbF3HUN5F/lVRTan2I+Wb2+WnMMbashS1hoVIfk/Bc" + "426M3xPy7JDsczXLUP0PAyhAc2+xiFnkI120l4IHT3JK53AArx7svpTxzLOvPffVIcodkyWXHeJeC6eQbD6Elt/wWbx5Ig1L" + "mJQP7egSoyqmzu0zpuIC0xVsmDdgnumSi7acTjp1mIsLHHLLaWPcpxaNeiwBQf/7FfKsGqJ6p+I+lzt/rdRDOR4apkwKRwnr" + "d5beSf38TJd7L8Zj1qXp36Q2rnvtKubSm7f451w9ZbKsqbMDrwoagNWsfG9kc0CR0wtraSRzBHD6GUfoapPbplOLsSyjzEvX" + "i48aG814wRmQk932cd8ETx1a9BIbN/bsOr6stttugYakpvD3b8ckhfXW7eXfpKywBBzu1s1KLrU=" +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES: Final = ( + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "I need to read `shop/validation.py` next to see how `validate_quantity` currently handles zero, negative, and", + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 27, + "totalTokenCount": 676, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 295, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "2C3RavD_A77W-sAP9brxyQk", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": " non-integer quantities.", + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 32, + "totalTokenCount": 681, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 295, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "2C3RavD_A77W-sAP9brxyQk", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "functionCall": { + "name": "read_file", + "args": { + "start_line": 1, + "end_line": 50, + "path": "shop/validation.py", + }, + "id": "call_654547", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 65, + "totalTokenCount": 714, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 295, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "2C3RavD_A77W-sAP9brxyQk", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "", + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 65, + "totalTokenCount": 714, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 295, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "2C3RavD_A77W-sAP9brxyQk", + } + ), +) + +ADAPTIVE_THINKING: Final[Mapping[str, JsonValue]] = {"type": "adaptive"} + +GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="thoughts_off", + litellm_request={**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "thinking": dict(ADAPTIVE_THINKING)}, + expected_provider_request=with_thinking_config( + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"thinkingLevel": "high", "includeThoughts": False} + ), + mock_provider_response=GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE, + expected_litellm_response={ + **GEMINI_MESSAGES_RESPONSE, + "id": "1C3RatW-Jtuc_uMPhaLroA0", + "model": "gemini/gemini-3.8-flash", + "content": [ + { + "type": "text", + "text": ( + "I need to inspect `shop/validation.py` next to examine how `validate_quantity` handles zero, " + "negative, and non-integer quantities." + ), + }, + { + "type": "tool_use", + "id": "call_412671", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 40}, + "provider_specific_fields": {"signature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE}, + }, + ], + "usage": {"input_tokens": 354, "output_tokens": 385}, + }, +) + + +def with_thinking(case: TranslationTestCase, scenario: str, field: str, value: JsonValue) -> TranslationTestCase: + return replace( + case, + scenario=scenario, + litellm_request={**case.litellm_request, "thinking": {**ADAPTIVE_THINKING, field: value}}, + ) + + +THINKING_ASKS: Final = ( + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_summarized", "display", "summarized"), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_updates", "display", "updates"), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_int", "display", 7), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_empty", "display", ""), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_list", "display", ["omitted"]), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_long", "display", LONG_SUMMARY), + with_thinking(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "display_omitted", "display", "omitted"), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "summary_auto", "summary", "auto"), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "summary_empty", "summary", ""), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "summary_long", "summary", LONG_SUMMARY), + with_thinking(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_int", "summary", 7), + with_thinking(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_list", "summary", ["auto"]), + with_thinking(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_null", "summary", None), +) + + +def test_gemini_keeps_thinking_and_billing_reasoning_without_thought_text_when_the_gateway_turns_it_off( + thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, thoughts_off, provider) + + +@pytest.mark.parametrize("case", THINKING_ASKS, ids=tuple(case.id for case in THINKING_ASKS)) +def test_a_thinking_display_or_summary_ask_brings_thought_text_back_over_the_gateway_default( + case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, thoughts_off, provider) + + +def test_display_omitted_alone_keeps_thought_text_on_the_default_gateway( + gateway: Gateway, provider: SharedProvider +) -> None: + assert_translation( + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_omitted_default", "display", "omitted"), gateway, provider + ) + + +def test_a_streamed_turn_keeps_thinking_and_billing_reasoning_without_thought_text( + thoughts_off: Gateway, provider: SharedProvider +) -> None: + provider.expect( + Reply( + content_type="text/event-stream", + chunks=tuple(encode_sse_frame(frame) for frame in GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES), + ) + ) + response: Final = thoughts_off.request( + "POST", "/v1/messages", {**GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.litellm_request, "stream": True} + ) + assert_provider_request( + provider.received(), + "/models/gemini-3.8-flash:streamGenerateContent?alt=sse", + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_headers, + GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.expected_provider_request, + ) + assert response.status_code == 200, response.text + events: Final = tuple(object_value(frame.data) for frame in parse_sse_frames(response.text)) + started: Final = tuple( + object_value(event["content_block"])["type"] for event in events if event.get("type") == "content_block_start" + ) + assert started == ("text", "tool_use") + deltas: Final = tuple( + object_value(event["delta"]) for event in events if event.get("type") == "content_block_delta" + ) + assert "".join(str(delta.get("text", "")) for delta in deltas) == ( + "I need to read `shop/validation.py` next to see how `validate_quantity` currently handles zero, negative, " + "and non-integer quantities." + ) + assert [object_value(event["usage"]) for event in events if event.get("type") == "message_delta"][-1] == { + "input_tokens": 354, + "output_tokens": 360, + } diff --git a/tests/unit/messages/__init__.py b/tests/integration/translation/responses/reasoning/__init__.py similarity index 100% rename from tests/unit/messages/__init__.py rename to tests/integration/translation/responses/reasoning/__init__.py diff --git a/tests/integration/translation/responses/reasoning/test_responses_reasoning_gemini.py b/tests/integration/translation/responses/reasoning/test_responses_reasoning_gemini.py new file mode 100644 index 00000000000..984f7d2a8c4 --- /dev/null +++ b/tests/integration/translation/responses/reasoning/test_responses_reasoning_gemini.py @@ -0,0 +1,400 @@ +from collections.abc import Mapping +from dataclasses import replace +from typing import Final +from unittest.mock import ANY + +import pytest +from integration._support.client import Gateway, list_value, object_value +from integration._support.process import graceful_stop_seconds +from integration._support.provider import SharedProvider +from integration._support.wire import Reply +from integration.translation.case import SSEFrame, TranslationTestCase +from integration.translation.gemini_thoughts import LONG_SUMMARY, thoughts_off, with_thinking_config +from integration.translation.responses.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE, GEMINI_RESPONSES_RESPONSE +from integration.translation.runner import ( + assert_provider_request, + assert_translation, + encode_sse_frame, + parse_sse_frames, +) +from pydantic import JsonValue + +__all__ = ("thoughts_off",) + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE: Final = ( + "EvQLCvELAWkUfRM8dwxvoWqyX5HLMaxcqWN9o4LPX6M6MZM/OD962KfMSvacFOxwOEb2B3rQM7QgoyT58YrItR4zpDSZJUxA" + "N4shekmtkSRpwrqE0rmH1v9AE8w5lR0Jewj8xmiiUXUrK0qvpyxhbVJRmIgdt9CqW+B4bJHYpp1VFAuHMk8eyBRoVNXakpsG" + "xjvHWGDomXbyJ7XuG1uHzvA5QtV2lgKUwRNvQqbFgbMhKN4dc8Imm2kwCLQONRN6wYIKJnYIi6ib+MhsFY+hZ6sPysMzVyrk" + "nGhc6pZwOKHbUjAraY/X98HO7UEL/q86LKUIvKTS4FfZ8keuyIQQBMdX68fdZPqurR4fI2dr5pFSDbaJWEVaE9XAY9wOXhOE" + "uC5SxY0U7nPU39eycNqAoAf+O7i/uHn1eCfLMhvAZyoH0z6vMfofnVn6L6alZtnRFQ3rEdJWEZHFmsLOk7G9SHlLAYKQjEdj" + "5aQHsKqBG1SFODPWSeKduI5w/kgFwgrw1ye1rnFTd+FDBs3DsQUlwh+jbPg8tr3KnCL5f2bxoGShUanPcGRAJ4b/IWGM3+wf" + "Whcx+3JY/OfuklkCnw86uRlcJiEjWd9XZ+jAby3ECZxBuXsHSE+Ynzpd0mp5g/S1orzqlmUJ/PlOcbDR/AsFduZvZwG0n4hq" + "hbXIIPszIpEyGDhkC//vxXsWdo5gLcsnEoUlLCESN3Sc12goKneId6u8pVWfzvHwRa0pK0LIWo4GcWZepZpEoYARlUZN6Dat" + "FNJ0nFW1sGaNb4CEYdSSp3m+GVobHpdz8qsGJxWrzzc6ZYmO7LbDJ25dJtyCl9QcfyzGbkFVDdAaTBEM23EdyZMFc2IIT24K" + "OVx5esCx+VjDZOmOzue/ZR0xrfaPXVa0CMuzGA8WwiphpWTbD8xZUnqEUZQNwKrjU91+SIhaBKcTVM8F1Wj0t9vlQkkMlzth" + "bJxIw+Ccu51gh/s8/awIDsBwC2Ax6eYIhUoWxUWYcphX7r3lBKhIuzv7FyO8tmaRkIOy6Eb1lXHB4fHAIOimZ/n4ClgK/KOZ" + "XVLosaGMKsmjYAdrxxPaNqs+say+pQ08bec/c+mqWeDbTOJ2JrCKM/Q+EcyD+82QqWlO68boeJyyt0fIvTCAdn57xN5Th7r1" + "WGhT9YHeSB1Y77P9FaJLQqAOjr+yogG7suiCP4QHg1arBlVzvaxAdnlBeXSv6pbQJuuvsFV9+SfwKHVXaBv68bcjSYKm8YZa" + "vfK57jweJUxvknCUpzrDRBdVZX28H60eQ3X6S0cV/696AO8NHoKuMv9zsCyu4WXWnpYEObtVD4Sn5hdSn1tnOkyUTLmN+bHV" + "sRGe5WTmrgnTKbrGDqpeqFpYLeMftkKWX+DHZj/LvZiNH1pHJid4nbEvcUKHciculasWiLZXDhEDKA1uzvSeG7CHjtHziqRi" + "b0dY7pQ0iEJUrjIwMVPnFZoyrywYZVPYgk+jna9FJypscuCz5XzdV+N2b2yrKryU6NhOQ5u/lzMNEFsyYxuXp6LZ5CVgdH1T" + "k+1rUf3sWPTO+CyDj+GzWvo5ljP7HihqyitafZJeKOSfaPfrJpBq+R8jnwDOVLCS1lnkhsof12PFoQHAE0YzGDtVscrstW4O" + "OaEdtNyNh30fhwwaZ/XGM4O9yJ8l7CgUvrgY/t8/f1hB/ZIc980HGs3G/bmVBYTxeZ18rqanrURtVuNtveyGief6eiEVIFjo" + "HWkAFXM+sn1zsJI6l4N9TPMbkoPxLCw2KaOaI/7mQjsH6VktpGuuqpHw3S5tiw3OGwh0lEwydKT0bk4hUhWqoBHqFpCrmKqA" + "g3z0kEHxQjQ1j8GmHusdPp4odrCWphw44m1VeWYjF4VjQ16Vik3SSrC8g3v9sHhtKM+3RL6z8zzukqgcQkXf7RXWvmFttckj" + "yYWM+vfQm/UvCzQ4dKvFHMqgmwuyEISPO9CD68/518wNpbQD1iB/RnxKwyYoD4u9YQ9TzXCzk2eP0NBB6pRmGg5nBo2SP0PB" + "BRZqlxnFX4XfQiS9pxNL" +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "To see how zero, negative, and non-integer quantities are currently handled and rejected by the validation logic, the next file I need to inspect is `shop/validation.py`.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "path": "shop/validation.py", + "start_line": 1, + "end_line": 50, + }, + "id": "call_584229", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + "finishMessage": "Model generated function call(s).", + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 70, + "totalTokenCount": 790, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 366, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "2y3RasgqkpyMxw-RxpeQBg", +} + +GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE: Final = ( + "EvYSCvMSAWkUfROtzl1tHuKrIElCi7xsI6n+7cPuqsM/qZSYIpK3gONKCCo0E4moadzMQ3sN09ll3gRjbcBg6rKSck+JTBhd" + "mrLsP/sqfxiYfMx4lWzDeq/Pw/IPhT2g5/2EvSU4FDsUczorTsD55gKvgRT8v05YTFMGF/np8r3VT52Q2yAzwRK+31qn8pYS" + "c25if5coTZJrvmJZDlZ9/64aibghRps9zN07yRVXpUjtLyKMaTDH0loGgiQn0I7Cxfhw9oUyZels7Un7pB7omA4RKW5xp5rx" + "Y+WgSeAEyyxNcPIkRvaxHaY1p2QaR3m17a0L7joJYNB+gNg7R/4Mfwpr8WuSOyLfBaKTuC5IiYWLaRm0ibASYzd59pPUyArt" + "mIx7pPfWxAr8/FT8TVQTSSeHXF88x/xJ2WSTSzQ3+FPw6Wm/LpfeGwQPX/g1+8f+KVcj1qAwT/tUobqOh7DdWl68CnNFReX4" + "SwBqRjTzZ276P7l15fhWmnQkiWXmLVsufJi5VXidEVrAsuEyVe7JmpExi3Y+0pn82D9x7lyv/TrOCVQgKwC0Zoq9NfuQxWa0" + "RvqmAv6toMfI8owTyzVoItnoGdhcYamwr7AyECKTuqnzMuylzOx5ttZnRc12/Ey2GQR3UxhyxMNLk01/UDicLnEXFl8Zn56i" + "rj51atVum7mJc48qekcKXQ5Lc3/t277v1d+eLmuiHebshpQwg+qQPHnnfky0WL98w343XDb29wwK3+veqQSjrffv5rQY9Tvq" + "3PMuN1pjU7Z3rL6fZmt5UrhqQQnTf4dA7X7n/aLkpA8SpY6ICDDr2x1z0yW5kMnHfxxk+N+j5eWwoTJuxxy94rLaYj3CKT57" + "XhMthQAoF1DGNpGxhbfC/tVuFE9FkxEBiYXqDz5fy2AU+S2zatT7/9mkFaHLY6hkOCS48X9WBVobxsQZZUnkbxplCzg4Huin" + "gPnl1bUDEprLHfTYgZHcnn3FgWFPwzxNBebrc4xjqXciMgReu4kOLP18JZbIJSjhP/RZAnQJEEOI8OrVGZlmiC7/JJx+8PTX" + "XYMTKtiVU/R3gOhOeVNN9WXruUoQzuy+JfJGg8ZQ2ni0Ugkx0sLO5rDsTOZ67bA+B99FQnQcBLyb40EPTOzrLrcM9wzXL+YZ" + "JUwNtdaQGxh/aUeKex/8of7GUH9ZsCZW0YN5xuEs8e2Q+vi2Msluckvki/wcyx7oOxmjw8bEbUkPBFq/ZEOcQqtetPLBzNPt" + "smJ4EfOU/kwBBLxplt6JrF+0ILeWJuH7OGpvpseXEosLpJ9HV4oIHL9co8EWGH/oAUzcuvhymR8FRziVE2IgVkreryIfHAs/" + "umsJgk1dXZk4A/R7KtjHKau0yxLAYJh6jWLvLsLv5N44VfKCQtyWFdUOSGbqn6E4YuOYVc+x+TGQcSMniobTekwm+UtqzVH8" + "mmtd8nwzwVkDgqZEAHrrnJOALLPeIbPuQ0/zjuUJmQovnnSqUyyCxQDpaxrlLD3iIZp9/9hPTY+lYHlam7VfChpQ7dPQ075t" + "ECiPA5TgMzh1qNQBF+rgKtrcrPo1ZDFqlseAybmSKqnCdvD5vdfdeVL8cmQIQao9qHDVh/aR0uwOsVKKzavPIYH2mhomg/cz" + "XB1nJG2zXB99g08Bd14c/zHHPNeqkqmwURCTxzPh8aeoRPcC8uL/eyawL66FfyUo40og9uridBDA4qx4ZMI+P1dBKFnMPM7o" + "Blpo+Sg9m6XZNPNVbDqRKKmX9OWCng7/PAKX+wa8EzRetNbyh73eMviMZaiN9YAyoKiS+nR5FqrPgNcFMqa8dHDEyY6O3Vuh" + "5Er5R3iGNACXMQvLRDVy/1UmFY1mzwew9lrDcUkZ0MNKBrbtDC87XEFGMejqvoJgF5Yl+q1mJE6gMBDTo6UEP1SOctxHyX23" + "mvqhCK9wOvK91UGnNyd+8ZQtQPP/L99d2y7hHmE6VEA/nUa2qkkO99vEQmJV3ilax+wZknABwk5sFdh7vWLwhEvxkdGYUh6c" + "HnA70VX5Kh0twJPKwE7gNASf/rvE5kYdrei5RtZNKMpVulXMPVho3eRBW6Qf6Ka0JjDZF5Dr8IXXsjP2LFTmnLNcnxu0FXx8" + "4xRuw1+yN3cP5R8LFlSGJmqPf5kFwAwpzT5FAXETiBaj7lg3WyG+OaDpab2NG+BvHbdfrDAuFKSsTDZ3Vt58fj6fuhS8tFd4" + "YzWL6X30pj+r7QRAefk+2CQDKM1TnbAJbyBcG2DVfeWpgmNjtzoV+RXNykIFIy6rSeOnUBAJCXwbAe7Scp5w/lx03a5aooRU" + "CP7x7CAmAU+YuGerG7A/3bbWDqN4EHGawkloK9UXmlOVX7uzLy0q3fJZb/6pSvPX+DR2rjFq0CSKnNFyBXYQRHx+DekAHWJS" + "VGUw7RLAPd4DE2HRiYTss7vtXalkP0hQ/vnt8lS9bjyBCi8ly+4FjFpo/SeKx0TKFulbFp1GPpsNERT+7kEkTsNZQHXWXPRB" + "PH0qmpFWZjI98fzWIGXGF376X0bcg6g67C/K/+YtH+QIU82ScAM0ozzQ1uzb0ci6wnM4JjNKA3ctVX6zuPR20Tb7VLlq5rwR" + "coOKjGQmBGbUiW5aKAAsMnIDhxf2lPCcwwb51C0Zb6hfOG+ka7LLcstgq4fmXHQ11BbTFu9ctAhy0E5V3aYSJAilb6VXZdD4" + "kSl3I0D2Xc/pTLPAzsxg6mEHL7J/DD1lRvEtSh4U34RCwxMk+pQbfEo6y/lkIVhXuE8p7d4mhPtsDJkls/3vnl+rR5ycebjH" + "g/zC01CjHhbVglb+HVEC5FMTrBNy/bweTE6c7qNHV6EWRN64nqvs4IW/jEVoaI2qt9WAJZ1br21rCE01iQagQUs1/ZH2KpQB" + "jTe3ju/BotwiE4TdqHcDL/VP62yq9KLeiBNXY2ctb3NwJzzfLrM2KHv2b4YmWuZ/VLABBfh9r3OzSmnFBvmWsZjZOHHEQXMI" + "hI+WDpRGB54CWmnWs6YzQK70IKrZUtq4BfPQJGf8OqO/ZeCJA4kinQbCaTTqQOJLsONKbPTqm8ImXz/QiML+KFvY3Z3ReW3T" + "NsJ3CGW7NRY1cfBIkhHuKHxPK1Ua5H3b7yP+h8OKrbqyjYPw2kE56TOH75HnIp/Ib6w+7uZpMVIbM4fYCRWlvQ7cUtaRq4fk" + "qqgHykE8bJYoU9eJFmPJNj4gt051sXJoslr2C6mvHYdSZgKxNa1ksd95cPKS3yURYQ==" +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES: Final = ( + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "To understand how zero, negative, and non-integer quantities are checked during validation, I need to inspect", + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 21, + "totalTokenCount": 847, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 472, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "3S3Rat6pNonj_uMPv9i78Qw", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": " `shop/validation.py` next.", + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 30, + "totalTokenCount": 856, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 472, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "3S3Rat6pNonj_uMPv9i78Qw", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "functionCall": { + "name": "read_file", + "args": { + "end_line": 50, + "path": "shop/validation.py", + "start_line": 1, + }, + "id": "call_5018", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 63, + "totalTokenCount": 889, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 472, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "3S3Rat6pNonj_uMPv9i78Qw", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "", + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 63, + "totalTokenCount": 889, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 472, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "3S3Rat6pNonj_uMPv9i78Qw", + } + ), +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="thoughts_off", + expected_provider_request=with_thinking_config( + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"thinkingLevel": "high", "includeThoughts": False} + ), + mock_provider_response=GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE, + expected_litellm_response={ + **GEMINI_RESPONSES_RESPONSE, + "model": "gemini/gemini-3.8-flash", + "output": [ + { + "type": "message", + "id": ANY, + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": ( + "To see how zero, negative, and non-integer quantities are currently handled and rejected " + "by the validation logic, the next file I need to inspect is `shop/validation.py`." + ), + "annotations": [], + }, + ], + "phase": None, + }, + { + "arguments": '{"path": "shop/validation.py", "start_line": 1, "end_line": 50}', + "call_id": "call_584229__thought__" + GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + "name": "read_file", + "type": "function_call", + "id": "fc_call_584229__thought__" + GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + "namespace": None, + "status": "completed", + "provider_specific_fields": {"thought_signature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE}, + }, + ], + "usage": { + "input_tokens": 354, + "input_tokens_details": { + "audio_tokens": None, + "cached_tokens": 0, + "cached_tokens_details": None, + "image_tokens": None, + "text_tokens": 354, + "video_tokens": None, + }, + "output_tokens": 436, + "output_tokens_details": {"audio_tokens": None, "reasoning_tokens": 366, "text_tokens": 70}, + "total_tokens": 790, + "cost": None, + }, + }, +) + + +def with_summary(case: TranslationTestCase, scenario: str, summary: JsonValue) -> TranslationTestCase: + reasoning: Final[Mapping[str, JsonValue]] = {"effort": "high", "summary": summary} + return replace( + case, + scenario=scenario, + litellm_request={**case.litellm_request, "reasoning": dict(reasoning)}, + expected_litellm_response={**case.expected_litellm_response, "reasoning": dict(reasoning)}, + ) + + +SUMMARY_ASKS: Final = ( + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_auto", "auto"), + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_concise", "concise"), + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_empty", ""), + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_long", LONG_SUMMARY), + with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_int", 7), + with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_list", ["auto"]), + with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_null", None), +) + + +def test_gemini_keeps_thinking_and_billing_reasoning_without_thought_text_when_the_gateway_turns_it_off( + thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, thoughts_off, provider) + + +@pytest.mark.parametrize("case", SUMMARY_ASKS, ids=tuple(case.id for case in SUMMARY_ASKS)) +def test_a_reasoning_summary_ask_brings_thought_text_back_over_the_gateway_default( + case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, thoughts_off, provider) + + +def test_a_streamed_turn_keeps_thinking_and_billing_reasoning_without_thought_text( + thoughts_off: Gateway, provider: SharedProvider +) -> None: + provider.expect( + Reply( + content_type="text/event-stream", + chunks=tuple(encode_sse_frame(frame) for frame in GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES), + ) + ) + response: Final = thoughts_off.request( + "POST", "/v1/responses", {**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "stream": True} + ) + assert_provider_request( + provider.received(), + "/models/gemini-3.8-flash:streamGenerateContent?alt=sse", + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_headers, + GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.expected_provider_request, + ) + assert response.status_code == 200, response.text + events: Final = tuple( + object_value(frame.data) for frame in parse_sse_frames(response.text) if frame.data != "[DONE]" + ) + assert [event["type"] for event in events if str(event["type"]).startswith("response.reasoning")] == [] + assert "".join(str(event["delta"]) for event in events if event["type"] == "response.output_text.delta") == ( + "To understand how zero, negative, and non-integer quantities are checked during validation, I need to " + "inspect `shop/validation.py` next." + ) + completed: Final = object_value( + next(event["response"] for event in events if event["type"] == "response.completed") + ) + assert [object_value(item)["type"] for item in list_value(completed["output"])] == ["message", "function_call"] + assert object_value(completed["usage"])["output_tokens_details"] == { + "reasoning_tokens": 472, + "text_tokens": 63, + } 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/test_litellm_rust/embeddings/__init__.py b/tests/test_litellm_rust/embeddings/__init__.py new file mode 100644 index 00000000000..e69de29bb2d 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.py b/tests/unit/caching/test_caching.py index b84529de0d6..91083921130 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -291,6 +291,37 @@ def test_exact_cache_key_includes_anthropic_messages_params(anthropic_param): assert baseline != cache.get_cache_key(model="claude-sonnet-4-5", messages=messages, **anthropic_param) +@pytest.mark.parametrize( + "summary_ask", + [ + {"reasoning_summary": "auto"}, + {"reasoningSummary": "auto"}, + {"extra_body": {"reasoning_summary": "auto"}}, + {"extra_body": {"reasoningSummary": "auto"}}, + ], + ids=["top_level", "top_level_camel", "extra_body", "extra_body_camel"], +) +def test_exact_cache_key_includes_the_reasoning_summary_alias(summary_ask: dict[str, object]) -> None: + """The AI SDK summary alias decides whether thought summaries come back, wherever the request + carries it, so a reply cached without it must not answer a request that asks for one.""" + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + messages: Final = [{"role": "user", "content": "how much is the ball?"}] + no_ask: Final = cache.get_cache_key(model="gemini/gemini-3.8-flash", messages=messages, reasoning_effort="medium") + ask: Final = cache.get_cache_key( + model="gemini/gemini-3.8-flash", messages=messages, reasoning_effort="medium", **summary_ask + ) + assert no_ask != ask + assert ask == cache.get_cache_key( + model="gemini/gemini-3.8-flash", messages=messages, reasoning_effort="medium", **summary_ask + ) + assert ask != cache.get_cache_key( + model="gemini/gemini-3.8-flash", + messages=messages, + reasoning_effort="medium", + reasoning_summary="detailed", + ) + + @pytest.mark.asyncio async def test_embedding_cache_skips_write_when_one_input_yields_many_embeddings(monkeypatch): """A cross-encoder behind /embeddings returns one score per document for a single @@ -561,10 +592,7 @@ def test_caching_v2(): # test in memory cache litellm.cache = None # disable cache litellm.success_callback = [] litellm._async_success_callback = [] - if ( - response2["choices"][0]["message"]["content"] - != response1["choices"][0]["message"]["content"] - ): + if response2["choices"][0]["message"]["content"] != response1["choices"][0]["message"]["content"]: print(f"response1: {response1}") print(f"response2: {response2}") pytest.fail(f"Error occurred:") @@ -595,10 +623,7 @@ def test_caching_with_ttl(): litellm.cache = None # disable cache litellm.success_callback = [] litellm._async_success_callback = [] - assert ( - response2["choices"][0]["message"]["content"] - != response1["choices"][0]["message"]["content"] - ) + assert response2["choices"][0]["message"]["content"] != response1["choices"][0]["message"]["content"] except Exception as e: print(f"error occurred: {traceback.format_exc()}") pytest.fail(f"Error occurred: {e}") @@ -632,9 +657,7 @@ def test_caching_with_default_ttl(): def test_caching_with_models_v2(): - messages = [ - {"role": "user", "content": "who is ishaan CTO of litellm from litellm 2023"} - ] + messages = [{"role": "user", "content": "who is ishaan CTO of litellm from litellm 2023"}] litellm.cache = Cache() print("test2 for caching") litellm.set_verbose = True @@ -657,18 +680,12 @@ def test_caching_with_models_v2(): litellm.cache = None litellm.success_callback = [] litellm._async_success_callback = [] - if ( - response3["choices"][0]["message"]["content"] - == response2["choices"][0]["message"]["content"] - ): + if response3["choices"][0]["message"]["content"] == response2["choices"][0]["message"]["content"]: # if models are different, it should not return cached response print(f"response2: {response2}") print(f"response3: {response3}") pytest.fail(f"Error occurred:") - if ( - response1["choices"][0]["message"]["content"] - != response2["choices"][0]["message"]["content"] - ): + if response1["choices"][0]["message"]["content"] != response2["choices"][0]["message"]["content"]: print(f"response1: {response1}") print(f"response2: {response2}") pytest.fail(f"Error occurred:") @@ -708,9 +725,7 @@ def test_get_cache_key(): cache_key = cache_instance.get_cache_key( **{ "model": "gpt-3.5-turbo", - "messages": [ - {"role": "user", "content": "write a one sentence poem about: 7510"} - ], + "messages": [{"role": "user", "content": "write a one sentence poem about: 7510"}], "max_tokens": 40, "temperature": 0.2, "stream": True, @@ -721,9 +736,7 @@ def test_get_cache_key(): cache_key_2 = cache_instance.get_cache_key( **{ "model": "gpt-3.5-turbo", - "messages": [ - {"role": "user", "content": "write a one sentence poem about: 7510"} - ], + "messages": [{"role": "user", "content": "write a one sentence poem about: 7510"}], "max_tokens": 40, "temperature": 0.2, "stream": True, @@ -736,9 +749,9 @@ def test_get_cache_key(): # Hexadecimal representation of the hash hash_hex = hash_object.hexdigest() assert cache_key == hash_hex - assert ( - cache_key_2 == hash_hex - ), f"{cache_key} != {cache_key_2}. The same kwargs should have the same cache key across runs" + assert cache_key_2 == hash_hex, ( + f"{cache_key} != {cache_key_2}. The same kwargs should have the same cache key across runs" + ) embedding_cache_key = cache_instance.get_cache_key( **{ @@ -756,15 +769,13 @@ def test_get_cache_key(): print(embedding_cache_key) - embedding_cache_key_str = ( - "model: azure/text-embedding-ada-002input: ['hi who is ishaan']" - ) + embedding_cache_key_str = "model: azure/text-embedding-ada-002input: ['hi who is ishaan']" hash_object = hashlib.sha256(embedding_cache_key_str.encode()) # Hexadecimal representation of the hash hash_hex = hash_object.hexdigest() - assert ( - embedding_cache_key == hash_hex - ), f"{embedding_cache_key} != 'model: azure/text-embedding-ada-002input: ['hi who is ishaan']'. The same kwargs should have the same cache key across runs" + assert embedding_cache_key == hash_hex, ( + f"{embedding_cache_key} != 'model: azure/text-embedding-ada-002input: ['hi who is ishaan']'. The same kwargs should have the same cache key across runs" + ) # Proxy - embedding cache, test if embedding key, gets model_group and not model embedding_cache_key_2 = cache_instance.get_cache_key( @@ -817,9 +828,7 @@ def test_get_cache_key(): ) print(embedding_cache_key_2) - embedding_cache_key_str_2 = ( - "model: EMBEDDING_MODEL_GROUPinput: ['hi who is ishaan']" - ) + embedding_cache_key_str_2 = "model: EMBEDDING_MODEL_GROUPinput: ['hi who is ishaan']" hash_object = hashlib.sha256(embedding_cache_key_str_2.encode()) # Hexadecimal representation of the hash hash_hex = hash_object.hexdigest() @@ -852,9 +861,7 @@ def test_redis_caching_multiple_namespaces(): # Mock the Redis client creation from the _redis module with ( patch("litellm._redis.get_redis_client") as mock_get_redis_client, - patch( - "litellm._redis.get_redis_connection_pool" - ) as mock_get_redis_connection_pool, + patch("litellm._redis.get_redis_connection_pool") as mock_get_redis_connection_pool, ): # Create a mock Redis client that simulates real Redis behavior mock_redis_client = MagicMock() @@ -927,18 +934,10 @@ def test_redis_caching_multiple_namespaces(): mock_response="Response without namespace", ) - print( - f"Response 1 type: {type(response_1)} - ID: {getattr(response_1, 'id', 'N/A')}" - ) - print( - f"Response 2 type: {type(response_2)} - ID: {getattr(response_2, 'id', 'N/A')}" - ) - print( - f"Response 3 type: {type(response_3)} - Cache hit: {isinstance(response_3, str)}" - ) - print( - f"Response 4 type: {type(response_4)} - ID: {getattr(response_4, 'id', 'N/A')}" - ) + print(f"Response 1 type: {type(response_1)} - ID: {getattr(response_1, 'id', 'N/A')}") + print(f"Response 2 type: {type(response_2)} - ID: {getattr(response_2, 'id', 'N/A')}") + print(f"Response 3 type: {type(response_3)} - Cache hit: {isinstance(response_3, str)}") + print(f"Response 4 type: {type(response_4)} - ID: {getattr(response_4, 'id', 'N/A')}") print(f"Redis storage keys: {list(redis_storage.keys())}") @@ -947,10 +946,7 @@ def test_redis_caching_multiple_namespaces(): namespace_1_keys = [k for k in cache_keys if k.startswith(f"{namespace_1}:")] namespace_2_keys = [k for k in cache_keys if k.startswith(f"{namespace_2}:")] no_namespace_keys = [ - k - for k in cache_keys - if not k.startswith(f"{namespace_1}:") - and not k.startswith(f"{namespace_2}:") + k for k in cache_keys if not k.startswith(f"{namespace_1}:") and not k.startswith(f"{namespace_2}:") ] print(f"Namespace 1 keys: {namespace_1_keys}") @@ -963,9 +959,7 @@ def test_redis_caching_multiple_namespaces(): assert len(no_namespace_keys) > 0, "Should have cache keys for no namespace" # The main test: response 3 should be a cache hit (string) because it uses same namespace as response 1 - assert isinstance( - response_3, str - ), "Response 3 should be a cache hit (string) for same namespace" + assert isinstance(response_3, str), "Response 3 should be a cache hit (string) for same namespace" # response 1 & 2 should be ModelResponse objects (cache misses) assert hasattr(response_1, "id"), "Response 1 should be a ModelResponse object" @@ -973,14 +967,14 @@ def test_redis_caching_multiple_namespaces(): assert hasattr(response_4, "id"), "Response 4 should be a ModelResponse object" # response 1 & 2 should have different IDs (different namespaces) - assert ( - response_1.id != response_2.id - ), f"Expected different response ID for different namespace. Got {response_1.id} and {response_2.id}" + assert response_1.id != response_2.id, ( + f"Expected different response ID for different namespace. Got {response_1.id} and {response_2.id}" + ) # response 1 & 4 should have different IDs (different namespaces) - assert ( - response_1.id != response_4.id - ), f"Expected different response ID for no namespace vs namespaced. Got {response_1.id} and {response_4.id}" + assert response_1.id != response_4.id, ( + f"Expected different response ID for no namespace vs namespaced. Got {response_1.id} and {response_4.id}" + ) _TOOL_TURN_ITEM: Final = {"role": "user", "content": "hi"} @@ -1305,12 +1299,8 @@ async def test_embedding_caching_individual_items(monkeypatch: pytest.MonkeyPatc monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL)) item: Final = f"individual embedding {uuid.uuid4().hex}" - first: Final = await aembedding( - model="text-embedding-3-small", input=item, caching=True, mock_response="0.1,0.2" - ) - second: Final = await aembedding( - model="text-embedding-3-small", input=item, caching=True, mock_response="0.3,0.4" - ) + first: Final = await aembedding(model="text-embedding-3-small", input=item, caching=True, mock_response="0.1,0.2") + second: Final = await aembedding(model="text-embedding-3-small", input=item, caching=True, mock_response="0.3,0.4") assert second.data[0].embedding == first.data[0].embedding assert second._hidden_params["cache_hit"] is True @@ -1448,9 +1438,7 @@ async def test_caching_with_cache_controls(monkeypatch: pytest.MonkeyPatch) -> N @pytest.mark.asyncio -async def test_embedding_caching_redis_ttl( - monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None -) -> None: +async def test_embedding_caching_redis_ttl(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None: mock_pipeline: Final = AsyncMock() mock_set: Final = AsyncMock() mock_pipeline.__aenter__.return_value.set = mock_set @@ -1506,18 +1494,14 @@ async def test_logging_turn_off_message_logging_streaming( assert cached_response["choices"][0]["message"]["content"] == "cached" -def test_redis_caching_default_ttl( - monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None -) -> None: +def test_redis_caching_default_ttl(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None: monkeypatch.setattr(litellm, "default_redis_ttl", 120) cache: Final = RedisCache(host="127.0.0.1", port="6379") assert cache.default_ttl == 120 -def test_redis_caching_llm_caching_ttl( - monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None -) -> None: +def test_redis_caching_llm_caching_ttl(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None: monkeypatch.setattr(litellm, "default_redis_ttl", 120) cache: Final = RedisCache(host="127.0.0.1", port="6379") set_cache: Final = MagicMock() @@ -1529,9 +1513,7 @@ def test_redis_caching_llm_caching_ttl( @pytest.mark.asyncio -async def test_redis_caching_ttl_pipeline( - monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None -) -> None: +async def test_redis_caching_ttl_pipeline(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None: monkeypatch.setattr(litellm, "default_redis_ttl", 120) cache: Final = RedisCache(host="127.0.0.1", port="6379") pipeline: Final = AsyncMock() @@ -1551,9 +1533,7 @@ async def test_redis_caching_ttl_pipeline( @pytest.mark.asyncio -async def test_redis_caching_ttl_sadd( - monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None -) -> None: +async def test_redis_caching_ttl_sadd(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None: monkeypatch.setattr(litellm, "default_redis_ttl", 120) cache: Final = RedisCache(host="127.0.0.1", port="6379") redis_client: Final = AsyncMock() 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/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_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 9bd41f4a602..614ff97553a 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -2790,7 +2790,7 @@ class TestLoggingOnlyApplyGuardrail: assert out_kwargs["messages"] == kwargs["messages"] assert ( out_kwargs["standard_logging_object"]["guardrail_information"][0]["guardrail_status"] - == "guardrail_intervened" + == "guardrail_flagged" ) @pytest.mark.asyncio 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/interactions/test_background_cost_polling.py b/tests/unit/interactions/test_background_cost_polling.py index 9dc71a4fa1b..d55994dfb19 100644 --- a/tests/unit/interactions/test_background_cost_polling.py +++ b/tests/unit/interactions/test_background_cost_polling.py @@ -1,11 +1,14 @@ import asyncio +import json import time -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from itertools import islice from typing import Optional +import httpx import pytest +import litellm from litellm.interactions.background_cost_polling import ( _create_context, _poll_intervals, @@ -21,6 +24,7 @@ from litellm.interactions.background_cost_polling import ( ) from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.types.interactions import InteractionsAPIResponse USAGE_BLOCK = { @@ -74,6 +78,7 @@ def _context( logging_obj: LitellmLogging, timeout_seconds: float = 1.0, store: Optional[InMemoryBackgroundSettlementStore] = None, + resumed: bool = False, ) -> BackgroundInteractionPollContext: return BackgroundInteractionPollContext( interaction_id="interactions/bg-abc", @@ -82,7 +87,9 @@ def _context( initial_interval_seconds=0.001, max_interval_seconds=0.002, timeout_seconds=timeout_seconds, + outcome_write_timeout_seconds=0.01, store=store if store is not None else InMemoryBackgroundSettlementStore(), + resumed=resumed, ) @@ -255,6 +262,180 @@ async def test_poller_retries_after_fetch_error_and_still_bills(): assert logging_obj.model_call_details["response_cost"] > 0 +def _provider_error_response(status_code: int) -> httpx.Response: + return httpx.Response(status_code, request=httpx.Request("GET", "https://generativelanguage.googleapis.com")) + + +@pytest.mark.parametrize( + "error", + [ + litellm.BadRequestError(message="invalid interaction id", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.AuthenticationError(message="API key not valid", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.PermissionDeniedError( + message="permission denied", + model="gemini-2.5-flash", + llm_provider="gemini", + response=_provider_error_response(403), + ), + litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini"), + ], +) +@pytest.mark.asyncio +async def test_a_resumed_poll_stops_fetching_at_the_first_error_no_retry_changes_and_releases_the_reservation(error): + """ + The regression: a resumed fetch the provider refuses for good, like an + unknown id or a key rotated away, was retried every interval until the + poll timed out. + """ + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, calls = _fetch_sequence(error) + + outcome = await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=0.05, resumed=True), fetch_interaction=fetch + ) + + assert outcome == "unsettled" + assert len(calls) == 1 + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_a_resumed_poll_the_provider_refuses_leaves_the_bill_to_the_replica_still_polling_it(): + """ + The regression: a replica that booted with credentials unable to read the + interaction settled it as unsettled on its first 404, and the replica that + created it, still polling, then found it settled and never billed it. + """ + store = InMemoryBackgroundSettlementStore() + creator_logging_obj = _logging_obj() + resumed_fetched = asyncio.Event() + + async def creator_fetch(context): + if resumed_fetched.is_set(): + return _response("completed", with_usage=True) + return _response("in_progress", with_usage=False) + + async def resumed_fetch(context): + resumed_fetched.set() + raise litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini") + + creator_outcome, resumed_outcome = await asyncio.gather( + poll_and_log_background_interaction_cost( + _context(creator_logging_obj, store=store), fetch_interaction=creator_fetch + ), + poll_and_log_background_interaction_cost( + _context(_logging_obj(), store=store, resumed=True), fetch_interaction=resumed_fetch + ), + ) + + assert (creator_outcome, resumed_outcome) == ("billed", None) + assert creator_logging_obj.model_call_details["response_cost"] > 0 + + +@pytest.mark.parametrize( + "error", + [ + litellm.RateLimitError(message="quota exceeded", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.ServiceUnavailableError(message="overloaded", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.InternalServerError(message="internal error", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.Timeout(message="timed out", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.BadRequestError(message="invalid interaction id", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.AuthenticationError(message="API key not valid", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini"), + ], +) +@pytest.mark.asyncio +async def test_the_creating_poll_retries_through_any_provider_error(error): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence(error, _response("completed", with_usage=True)) + + outcome = await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert outcome == "billed" + assert len(calls) == 2 + + +class _OutcomeWritesFail: + def __init__(self): + self.store = InMemoryBackgroundSettlementStore() + + async def register(self, pending): + await self.store.register(pending) + + async def pending(self, interaction_id): + return await self.store.pending(interaction_id) + + async def is_claimed(self, interaction_id): + return await self.store.is_claimed(interaction_id) + + async def claim(self, interaction_id): + return await self.store.claim(interaction_id) + + async def record_outcome(self, interaction_id, outcome): + raise RuntimeError("database unavailable") + + async def unsettled(self): + return await self.store.unsettled() + + +@pytest.mark.parametrize( + "last_status, timeout_seconds", + [("failed", 1.0), ("in_progress", 0.01)], +) +@pytest.mark.asyncio +async def test_a_release_it_could_not_record_leaves_the_reservation_to_the_settler_that_takes_over( + last_status, timeout_seconds +): + """ + The regression: a settler released the reservation and then failed to + record it, so the replica that took over the lapsed claim released the + same reservation again and the key's budget lost that amount twice. + """ + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence(_response(last_status, with_usage=False)) + + await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=timeout_seconds, store=_OutcomeWritesFail()), + fetch_interaction=fetch, + ) + + assert reservation["finalized"] is False + + +class _FirstOutcomeWriteFails(_OutcomeWritesFail): + def __init__(self): + super().__init__() + self.failed_once = False + + async def record_outcome(self, interaction_id, outcome): + if not self.failed_once: + self.failed_once = True + raise RuntimeError("connection reset") + await self.store.record_outcome(interaction_id, outcome) + + +@pytest.mark.asyncio +async def test_a_release_whose_first_outcome_write_failed_still_frees_the_reservation(): + """ + The regression: one failed outcome write left the key's budget + reservation held until a later boot or delete took the row over, and an + active key's spend counter never expired it in the meantime. + """ + reservation = _reservation() + store = _FirstOutcomeWriteFails() + fetch, _ = _fetch_sequence(_response("failed", with_usage=False)) + + outcome = await poll_and_log_background_interaction_cost( + _context(_logging_obj_with_reservation(reservation), store=store), fetch_interaction=fetch + ) + + assert outcome == "released" + assert reservation["finalized"] is True + assert await store.store.unsettled() == () + + @pytest.mark.asyncio async def test_schedule_creates_poll_task_for_in_progress_create(): logging_obj = _logging_obj() @@ -302,6 +483,59 @@ async def test_schedule_registers_an_agent_only_create_that_names_no_model(): assert _rebuild_logging_obj(pending.create_context).model is None +@pytest.mark.asyncio +async def test_schedule_registers_a_create_whose_tags_are_a_string_and_keeps_them_as_sent(): + """ + The regression: storing the create's context read string tags as a list + and raised, so a create sent with a tags string answered with an error + after the provider had already started the interaction. + """ + logging_obj = _logging_obj(litellm_params={"metadata": _create_metadata(tags="team-a,nightly-batch")}) + store = InMemoryBackgroundSettlementStore() + + task = await maybe_schedule_background_interaction_cost_polling( + response=_response("in_progress", with_usage=False), + create_kwargs={"litellm_logging_obj": logging_obj}, + custom_llm_provider="gemini", + store=store, + ) + + assert isinstance(task, asyncio.Task) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + pending = await store.pending("interactions/bg-abc") + assert pending is not None + rebuilt = _rebuild_logging_obj(pending.create_context) + assert get_litellm_metadata_from_kwargs(kwargs=rebuilt.model_call_details)["tags"] == "team-a,nightly-batch" + + +@pytest.mark.asyncio +async def test_delete_fetches_with_the_vertex_project_location_and_credentials_the_request_used(): + store = InMemoryBackgroundSettlementStore() + await _create_on_a_replica_that_then_dies(_logging_obj(litellm_params={"metadata": _create_metadata()}), store) + fetch, captured = _capturing_fetch(_response("completed", with_usage=True)) + + outcome = await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + delete_kwargs={ + "vertex_project": "request-project", + "vertex_location": "europe-west4", + "vertex_credentials": {"type": "service_account", "project_id": "request-project"}, + }, + fetch_interaction=fetch, + store=store, + ) + + assert outcome == "billed" + credentials = captured[0].credentials + assert (credentials.vertex_project, credentials.vertex_location, credentials.vertex_credentials) == ( + "request-project", + "europe-west4", + {"type": "service_account", "project_id": "request-project"}, + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( "response,create_kwargs", @@ -781,6 +1015,105 @@ async def test_restart_resumes_only_the_rows_no_replica_claimed(): assert await store.is_claimed("interactions/bg-orphaned") +@pytest.mark.asyncio +async def test_a_restart_long_after_the_create_still_retries_a_failed_fetch_before_giving_up(): + """ + The regression: a resumed poll kept only what was left of the original + window, so a row older than the window got a single fetch, and one + transient error right after the restart ended it unbilled. + """ + store = InMemoryBackgroundSettlementStore() + await store.register( + PendingBackgroundInteraction( + interaction_id="interactions/bg-abc", + custom_llm_provider="gemini", + create_context=_create_context(_logging_obj(litellm_params={"metadata": _create_metadata()}), "gemini"), + created_at=datetime.now(timezone.utc) - timedelta(seconds=FAST_SCHEDULE.timeout_seconds * 100), + ) + ) + fetch, calls = _fetch_sequence( + RuntimeError("transient network error"), + _response("completed", with_usage=True), + ) + + (resumed,) = await resume_unsettled_background_interactions(store, fetch, schedule=FAST_SCHEDULE) + + assert await asyncio.wait_for(resumed, timeout=5) == "billed" + assert len(calls) == 2 + + +def _create_through_the_proxy_router(tags: list) -> LitellmLogging: + logging_obj = LitellmLogging( + model="gemini-2.5-flash", + messages=[], + stream=False, + call_type="acreate_interaction", + start_time=time.time(), + litellm_call_id="bg-interactions-call-id", + function_id="bg-interactions-fn-id", + ) + logging_obj.update_from_kwargs( + kwargs={ + "metadata": _create_metadata( + tags=tags, spend_logs_metadata={"job": "nightly"}, model_group="client-chosen-group" + ), + "litellm_metadata": {"model_info": {"id": "deployment-1"}, "model_group": "gemini-background"}, + }, + model="gemini-2.5-flash", + optional_params={}, + litellm_params={"litellm_call_id": "bg-interactions-call-id"}, + custom_llm_provider="gemini", + ) + return logging_obj + + +def _create_with_header_tags(tags: list) -> LitellmLogging: + return _logging_obj( + litellm_params={ + "metadata": _create_metadata(tags=tags, spend_logs_metadata={"job": "nightly"}), + "proxy_server_request": {"headers": {"user-agent": "claude-cli/2.1.0"}}, + } + ) + + +async def _billed_spend_row(logging_obj: LitellmLogging) -> dict: + response = _response("completed", with_usage=True) + fetch, _ = _fetch_sequence(response) + assert await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) == "billed" + settled_at = datetime.now(timezone.utc) + row = get_logging_payload( + kwargs={**logging_obj.model_call_details, "completion_start_time": settled_at}, + response_obj=response, + start_time=settled_at, + end_time=settled_at, + ) + return { + "request_tags": json.loads(row["request_tags"]), + "model_id": row["model_id"], + "model_group": row["model_group"], + "spend_logs_metadata": json.loads(row["metadata"]).get("spend_logs_metadata"), + } + + +@pytest.mark.parametrize("create", [_create_through_the_proxy_router, _create_with_header_tags]) +@pytest.mark.asyncio +async def test_a_bill_settled_off_the_creating_worker_matches_the_creating_worker_spend_row(create): + """ + The regressions: the stored create context kept only the router's metadata + on the proxy path, so a row billed after a restart or by another replica + lost the request tags that tag-based spend reports group by; then it let + the client's request metadata override the router's, so a client-sent + model group relabeled the bill another replica wrote. + """ + creator_row = await _billed_spend_row(create(["team-a", "nightly-batch"])) + stored = _create_context(create(["team-a", "nightly-batch"]), "gemini") + + rebuilt_row = await _billed_spend_row(_rebuild_logging_obj(stored)) + + assert "team-a" in creator_row["request_tags"] + assert rebuilt_row == creator_row + + class _ClaimAnswersOnlyAfterTheLastFetch: def __init__(self): self.store = InMemoryBackgroundSettlementStore() @@ -809,8 +1142,8 @@ class _ClaimAnswersOnlyAfterTheLastFetch: async def record_outcome(self, interaction_id, outcome): return None - async def unclaimed(self): - return await self.store.unclaimed() + async def unsettled(self): + return await self.store.unsettled() @pytest.mark.asyncio @@ -844,7 +1177,7 @@ class _DownStore: async def record_outcome(self, interaction_id, outcome): raise RuntimeError("database unavailable") - async def unclaimed(self): + async def unsettled(self): raise RuntimeError("database unavailable") @@ -868,8 +1201,8 @@ class _RegistersThenRaises: async def record_outcome(self, interaction_id, outcome): return None - async def unclaimed(self): - return await self.store.unclaimed() + async def unsettled(self): + return await self.store.unsettled() @pytest.mark.asyncio @@ -898,7 +1231,7 @@ async def test_create_whose_registration_raised_after_landing_still_claims_the_s assert outcome == "billed" assert await store.is_claimed("interactions/bg-abc") - assert await store.unclaimed() == () + assert await store.unsettled() == () task.cancel() with pytest.raises(asyncio.CancelledError): await task @@ -907,7 +1240,7 @@ async def test_create_whose_registration_raised_after_landing_still_claims_the_s @pytest.mark.asyncio async def test_delete_on_a_worker_that_resumed_the_poll_fails_when_it_cannot_fetch(): """ - After a restart every worker resumes the unclaimed rows, so none of them + After a restart every worker resumes the unsettled rows, so none of them is the creator whose delete may release and delete on a failed fetch. A resumed worker's delete fails like any other replica's, and its own poll still bills the interaction once it completes. @@ -965,9 +1298,9 @@ class _LandsThenGoesDown: async def record_outcome(self, interaction_id, outcome): return None - async def unclaimed(self): + async def unsettled(self): self._answer() - return await self.store.unclaimed() + return await self.store.unsettled() def _answer(self): if self.down: @@ -992,7 +1325,7 @@ class _TableLessStore: async def record_outcome(self, interaction_id, outcome): raise RuntimeError("the settlement table does not exist") - async def unclaimed(self): + async def unsettled(self): raise RuntimeError("the settlement table does not exist") @@ -1027,7 +1360,7 @@ async def test_create_whose_registration_and_read_back_both_failed_bills_once_th assert (while_down, recovered) == (None, "billed") assert len(calls) == 2 assert await store.is_claimed("interactions/bg-abc") - assert await store.unclaimed() == () + assert await store.unsettled() == () assert await resume_unsettled_background_interactions(store, poll_fetch, schedule=FAST_SCHEDULE) == () task.cancel() with pytest.raises(asyncio.CancelledError): 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 320a4eff371..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(): @@ -593,6 +596,422 @@ _BEDROCK_MANTLE_WEB_SEARCH_MODELS = ( _BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012 + + +def _openai_responses_response(model, output, usage=None, tools=None): + return ResponsesAPIResponse.model_validate( + { + "id": "resp_1", + "created_at": 1754900000, + "model": model, + "object": "response", + "status": "completed", + "output": output, + "usage": usage or {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + **({"tools": tools} if tools is not None else {}), + } + ) + + +_ASSISTANT_MESSAGE_OUTPUT_ITEM = { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "done", "annotations": []}], +} + +_GPT_IMAGE_1_HIGH_1024_COST_KEY = "high/1024-x-1024/gpt-image-1" +_GPT_IMAGE_1_TOOLS = [{"type": "image_generation", "model": "gpt-image-1"}] + + +@pytest.mark.parametrize("custom_llm_provider", ["openai", None]) +def test_responses_image_generation_call_billed_as_tool_usage_cost(local_model_cost_map, custom_llm_provider): + """A completed image_generation_call bills at the tools[].model rate for its quality/size, provider given or inferred.""" + expected_image_cost = litellm.model_cost[_GPT_IMAGE_1_HIGH_1024_COST_KEY]["input_cost_per_image"] + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "1024x1024", + "result": "AAAA", + }, + dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM), + ], + tools=_GPT_IMAGE_1_TOOLS, + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider=custom_llm_provider, + standard_built_in_tools_params=None, + ) + + assert cost > 0 + assert cost == pytest.approx(expected_image_cost) + + +def test_responses_image_generation_call_without_tool_model_bills_zero(local_model_cost_map): + """With no tools[].model to price against, the image tool part bills $0 instead of guessing a model or raising.""" + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "1024x1024", + "result": "AAAA", + } + ], + usage=_image_gen_token_usage(10, 0, 50, 0), + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert cost == 0.0 + + +def test_responses_web_search_and_image_generation_costs_are_additive(local_model_cost_map): + """A response billed for web search must still also bill its image_generation_call items.""" + model = "gpt-4o-search-preview" + image_cost = litellm.model_cost[_GPT_IMAGE_1_HIGH_1024_COST_KEY]["input_cost_per_image"] + image_item = { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "1024x1024", + "result": "AAAA", + } + web_search_item = {"type": "web_search_call", "id": "ws_1", "status": "completed"} + + combined = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=_openai_responses_response(model, [web_search_item, image_item], tools=_GPT_IMAGE_1_TOOLS), + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + web_search_only = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=_openai_responses_response(model, [web_search_item]), + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + image_only = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=_openai_responses_response(model, [image_item], tools=_GPT_IMAGE_1_TOOLS), + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert image_only == pytest.approx(image_cost) + assert web_search_only > 0 + assert combined == pytest.approx(web_search_only + image_only) + + +def test_responses_incomplete_image_generation_call_not_billed(local_model_cost_map): + """A failed image_generation_call produced no billable image, so it must cost $0.""" + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "failed", + "quality": "high", + "size": "1024x1024", + "result": None, + }, + dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM), + ], + tools=_GPT_IMAGE_1_TOOLS, + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert cost == 0.0 + + +def test_completion_cost_includes_responses_image_generation_tool_cost(local_model_cost_map): + """The image tool fee must flow through completion_cost on top of the token-only baseline.""" + image_item = { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "1024x1024", + "result": "AAAA", + } + response_with_image = _openai_responses_response( + "gpt-5", [image_item, dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM)], tools=_GPT_IMAGE_1_TOOLS + ) + response_without_image = _openai_responses_response("gpt-5", [dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM)]) + + cost_with_image = litellm.completion_cost( + completion_response=response_with_image, + model="gpt-5", + custom_llm_provider="openai", + call_type="aresponses", + ) + cost_without_image = litellm.completion_cost( + completion_response=response_without_image, + model="gpt-5", + custom_llm_provider="openai", + call_type="aresponses", + ) + + assert cost_with_image > cost_without_image + assert cost_with_image - cost_without_image == pytest.approx( + litellm.model_cost[_GPT_IMAGE_1_HIGH_1024_COST_KEY]["input_cost_per_image"] + ) + + +def test_responses_usage_tool_usage_web_search_billed_without_output_item(local_model_cost_map): + """usage.tool_usage.web_search.num_requests bills web search even when no web_search_call item is present.""" + model = "gpt-5.4-mini" + per_call = litellm.get_model_info(model)["search_context_cost_per_query"]["search_context_size_medium"] + + for num_requests in (1, 2): + response = _openai_responses_response( + model, + [dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM)], + usage={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "tool_usage": {"web_search": {"num_requests": num_requests}}, + }, + ) + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + assert cost == pytest.approx(num_requests * per_call) + + +def _image_gen_token_usage(input_text, input_image, output_image, output_text): + return { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "tool_usage": { + "image_gen": { + "input_tokens": input_text + input_image, + "output_tokens": output_image + output_text, + "total_tokens": input_text + input_image + output_image + output_text, + "input_tokens_details": {"image_tokens": input_image, "text_tokens": input_text}, + "output_tokens_details": {"image_tokens": output_image, "text_tokens": output_text}, + } + }, + } + + +def _expected_image_gen_token_cost(model_info, input_text, input_image, output_image, output_text): + return ( + input_text * (model_info.get("input_cost_per_token") or 0) + + input_image * (model_info.get("input_cost_per_image_token") or 0) + + output_image * (model_info.get("output_cost_per_image_token") or 0) + + output_text * (model_info.get("output_cost_per_token") or 0) + ) + + +def test_responses_image_tool_model_from_tools_bills_token_usage(local_model_cost_map): + """The image tool's tools[].model is used and usage.tool_usage.image_gen tokens bill at that model's rates.""" + tool_model = "gpt-image-2" + model_info = litellm.get_model_info(tool_model, custom_llm_provider="openai") + tools = [{"type": "image_generation", "model": tool_model, "quality": "low", "size": "1024x1024"}] + output = [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "low", + "size": "1024x1024", + "result": "AAAA", + }, + dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM), + ] + + response = _openai_responses_response("gpt-5", output, usage=_image_gen_token_usage(10, 0, 50, 0), tools=tools) + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + assert cost > 0 + assert cost == pytest.approx(_expected_image_gen_token_cost(model_info, 10, 0, 50, 0)) + + response_more_tokens = _openai_responses_response( + "gpt-5", output, usage=_image_gen_token_usage(10, 0, 80, 0), tools=tools + ) + cost_more = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response_more_tokens, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + assert cost_more != cost + assert cost_more == pytest.approx(_expected_image_gen_token_cost(model_info, 10, 0, 80, 0)) + + +def test_responses_zero_image_gen_tokens_fall_back_to_per_image_pricing(local_model_cost_map): + """An all-zero image_gen usage block keeps the per-image path for the tool's model/quality/size.""" + tool_model = "gpt-image-1" + quality = "low" + size = "1024x1024" + tools = [{"type": "image_generation", "model": tool_model, "quality": quality, "size": size}] + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": quality, + "size": size, + "result": "AAAA", + } + ], + usage=_image_gen_token_usage(0, 0, 0, 0), + tools=tools, + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + from litellm.cost_calculator import default_image_cost_calculator + + assert cost == pytest.approx( + default_image_cost_calculator( + model=tool_model, + custom_llm_provider="openai", + quality=quality, + n=1, + size=size, + ) + ) + assert cost > 0 + + +def test_responses_auto_size_image_generation_call_billed_at_default_size(local_model_cost_map): + """An image_generation_call with size "auto" bills at the default size instead of erroring to $0.""" + from litellm.cost_calculator import default_image_cost_calculator + + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "auto", + "result": "AAAA", + } + ], + tools=_GPT_IMAGE_1_TOOLS, + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert cost == pytest.approx( + default_image_cost_calculator( + model="gpt-image-1", + custom_llm_provider="openai", + quality="high", + n=1, + size=None, + ) + ) + assert cost > 0 + + +def test_responses_image_gen_total_without_token_details_falls_back_to_per_image(local_model_cost_map): + """A positive image_gen total with no token details falls back to per-image pricing, not $0.""" + tool_model = "gpt-image-1" + quality = "low" + size = "1024x1024" + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": quality, + "size": size, + "result": "AAAA", + } + ], + usage={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "tool_usage": {"image_gen": {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10}}, + }, + tools=[{"type": "image_generation", "model": tool_model, "quality": quality, "size": size}], + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + from litellm.cost_calculator import default_image_cost_calculator + + assert cost == pytest.approx( + default_image_cost_calculator( + model=tool_model, + custom_llm_provider="openai", + quality=quality, + n=1, + size=size, + ) + ) + assert cost > 0 @pytest.mark.parametrize( ("model_key", "web_search_requests", "medium_cost_per_query", "expected_cost"), [ @@ -654,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/llm_response_utils/test_get_headers.py b/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py index f52e837136a..1eccf5010ec 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py @@ -8,6 +8,7 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( _get_llm_provider_headers, get_response_headers, get_provider_request_id, + provider_request_id_headers, ) from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -89,6 +90,31 @@ def test_invalid_provider_request_ids_remain_absent(headers: object) -> None: assert get_provider_request_id(headers) is None +@pytest.mark.parametrize( + ("headers", "expected"), + ( + ( + {"X-Amzn-RequestId": "amzn-1", "retry-after": "7", "content-type": "text/event-stream"}, + {"x-amzn-requestid": "amzn-1"}, + ), + ( + {"x-request-id": "req_1", "cf-ray": "ray-1", "x-ratelimit-remaining-requests": "0"}, + {"x-request-id": "req_1", "cf-ray": "ray-1"}, + ), + ({"retry-after": "7"}, None), + ({"x-request-id": 7, "request-id": None, "x-amzn-requestid": "amzn-2"}, {"x-amzn-requestid": "amzn-2"}), + ({"x-amzn-requestid": "amzn-\u00fc", "x-request-id": "req_2"}, {"x-request-id": "req_2"}), + ({}, None), + (None, None), + ), +) +def test_provider_request_id_headers_keeps_only_request_id_headers( + headers: dict[str, object] | None, expected: dict[str, str] | None +) -> None: + result = provider_request_id_headers(headers) + assert (None if result is None else dict(result)) == expected + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index dba1730208d..bf977325a85 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -3,10 +3,12 @@ import json import logging import os import re +import sys import time -from collections.abc import Callable +import types +from collections.abc import Callable, Mapping, Sequence from copy import deepcopy -from typing import Final, List +from typing import Final, List, TypeAlias from unittest.mock import MagicMock, patch import httpx @@ -16,6 +18,7 @@ import respx import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( BEDROCK_DOCUMENT_PLACEHOLDER_TEXT, + _EMPTY_TEXT_PLACEHOLDER, BedrockConverseMessagesProcessor, BedrockImageProcessor, _bedrock_converse_messages_pt, @@ -68,6 +71,17 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.llms.vertex_ai.gemini.transformation import gemini_convert_messages_with_history from litellm.types.llms.bedrock import BedrockConverseReasoningContentBlock from litellm.types.llms.openai import AllMessageValues +from litellm.litellm_core_utils.prompt_templates.factory import response_schema_prompt +from litellm.llms.anthropic.completion.transformation import AnthropicTextConfig +from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import AmazonInvokeConfig +from litellm.llms.codestral.completion.handler import CodestralTextCompletion +from litellm.llms.ollama.completion.transformation import OllamaConfig +from litellm.llms.petals.completion import handler as petals_handler +from litellm.llms.predibase.chat.transformation import PredibaseConfig +from litellm.llms.sagemaker.completion.transformation import SagemakerConfig +from litellm.llms.vllm.completion import handler as vllm_handler +from litellm.llms.watsonx.common_utils import _convert_watsonx_messages_core +from litellm.types.utils import ModelResponse, TextCompletionResponse READ_FILE: Final[list[object]] = [{"name": "read_file", "description": "Read the contents of a file."}] @@ -628,6 +642,64 @@ async def test_bedrock_converse_replays_reasoning_to_claude_only_as_base64_of_cl assert assistant_content == [*sent, {"text": "hello"}] +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize( + "model, final_text", + [("us.anthropic.claude-sonnet-4-6", "Leaves fall"), ("us.amazon.nova-pro-v1:0", "Leaves fall \n")], + ids=["claude", "nova"], +) +async def test_bedrock_converse_strips_trailing_whitespace_from_a_final_assistant_turn_for_claude_only( + use_async: bool, model: str, final_text: str +): + """Claude on Converse 400s on a final assistant turn ending in whitespace, which a cut stream's prefill often has.""" + messages: Final = [ + {"role": "user", "content": "Write a haiku."}, + {"role": "assistant", "content": "Sure. "}, + {"role": "user", "content": "About autumn."}, + {"role": "assistant", "content": "Leaves fall \n"}, + ] + + result: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model=model, llm_provider="bedrock" + ) + if use_async + else _bedrock_converse_messages_pt(messages=messages, model=model, llm_provider="bedrock") + ) + + assert [m["content"] for m in result if m["role"] == "assistant"] == [ + [{"text": "Sure. "}], + [{"text": final_text}], + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_bedrock_converse_keeps_whitespace_between_text_blocks_of_a_final_claude_assistant_turn( + use_async: bool, +): + messages: Final = [ + {"role": "user", "content": "Write a haiku."}, + { + "role": "assistant", + "content": [{"type": "text", "text": "Leaves fall "}, {"type": "text", "text": "softly \n"}], + }, + ] + + result: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model="us.anthropic.claude-sonnet-4-6", llm_provider="bedrock" + ) + if use_async + else _bedrock_converse_messages_pt( + messages=messages, model="us.anthropic.claude-sonnet-4-6", llm_provider="bedrock" + ) + ) + + assert result[-1]["content"] == [{"text": "Leaves fall "}, {"text": "softly"}] + + def test_translate_thinking_blocks_keeps_dropping_redacted_data_when_called_the_old_way(): """Outside callers that only pass the thinking blocks keep the old behavior, with no redactedContent sent.""" assert BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( @@ -795,6 +867,138 @@ def test_anthropic_messages_pt_keeps_signed_thinking_block(): assert thinking_blocks[0]["thinking"] == "genuine anthropic reasoning" +EMPTY_USER_MESSAGES: Final = pytest.mark.parametrize( + "empty_user_message", + [{"role": "user", "content": None}, {"role": "user", "content": []}, {"role": "user"}], + ids=["null", "empty-list", "content-key-dropped"], +) + +Message: TypeAlias = Mapping[str, object] + + +def _convert(messages: Sequence[Message]) -> Sequence[Message]: + return anthropic_messages_pt(messages=list(messages), model="claude-opus-5", llm_provider="vertex_ai") + + +@EMPTY_USER_MESSAGES +def test_anthropic_messages_pt_preserves_trailing_user_message_with_empty_content(empty_user_message: Message) -> None: + result: Final = _convert( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + empty_user_message, + ] + ) + + assert [message["role"] for message in result] == ["user", "assistant", "user"] + assert result[-1]["content"] == [{"type": "text", "text": _EMPTY_TEXT_PLACEHOLDER}] + + +@EMPTY_USER_MESSAGES +@pytest.mark.parametrize( + "cache_control", [{}, {"cache_control": {"type": "ephemeral"}}], ids=["no-cache-control", "with-cache-control"] +) +def test_anthropic_messages_pt_empty_user_message_beside_tool_result_keeps_tool_result_only( + empty_user_message: Message, cache_control: Message +) -> None: + result: Final = _convert( + [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "toolu_1", "type": "function", "function": {"name": "f", "arguments": "{}"}}], + }, + {**empty_user_message, **cache_control}, + {"role": "tool", "tool_call_id": "toolu_1", "content": "ok"}, + ] + ) + + assert [message["role"] for message in result] == ["user", "assistant", "user"] + assert [block["type"] for block in result[-1]["content"]] == ["tool_result"] + + +@EMPTY_USER_MESSAGES +def test_anthropic_messages_pt_empty_user_message_keeps_cache_control_like_an_empty_string( + empty_user_message: Message, +) -> None: + def last_turn(last_user_message: Message) -> Message: + return _convert( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {**last_user_message, "cache_control": {"type": "ephemeral"}}, + ] + )[-1] + + assert last_turn(empty_user_message)["content"][0]["cache_control"] == {"type": "ephemeral"} + assert last_turn(empty_user_message) == last_turn({"role": "user", "content": ""}) + + +TOOL_CALL_AND_RESULT: Final[tuple[Message, ...]] = ( + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "toolu_1", "type": "function", "function": {"name": "f", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "toolu_1", "content": "ok"}, +) + + +@EMPTY_USER_MESSAGES +@pytest.mark.parametrize( + "messages_before", + [ + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "q"}, + ], + [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}, {"role": "user", "content": []}], + TOOL_CALL_AND_RESULT, + ], + ids=["after-user-text", "after-another-empty-message", "after-tool-result"], +) +def test_anthropic_messages_pt_empty_user_message_beside_other_content_sends_nothing_extra( + empty_user_message: Message, messages_before: Sequence[Message] +) -> None: + assert _convert([*messages_before, empty_user_message]) == _convert(messages_before) + + +@EMPTY_USER_MESSAGES +def test_anthropic_messages_pt_empty_user_message_with_cache_control_after_text_keeps_its_breakpoint( + empty_user_message: Message, +) -> None: + def last_turn(last_user_message: Message) -> Message: + return _convert( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "long document"}, + {**last_user_message, "cache_control": {"type": "ephemeral"}}, + ] + )[-1] + + assert last_turn(empty_user_message)["content"][-1]["cache_control"] == {"type": "ephemeral"} + assert last_turn(empty_user_message) == last_turn({"role": "user", "content": ""}) + + +def test_anthropic_messages_pt_keeps_every_breakpoint_of_a_long_run_of_empty_user_messages() -> None: + result: Final = _convert( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + *[{"role": "user", "content": [], "cache_control": {"type": "ephemeral"}}] * 10_000, + ] + ) + + assert ( + result[-1]["content"] + == [{"type": "text", "text": _EMPTY_TEXT_PLACEHOLDER, "cache_control": {"type": "ephemeral"}}] * 10_000 + ) + + def test_convert_to_azure_openai_messages(): """Test coverting image_url to azure_openai spec""" @@ -4933,7 +5137,9 @@ def test_bedrock_converse_messages_pt_user_message_without_content_adds_no_block assert _bedrock_converse_messages_pt( messages=with_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock" - ) == _bedrock_converse_messages_pt(messages=without_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock") + ) == _bedrock_converse_messages_pt( + messages=without_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock" + ) @pytest.mark.asyncio @@ -5063,6 +5269,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -5115,6 +5322,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -5135,6 +5343,7 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -5157,6 +5366,7 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_function_call_non_openai_model(): try: @@ -5185,6 +5395,7 @@ def test_function_call_non_openai_model(): print(e) pass + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_empty_content(): """ @@ -5203,6 +5414,7 @@ def test_empty_content(): litellm_call_id=str(uuid.uuid4()), ) + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_signature_removal_for_non_gemini(): """ @@ -5251,6 +5463,7 @@ def test_thought_signature_removal_for_non_gemini(): assert THOUGHT_SIGNATURE_SEPARATOR not in processed_messages[1]["tool_calls"][0]["id"] assert THOUGHT_SIGNATURE_SEPARATOR not in processed_messages[2]["tool_call_id"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_signature_preserved_for_gemini(): """ @@ -5297,6 +5510,7 @@ def test_thought_signature_preserved_for_gemini(): assert THOUGHT_SIGNATURE_SEPARATOR in processed_messages[1]["tool_calls"][0]["id"] assert THOUGHT_SIGNATURE_SEPARATOR in processed_messages[2]["tool_call_id"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_signature_removal_with_multiple_tool_calls(): """ @@ -5351,6 +5565,7 @@ def test_thought_signature_removal_with_multiple_tool_calls(): assert processed_messages[2]["tool_call_id"] == "call_1" assert processed_messages[3]["tool_call_id"] == "call_2" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_messages_without_tool_calls_unchanged(): """ @@ -7727,3 +7942,334 @@ def test_anthropic_messages_pt_merges_a_long_assistant_run_in_linear_time(): growth: Final = _merged_assistant_run_cost(4000) / _merged_assistant_run_cost(1000) assert growth < 8.0, f"cost grew {growth:.1f}x for 4x the assistant run" + + +TEMPLATE_MODEL = "partial-template-model" +TEMPLATE_ROLES = {"user": {"pre_message": "[INST] ", "post_message": " [/INST]"}} +TEMPLATE_MESSAGES = [{"role": "user", "content": "hi"}] +TEMPLATE_RENDERED_MESSAGES = "[INST] hi [/INST]" +TEMPLATE_BARE_MESSAGES = "hi" + +TEMPLATE_RESPONSE_SCHEMA = {"type": "object"} +TEMPLATE_RENDERED_SCHEMA = f"[INST] {TEMPLATE_RESPONSE_SCHEMA} [/INST]" +TEMPLATE_BARE_SCHEMA = str(TEMPLATE_RESPONSE_SCHEMA) + +TEMPLATE_INITIAL = "" +TEMPLATE_FINAL = "" + +PARTIAL_TEMPLATE = {"roles": TEMPLATE_ROLES} +FULL_TEMPLATE = {"roles": TEMPLATE_ROLES, "initial_prompt_value": TEMPLATE_INITIAL, "final_prompt_value": TEMPLATE_FINAL} +ROLELESS_TEMPLATE = {"initial_prompt_value": TEMPLATE_INITIAL, "final_prompt_value": TEMPLATE_FINAL} + + +class _PromptCaptured(Exception): + def __init__(self, prompt): + super().__init__(prompt) + self.prompt = prompt + + +class _CapturingLogging: + def pre_call(self, input, api_key, additional_args=None, **kwargs): + raise _PromptCaptured(input) + + def post_call(self, *args, **kwargs): + pass + + +def _captured_prompt(call): + with pytest.raises(_PromptCaptured) as excinfo: + call() + return excinfo.value.prompt + + +def _stub_vllm_import(monkeypatch): + class _SamplingParams: + def __init__(self, **kwargs): + self.kwargs = kwargs + + class _LLM: + def __init__(self, model): + self.model = model + + def generate(self, prompts, sampling_params): + raise _PromptCaptured(prompts) + + stub = types.ModuleType("vllm") + stub.LLM = _LLM + stub.SamplingParams = _SamplingParams + + monkeypatch.setitem(sys.modules, "vllm", stub) + monkeypatch.setattr(vllm_handler, "llm", None) + + +def _render_ollama(template, monkeypatch, messages=TEMPLATE_MESSAGES): + request = OllamaConfig().transform_request( + model=TEMPLATE_MODEL, + messages=messages, + optional_params={}, + litellm_params={"custom_prompt_dict": {TEMPLATE_MODEL: template}}, + headers={}, + ) + return request["prompt"] + + +def _render_predibase(template, monkeypatch): + request = PredibaseConfig().transform_request( + model=TEMPLATE_MODEL, + messages=TEMPLATE_MESSAGES, + optional_params={}, + litellm_params={"custom_prompt_dict": {TEMPLATE_MODEL: template}}, + headers={}, + ) + return request["inputs"] + + +def _render_anthropic_text(template, monkeypatch): + monkeypatch.setattr(litellm, "custom_prompt_dict", {TEMPLATE_MODEL: template}) + return AnthropicTextConfig()._get_anthropic_text_prompt_from_messages( + messages=TEMPLATE_MESSAGES, model=TEMPLATE_MODEL + ) + + +def _render_response_schema_prompt(template, monkeypatch): + monkeypatch.setattr(litellm, "custom_prompt_dict", {f"{TEMPLATE_MODEL}/response_schema_prompt": template}) + return response_schema_prompt( + model=TEMPLATE_MODEL, response_schema=TEMPLATE_RESPONSE_SCHEMA, llm_provider="vertex_ai" + ) + + +def _render_petals(template, monkeypatch): + monkeypatch.setattr(litellm, "custom_prompt_dict", {TEMPLATE_MODEL: template}) + return _captured_prompt( + lambda: petals_handler.completion( + model=TEMPLATE_MODEL, + messages=TEMPLATE_MESSAGES, + api_base="http://petals.invalid", + model_response=ModelResponse(), + print_verbose=lambda *args, **kwargs: None, + encoding=None, + logging_obj=_CapturingLogging(), + optional_params={}, + ) + ) + + +def _render_codestral(template, monkeypatch): + return _captured_prompt( + lambda: CodestralTextCompletion().completion( + model=TEMPLATE_MODEL, + messages=TEMPLATE_MESSAGES, + api_base="http://codestral.invalid", + custom_prompt_dict={TEMPLATE_MODEL: template}, + model_response=TextCompletionResponse(), + print_verbose=lambda *args, **kwargs: None, + encoding=None, + api_key="sk-not-used", + logging_obj=_CapturingLogging(), + optional_params={}, + timeout=1.0, + ) + ) + + +def _render_vllm_completion(template, monkeypatch): + _stub_vllm_import(monkeypatch) + return _captured_prompt( + lambda: vllm_handler.completion( + model=TEMPLATE_MODEL, + messages=TEMPLATE_MESSAGES, + model_response=ModelResponse(), + print_verbose=lambda *args, **kwargs: None, + encoding=None, + logging_obj=_CapturingLogging(), + optional_params={}, + custom_prompt_dict={TEMPLATE_MODEL: template}, + ) + ) + + +def _render_vllm_batch(template, monkeypatch): + _stub_vllm_import(monkeypatch) + prompts = _captured_prompt( + lambda: vllm_handler.batch_completions( + model=TEMPLATE_MODEL, + messages=[TEMPLATE_MESSAGES], + optional_params={}, + custom_prompt_dict={TEMPLATE_MODEL: template}, + ) + ) + assert len(prompts) == 1 + return prompts[0] + + +def _render_bedrock_invoke(template, monkeypatch): + prompt, _ = AmazonInvokeConfig().convert_messages_to_prompt( + model=TEMPLATE_MODEL, messages=TEMPLATE_MESSAGES, provider="meta", custom_prompt_dict={TEMPLATE_MODEL: template} + ) + return prompt + + +def _render_sagemaker(template, monkeypatch): + return SagemakerConfig()._transform_prompt( + model=TEMPLATE_MODEL, + messages=TEMPLATE_MESSAGES, + custom_prompt_dict={TEMPLATE_MODEL: template}, + hf_model_name=None, + ) + + +def _render_watsonx(template, monkeypatch): + return _convert_watsonx_messages_core( + model=TEMPLATE_MODEL, + messages=TEMPLATE_MESSAGES, + provider="ibm", + custom_prompt_dict={TEMPLATE_MODEL: template}, + apply_template_fn=None, + ) + + +TEMPLATE_CALL_SITES = [ + pytest.param(_render_ollama, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="ollama"), + pytest.param(_render_petals, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="petals"), + pytest.param(_render_vllm_completion, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="vllm-completion"), + pytest.param(_render_vllm_batch, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="vllm-batch-completions"), + pytest.param(_render_predibase, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="predibase"), + pytest.param(_render_codestral, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="codestral"), + pytest.param(_render_anthropic_text, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="anthropic-text"), + pytest.param( + _render_response_schema_prompt, TEMPLATE_RENDERED_SCHEMA, TEMPLATE_BARE_SCHEMA, id="response-schema-prompt" + ), + pytest.param(_render_bedrock_invoke, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="bedrock-invoke"), + pytest.param(_render_sagemaker, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="sagemaker"), + pytest.param(_render_watsonx, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="watsonx"), +] + + +@pytest.mark.parametrize(("render", "rendered_body", "bare_body"), TEMPLATE_CALL_SITES) +def test_template_without_the_optional_prompt_values_renders(render, rendered_body, bare_body, monkeypatch): + assert render(PARTIAL_TEMPLATE, monkeypatch) == rendered_body + + +@pytest.mark.parametrize(("render", "rendered_body", "bare_body"), TEMPLATE_CALL_SITES) +def test_template_with_the_optional_prompt_values_still_applies_them(render, rendered_body, bare_body, monkeypatch): + assert render(FULL_TEMPLATE, monkeypatch) == TEMPLATE_INITIAL + rendered_body + TEMPLATE_FINAL + + +@pytest.mark.parametrize(("render", "rendered_body", "bare_body"), TEMPLATE_CALL_SITES) +def test_template_without_roles_renders_messages_unwrapped(render, rendered_body, bare_body, monkeypatch): + assert render(ROLELESS_TEMPLATE, monkeypatch) == TEMPLATE_INITIAL + bare_body + TEMPLATE_FINAL + + +def test_template_wraps_only_the_roles_it_defines_across_a_conversation(monkeypatch): + messages = [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "sky colour?"}, + ] + + rendered = _render_ollama(FULL_TEMPLATE, monkeypatch, messages=messages) + + assert rendered == "Be terse.[INST] hi [/INST]hello[INST] sky colour? [/INST]" + + +@pytest.mark.parametrize( + "template", + [ + pytest.param({"roles": {**TEMPLATE_ROLES, "tool": None}}, id="null-role-entry"), + pytest.param({"roles": {"user": {**TEMPLATE_ROLES["user"], "priority": 1}}}, id="extra-role-key"), + pytest.param({"roles": {**TEMPLATE_ROLES, "tool": {"pre_message": 5}}}, id="unused-role-wrong-type"), + pytest.param({"roles": {**TEMPLATE_ROLES, 1: {"pre_message": "x"}}}, id="unused-non-string-role"), + pytest.param({"roles": TEMPLATE_ROLES, "initial_prompt_value": None, "bos_token": 7}, id="null-and-unread-keys"), + ], +) +def test_template_shapes_main_rendered_keep_rendering(template, monkeypatch): + assert _render_ollama(template, monkeypatch) == TEMPLATE_RENDERED_MESSAGES + + +@pytest.mark.parametrize( + "template", + [ + pytest.param({"roles": {"user": "[INST]"}}, id="string-role-entry"), + pytest.param({"roles": {"user": ["[INST]"]}}, id="list-role-entry"), + pytest.param({"roles": ["user"]}, id="list-roles"), + pytest.param({"roles": '{"user": {"pre_message": "[INST]"}}'}, id="json-string-roles"), + ], +) +def test_roles_that_are_not_a_mapping_render_unwrapped_with_a_warning(template, monkeypatch, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + rendered = _render_ollama({**template, "final_prompt_value": TEMPLATE_FINAL}, monkeypatch) + + assert rendered == TEMPLATE_BARE_MESSAGES + TEMPLATE_FINAL + assert "not a mapping" in caplog.text + + +def test_null_role_entry_renders_that_role_unwrapped_without_a_warning(monkeypatch, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + rendered = _render_ollama({"roles": {"user": None}}, monkeypatch) + + assert rendered == TEMPLATE_BARE_MESSAGES + assert "not a mapping" not in caplog.text + + +@pytest.mark.parametrize( + ("template", "field"), + [ + pytest.param({"initial_prompt_value": 5}, "initial_prompt_value", id="initial-prompt-value"), + pytest.param({"final_prompt_value": ["x"]}, "final_prompt_value", id="final-prompt-value"), + pytest.param({"roles": {"user": {"pre_message": 5}}}, "roles.user.pre_message", id="pre-message"), + ], +) +def test_wrong_type_template_field_fails_as_a_server_error_naming_it(template, field): + with pytest.raises(litellm.APIConnectionError) as excinfo: + litellm.completion( + model=f"ollama/{TEMPLATE_MODEL}", + messages=TEMPLATE_MESSAGES, + api_base="http://ollama.invalid", + **template, + ) + + assert excinfo.value.status_code == 500 + assert f"{field}: Input should be a valid string" in excinfo.value.message + + +@pytest.mark.parametrize( + "template", + [ + pytest.param({"bos_token": ""}, id="bos-token-only"), + pytest.param({"eos_token": ""}, id="eos-token-only"), + pytest.param({"roles": '{"user": {"pre_message": "[INST]"}}'}, id="json-string-roles-only"), + ], +) +def test_template_that_sets_nothing_the_prompt_uses_fails_as_a_server_error(template): + with pytest.raises(litellm.APIConnectionError) as excinfo: + litellm.completion( + model=f"ollama/{TEMPLATE_MODEL}", + messages=TEMPLATE_MESSAGES, + api_base="http://ollama.invalid", + **template, + ) + + assert excinfo.value.status_code == 500 + assert "is invalid: it sets none of roles, initial_prompt_value, or final_prompt_value" in excinfo.value.message + + +def test_wrong_type_template_field_in_requests_leaves_the_deployments_in_rotation(): + router = litellm.Router( + model_list=[ + { + "model_name": "template-group", + "litellm_params": {"model": f"ollama/{TEMPLATE_MODEL}", "api_base": "http://ollama.invalid"}, + } + for _ in range(2) + ], + num_retries=0, + allowed_fails=0, + ) + + for _ in range(2): + with pytest.raises(litellm.APIConnectionError): + router.completion(model="template-group", messages=TEMPLATE_MESSAGES, initial_prompt_value=5) + response = router.completion(model="template-group", messages=TEMPLATE_MESSAGES, mock_response="Blue") + + assert response.choices[0].message.content == "Blue" diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py index 6bc1178f64b..91d7acd331e 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py @@ -29,11 +29,8 @@ def _texts(message: dict) -> list[str]: return [block["text"] for block in message["content"]] -SENDS_NOTHING = pytest.mark.parametrize( - "empty_content", - [[], None, [{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}]], - ids=["empty-list", "none", "unsupported-part-only"], -) +EMPTY_USER_CONTENT: Final = pytest.mark.parametrize("empty_content", [[], None], ids=["empty-list", "none"]) +UNSUPPORTED_PART_ONLY: Final = ({"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}},) def test_split_leading_system_run_keeps_later_system_messages_in_the_conversation(): @@ -213,9 +210,9 @@ def test_flagged_placement_of_an_earlier_run_does_not_move_when_later_turns_are_ assert _roles(placed_n_plus_one) == ["user", "assistant", "user", "assistant", "user"] -@SENDS_NOTHING -def test_flagged_placement_converts_a_run_whose_preceding_user_turn_sends_nothing(empty_content): - placed = place_mid_conversation_system( +@EMPTY_USER_CONTENT +def test_flagged_placement_anchors_a_run_on_a_preceding_empty_user_turn(empty_content: list[object] | None) -> None: + placed: Final = place_mid_conversation_system( [ {"role": "user", "content": empty_content}, {"role": "system", "content": "reminder"}, @@ -225,18 +222,48 @@ def test_flagged_placement_converts_a_run_whose_preceding_user_turn_sends_nothin supports_mid_conversation_system=True, ) + assert _roles(placed) == ["user", "system", "assistant", "user"] + assert _texts(placed[1]) == ["reminder"] + + +@EMPTY_USER_CONTENT +def test_flagged_placement_anchors_a_run_on_a_following_empty_user_turn(empty_content: list[object] | None) -> None: + placed: Final = place_mid_conversation_system( + [ + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "a1"}, + {"role": "system", "content": "reminder"}, + {"role": "user", "content": empty_content}, + ], + supports_mid_conversation_system=True, + ) + + assert _roles(placed) == ["user", "assistant", "user", "system"] + assert _texts(placed[3]) == ["reminder"] + + +def test_flagged_placement_converts_a_run_whose_preceding_user_turn_sends_nothing() -> None: + placed: Final = place_mid_conversation_system( + [ + {"role": "user", "content": list(UNSUPPORTED_PART_ONLY)}, + {"role": "system", "content": "reminder"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "q2"}, + ], + supports_mid_conversation_system=True, + ) + assert _roles(placed) == ["user", "user", "assistant", "user"] assert _texts(placed[1]) == [CONVERTED_SYSTEM_NOTE, "reminder"] -@SENDS_NOTHING -def test_flagged_placement_converts_a_run_whose_following_user_turn_sends_nothing(empty_content): - placed = place_mid_conversation_system( +def test_flagged_placement_converts_a_run_whose_following_user_turn_sends_nothing() -> None: + placed: Final = place_mid_conversation_system( [ {"role": "user", "content": "q1"}, {"role": "assistant", "content": "a1"}, {"role": "system", "content": "reminder"}, - {"role": "user", "content": empty_content}, + {"role": "user", "content": list(UNSUPPORTED_PART_ONLY)}, ], supports_mid_conversation_system=True, ) 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_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 4347a6639bd..93f700b1c2a 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -1019,6 +1019,87 @@ async def test_bridged_stream_mid_stream_fallback_error_is_rebuilt_around_the_pr assert excinfo.value.generated_content == "" +@pytest.mark.asyncio +@pytest.mark.parametrize("merge_reasoning", [True, False], ids=["merged_into_content", "kept_as_reasoning"]) +async def test_mid_stream_break_after_reasoning_disqualifies_continuation(logging_obj: Logging, merge_reasoning: bool): + import httpx + + from litellm.exceptions import MidStreamFallbackError + + class _ReasoningThenCut: + def __init__(self) -> None: + self.chunks = iter( + ( + ModelResponseStream(choices=[{"index": 0, "delta": {"reasoning_content": "Let me think"}}]), + ModelResponseStream(choices=[{"index": 0, "delta": {"content": "The answer"}}]), + ) + ) + + def __aiter__(self) -> "_ReasoningThenCut": + return self + + async def __anext__(self) -> ModelResponseStream: + next_chunk: Final = next(self.chunks, None) + if next_chunk is None: + raise httpx.ReadTimeout("connection dropped") + return next_chunk + + response: Final = CustomStreamWrapper( + completion_stream=_ReasoningThenCut(), model="deepseek-reasoner", logging_obj=logging_obj + ) + response.merge_reasoning_content_in_choices = merge_reasoning + + with pytest.raises(MidStreamFallbackError) as excinfo: + async for _ in response: + pass + + assert excinfo.value.is_pre_first_chunk is False + assert excinfo.value.emitted_disqualifying_content is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "outer_delta,inner_flag,expected", + [ + ({"tool_calls": [{"index": 0}]}, False, True), + ({"content": "Let me check."}, True, True), + ({"content": "Let me check."}, False, False), + ], + ids=["bridge_wrapper_sent_a_tool_call", "inner_stream_sent_a_tool_call", "plain_text_only"], +) +async def test_bridged_stream_rebuild_keeps_the_continuation_disqualifier( + logging_obj: Logging, outer_delta: dict[str, object], inner_flag: bool, expected: bool +): + from litellm.exceptions import InternalServerError, MidStreamFallbackError + + provider_error: Final = InternalServerError(message="upstream 500", llm_provider="azure", model="gpt-6.1-sol") + inner_error: Final = MidStreamFallbackError( + message=str(provider_error), + model="gpt-6.1-sol", + llm_provider="azure", + original_exception=provider_error, + is_pre_first_chunk=False, + emitted_disqualifying_content=inner_flag, + ) + + async def _raise_inner_error(**kwargs: object) -> NoReturn: + raise inner_error + + response: Final = CustomStreamWrapper( + completion_stream=None, + model="gpt-6.1-sol", + logging_obj=logging_obj, + custom_llm_provider="azure", + make_call=_raise_inner_error, + ) + response._accumulate_streamed_delta(outer_delta) + + with pytest.raises(MidStreamFallbackError) as excinfo: + await response.__anext__() + + assert excinfo.value.emitted_disqualifying_content is expected + + @pytest.mark.asyncio @pytest.mark.parametrize("sync_mode", [True, False]) async def test_streaming_rate_limit_midstream_fallback_keeps_provider_headers(logging_obj: Logging, sync_mode: bool): @@ -6062,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 @@ -6230,6 +6245,59 @@ class TestStableStreamingResponseId: assert wrapper.model_response_creator().id == "chatcmpl-from-provider" +class TestContinuationDisqualifiers: + """The mid-stream continuation eligibility hinges on classifying which deltas + carry output a text-only prefill cannot represent.""" + + @pytest.mark.parametrize( + "field", + [ + "tool_calls", + "function_call", + "reasoning_content", + "thinking_blocks", + "reasoning_items", + "audio", + "images", + "annotations", + ], + ) + def test_disqualifying_fields_flagged(self, field): + # reasoning_content included: it reaches the caller as visible reasoning a + # text-only prefill cannot carry, so a reasoning fallback would re-derive it + assert CustomStreamWrapper._delta_disqualifies_continuation({field: [{"x": 1}]}) is True + + @pytest.mark.parametrize("delta", [{"content": "hi"}, {}, {"role": "assistant"}]) + def test_plain_text_not_flagged(self, delta): + assert CustomStreamWrapper._delta_disqualifies_continuation(delta) is False + + def test_non_mapping_delta_is_safe(self): + assert CustomStreamWrapper._delta_disqualifies_continuation(object()) is False + + def test_accumulate_grows_text_and_leaves_flag_clear_for_plain_text(self): + wrapper = object.__new__(CustomStreamWrapper) + wrapper.response_uptil_now = "" + wrapper._emitted_disqualifying_content = False + + wrapper._accumulate_streamed_delta({"content": "Hel"}) + wrapper._accumulate_streamed_delta({"content": "lo"}) + + assert wrapper.response_uptil_now == "Hello" + assert wrapper._emitted_disqualifying_content is False + + def test_accumulate_latches_flag_on_disqualifying_delta(self): + wrapper = object.__new__(CustomStreamWrapper) + wrapper.response_uptil_now = "" + wrapper._emitted_disqualifying_content = False + + wrapper._accumulate_streamed_delta({"content": "Hi"}) + wrapper._accumulate_streamed_delta({"tool_calls": [{"index": 0}]}) + wrapper._accumulate_streamed_delta({"content": "there"}) + + assert wrapper.response_uptil_now == "Hithere" + assert wrapper._emitted_disqualifying_content is True + + @pytest.mark.asyncio async def test_async_stream_without_usage_counts_tokens_off_the_event_loop(): from tests.large_text import text 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/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 0c093700822..724e137c835 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6763,15 +6763,11 @@ def test_chat_flagged_model_converts_system_with_no_following_user_turn(local_mo assert _texts(result["messages"][2]) == [CONVERTED_SYSTEM_NOTE, REMINDER_TEXT] -@pytest.mark.parametrize( - "empty_content", - [[], None, [{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}]], - ids=["empty-list", "none", "unsupported-part-only"], -) -def test_chat_flagged_model_converts_a_system_behind_a_user_turn_that_sends_nothing( - local_model_cost_map, empty_content -): - messages = [ +@pytest.mark.parametrize("empty_content", [[], None], ids=["empty-list", "none"]) +def test_chat_flagged_model_anchors_a_system_on_an_empty_user_turn( + local_model_cost_map: None, empty_content: list[object] | None +) -> None: + messages: Final = [ {"role": "system", "content": "You are terse."}, {"role": "user", "content": empty_content}, {"role": "system", "content": REMINDER_TEXT}, @@ -6779,7 +6775,22 @@ def test_chat_flagged_model_converts_a_system_behind_a_user_turn_that_sends_noth {"role": "user", "content": "Second question"}, ] - result = _chat_request(AnthropicConfig(), FLAGGED_CLAUDE, messages) + result: Final = _chat_request(AnthropicConfig(), FLAGGED_CLAUDE, messages) + + assert [m["role"] for m in result["messages"]] == ["user", "system", "assistant", "user"] + assert _texts(result["messages"][1]) == [REMINDER_TEXT] + + +def test_chat_flagged_model_converts_a_system_behind_a_user_turn_that_sends_nothing(local_model_cost_map: None) -> None: + messages: Final = [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": [{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}]}, + {"role": "system", "content": REMINDER_TEXT}, + {"role": "assistant", "content": "First answer"}, + {"role": "user", "content": "Second question"}, + ] + + result: Final = _chat_request(AnthropicConfig(), FLAGGED_CLAUDE, messages) assert [m["role"] for m in result["messages"]] == ["user", "assistant", "user"] assert _texts(result["messages"][0]) == [CONVERTED_SYSTEM_NOTE, REMINDER_TEXT] 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_handler_reasoning_effort_normalization.py b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py index 0dfe4a93649..4b96d2e7991 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py @@ -6,13 +6,17 @@ regression they guard is the one a caller sees: a tier the proxy advertises has leaves the adapter, in the shape the target expects. """ +from collections.abc import Mapping from typing import Final import pytest +import litellm from litellm.llms.anthropic.pass_through.adapters.handler import ( LiteLLMMessagesToCompletionTransformationHandler, ) +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig +from litellm.utils import reasoning_effort_with_summary_alias MESSAGES = [{"role": "user", "content": "hello"}] @@ -38,14 +42,14 @@ def _reasoning_effort_sent(model: str, provider: str, reasoning_effort: object) return completion_kwargs.get("reasoning_effort") -def _reasoning_effort_sent_for_thinking( +def _completion_kwargs_for_thinking( model: str, provider: str | None, thinking: dict[str, object], *, tools: list[dict[str, object]] | None = None, api_base: str | None = None, -) -> object: +) -> Mapping[str, object]: extra_kwargs: Final = { key: value for key, value in (("custom_llm_provider", provider), ("api_base", api_base)) if value is not None } @@ -66,7 +70,20 @@ def _reasoning_effort_sent_for_thinking( output_format=None, extra_kwargs=extra_kwargs, ) - return completion_kwargs.get("reasoning_effort") + return completion_kwargs + + +def _reasoning_effort_sent_for_thinking( + model: str, + provider: str | None, + thinking: dict[str, object], + *, + tools: list[dict[str, object]] | None = None, + api_base: str | None = None, +) -> object: + return _completion_kwargs_for_thinking(model, provider, thinking, tools=tools, api_base=api_base).get( + "reasoning_effort" + ) SUMMARIZED_THINKING = {"type": "enabled", "budget_tokens": 4096, "summary": "auto"} @@ -78,6 +95,66 @@ MULTIPLY_TOOL = { } +class TestAThinkingSummaryAskReachesGemini: + """Gemini maps a summary ask to ``includeThoughts``, so ``gemini_include_thoughts_default`` only + decides for a ``/v1/messages`` request that asks for nothing, the same as on chat completions.""" + + @pytest.mark.parametrize( + "model, provider", + [("gemini/gemini-3.1-pro-preview", "gemini"), ("vertex_ai/gemini-2.5-flash", "vertex_ai")], + ) + @pytest.mark.parametrize( + "thinking, sends_thoughts", + [ + (SUMMARIZED_THINKING, True), + ({**PLAIN_THINKING, "display": "summarized"}, True), + ({**PLAIN_THINKING, "display": "updates"}, True), + ({**PLAIN_THINKING, "display": "omitted"}, False), + (PLAIN_THINKING, False), + ({"type": "disabled", "display": "summarized"}, False), + ], + ) + def test_the_request_ask_beats_the_gateway_default( + self, + local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, + model: str, + provider: str, + thinking: dict[str, object], + sends_thoughts: bool, + ) -> None: + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False) + + sent: Final = _completion_kwargs_for_thinking(model, provider, thinking) + gemini: Final = VertexGeminiConfig() + reasoning_effort: Final = reasoning_effort_with_summary_alias( + sent.get("reasoning_effort"), sent, accepts_summary=gemini.accepts_reasoning_summary_in_reasoning_effort + ) + mapped: Final = gemini.map_openai_params( + {"reasoning_effort": reasoning_effort}, {}, model.split("/", 1)[1], True + ) + + assert mapped["thinkingConfig"]["includeThoughts"] is sends_thoughts + assert not isinstance(sent.get("reasoning_effort"), dict) + + def test_auto_summary_leaves_the_gateway_default_in_charge( + self, local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False) + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") + + sent: Final = _reasoning_effort_sent_for_thinking("gemini/gemini-3.1-pro-preview", "gemini", PLAIN_THINKING) + + assert sent == "high" + + def test_a_display_ask_leaves_other_chat_targets_on_the_plain_tier(self, local_model_cost_map: None) -> None: + sent: Final = _reasoning_effort_sent_for_thinking( + "databricks/databricks-qwen35-122b-a10b", "databricks", {**PLAIN_THINKING, "display": "summarized"} + ) + + assert sent == "high" + + class TestTheSummaryWrappingOnlyRidesTheResponsesBridge: """Only the Responses API takes ``reasoning_effort`` as a dict. Databricks answered the wrapped ``{"effort", "summary"}`` with ``field 'reasoning_effort' expects input with json type 'string' @@ -100,7 +177,9 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge: ) -> None: monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") - sent = _reasoning_effort_sent_for_thinking("databricks/databricks-qwen35-122b-a10b", "databricks", PLAIN_THINKING) + sent: Final = _reasoning_effort_sent_for_thinking( + "databricks/databricks-qwen35-122b-a10b", "databricks", PLAIN_THINKING + ) assert sent == "high" @@ -113,7 +192,7 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge: ], ) def test_a_bridged_target_keeps_the_summary(self, local_model_cost_map: None, model: str, provider: str) -> None: - sent = _reasoning_effort_sent_for_thinking(model, provider, SUMMARIZED_THINKING) + sent: Final = _reasoning_effort_sent_for_thinking(model, provider, SUMMARIZED_THINKING) assert sent == {"effort": "high", "summary": "auto"} @@ -122,7 +201,7 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge: ) -> None: monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") - sent = _reasoning_effort_sent_for_thinking("azure/responses/gpt-5-mini", "azure", PLAIN_THINKING) + sent: Final = _reasoning_effort_sent_for_thinking("azure/responses/gpt-5-mini", "azure", PLAIN_THINKING) assert sent == {"effort": "high", "summary": "detailed"} @@ -139,21 +218,21 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge: """``completion()`` keeps a gpt-5.5 deployment with function tools on Foundry's chat route when its ``api_base`` is a Foundry OpenAI host, and bridges it to Responses when the base makes it an Azure OpenAI deployment. The adapter has to read the same ``api_base`` to land on the same call.""" - sent = _reasoning_effort_sent_for_thinking( + sent: Final = _reasoning_effort_sent_for_thinking( "azure_ai/gpt-5.5", "azure_ai", SUMMARIZED_THINKING, tools=[MULTIPLY_TOOL], api_base=api_base ) assert sent == expected def test_a_provider_resolved_from_the_api_base_gets_the_plain_tier(self, local_model_cost_map: None) -> None: - sent = _reasoning_effort_sent_for_thinking( + sent: Final = _reasoning_effort_sent_for_thinking( "kimi-k3", None, SUMMARIZED_THINKING, api_base="https://api.together.xyz/v1" ) assert sent == "high" def test_a_chained_gateway_keeps_the_dict_for_its_own_bridge(self, local_model_cost_map: None) -> None: - sent = _reasoning_effort_sent_for_thinking("litellm_proxy/gpt-5.4", "litellm_proxy", SUMMARIZED_THINKING) + sent: Final = _reasoning_effort_sent_for_thinking("litellm_proxy/gpt-5.4", "litellm_proxy", SUMMARIZED_THINKING) assert sent == {"effort": "high", "summary": "auto"} @@ -192,7 +271,7 @@ class TestTheNormalizedTierIsTheTierSent: def test_the_dict_form_normalizes_effort_and_keeps_its_siblings( self, local_model_cost_map, model, provider, expected ): - sent = _reasoning_effort_sent(model, provider, {"effort": "max", "summary": "detailed"}) + sent: Final = _reasoning_effort_sent(model, provider, {"effort": "max", "summary": "detailed"}) assert sent == {"effort": expected, "summary": "detailed"} @@ -215,8 +294,8 @@ class TestTheNormalizedTierIsTheTierSent: ], ) def test_an_azure_hosted_astra_deployment_drops_to_the_tier_it_accepts( - self, local_model_cost_map, model, provider - ): + self, local_model_cost_map: None, model: str, provider: str + ) -> None: """The deployment answers ``max`` with a 400 naming ``none`` through ``xhigh``, so the rows say so and the adapter sends the tier below instead of the rejected one.""" assert _reasoning_effort_sent(model, provider, "max") == "xhigh" 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/anthropic/pass_through/test_convertible_blocks.py b/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py new file mode 100644 index 00000000000..b5301b413ff --- /dev/null +++ b/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py @@ -0,0 +1,192 @@ +import asyncio +import json +from collections.abc import Mapping +from typing import Final + +import pytest +from pydantic import JsonValue, TypeAdapter + +from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter +from litellm.llms.anthropic.pass_through.convertible_blocks import ( + unconvertible_anthropic_blocks, + unconvertible_blocks_error, +) +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) +from litellm.llms.vertex_ai.count_tokens.context import gemini_count_contents +from litellm.types.llms.anthropic import AllAnthropicPassThroughMessageValues + +_JSON_OBJECTS: Final = TypeAdapter(list[dict[str, JsonValue]]) + +_PDF: Final = {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"} +_PNG: Final = {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="} + +_CONVERTED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object]], ...]] = ( + ("user", {"type": "text", "text": "hi"}), + ("user", {"type": "image", "source": _PNG}), + ("user", {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}), + ("user", {"type": "document", "source": _PDF}), + ("user", {"type": "document", "source": {"type": "url", "url": "https://example.com/a.pdf"}}), + ("user", {"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}), + ("assistant", {"type": "text", "text": "hi"}), + ("assistant", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}), + ("assistant", {"type": "thinking", "thinking": "Let me think.", "signature": "sig"}), +) + +_DROPPED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object], str], ...]] = ( + ("user", {"type": "search_result", "source": "https://example.com", "title": "t", "content": []}, "search_result"), + ("user", {"type": "document", "source": {"type": "text", "data": "notes"}}, "document with a text source"), + ("user", {"type": "image", "source": {"type": "file", "file_id": "file_01"}}, "image with a file source"), + ("user", {"type": "image", "source": {"type": "base64", "data": ""}}, "image with an empty base64 source"), + ("user", {"type": "document"}, "document without a source"), + ("user", {"text": "untyped"}, "a block without a type"), + ("user", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}, "tool_use"), + ( + "assistant", + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {}}, + "server_tool_use", + ), +) + + +def _adapter_messages(role: str, block: Mapping[str, object]) -> list[AllAnthropicPassThroughMessageValues]: + return [{"role": role, "content": [dict(block)]}] # pyright: ignore[reportReturnType] # sends block shapes the TypedDicts leave out on purpose + + +def _converted_by_both_adapters(role: str, block: Mapping[str, object]) -> tuple[bool, bool]: + messages: Final = _adapter_messages(role, block) + chat: Final = _JSON_OBJECTS.validate_python( + LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=messages) + ) + responses: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input(messages) + chat_kept: Final = any( + message.get("content") or message.get("tool_calls") or message.get("thinking_blocks") for message in chat + ) + return chat_kept, bool(responses) + + +@pytest.mark.parametrize(("role", "block"), _CONVERTED_BLOCKS) +def test_blocks_both_conversions_keep_are_accepted(role: str, block: Mapping[str, object]) -> None: + assert _converted_by_both_adapters(role, block) == (True, True) + assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == () + + +@pytest.mark.parametrize(("role", "block", "label"), _DROPPED_BLOCKS) +def test_blocks_both_conversions_drop_are_named(role: str, block: Mapping[str, object], label: str) -> None: + assert _converted_by_both_adapters(role, block) == (False, False) + assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == (label,) + + +def test_labels_are_listed_once_in_order_and_string_content_is_accepted() -> None: + messages: Final = [ + {"role": "user", "content": "plain text"}, + {"role": "user", "content": [{"type": "search_result"}, {"type": "container_upload"}]}, + {"role": "assistant", "content": [{"type": "search_result"}]}, + ] + assert unconvertible_anthropic_blocks(messages) == ("container_upload", "search_result") + + +def test_redacted_thinking_is_accepted_like_other_replayed_thinking() -> None: + messages: Final = [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "opaque"}]}] + assert unconvertible_anthropic_blocks(messages) == () + + +def test_messages_that_are_not_json_are_left_to_the_count() -> None: + messages: Final = [{"role": "user", "content": [{"type": "search_result", "source": object()}]}] + assert unconvertible_anthropic_blocks(messages) == () + + +def _tool_turn(inner: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + return ( + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": [dict(inner)]}]}, + ) + + +def _sent_by_each_counter(inner: Mapping[str, object]) -> tuple[str, str]: + messages: Final = _tool_turn(inner) + gemini: Final = asyncio.run(gemini_count_contents(messages, "anthropic", "gemini-2.5-flash", "gemini", {})) + openai: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input( + [dict(message) for message in messages] # pyright: ignore[reportArgumentType] # sends block shapes the TypedDicts leave out on purpose + ) + return json.dumps(gemini, default=str), json.dumps(openai, default=str) + + +_KEPT_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str], ...]] = ( + ({"type": "text", "text": "forty-two"}, "forty-two"), + ({"type": "image", "source": _PNG}, _PNG["data"]), + ({"type": "document", "source": _PDF}, _PDF["data"]), +) + +_DROPPED_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str, str], ...]] = ( + ( + { + "type": "search_result", + "source": "https://example.com", + "title": "t", + "content": [{"type": "text", "text": "search-snippet"}], + }, + "search-snippet", + "search_result inside a tool_result", + ), + ( + {"type": "document", "source": {"type": "text", "data": "text-notes"}}, + "text-notes", + "document with a text source inside a tool_result", + ), + ({"type": "tool_reference", "tool_name": "find_docs"}, "find_docs", "tool_reference inside a tool_result"), +) + + +@pytest.mark.parametrize(("inner", "marker"), _KEPT_INSIDE_TOOL_RESULT) +def test_tool_result_content_both_counters_send_is_accepted(inner: Mapping[str, object], marker: str) -> None: + gemini, openai = _sent_by_each_counter(inner) + + assert marker in gemini + assert marker in openai + assert unconvertible_anthropic_blocks(_tool_turn(inner)) == () + + +@pytest.mark.parametrize(("inner", "marker", "label"), _DROPPED_INSIDE_TOOL_RESULT) +def test_tool_result_content_both_counters_drop_is_named(inner: Mapping[str, object], marker: str, label: str) -> None: + gemini, openai = _sent_by_each_counter(inner) + + assert marker not in gemini + assert marker not in openai + assert unconvertible_anthropic_blocks(_tool_turn(inner)) == (label,) + + +def _replayed_web_search(encrypted_content: str) -> tuple[Mapping[str, object], ...]: + return ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": encrypted_content, + "snippet": "LiteLLM ships a weekly release candidate.", + } + ], + }, + ], + }, + ) + + +def test_a_replayed_web_search_inference_sends_as_text_is_accepted() -> None: + assert unconvertible_blocks_error(_replayed_web_search(""), "anthropic") is None + + +def test_an_encrypted_web_search_inference_cannot_send_is_named() -> None: + error: Final = unconvertible_blocks_error(_replayed_web_search("opaque"), "anthropic") + assert error is not None + assert error.endswith("server_tool_use, web_search_tool_result") diff --git a/tests/unit/llms/anthropic/test_count_tokens_oauth.py b/tests/unit/llms/anthropic/test_count_tokens_oauth.py index 7f717091754..16ec513916c 100644 --- a/tests/unit/llms/anthropic/test_count_tokens_oauth.py +++ b/tests/unit/llms/anthropic/test_count_tokens_oauth.py @@ -197,6 +197,30 @@ class TestCountTokensUsesWorkloadIdentity: assert ANTHROPIC_TOKEN_COUNTING_BETA_VERSION in betas assert ANTHROPIC_OAUTH_BETA_HEADER not in betas + @pytest.mark.asyncio + async def test_a_global_anthropic_key_counts_like_chat_would_authenticate( + self, monkeypatch, httpx_transport_clients + ): + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "anthropic_key", FAKE_REGULAR_KEY) + + with respx.mock(assert_all_called=True) as router: + route = router.post("https://api.anthropic.com/v1/messages/count_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 13}) + ) + result = await AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}}, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.total_tokens == 13 + assert route.calls.last.request.headers["x-api-key"] == FAKE_REGULAR_KEY + @pytest.mark.asyncio async def test_a_failed_mint_degrades_like_an_anthropic_error(self, monkeypatch): from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module diff --git a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py index 9875cf439cf..f753fb6b8f9 100644 --- a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -753,3 +753,33 @@ def test_azure_deployment_id_selects_chat_completions_endpoint(respx_mock: MockR assert route.call_count == 1 assert route.calls[0].request.url.path == "/openai/deployments/deployment-gpt-4o/chat/completions" assert response.choices[0].message.content == "Hello" + + +@pytest.mark.respx(assert_all_called=True) +def test_responses_api_version_leaves_the_chat_route_and_body_untouched(respx_mock: MockRouter) -> None: + url: Final = "https://example-resource.openai.azure.com/openai/deployments/gpt-5.4-nano/chat/completions" + payload: Final = { + "id": "chatcmpl-dated", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-5.4-nano", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, + } + route: Final = respx_mock.post(url__startswith=url).mock(return_value=httpx.Response(200, json=payload)) + messages: Final = [{"role": "user", "content": "Hello"}] + + response: Final = litellm.completion( + model="azure/gpt-5.4-nano", + messages=messages, + api_base="https://example-resource.openai.azure.com", + api_key="azure-test-key", + api_version="2024-12-01-preview", + responses_api_version="preview", + ) + request: Final = route.calls[0].request + + assert route.call_count == 1 + assert request.url.params["api-version"] == "2024-12-01-preview" + assert json.loads(request.content) == {"model": "gpt-5.4-nano", "messages": messages} + assert response.choices[0].message.content == "Hello" diff --git a/tests/unit/llms/azure/response/test_azure_transformation.py b/tests/unit/llms/azure/response/test_azure_transformation.py index fea5304a689..9b4c618c78b 100644 --- a/tests/unit/llms/azure/response/test_azure_transformation.py +++ b/tests/unit/llms/azure/response/test_azure_transformation.py @@ -1,11 +1,15 @@ import json +from collections.abc import Callable from copy import deepcopy from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest +import respx import litellm +from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.azure.responses.o_series_transformation import ( @@ -282,16 +286,25 @@ class TestAzureResponsesAPIConfig: == "https://litellm8397336933.openai.azure.com/openai/v1/responses?api-version=latest" ) - # Test with date-based version - should use openai/responses - result_date = self.config.get_complete_url( + result_dated_preview = self.config.get_complete_url( api_base=base_url, - litellm_params={"api_version": "2025-01-01"}, + litellm_params={"api_version": "2025-04-01-preview"}, ) assert ( - result_date - == "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2025-01-01" + result_dated_preview + == "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2025-04-01-preview" ) + @pytest.mark.parametrize("api_version", ["2024-10-21", "2024-06-01", "2025-01-01"]) + def test_dated_ga_api_version_routes_responses_like_an_unset_one(self, api_version: str): + base_url: Final = "https://litellm8397336933.openai.azure.com" + + dated: Final = self.config.get_complete_url(api_base=base_url, litellm_params={"api_version": api_version}) + unset: Final = self.config.get_complete_url(api_base=base_url, litellm_params={}) + + assert dated == unset + assert httpx.URL(dated).params["api-version"] == AZURE_DEFAULT_RESPONSES_API_VERSION + def test_azure_get_complete_url_with_default_api_version(self): """Test Azure get_complete_url uses default API version when none is provided""" from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION @@ -320,6 +333,64 @@ class TestAzureResponsesAPIConfig: ) assert result_none_version == expected_url + @pytest.mark.parametrize( + ("api_version", "responses_api_version", "expected_path"), + [ + ("2024-10-21", "2025-04-01-preview", "/openai/responses?api-version=2025-04-01-preview"), + ("2025-04-01-preview", "v1", "/openai/v1/responses?api-version=v1"), + (None, "2025-04-01-preview", "/openai/responses?api-version=2025-04-01-preview"), + ], + ) + def test_responses_api_version_wins_over_api_version( + self, api_version: str | None, responses_api_version: str, expected_path: str + ): + base_url: Final = "https://litellm8397336933.openai.azure.com" + + url: Final = self.config.get_complete_url( + api_base=base_url, + litellm_params={"api_version": api_version, "responses_api_version": responses_api_version}, + ) + + assert url == f"{base_url}{expected_path}" + + @pytest.mark.parametrize("responses_api_version", [None, "v1"]) + def test_api_version_in_api_base_query_wins_over_a_dated_ga_api_version(self, responses_api_version: str | None): + api_base: Final = "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2025-04-01-preview" + + url: Final = self.config.get_complete_url( + api_base=api_base, + litellm_params={"api_version": "2024-10-21", "responses_api_version": responses_api_version}, + ) + + assert url == api_base + + def test_api_version_in_env_api_base_query_wins_over_a_dated_ga_api_version(self, monkeypatch: pytest.MonkeyPatch): + env_api_base: Final = ( + "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2025-04-01-preview" + ) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.setenv("AZURE_API_BASE", env_api_base) + + url: Final = self.config.get_complete_url(api_base=None, litellm_params={"api_version": "2024-10-21"}) + + assert url == env_api_base + + @pytest.mark.parametrize( + ("api_base", "expected_path"), + [ + ("https://apim.example.com/openai-gw/team", "/openai-gw/team/openai/v1/responses"), + ("https://apim.example.com/openai/team-a", "/openai/team-a/openai/v1/responses"), + ], + ) + @pytest.mark.parametrize("litellm_params", [{"api_version": "2024-10-21"}, {"responses_api_version": "v1"}]) + def test_v1_route_only_rewrites_the_responses_route_segment_of_a_gateway_api_base( + self, api_base: str, expected_path: str, litellm_params: dict[str, str] + ): + url: Final = httpx.URL(self.config.get_complete_url(api_base=api_base, litellm_params=litellm_params)) + + assert url.host == "apim.example.com" + assert url.path == expected_path + def test_azure_cancel_response_api_request(self): """Test Azure cancel response API request transformation""" from litellm.types.router import GenericLiteLLMParams @@ -684,6 +755,207 @@ def test_azure_responses_gpt6_astra_rejects_temperature_while_reasoning(local_mo ) +_COMPLETED_RESPONSE: Final = { + "id": "resp_dated_chat_deployment", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "gpt-5.4-nano", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": None, + "temperature": None, + "tool_choice": "auto", + "tools": [], + "top_p": None, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": None, + "truncation": None, + "user": None, +} + + +_AZURE_RESOURCE: Final = "https://example-resource.openai.azure.com" + + +def _unset_api_version_responses_url() -> httpx.URL: + return httpx.URL(AzureOpenAIResponsesAPIConfig().get_complete_url(api_base=_AZURE_RESOURCE, litellm_params={})) + + +def _azure_responses_create(api_version: str) -> str | None: + response: Final = litellm.responses( + model="azure/gpt-5.4-nano", + input="hi", + api_base=_AZURE_RESOURCE, + api_key="azure-test-key", + api_version=api_version, + ) + return response.output[0].content[0].text + + +def _azure_chat_bridged_to_responses(api_version: str) -> str | None: + response: Final = litellm.completion( + model="azure/responses/gpt-5.4-nano", + messages=[{"role": "user", "content": "hi"}], + api_base=_AZURE_RESOURCE, + api_key="azure-test-key", + api_version=api_version, + ) + return response.choices[0].message.content + + +@pytest.mark.parametrize("call_responses", [_azure_responses_create, _azure_chat_bridged_to_responses]) +def test_dated_ga_api_version_sends_responses_to_the_unset_api_version_url( + respx_mock: respx.MockRouter, call_responses: Callable[[str], str | None] +) -> None: + route: Final = respx_mock.post(host="example-resource.openai.azure.com").mock( + return_value=httpx.Response(200, json=_COMPLETED_RESPONSE) + ) + + text: Final = call_responses("2024-10-21") + + assert text == "Done." + assert route.call_count == 1 + assert route.calls[0].request.url == _unset_api_version_responses_url() + + +_LIFECYCLE_BODIES: Final = { + "get": _COMPLETED_RESPONSE, + "delete": {"id": "resp_dated", "object": "response", "deleted": True}, + "cancel": {**_COMPLETED_RESPONSE, "status": "cancelled"}, + "input_items": {"object": "list", "data": [], "first_id": None, "last_id": None, "has_more": False}, +} + + +def _call_lifecycle(operation: str) -> object: + kwargs: Final = { + "response_id": "resp_dated", + "custom_llm_provider": "azure", + "api_base": _AZURE_RESOURCE, + "api_key": "azure-test-key", + "api_version": "2024-10-21", + } + match operation: + case "get": + return litellm.get_responses(**kwargs) + case "delete": + return litellm.delete_responses(**kwargs) + case "cancel": + return litellm.cancel_responses(**kwargs) + case _: + return litellm.list_input_items(**kwargs) + + +@pytest.mark.parametrize( + ("operation", "path_suffix"), + [("get", ""), ("delete", ""), ("cancel", "/cancel"), ("input_items", "/input_items")], +) +def test_dated_ga_api_version_sends_response_lifecycle_calls_to_the_unset_api_version_route( + respx_mock: respx.MockRouter, operation: str, path_suffix: str +) -> None: + route: Final = respx_mock.route(host="example-resource.openai.azure.com").mock( + return_value=httpx.Response(200, json=_LIFECYCLE_BODIES[operation]) + ) + + _call_lifecycle(operation) + + unset_url: Final = _unset_api_version_responses_url() + sent_url: Final = route.calls[0].request.url + assert route.call_count == 1 + assert sent_url.path == f"{unset_url.path}/resp_dated{path_suffix}" + assert sent_url.params["api-version"] == unset_url.params["api-version"] + + +def test_chat_completion_keeps_the_dated_ga_api_version(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post( + f"{_AZURE_RESOURCE}/openai/deployments/gpt-5.4-nano/chat/completions", params={"api-version": "2024-10-21"} + ).mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-dated", + "object": "chat.completion", + "created": 1734366691, + "model": "gpt-5.4-nano", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Done."}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ) + ) + + response: Final = litellm.completion( + model="azure/gpt-5.4-nano", + messages=[{"role": "user", "content": "hi"}], + api_base=_AZURE_RESOURCE, + api_key="azure-test-key", + api_version="2024-10-21", + ) + + assert response.choices[0].message.content == "Done." + assert route.call_count == 1 + + +@pytest.mark.respx(assert_all_called=True) +def test_responses_api_version_routes_a_dated_chat_deployment_to_its_own_responses_version( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post( + "https://example-resource.openai.azure.com/openai/responses", params={"api-version": "2025-04-01-preview"} + ).mock(return_value=httpx.Response(200, json=_COMPLETED_RESPONSE)) + + response: Final = litellm.responses( + model="azure/gpt-5.4-nano", + input="hi", + api_base="https://example-resource.openai.azure.com", + api_key="azure-test-key", + api_version="2024-10-21", + responses_api_version="2025-04-01-preview", + ) + request_body: Final = json.loads(route.calls[0].request.content) + + assert route.call_count == 1 + assert response.output[0].content[0].text == "Done." + assert request_body["model"] == "gpt-5.4-nano" + assert "responses_api_version" not in request_body + assert "api_version" not in request_body + + +@pytest.mark.respx(assert_all_called=True) +def test_responses_api_version_reaches_the_url_when_chat_completion_bridges_to_responses( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post( + "https://example-resource.openai.azure.com/openai/v1/responses", params={"api-version": "v1"} + ).mock(return_value=httpx.Response(200, json=_COMPLETED_RESPONSE)) + + response: Final = litellm.completion( + model="azure/responses/gpt-5.4-nano", + messages=[{"role": "user", "content": "hi"}], + api_base="https://example-resource.openai.azure.com", + api_key="azure-test-key", + api_version="2024-10-21", + responses_api_version="v1", + ) + request_body: Final = json.loads(route.calls[0].request.content) + + assert route.call_count == 1 + assert response.choices[0].message.content == "Done." + assert "responses_api_version" not in request_body + + def test_azure_responses_sends_the_deployment_name_when_azure_ai_prefix_survives_provider_remap(): request = AzureOpenAIResponsesAPIConfig().transform_responses_api_request( model="azure_ai/gpt-5.4-nano", diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py index 29380c6fc53..86ad6c8d9d9 100644 --- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -186,3 +186,27 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_ "status_code": 500, } ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("set_base", ("env", "global"), ids=("AZURE_API_BASE", "litellm.api_base")) +async def test_count_uses_the_shared_api_base_chat_would_use(monkeypatch: pytest.MonkeyPatch, set_base: str) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("AZURE_AI_API_BASE", raising=False) + monkeypatch.delenv("AZURE_API_BASE", raising=False) + if set_base == "env": + monkeypatch.setenv("AZURE_API_BASE", API_BASE) + else: + monkeypatch.setattr(litellm, "api_base", API_BASE) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key"}}, + request_model="azure_ai/claude-sonnet-4-5", + ) + assert result is not None + assert result.total_tokens == 6 + assert route.called 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/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/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index 0609f122949..069438e25c0 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -16,6 +16,7 @@ import litellm from litellm import AmazonInvokeConfig from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock.chat.invoke_handler import ( AmazonOpenAICompatibleStreamDecoder, @@ -864,14 +865,16 @@ async def test_async_invoke_streaming_non_200_forwards_bedrock_response_headers( def _event_stream_frame(event_type: str, payload: bytes) -> bytes: + return _typed_event_stream_frame( + {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, payload + ) + + +def _typed_event_stream_frame(frame_headers: Mapping[str, str], payload: bytes) -> bytes: def header(name: str, value: str) -> bytes: return bytes([len(name)]) + name.encode() + bytes([7]) + struct.pack(">H", len(value)) + value.encode() - headers: Final = ( - header(":event-type", event_type) - + header(":content-type", "application/json") - + header(":message-type", "event") - ) + headers: Final = b"".join(header(name, value) for name, value in frame_headers.items()) prelude: Final = struct.pack(">II", 12 + len(headers) + len(payload) + 4, len(headers)) body: Final = prelude + struct.pack(">I", binascii.crc32(prelude)) + headers + payload return body + struct.pack(">I", binascii.crc32(body)) @@ -1367,6 +1370,87 @@ def test_iter_bytes_keeps_a_stream_whose_unknown_event_sits_beside_known_frames( assert "stop" in finish_reasons +_STREAM_FAILURE_BODIES: Final = ( + pytest.param( + _converse_event_frame("internalServerException", {"message": _UPSTREAM_REJECTION}), 500, id="exception-event" + ), + pytest.param( + _typed_event_stream_frame( + { + ":exception-type": "throttlingException", + ":content-type": "application/json", + ":message-type": "exception", + }, + json.dumps({"message": _UPSTREAM_REJECTION}).encode(), + ), + 429, + id="exception-message", + ), + pytest.param(_unknown_event_frame(), 502, id="no-known-event"), + pytest.param(b"", 502, id="no-events"), + pytest.param(_CONVERSE_CONTENT_FRAMES[0] + _truncated_frame(), 502, id="undecoded-tail"), +) + + +def _stream_headers(request_id: str | None) -> httpx.Headers: + request_id_header: Final = {} if request_id is None else {"x-amzn-RequestId": request_id} + return httpx.Headers( + {"content-type": "application/vnd.amazon.eventstream", "retry-after": "7", **request_id_header} + ) + + +def _assert_carries_only_the_request_id(error: BedrockError, status: int, request_id: str | None) -> None: + assert error.status_code == status + assert _get_provider_request_id(error) == request_id + if request_id is None: + assert error.headers is None + return + assert dict(error.headers or {}) == {"x-amzn-requestid": request_id} + assert dict(error.response.headers) == {"x-amzn-requestid": request_id} + + +@pytest.mark.parametrize("request_id", ("req-stream-1", None)) +@pytest.mark.parametrize(("body", "status"), _STREAM_FAILURE_BODIES) +def test_iter_bytes_stream_error_carries_only_the_provider_request_id_header( + body: bytes, status: int, request_id: str | None +) -> None: + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + + with pytest.raises(BedrockError) as exc_info: + list(decoder.iter_bytes(iter([body]), response_headers=_stream_headers(request_id))) + + _assert_carries_only_the_request_id(exc_info.value, status, request_id) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("body", "status"), _STREAM_FAILURE_BODIES) +async def test_aiter_bytes_stream_error_carries_only_the_provider_request_id_header(body: bytes, status: int) -> None: + async def _chunks() -> AsyncIterator[bytes]: + yield body + + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + + with pytest.raises(BedrockError) as exc_info: + _ = [chunk async for chunk in decoder.aiter_bytes(_chunks(), response_headers=_stream_headers("req-stream-2"))] + + _assert_carries_only_the_request_id(exc_info.value, status, "req-stream-2") + + +@pytest.mark.parametrize("request_id", ("req-\u00fc-1", "req-\udcfc-1"), ids=("latin-1", "surrogate-escaped")) +def test_iter_bytes_streams_and_fails_normally_with_a_non_ascii_provider_request_id(request_id: str) -> None: + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + headers: Final = {"content-type": "application/vnd.amazon.eventstream", "x-amzn-RequestId": request_id} + + chunks: Final = list(decoder.iter_bytes(iter(_CONVERSE_CONTENT_FRAMES), response_headers=headers)) + with pytest.raises(BedrockError) as exc_info: + list(decoder.iter_bytes(iter([_unknown_event_frame()]), response_headers=headers)) + + texts: Final = [chunk.choices[0].delta.content for chunk in chunks if isinstance(chunk, ModelResponseStream)] + assert "".join(text or "" for text in texts) == "hi" + assert exc_info.value.status_code == 502 + assert exc_info.value.headers is None + + def _assert_exception_event_surfaced_with_its_modeled_status(error: BaseException, event_type: str) -> None: assert not isinstance(error, litellm.BadGatewayError) assert getattr(error, "status_code", None) == _modeled_status(event_type) 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/test_web_identity_session_policy.py b/tests/unit/llms/bedrock/test_web_identity_session_policy.py index cbb69b4ceed..332e1258e9c 100644 --- a/tests/unit/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/unit/llms/bedrock/test_web_identity_session_policy.py @@ -176,6 +176,44 @@ class TestClaudePlatformActionsCovered: ) +class TestClaudePlatformOutboundIdentityCovered: + """Claude Platform on AWS calls ``sts:GetWebIdentityToken`` with the + caller's session before serving inference, and tags the token it + requests (``sts:TagGetWebIdentityToken``). Without either in the ceiling, + every claude_platform request via OIDC auth denies with "no session + policy allows the sts: action" even when the role's identity + policy grants it.""" + + @pytest.mark.parametrize("action", ["sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"]) + def test_outbound_identity_action_present(self, action: str): + assert action in _granted_actions(_captured_policy()), ( + f"{action} missing from session policy, so bedrock/claude_platform/* requests will 403 on OIDC auth" + ) + + @pytest.mark.parametrize("action", ["sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"]) + def test_outbound_identity_action_is_granted_only_via_claude_platform(self, action: str): + granting: Final = tuple(stmt for stmt in _captured_policy()["Statement"] if action in stmt["Action"]) + assert granting, f"{action} missing from session policy" + assert all( + stmt["Condition"].get("StringEquals") == {"aws:CalledViaLast": "aws-external-anthropic.amazonaws.com"} + and stmt["Condition"]["Bool"] == {"aws:SecureTransport": "true"} + for stmt in granting + ), f"{action} must only be usable over TLS when Claude Platform on AWS calls STS for the session" + + def test_outbound_identity_statement_is_scoped(self): + policy = _captured_policy() + stmt = _statement_by_sid(policy, "ClaudePlatformOutboundIdentityLiteLLM") + assert stmt["Effect"] == "Allow" + assert tuple(stmt["Action"]) == ("sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"), ( + "the outbound identity statement must grant only the two outbound identity actions" + ) + + def test_no_sts_wildcard(self): + assert "sts:*" not in _granted_actions(_captured_policy()), ( + "session policy must not grant sts:*, only the two outbound identity actions are needed" + ) + + class TestBedrockMantleActionsCovered: """LIT-3859: bedrock_mantle inference authorizes against the ``bedrock-mantle`` action namespace, so the session-policy ceiling @@ -306,6 +344,15 @@ class TestPolicyTransportConditions: "to keep parity with the bedrock statement" ) + def test_only_the_outbound_identity_statement_requires_a_claude_platform_caller(self): + directly_called: Final = tuple( + stmt for stmt in _captured_policy()["Statement"] if stmt["Sid"] != "ClaudePlatformOutboundIdentityLiteLLM" + ) + assert directly_called + assert all("StringEquals" not in stmt["Condition"] for stmt in directly_called), ( + "statements LiteLLM calls directly must not require aws:CalledViaLast, or OIDC auth 403s on those routes" + ) + _STS_SESSION_POLICY_PLAINTEXT_LIMIT: Final = 2048 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/dashscope/test_dashscope_video_transformation.py b/tests/unit/llms/dashscope/test_dashscope_video_transformation.py index 65eefcab1af..15b77aaee3d 100644 --- a/tests/unit/llms/dashscope/test_dashscope_video_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_video_transformation.py @@ -194,8 +194,9 @@ class TestDashScopeVideoInputReference: lambda tmp_path: ("first.png", PNG_BYTES, "image/png"), lambda tmp_path: ("first.png", io.BytesIO(PNG_BYTES)), lambda tmp_path: _write_png(tmp_path), + lambda tmp_path: ("first.png", ("inner.png", PNG_BYTES)), ], - ids=["filename-bytes-content-type", "filename-file", "path"], + ids=["filename-bytes-content-type", "filename-file", "path", "nested-tuple"], ) def test_file_tuple_and_path_references_are_encoded(self, make_reference, tmp_path): """Regression: tuples and paths are valid FileTypes but raised before reaching DashScope.""" diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index 650517c92ca..1e353a7a82a 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -461,6 +461,21 @@ def test_transform_messages_sanitizes_empty_content(): assert result[1]["content"] == "Hi" +def test_transform_messages_drops_the_prefill_prefix_flag(): + """Databricks rejects unknown message keys, so a continuation prefill reaches Claude without ``prefix``.""" + messages = [ + {"role": "user", "content": "How do transformers work?"}, + {"role": "assistant", "content": "Transformers use", "prefix": True}, + ] + + result = DatabricksConfig().transform_messages( + messages=messages, model="databricks-claude-sonnet-4-6", is_async=False + ) + + assert result[-1] == {"role": "assistant", "content": "Transformers use"} + assert messages[-1]["prefix"] is True + + def test_transform_request_preserves_unity_model_service_name(): config = DatabricksConfig() result = config.transform_request( 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/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 9a47d01ec46..e54023c6596 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -5,6 +5,7 @@ import httpx import pytest from pydantic import JsonValue, TypeAdapter +import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.common_utils import GeminiModelInfo, GoogleAIStudioTokenCounter from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler @@ -100,6 +101,58 @@ async def test_google_ai_studio_count_tokens_rejects_image_system_without_reques assert peer.request_body is None +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_rejects_an_image_url_it_refuses_to_fetch_without_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + peer: Final = _GeminiCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}], + } + ], + contents=None, + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "http://127.0.0.1:9/image.png" in (result.error_message or "") + assert peer.request_body is None + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_answers_with_the_status_google_rejected_it_with() -> None: + rejection: Final = '{"error": {"code": 400, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}' + client: Final = AsyncHTTPHandler( + transport=httpx.MockTransport(lambda request: httpx.Response(400, text=rejection, request=request)) + ) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=None, + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "Invalid function name" in (result.error_message or "") + + @pytest.mark.asyncio @pytest.mark.parametrize( ("system_instruction", "tools", "expected_body"), @@ -129,9 +182,7 @@ async def test_gemini_count_tokens_sends_optional_request_fields( expected_body: dict[str, JsonValue], ) -> None: peer: Final = _GeminiCountTokensPeer() - token_counter: Final = GeminiCountTokensHandler( - client=AsyncHTTPHandler(transport=httpx.MockTransport(peer)) - ) + token_counter: Final = GeminiCountTokensHandler(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer))) result: Final = await token_counter.acount_tokens( contents=[{"role": "user", "parts": [{"text": "hello"}]}], @@ -177,18 +228,10 @@ class TestGeminiModelInfo: # Test edge cases where model names end with characters from "models/" # These would be incorrectly processed if using strip("models/") instead of replace("models/", "") models = [ - { - "name": "models/gemini-1.5-pro" - }, # ends with 'o' - would become "gemini-1.5-pr" with strip() - { - "name": "models/test-model" - }, # ends with 'l' - would become "gemini/test-mode" with strip() - { - "name": "models/custom-models" - }, # ends with 's' - would become "gemini/custom-model" with strip() - { - "name": "models/demo" - }, # ends with 'o' - would become "gemini/dem" with strip() + {"name": "models/gemini-1.5-pro"}, # ends with 'o' - would become "gemini-1.5-pr" with strip() + {"name": "models/test-model"}, # ends with 'l' - would become "gemini/test-mode" with strip() + {"name": "models/custom-models"}, # ends with 's' - would become "gemini/custom-model" with strip() + {"name": "models/demo"}, # ends with 'o' - would become "gemini/dem" with strip() ] result = gemini_model_info.process_model_name(models) @@ -239,16 +282,10 @@ class TestGoogleAIStudioTokenCounter: token_counter = GoogleAIStudioTokenCounter() # Test with gemini provider - should return True - assert ( - token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) - is True - ) + assert token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) is True # Test with other providers - should return False - assert ( - token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) - is False - ) + assert token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) is False assert token_counter.should_use_token_counting_api("anthropic") is False assert token_counter.should_use_token_counting_api("vertex_ai") is False @@ -320,9 +357,7 @@ class TestGoogleAIStudioTokenCounter: "functionResponse": { "id": "read_many_files-1757526647518-730a691aac11c", # This should be removed "name": "read_many_files", - "response": { - "output": "No files matching the criteria were found or all were skipped." - }, + "response": {"output": "No files matching the criteria were found or all were skipped."}, } } ], @@ -331,9 +366,7 @@ class TestGoogleAIStudioTokenCounter: ] # Clean the contents - cleaned_contents = token_counter._clean_contents_for_gemini_api( - contents_with_id - ) + cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_with_id) # Verify the 'id' field was removed function_response = cleaned_contents[1]["parts"][0]["functionResponse"] @@ -342,8 +375,7 @@ class TestGoogleAIStudioTokenCounter: assert "response" in function_response assert function_response["name"] == "read_many_files" assert ( - function_response["response"]["output"] - == "No files matching the criteria were found or all were skipped." + function_response["response"]["output"] == "No files matching the criteria were found or all were skipped." ) def test_clean_contents_for_gemini_api_preserves_other_fields(self): @@ -359,9 +391,7 @@ class TestGoogleAIStudioTokenCounter: ] # Clean the contents - cleaned_contents = token_counter._clean_contents_for_gemini_api( - contents_without_function_response - ) + cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_without_function_response) # Verify the contents are unchanged assert cleaned_contents == contents_without_function_response 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/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/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py index 976ecff01ea..befdc5ee9fd 100644 --- a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -388,3 +388,44 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_ "status_code": 500, } ) + + +@pytest.mark.asyncio +async def test_count_uses_the_global_openai_key_chat_would_use(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "openai_key", "sk-from-litellm-settings") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 9})) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.total_tokens == 9 + assert route.calls.last.request.headers["authorization"] == "Bearer sk-from-litellm-settings" + + +@pytest.mark.asyncio +async def test_count_sends_the_global_key_only_to_the_global_api_base(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "openai_key", "sk-private-gateway") + monkeypatch.setattr(litellm, "api_base", "https://gateway.internal.example/v1") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with respx.mock: + route: Final = respx.post("https://gateway.internal.example/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 4}) + ) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.total_tokens == 4 + assert route.calls.last.request.headers["authorization"] == "Bearer sk-private-gateway" diff --git a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py index 0da28e7207e..f85e4db447b 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py @@ -1434,11 +1434,11 @@ class TestAzureResponsesAPIConfig: # Test with date-based version - should use openai/responses result_date = self.config.get_complete_url( api_base=base_url, - litellm_params={"api_version": "2025-01-01"}, + litellm_params={"api_version": "2025-04-01-preview"}, ) assert ( result_date - == "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2025-01-01" + == "https://litellm8397336933.openai.azure.com/openai/responses?api-version=2025-04-01-preview" ) def test_azure_transform_then_normalize_strips_custom_tool_call_namespace(self): diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index e37fd4b6ae5..d7ca9319f3a 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -10,7 +10,7 @@ from unittest.mock import Mock, patch import httpx import pytest import respx -from openai import AsyncOpenAI, OpenAI +from openai import AsyncOpenAI, AzureOpenAI, OpenAI from openai._types import Response as SDKResponse from pydantic import BaseModel, ConfigDict, JsonValue @@ -221,6 +221,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( diff --git a/tests/unit/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py index 3c7dff34296..831df30c98a 100644 --- a/tests/unit/llms/openai/test_openai_common_utils.py +++ b/tests/unit/llms/openai/test_openai_common_utils.py @@ -442,15 +442,26 @@ def _sdk_api_client( ) if api == "image_variations": variation_factory: Final = OpenAIImageVariationsHandler() - params: Final = { - "api_key": "transport-only", - "base_url": kwargs["api_base"], - "timeout": timeout, - } return ( - variation_factory.get_async_client(client=client, init_client_params=params) + variation_factory.get_async_client( + client=client, + api_key="transport-only", + api_base="https://sdk-default.example/v1", + timeout=timeout, + max_retries=0, + organization=None, + litellm_params={}, + ) if is_async - else variation_factory.get_sync_client(client=client, init_client_params=params) + else variation_factory.get_sync_client( + client=client, + api_key="transport-only", + api_base="https://sdk-default.example/v1", + timeout=timeout, + max_retries=0, + organization=None, + litellm_params={}, + ) ) if api == "azure_gateway": return BaseAzureLLM()._init_azure_client_for_cloudflare_ai_gateway( diff --git a/tests/unit/llms/openai/test_openai_workload_identity.py b/tests/unit/llms/openai/test_openai_workload_identity.py index d0cf52d0e28..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 @@ -648,3 +988,1017 @@ class TestClientsideBaseOverride: ) is None ) + + +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 + async def test_count_authenticates_with_the_workload_identity_bearer( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + mock_token_exchange() + count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 12}) + ) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.error is False + assert result.total_tokens == 12 + assert count_route.calls.last.request.headers["authorization"] == "Bearer exchanged-bearer-token" + + @pytest.mark.asyncio + @respx.mock + async def test_failed_mint_becomes_a_count_error_response( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + respx.post(TOKEN_EXCHANGE_URL).mock(return_value=httpx.Response(400, json={"error": "invalid_grant"})) + count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens") + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.error is True + assert result.status_code == 401 + assert result.total_tokens == 0 + assert not count_route.called 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/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/soniox/audio_transcription/test_soniox_audio_transcription_handler.py b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py index d2d7d2247f1..9c70ebe7ece 100644 --- a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py +++ b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py @@ -339,6 +339,65 @@ class TestPollLimitsClamping: a worker on tight poll loops. """ + def test_should_raise_400_when_form_bool_param_is_invalid(self): + handler = SonioxAudioTranscriptionHandler() + with pytest.raises(SonioxException) as exc_info: + handler._prepare( + audio_file=None, + optional_params={ + "enable_speaker_diarization": "maybe", + "audio_url": "https://example.com/a.wav", + }, + litellm_params={}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert exc_info.value.status_code == 400 + assert "enable_speaker_diarization" in str(exc_info.value) + + def test_should_raise_400_when_poll_interval_is_not_a_number(self): + with pytest.raises(SonioxException) as exc_info: + SonioxAudioTranscriptionHandler()._prepare( + audio_file=None, + optional_params={"soniox_polling_interval": "abc", "audio_url": "https://example.com/a.wav"}, + litellm_params={}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert exc_info.value.status_code == 400 + assert "soniox_polling_interval" in str(exc_info.value) + + def test_should_use_default_poll_interval_for_bad_value_when_drop_params_is_set(self): + from litellm.llms.soniox.common_utils import SONIOX_DEFAULT_POLL_INTERVAL + + _, _, _, opts = SonioxAudioTranscriptionHandler()._prepare( + audio_file=None, + optional_params={"soniox_polling_interval": "abc", "audio_url": "https://example.com/a.wav"}, + litellm_params={"drop_params": True}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert opts["poll_interval"] == SONIOX_DEFAULT_POLL_INTERVAL + + @pytest.mark.parametrize("raw, expected", [(10.0, 10), (10, 10), ("10", 10)]) + def test_should_keep_numeric_max_polling_attempts(self, raw: object, expected: int): + _, _, _, opts = SonioxAudioTranscriptionHandler()._prepare( + audio_file=None, + optional_params={"soniox_max_polling_attempts": raw, "audio_url": "https://example.com/a.wav"}, + litellm_params={}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert opts["max_attempts"] == expected + def test_should_clamp_poll_interval_to_minimum(self): from litellm.llms.soniox.common_utils import SONIOX_MIN_POLL_INTERVAL @@ -707,6 +766,112 @@ class TestCleanupNormalization: assert "https://api.soniox.com/v1/transcriptions/tx_1" in deletes + def test_should_skip_non_string_cleanup_entries(self): + _, _, _, opts = SonioxAudioTranscriptionHandler()._prepare( + audio_file=None, + optional_params={"soniox_cleanup": ["transcription", 1, None], "audio_url": "https://example.com/a.wav"}, + litellm_params={}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert opts["cleanup"] == ("transcription",) + + @pytest.mark.parametrize( + "language_hints, expected", + [("en,es", ["en", "es"]), (["es"], ["es"]), (None, ["de"])], + ) + def test_should_use_language_only_without_language_hints(self, language_hints, expected, monkeypatch): + import litellm + + monkeypatch.setattr("time.sleep", lambda *_: None) + client = _MockSyncClient( + { + "POST https://api.soniox.com/v1/transcriptions": [_make_response({"id": "tx_1"})], + "GET https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({"status": "completed"})], + "GET https://api.soniox.com/v1/transcriptions/tx_1/transcript": [ + _make_response({"text": "hallo", "tokens": []}) + ], + "DELETE https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({})], + } + ) + litellm.transcription( + model="soniox/stt-async-v4", + file=("a.wav", b"RIFF", "audio/wav"), + language="de", + api_key="sk-test", + client=client, + audio_url="https://example.com/a.wav", + **({} if language_hints is None else {"language_hints": language_hints}), + ) + body = next(c for c in client.calls if c["method"] == "POST")["json"] + assert list(body["language_hints"]) == expected + assert "language" not in body + + +class TestProxyFormFields: + def test_should_send_soniox_types_for_fields_the_proxy_parsed_from_an_sdk_form(self, monkeypatch): + import litellm + from starlette.requests import Request + + from litellm.proxy.common_utils.http_parsing_utils import get_form_data + + monkeypatch.setattr("time.sleep", lambda *_: None) + multipart = httpx.Request( + "POST", + "http://proxy/v1/audio/transcriptions", + data={ + "model": "soniox/stt-async-v4", + "audio_url": "https://example.com/a.wav", + "context[terms][]": ["Celebrex"], + "context[general][][key]": "domain", + "context[general][][value]": "Healthcare", + "translation[type]": "one_way", + "translation[target_language]": "de", + "enable_speaker_diarization": "true", + "language_hints[]": ["en", "es"], + }, + files={"file": ("a.wav", b"RIFF", "audio/wav")}, + ) + multipart_body = multipart.read() + + async def receive(): + return {"type": "http.request", "body": multipart_body, "more_body": False} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/audio/transcriptions", + "headers": [(b"content-type", multipart.headers["content-type"].encode())], + }, + receive, + ) + form = asyncio.run(get_form_data(request)) + client = _MockSyncClient( + { + "POST https://api.soniox.com/v1/transcriptions": [_make_response({"id": "tx_1"})], + "GET https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({"status": "completed"})], + "GET https://api.soniox.com/v1/transcriptions/tx_1/transcript": [ + _make_response({"text": "hello", "tokens": []}) + ], + "DELETE https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({})], + } + ) + litellm.transcription( + file=("a.wav", b"RIFF", "audio/wav"), + api_key="sk-test", + client=client, + **{key: value for key, value in form.items() if key != "file"}, + ) + body = json.loads(json.dumps(next(c for c in client.calls if c["method"] == "POST")["json"])) + assert body["context"] == {"terms": ["Celebrex"], "general": [{"key": "domain", "value": "Healthcare"}]} + assert body["translation"] == {"type": "one_way", "target_language": "de"} + assert body["enable_speaker_diarization"] is True + assert body["language_hints"] == ["en", "es"] + + class TestErrorResponses: def test_should_raise_on_4xx_during_create_with_json_error(self, monkeypatch): monkeypatch.setattr("time.sleep", lambda *_: None) @@ -780,6 +945,51 @@ class TestPassthroughBodyBuilding: post_call = next(c for c in client.calls if c["method"] == "POST") assert "context" not in post_call["json"] + def test_should_send_form_string_params_to_soniox_as_json_types(self, monkeypatch): + """Simulates the proxy path, where every multipart form field is a string.""" + monkeypatch.setattr("time.sleep", lambda *_: None) + responses = { + "POST https://api.soniox.com/v1/transcriptions": [ + _make_response({"id": "tx_1"}), + ], + "GET https://api.soniox.com/v1/transcriptions/tx_1": [ + _make_response({"status": "completed"}), + ], + "GET https://api.soniox.com/v1/transcriptions/tx_1/transcript": [ + _make_response({"text": "ok", "tokens": []}), + ], + "DELETE https://api.soniox.com/v1/transcriptions/tx_1": [ + _make_response({}), + ], + } + client = _MockSyncClient(responses) + form_params = { + "audio_url": "https://example.com/a.wav", + "context": '{"terms": ["Celebrex"], "general": [{"key": "domain", "value": "Healthcare"}]}', + "translation": '{"type": "one_way", "target_language": "es"}', + "language_hints": "en,es", + "language_hints_strict": "true", + "enable_speaker_diarization": "true", + "enable_language_identification": "false", + "soniox_polling_interval": "0.5", + } + SonioxAudioTranscriptionHandler().audio_transcriptions( + audio_file=None, + optional_params=form_params, + litellm_params={}, + atranscription=False, + **_common_call_kwargs(client), + ) + body = next(c for c in client.calls if c["method"] == "POST")["json"] + assert body["context"] == {"terms": ["Celebrex"], "general": [{"key": "domain", "value": "Healthcare"}]} + assert body["translation"] == {"type": "one_way", "target_language": "es"} + assert list(body["language_hints"]) == ["en", "es"] + assert body["language_hints_strict"] is True + assert body["enable_speaker_diarization"] is True + assert body["enable_language_identification"] is False + assert "soniox_polling_interval" not in body + assert form_params["context"].startswith("{"), "caller's params must not be mutated" + class TestSecretRedaction: """Secret-bearing fields must be redacted before reaching logging callbacks. diff --git a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py index c99b28adc49..e3e1c9ab493 100644 --- a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py +++ b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py @@ -9,8 +9,16 @@ import pytest from litellm.llms.soniox.audio_transcription.transformation import ( SonioxAudioTranscriptionConfig, + SonioxInvalidFormParam, + decode_soniox_form_params, + raise_soniox_form_error, ) from litellm.llms.soniox.common_utils import SonioxException +from litellm.llms.soniox.types import ( + SonioxContext, + SonioxContextGeneralEntry, + SonioxTranslationTerm, +) from litellm.types.utils import TranscriptionResponse @@ -40,9 +48,9 @@ class TestMapOpenAIParams: model="stt-async-v4", drop_params=False, ) - assert result["language_hints"] == ["en"] + assert list(result["language_hints"]) == ["en"] - def test_should_prepend_language_to_existing_hints(self): + def test_should_keep_explicit_hints_over_language(self): cfg = SonioxAudioTranscriptionConfig() result = cfg.map_openai_params( non_default_params={"language": "en"}, @@ -50,7 +58,7 @@ class TestMapOpenAIParams: model="stt-async-v4", drop_params=False, ) - assert result["language_hints"] == ["en", "fr"] + assert list(result["language_hints"]) == ["fr"] def test_should_not_duplicate_language_already_in_hints(self): cfg = SonioxAudioTranscriptionConfig() @@ -60,7 +68,7 @@ class TestMapOpenAIParams: model="stt-async-v4", drop_params=False, ) - assert result["language_hints"] == ["en", "fr"] + assert list(result["language_hints"]) == ["en", "fr"] def test_should_passthrough_soniox_native_kwargs(self): cfg = SonioxAudioTranscriptionConfig() @@ -243,6 +251,21 @@ class TestTransformAudioTranscriptionResponse: assert "Speaker 1:" in resp.text assert "Speaker 2:" in resp.text + def test_should_keep_source_transcript_for_translation_without_tags(self): + cfg = SonioxAudioTranscriptionConfig() + payload = { + "transcript": { + "text": "Hello there.", + "tokens": [ + {"text": "Hello", "translation_status": "original"}, + {"text": " there.", "translation_status": "original"}, + {"text": "Hallo", "translation_status": "translation"}, + {"text": " zusammen.", "translation_status": "translation"}, + ], + } + } + assert cfg.build_response_from_payload(payload).text == "Hello there." + def test_should_set_language_when_all_tokens_share_one(self): cfg = SonioxAudioTranscriptionConfig() payload = { @@ -678,3 +701,175 @@ class TestGetErrorClass: err = cfg.get_error_class(error_message="boom", status_code=500, headers={}) assert isinstance(err, SonioxException) assert err.status_code == 500 + + +class TestDecodeSonioxFormParams: + """Proxy multipart form fields arrive as strings; Soniox needs real JSON types.""" + + def test_should_decode_json_string_context_with_all_four_sections(self): + context = { + "general": [{"key": "domain", "value": "Healthcare"}], + "text": "Follow-up visit for a patient on blood thinners.", + "terms": ["Celebrex", "Zyrtec"], + "translation_terms": [{"source": "Mr. Smith", "target": "Sr. Smith"}], + } + result = decode_soniox_form_params({"context": json.dumps(context)}) + assert result["context"] == context + + def test_should_keep_plain_text_context_as_string(self): + result = decode_soniox_form_params({"context": "medical conversation"}) + assert result["context"] == "medical conversation" + + def test_should_leave_sdk_typed_values_untouched(self): + context = SonioxContext( + general=[SonioxContextGeneralEntry(key="domain", value="Healthcare")], + text="Follow-up visit notes.", + terms=["Celebrex"], + translation_terms=[SonioxTranslationTerm(source="Mr. Smith", target="Sr. Smith")], + ) + translation = {"type": "one_way", "target_language": "es"} + result = decode_soniox_form_params( + { + "context": context, + "translation": translation, + "enable_speaker_diarization": True, + "language_hints": ["en", "es"], + } + ) + assert result["context"] is context + assert result["translation"] is translation + assert result["enable_speaker_diarization"] is True + assert result["language_hints"] == ["en", "es"] + + @pytest.mark.parametrize("raw", ["[inaudible] consultation", "{unbalanced", '{"terms": ["Celebrex"']) + def test_should_keep_bracket_prefixed_free_form_context_as_string(self, raw: str): + assert decode_soniox_form_params({"context": raw})["context"] == raw + + def test_should_decode_json_string_translation(self): + result = decode_soniox_form_params({"translation": '{"type": "one_way", "target_language": "es"}'}) + assert result["translation"] == {"type": "one_way", "target_language": "es"} + + @pytest.mark.parametrize( + "raw, expected", + [("true", True), ("True", True), ("1", True), ("false", False), ("FALSE", False), ("0", False)], + ) + def test_should_decode_boolean_strings(self, raw: str, expected: bool): + result = decode_soniox_form_params( + { + "enable_speaker_diarization": raw, + "enable_language_identification": raw, + "language_hints_strict": raw, + } + ) + assert result == { + "enable_speaker_diarization": expected, + "enable_language_identification": expected, + "language_hints_strict": expected, + } + + @pytest.mark.parametrize("raw, expected", [("on", True), ("yes", True), ("off", False), ("no", False)]) + def test_should_decode_html_style_boolean_strings(self, raw: str, expected: bool): + assert decode_soniox_form_params({"enable_speaker_diarization": raw})["enable_speaker_diarization"] is expected + + def test_should_return_error_value_on_non_boolean_string(self): + result = decode_soniox_form_params({"context": '{"general": []}', "enable_speaker_diarization": "maybe"}) + assert result == SonioxInvalidFormParam(key="enable_speaker_diarization", value="maybe", expected="a boolean") + + def test_should_drop_non_boolean_string_when_drop_params_is_set(self): + result = decode_soniox_form_params( + {"enable_speaker_diarization": "maybe", "language_hints_strict": "true"}, drop_params=True + ) + assert result == {"language_hints_strict": True} + + def test_should_return_error_value_on_non_numeric_polling_interval(self): + result = decode_soniox_form_params({"soniox_polling_interval": "abc"}) + assert result == SonioxInvalidFormParam(key="soniox_polling_interval", value="abc", expected="a number") + + def test_should_drop_non_numeric_polling_interval_when_drop_params_is_set(self): + assert decode_soniox_form_params({"soniox_polling_interval": "abc"}, drop_params=True) == {} + + def test_should_map_invalid_form_value_error_to_400(self): + with pytest.raises(SonioxException) as exc_info: + raise_soniox_form_error( + SonioxInvalidFormParam(key="enable_speaker_diarization", value="maybe", expected="a boolean") + ) + assert exc_info.value.status_code == 400 + assert "enable_speaker_diarization" in str(exc_info.value) + assert "maybe" in str(exc_info.value) + + @pytest.mark.parametrize( + "raw", ['["Celebrex", "Zyrtec"]', '{"note": "follow-up visit"}', '{"terms": "Celebrex"}', '"quoted"'] + ) + def test_should_keep_json_context_soniox_cannot_take_as_an_object_as_string(self, raw: str): + assert decode_soniox_form_params({"context": raw})["context"] == raw + + def test_should_nest_openai_sdk_bracketed_form_fields(self): + result = decode_soniox_form_params( + { + "context[terms]": ["Celebrex", "Zyrtec"], + "context[text]": "Follow-up visit.", + "translation[type]": "one_way", + "translation[target_language]": "de", + "client_reference_id": "ref-1", + } + ) + assert result == { + "context": {"terms": ["Celebrex", "Zyrtec"], "text": "Follow-up visit."}, + "translation": {"type": "one_way", "target_language": "de"}, + "client_reference_id": "ref-1", + } + + def test_should_prefer_explicit_field_over_bracketed_fields(self): + result = decode_soniox_form_params({"context": "plain text", "context[terms]": ["Celebrex"]}) + assert result == {"context": "plain text"} + + @pytest.mark.parametrize( + "raw, expected", + [('["en", "es"]', ["en", "es"]), ("en,es", ["en", "es"]), ("en, es ,", ["en", "es"]), ("en", ["en"])], + ) + def test_should_decode_language_hints_strings(self, raw: str, expected: list): + assert list(decode_soniox_form_params({"language_hints": raw})["language_hints"]) == expected + + @pytest.mark.parametrize("raw", ['{"type": "one_way"', "de", '["one_way"]']) + def test_should_return_error_value_on_translation_that_is_not_a_json_object(self, raw: str): + assert decode_soniox_form_params({"translation": raw}) == SonioxInvalidFormParam( + key="translation", value=raw, expected="a JSON object" + ) + + def test_should_return_error_value_on_malformed_json_language_hints(self): + assert decode_soniox_form_params({"language_hints": "[en, es]"}) == SonioxInvalidFormParam( + key="language_hints", value="[en, es]", expected="a JSON list" + ) + + def test_should_drop_malformed_json_params_when_drop_params_is_set(self): + result = decode_soniox_form_params( + {"translation": '{"type": "one_way"', "language_hints": "[en, es]", "language_hints_strict": "true"}, + drop_params=True, + ) + assert result == {"language_hints_strict": True} + + def test_should_rebuild_openai_sdk_list_of_object_context_fields(self): + result = decode_soniox_form_params( + { + "context[terms]": ["Celebrex"], + "context[general][][key]": "domain", + "context[general][][value]": "Healthcare", + "context[translation_terms][][source]": "Mr. Smith", + "context[translation_terms][][target]": "Sr. Smith", + } + ) + assert result == { + "context": { + "terms": ["Celebrex"], + "general": ({"key": "domain", "value": "Healthcare"},), + "translation_terms": ({"source": "Mr. Smith", "target": "Sr. Smith"},), + } + } + + def test_should_leave_bracketed_context_soniox_cannot_take_unnested(self): + params = {"context[terms]": ["Celebrex"], "context[note]": "follow-up"} + assert decode_soniox_form_params(params) == params + + def test_should_not_json_decode_unrelated_string_params(self): + result = decode_soniox_form_params({"client_reference_id": '["not", "json-decoded"]'}) + assert result["client_reference_id"] == '["not", "json-decoded"]' diff --git a/tests/unit/llms/test_predibase_transformation.py b/tests/unit/llms/test_predibase_transformation.py index d6425488be3..20d547c3b63 100644 --- a/tests/unit/llms/test_predibase_transformation.py +++ b/tests/unit/llms/test_predibase_transformation.py @@ -39,14 +39,9 @@ def test_predibase_transform_request_non_stream(): assert "stream" not in request_data["parameters"] -def test_predibase_transform_request_custom_prompt(monkeypatch): +def test_predibase_transform_request_custom_prompt(): config = PredibaseConfig() - monkeypatch.setattr( - "litellm.llms.predibase.chat.transformation.custom_prompt", - lambda **kwargs: "custom-prompt", - ) - request_data = config.transform_request( model="predibase-model", messages=[{"role": "user", "content": "hello"}], @@ -54,16 +49,16 @@ def test_predibase_transform_request_custom_prompt(monkeypatch): litellm_params={ "custom_prompt_dict": { "predibase-model": { - "roles": {}, - "initial_prompt_value": "", - "final_prompt_value": "", + "roles": {"user": {"pre_message": "", "post_message": ""}}, + "initial_prompt_value": "", + "final_prompt_value": "", } } }, headers={}, ) - assert request_data["inputs"] == "custom-prompt" + assert request_data["inputs"] == "hello" def test_predibase_get_complete_url_stream_and_non_stream(): 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 444920e1ecc..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,15 +1,31 @@ -from typing import Final +import asyncio +import base64 +import logging +import time +import uuid +from collections.abc import Generator +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Thread +from typing import Final, Literal +from unittest.mock import patch import pytest -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter +import litellm from litellm.llms.anthropic.pass_through.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, 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 @@ -19,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": [ { @@ -193,3 +210,461 @@ def test_invalid_tools_return_an_error_value(tools: JsonValue) -> None: assert isinstance(context, VertexGeminiCountContextError) assert "tools" in context.error_message + + +_ANTHROPIC_CONVERSATION: Final = ( + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize both"}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0="}}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + ], + }, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "x"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}]}, +) + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_anthropic_documents_images_and_tool_turns_become_gemini_parts( + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> None: + contents: Final = await gemini_count_contents( + _ANTHROPIC_CONVERSATION, "anthropic", "gemini-2.5-flash", custom_llm_provider, {} + ) + + assert contents == [ + { + "role": "user", + "parts": [ + {"text": "Summarize both"}, + {"inline_data": {"data": "JVBERi0=", "mime_type": "application/pdf"}}, + {"inline_data": {"data": "aW1hZ2U=", "mime_type": "image/png"}}, + ], + }, + {"role": "model", "parts": [{"function_call": {"name": "lookup", "args": {"query": "x"}}}]}, + {"role": "user", "parts": [{"function_response": {"name": "lookup", "response": {"content": "42"}}}]}, + ] + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_openai_messages_keep_their_chat_conversion(custom_llm_provider: Literal["vertex_ai", "gemini"]) -> None: + contents: Final = await gemini_count_contents( + ({"role": "user", "content": "hello"}, {"role": "assistant", "content": "hi"}), + "openai", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ) + + assert contents == [ + {"role": "user", "parts": [{"text": "hello"}]}, + {"role": "model", "parts": [{"text": "hi"}]}, + ] + + +_PNG_BYTES: Final = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" +) +_DOWNLOAD_SECONDS: Final = 0.5 + + +class _SlowPngHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + time.sleep(_DOWNLOAD_SECONDS) + self.send_response(200) + self.send_header("Content-Type", "image/png") + self.send_header("Content-Length", str(len(_PNG_BYTES))) + self.end_headers() + self.wfile.write(_PNG_BYTES) + + def log_message(self, format: str, *args: object) -> None: + return None + + +@contextmanager +def _slow_png_server() -> Generator[str, None, None]: + with ThreadingHTTPServer(("127.0.0.1", 0), _SlowPngHandler) as server: + thread: Final = Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + + +async def _event_loop_wake_delay(nap_seconds: float) -> float: + started: Final = time.monotonic() + await asyncio.sleep(nap_seconds) + return time.monotonic() - started + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_remote_image_downloads_without_blocking_the_event_loop( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + with _slow_png_server() as base_url: + image_url: Final = f"{base_url}/{uuid.uuid4().hex}.png" + wake_delay, contents = await asyncio.gather( + _event_loop_wake_delay(0.05), + gemini_count_contents( + ( + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe it"}, + {"type": "image", "source": {"type": "url", "url": image_url}}, + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ), + ) + + assert contents == [ + { + "role": "user", + "parts": [ + {"text": "Describe it"}, + {"inline_data": {"data": base64.b64encode(_PNG_BYTES).decode(), "mime_type": "image/png"}}, + ], + } + ] + assert wake_delay < _DOWNLOAD_SECONDS / 2 + + +class _MissingImageHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + self.send_response(404) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: + return None + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_image_that_cannot_be_fetched_returns_an_error_value( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + with ThreadingHTTPServer(("127.0.0.1", 0), _MissingImageHandler) as server: + Thread(target=server.serve_forever, daemon=True).start() + try: + contents = await gemini_count_contents( + ( + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "url", + "url": f"http://127.0.0.1:{server.server_address[1]}/{uuid.uuid4().hex}.png", + }, + } + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ) + finally: + server.shutdown() + + assert isinstance(contents, VertexGeminiCountContextError) + + +@pytest.mark.parametrize( + "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system") +) +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."}]}, + {"role": "user", "content": "latest question"}, + ), + 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) + + +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 = await inference_request_view( + ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + ), + 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"] + + +async def test_openai_view_is_the_request_unchanged() -> None: + messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},) + + 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." + + +async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None: + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value="image/png", + ) as lookup: + contents: Final = await gemini_count_contents( + ( + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "url", "url": "gs://bucket/image-without-extension"}} + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + "vertex_ai", + {"vertex_project": "deployment-project"}, + ) + + assert lookup.call_args.kwargs.get("vertex_project") == "deployment-project" + assert not isinstance(contents, VertexGeminiCountContextError) + assert contents[0]["parts"][0] == { + "file_data": {"mime_type": "image/png", "file_uri": "gs://bucket/image-without-extension"} + } + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +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." + messages: Final = ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + {"role": "user", "content": "Thanks."}, + ) + + 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/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 40ede74b98b..5c0e9f1a305 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1,6 +1,7 @@ import asyncio, importlib, os, uuid import json import re +from collections.abc import Awaitable, Callable, Mapping from copy import deepcopy from typing import Final, List, Literal, cast, get_args from unittest.mock import MagicMock, Mock, patch @@ -25,6 +26,7 @@ from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata from litellm.types.utils import ChoiceLogprobs, Usage from litellm.utils import _invalidate_model_cost_lowercase_map, CustomStreamWrapper from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.vertex_ai.gemini.transformation import ( gemini_convert_messages_with_history, ) @@ -38,62 +40,62 @@ from litellm.llms.vertex_ai.context_caching.transformation import ( import tempfile GEMINI_3_IMAGE_SIZE_MAPPINGS: Final[tuple[tuple[str, str, str], ...]] = ( - ('512x512', '1:1', '512'), - ('1024x1024', '1:1', '1K'), - ('2048x2048', '1:1', '2K'), - ('4096x4096', '1:1', '4K'), - ('256x1024', '1:4', '512'), - ('512x2048', '1:4', '1K'), - ('1024x4096', '1:4', '2K'), - ('2048x8192', '1:4', '4K'), - ('192x1536', '1:8', '512'), - ('384x3072', '1:8', '1K'), - ('768x6144', '1:8', '2K'), - ('1536x12288', '1:8', '4K'), - ('424x632', '2:3', '512'), - ('848x1264', '2:3', '1K'), - ('1696x2528', '2:3', '2K'), - ('3392x5056', '2:3', '4K'), - ('632x424', '3:2', '512'), - ('1264x848', '3:2', '1K'), - ('2528x1696', '3:2', '2K'), - ('5056x3392', '3:2', '4K'), - ('448x600', '3:4', '512'), - ('896x1200', '3:4', '1K'), - ('1792x2400', '3:4', '2K'), - ('3584x4800', '3:4', '4K'), - ('1024x256', '4:1', '512'), - ('2048x512', '4:1', '1K'), - ('4096x1024', '4:1', '2K'), - ('8192x2048', '4:1', '4K'), - ('600x448', '4:3', '512'), - ('1200x896', '4:3', '1K'), - ('2400x1792', '4:3', '2K'), - ('4800x3584', '4:3', '4K'), - ('464x576', '4:5', '512'), - ('928x1152', '4:5', '1K'), - ('1856x2304', '4:5', '2K'), - ('3712x4608', '4:5', '4K'), - ('576x464', '5:4', '512'), - ('1152x928', '5:4', '1K'), - ('2304x1856', '5:4', '2K'), - ('4608x3712', '5:4', '4K'), - ('1536x192', '8:1', '512'), - ('3072x384', '8:1', '1K'), - ('6144x768', '8:1', '2K'), - ('12288x1536', '8:1', '4K'), - ('384x688', '9:16', '512'), - ('768x1376', '9:16', '1K'), - ('1536x2752', '9:16', '2K'), - ('3072x5504', '9:16', '4K'), - ('688x384', '16:9', '512'), - ('1376x768', '16:9', '1K'), - ('2752x1536', '16:9', '2K'), - ('5504x3072', '16:9', '4K'), - ('792x336', '21:9', '512'), - ('1584x672', '21:9', '1K'), - ('3168x1344', '21:9', '2K'), - ('6336x2688', '21:9', '4K'), + ("512x512", "1:1", "512"), + ("1024x1024", "1:1", "1K"), + ("2048x2048", "1:1", "2K"), + ("4096x4096", "1:1", "4K"), + ("256x1024", "1:4", "512"), + ("512x2048", "1:4", "1K"), + ("1024x4096", "1:4", "2K"), + ("2048x8192", "1:4", "4K"), + ("192x1536", "1:8", "512"), + ("384x3072", "1:8", "1K"), + ("768x6144", "1:8", "2K"), + ("1536x12288", "1:8", "4K"), + ("424x632", "2:3", "512"), + ("848x1264", "2:3", "1K"), + ("1696x2528", "2:3", "2K"), + ("3392x5056", "2:3", "4K"), + ("632x424", "3:2", "512"), + ("1264x848", "3:2", "1K"), + ("2528x1696", "3:2", "2K"), + ("5056x3392", "3:2", "4K"), + ("448x600", "3:4", "512"), + ("896x1200", "3:4", "1K"), + ("1792x2400", "3:4", "2K"), + ("3584x4800", "3:4", "4K"), + ("1024x256", "4:1", "512"), + ("2048x512", "4:1", "1K"), + ("4096x1024", "4:1", "2K"), + ("8192x2048", "4:1", "4K"), + ("600x448", "4:3", "512"), + ("1200x896", "4:3", "1K"), + ("2400x1792", "4:3", "2K"), + ("4800x3584", "4:3", "4K"), + ("464x576", "4:5", "512"), + ("928x1152", "4:5", "1K"), + ("1856x2304", "4:5", "2K"), + ("3712x4608", "4:5", "4K"), + ("576x464", "5:4", "512"), + ("1152x928", "5:4", "1K"), + ("2304x1856", "5:4", "2K"), + ("4608x3712", "5:4", "4K"), + ("1536x192", "8:1", "512"), + ("3072x384", "8:1", "1K"), + ("6144x768", "8:1", "2K"), + ("12288x1536", "8:1", "4K"), + ("384x688", "9:16", "512"), + ("768x1376", "9:16", "1K"), + ("1536x2752", "9:16", "2K"), + ("3072x5504", "9:16", "4K"), + ("688x384", "16:9", "512"), + ("1376x768", "16:9", "1K"), + ("2752x1536", "16:9", "2K"), + ("5504x3072", "16:9", "4K"), + ("792x336", "21:9", "512"), + ("1584x672", "21:9", "1K"), + ("3168x1344", "21:9", "2K"), + ("6336x2688", "21:9", "4K"), ) @@ -314,9 +316,7 @@ def test_vertex_ai_response_json_schema_preserves_refs_for_gemini_2(): # $defs and $ref should be preserved (not unpacked) assert "response_json_schema" in transformed_request result_schema = transformed_request["response_json_schema"] - assert ( - "$defs" in result_schema - ), "responseJsonSchema should preserve $defs for Gemini 2.0+" + assert "$defs" in result_schema, "responseJsonSchema should preserve $defs for Gemini 2.0+" def test_vertex_ai_get_json_schema_preserves_refs_for_nested_pydantic(): @@ -396,22 +396,14 @@ def test_vertex_ai_response_json_schema_for_gemini_2(): # Types should be lowercase (standard JSON Schema format) assert transformed_request["response_json_schema"]["type"] == "object" - assert ( - transformed_request["response_json_schema"]["properties"]["name"]["type"] - == "string" - ) - assert ( - transformed_request["response_json_schema"]["properties"]["age"]["type"] - == "integer" - ) + assert transformed_request["response_json_schema"]["properties"]["name"]["type"] == "string" + assert transformed_request["response_json_schema"]["properties"]["age"]["type"] == "integer" # Should NOT have propertyOrdering (not needed for responseJsonSchema) assert "propertyOrdering" not in transformed_request["response_json_schema"] # additionalProperties should be preserved (supported by responseJsonSchema) - assert ( - transformed_request["response_json_schema"].get("additionalProperties") == False - ) + assert transformed_request["response_json_schema"].get("additionalProperties") == False def test_vertex_ai_response_schema_for_old_models(): @@ -648,7 +640,9 @@ def test_response_has_search_grounding_detection(): is False ) assert ( - VertexGeminiConfig._response_has_search_grounding({"candidates": [{"groundingMetadata": {"webSearchQueries": []}}]}) + VertexGeminiConfig._response_has_search_grounding( + {"candidates": [{"groundingMetadata": {"webSearchQueries": []}}]} + ) is False ) assert VertexGeminiConfig._response_has_search_grounding({"candidates": []}) is False @@ -789,9 +783,7 @@ def test_vertex_ai_url_context_tool_use_tokens_billed_as_input_tokens(): "candidates": [ { "urlContextMetadata": {"urlMetadata": []}, - "groundingMetadata": { - "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}] - }, + "groundingMetadata": {"groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]}, } ], "usageMetadata": UsageMetadata( @@ -823,9 +815,7 @@ def test_url_context_tool_use_tokens_are_priced_as_input(): "candidates": [ { "urlContextMetadata": {"urlMetadata": []}, - "groundingMetadata": { - "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}] - }, + "groundingMetadata": {"groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]}, } ], "usageMetadata": UsageMetadata( @@ -886,9 +876,7 @@ def test_missing_prompt_tokens_details_still_prices_the_full_prompt(): "candidates": [ { "urlContextMetadata": {"urlMetadata": []}, - "groundingMetadata": { - "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}] - }, + "groundingMetadata": {"groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]}, } ], "usageMetadata": UsageMetadata( @@ -925,9 +913,7 @@ def test_streaming_chunk_includes_reasoning_tokens(): "thoughtsTokenCount": 3, }, } - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=litellm_logging - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging) streaming_chunk = iterator.chunk_parser(chunk) assert streaming_chunk.usage is not None assert streaming_chunk.usage.prompt_tokens == 5 @@ -963,17 +949,12 @@ def test_streaming_chunk_includes_reasoning_content(): "usageMetadata": {}, } - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=litellm_logging - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging) streaming_chunk = iterator.chunk_parser(chunk) # The text content should be empty and reasoning_content should be populated assert streaming_chunk.choices[0].delta.content is None - assert ( - streaming_chunk.choices[0].delta.reasoning_content - == "I'm thinking through the problem..." - ) + assert streaming_chunk.choices[0].delta.reasoning_content == "I'm thinking through the problem..." def test_streaming_chunk_with_tool_calls_and_thought_includes_reasoning_content(): @@ -1018,24 +999,16 @@ def test_streaming_chunk_with_tool_calls_and_thought_includes_reasoning_content( }, } - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=litellm_logging - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging) streaming_chunk = iterator.chunk_parser(chunk) # Verify reasoning_content comes from the thought: true part - assert ( - streaming_chunk.choices[0].delta.reasoning_content - == "Let me think about how to get the time..." - ) + assert streaming_chunk.choices[0].delta.reasoning_content == "Let me think about how to get the time..." # Verify tool calls are also present assert streaming_chunk.choices[0].delta.tool_calls is not None assert len(streaming_chunk.choices[0].delta.tool_calls) == 1 - assert ( - streaming_chunk.choices[0].delta.tool_calls[0].function.name - == "get_current_time" - ) + assert streaming_chunk.choices[0].delta.tool_calls[0].function.name == "get_current_time" def test_streaming_chunk_with_tool_calls_no_thought_no_reasoning_content(): @@ -1077,9 +1050,7 @@ def test_streaming_chunk_with_tool_calls_no_thought_no_reasoning_content(): }, } - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=litellm_logging - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging) streaming_chunk = iterator.chunk_parser(chunk) # reasoning_content should be None - thoughtSignature alone does NOT mean reasoning @@ -1286,11 +1257,7 @@ def test_vertex_ai_usage_metadata_with_image_tokens_in_prompt(): # Verify the math: prompt_tokens = text + image # 533 = 6 (text) + 527 (image) - assert ( - result.prompt_tokens_details.text_tokens - + result.prompt_tokens_details.image_tokens - == result.prompt_tokens - ) + assert result.prompt_tokens_details.text_tokens + result.prompt_tokens_details.image_tokens == result.prompt_tokens def test_map_response_modalities_video(): @@ -1480,9 +1447,7 @@ def test_vertex_ai_streaming_usage_calculation(): } # Create iterator and parse chunk - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=MagicMock() - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock()) iterator.chunk_parser(chunk) # Verify calculate_usage was called with correct parameters @@ -1554,9 +1519,7 @@ def test_vertex_ai_streaming_usage_web_search_calculation(): } # Create iterator and parse chunk - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=MagicMock() - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock()) completed_response = iterator.chunk_parser(chunk) usage: Usage = completed_response.usage @@ -1775,7 +1738,9 @@ async def test_gemini_image_generation_async_stream(respx_mock: respx.MockRouter "modelVersion": "gemini-2.5-flash-image", }, { - "candidates": [{"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP", "index": 0}], + "candidates": [ + {"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP", "index": 0} + ], "usageMetadata": {"promptTokenCount": 13, "candidatesTokenCount": 1290, "totalTokenCount": 1303}, "modelVersion": "gemini-2.5-flash-image", }, @@ -2029,9 +1994,14 @@ def test_gemini_openai_web_search_tool_to_google_search(respx_mock: respx.MockRo GEMINI_GROUNDING_METADATA: Final = { "webSearchQueries": ["capital of France"], - "searchEntryPoint": {"renderedContent": "
"}, + "searchEntryPoint": {"renderedContent": '
'}, "groundingChunks": [ - {"web": {"uri": "https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9w", "title": "wikipedia.org"}} + { + "web": { + "uri": "https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9w", + "title": "wikipedia.org", + } + } ], "groundingSupports": [ { @@ -2188,9 +2158,7 @@ def test_vertex_ai_transform_parts(): # Test case 2: Tool call mode (is_function_call=False) - Single message with multiple tool calls parts_with_multiple_functions = [ - HttpxPartType( - functionCall={"name": "get_current_weather", "args": {"location": "Boston"}} - ), + HttpxPartType(functionCall={"name": "get_current_weather", "args": {"location": "Boston"}}), HttpxPartType( functionCall={ "name": "get_forecast", @@ -2220,9 +2188,7 @@ def test_vertex_ai_transform_parts(): # Test case 3: Simulating multiple messages - cumulative indexing across messages # First message with 2 tool calls (starting from index 0) first_message_parts = [ - HttpxPartType( - functionCall={"name": "get_weather", "args": {"location": "Boston"}} - ), + HttpxPartType(functionCall={"name": "get_weather", "args": {"location": "Boston"}}), HttpxPartType(functionCall={"name": "get_time", "args": {"timezone": "EST"}}), ] @@ -2239,9 +2205,7 @@ def test_vertex_ai_transform_parts(): # Second message with 1 tool call (continuing from previous index) second_message_parts = [ - HttpxPartType( - functionCall={"name": "send_email", "args": {"to": "user@example.com"}} - ), + HttpxPartType(functionCall={"name": "send_email", "args": {"to": "user@example.com"}}), ] function, tools, updated_idx = VertexGeminiConfig._transform_parts( @@ -2259,9 +2223,7 @@ def test_vertex_ai_transform_parts(): # Third message with 2 more tool calls (continuing from previous index) third_message_parts = [ - HttpxPartType( - functionCall={"name": "create_calendar_event", "args": {"title": "Meeting"}} - ), + HttpxPartType(functionCall={"name": "create_calendar_event", "args": {"title": "Meeting"}}), HttpxPartType(functionCall={"name": "set_reminder", "args": {"time": "10:00"}}), ] @@ -2308,9 +2270,7 @@ def test_vertex_ai_transform_parts(): assert updated_idx == 10 # Index should remain unchanged # Test case 6: Function call with empty args - parts_with_empty_args = [ - HttpxPartType(functionCall={"name": "simple_function", "args": {}}) - ] + parts_with_empty_args = [HttpxPartType(functionCall={"name": "simple_function", "args": {}})] function, tools, updated_idx = VertexGeminiConfig._transform_parts( parts=parts_with_empty_args, cumulative_tool_call_idx=0, is_function_call=True @@ -2326,13 +2286,9 @@ def test_vertex_ai_transform_parts(): # Test case 7: Mixed content with function calls - ensuring tool call IDs are unique mixed_parts = [ HttpxPartType(text="Before function call"), - HttpxPartType( - functionCall={"name": "function_a", "args": {"param": "value_a"}} - ), + HttpxPartType(functionCall={"name": "function_a", "args": {"param": "value_a"}}), HttpxPartType(text="Between function calls"), - HttpxPartType( - functionCall={"name": "function_b", "args": {"param": "value_b"}} - ), + HttpxPartType(functionCall={"name": "function_b", "args": {"param": "value_b"}}), HttpxPartType(text="After function calls"), ] @@ -2524,25 +2480,17 @@ def test_vertex_ai_tool_call_id_format(): tool_id = tool["id"] # Should start with 'call_' - assert tool_id.startswith( - "call_" - ), f"ID should start with 'call_', got: {tool_id}" + assert tool_id.startswith("call_"), f"ID should start with 'call_', got: {tool_id}" # Should have exactly 33 total characters (call_ + 28 hex chars) - assert ( - len(tool_id) == 33 - ), f"ID should be 33 characters long, got {len(tool_id)}: {tool_id}" + assert len(tool_id) == 33, f"ID should be 33 characters long, got {len(tool_id)}: {tool_id}" # The part after 'call_' should be 28 hex characters hex_part = tool_id[5:] # Remove 'call_' prefix - assert ( - len(hex_part) == 28 - ), f"Hex part should be 28 characters, got {len(hex_part)}: {hex_part}" + assert len(hex_part) == 28, f"Hex part should be 28 characters, got {len(hex_part)}: {hex_part}" # Should only contain valid hex characters - assert re.match( - r"^[0-9a-f]{28}$", hex_part - ), f"Should contain only lowercase hex chars, got: {hex_part}" + assert re.match(r"^[0-9a-f]{28}$", hex_part), f"Should contain only lowercase hex chars, got: {hex_part}" # Verify IDs are unique assert tools[0]["id"] != tools[1]["id"], "Tool call IDs should be unique" @@ -2559,9 +2507,7 @@ def test_vertex_ai_tool_call_id_format(): ids_generated.add(test_tools[0]["id"]) # All generated IDs should be unique - assert ( - len(ids_generated) == 10 - ), f"All 10 IDs should be unique, got {len(ids_generated)} unique IDs" + assert len(ids_generated) == 10, f"All 10 IDs should be unique, got {len(ids_generated)} unique IDs" def test_vertex_ai_map_google_maps_tool_simple(): @@ -2675,19 +2621,16 @@ def test_vertex_ai_penalty_parameters_validation(): for model, should_support in test_cases: # Test _supports_penalty_parameters method - assert ( - v._supports_penalty_parameters(model) == should_support - ), f"Model {model} penalty support should be {should_support}" + assert v._supports_penalty_parameters(model) == should_support, ( + f"Model {model} penalty support should be {should_support}" + ) # Test get_supported_openai_params method supported_params = v.get_supported_openai_params(model) - has_penalty_params = ( - "frequency_penalty" in supported_params - and "presence_penalty" in supported_params + has_penalty_params = "frequency_penalty" in supported_params and "presence_penalty" in supported_params + assert has_penalty_params == should_support, ( + f"Model {model} should {'include' if should_support else 'exclude'} penalty params in supported list" ) - assert ( - has_penalty_params == should_support - ), f"Model {model} should {'include' if should_support else 'exclude'} penalty params in supported list" # Test parameter mapping for unsupported model model = "gemini-2.5-pro-preview-06-05" @@ -2707,12 +2650,8 @@ def test_vertex_ai_penalty_parameters_validation(): ) # Penalty parameters should be filtered out for unsupported models - assert ( - "frequency_penalty" not in result - ), "frequency_penalty should be filtered out for unsupported model" - assert ( - "presence_penalty" not in result - ), "presence_penalty should be filtered out for unsupported model" + assert "frequency_penalty" not in result, "frequency_penalty should be filtered out for unsupported model" + assert "presence_penalty" not in result, "presence_penalty should be filtered out for unsupported model" # Other parameters should still be included assert "temperature" in result, "temperature should still be included" @@ -2741,18 +2680,18 @@ def test_vertex_ai_gemini_3_penalty_parameters_unsupported(): for model in gemini_3_models: # Test _supports_penalty_parameters method - assert ( - v._supports_penalty_parameters(model) == False - ), f"Gemini 3 model {model} should not support penalty parameters" + assert v._supports_penalty_parameters(model) == False, ( + f"Gemini 3 model {model} should not support penalty parameters" + ) # Test get_supported_openai_params method supported_params = v.get_supported_openai_params(model) - assert ( - "frequency_penalty" not in supported_params - ), f"frequency_penalty should not be in supported params for {model}" - assert ( - "presence_penalty" not in supported_params - ), f"presence_penalty should not be in supported params for {model}" + assert "frequency_penalty" not in supported_params, ( + f"frequency_penalty should not be in supported params for {model}" + ) + assert "presence_penalty" not in supported_params, ( + f"presence_penalty should not be in supported params for {model}" + ) # Test parameter mapping - penalty params should be filtered out non_default_params = { @@ -2771,36 +2710,28 @@ def test_vertex_ai_gemini_3_penalty_parameters_unsupported(): ) # Penalty parameters should be filtered out for Gemini 3 models - assert ( - "frequency_penalty" not in result - ), f"frequency_penalty should be filtered out for Gemini 3 model {model}" - assert ( - "presence_penalty" not in result - ), f"presence_penalty should be filtered out for Gemini 3 model {model}" + assert "frequency_penalty" not in result, f"frequency_penalty should be filtered out for Gemini 3 model {model}" + assert "presence_penalty" not in result, f"presence_penalty should be filtered out for Gemini 3 model {model}" # Other parameters should still be included - assert ( - "temperature" in result - ), f"temperature should still be included for Gemini 3 model {model}" - assert ( - "max_output_tokens" in result - ), f"max_output_tokens should still be included for Gemini 3 model {model}" + assert "temperature" in result, f"temperature should still be included for Gemini 3 model {model}" + assert "max_output_tokens" in result, f"max_output_tokens should still be included for Gemini 3 model {model}" assert result["temperature"] == 0.7 assert result["max_output_tokens"] == 100 # Test that non-Gemini 3 models still support penalty parameters (if they're not in the unsupported list) non_gemini_3_model = "gemini-2.5-pro" - assert ( - v._supports_penalty_parameters(non_gemini_3_model) == True - ), f"Non-Gemini 3 model {non_gemini_3_model} should support penalty parameters" + assert v._supports_penalty_parameters(non_gemini_3_model) == True, ( + f"Non-Gemini 3 model {non_gemini_3_model} should support penalty parameters" + ) supported_params = v.get_supported_openai_params(non_gemini_3_model) - assert ( - "frequency_penalty" in supported_params - ), f"frequency_penalty should be in supported params for {non_gemini_3_model}" - assert ( - "presence_penalty" in supported_params - ), f"presence_penalty should be in supported params for {non_gemini_3_model}" + assert "frequency_penalty" in supported_params, ( + f"frequency_penalty should be in supported params for {non_gemini_3_model}" + ) + assert "presence_penalty" in supported_params, ( + f"presence_penalty should be in supported params for {non_gemini_3_model}" + ) def test_vertex_ai_annotation_streaming_events(): @@ -2823,14 +2754,10 @@ def test_vertex_ai_annotation_streaming_events(): chunk_with_annotations = { "candidates": [ { - "content": { - "parts": [{"text": "The weather in San Francisco today is clear."}] - }, + "content": {"parts": [{"text": "The weather in San Francisco today is clear."}]}, "groundingMetadata": { "webSearchQueries": ["weather San Francisco today"], - "searchEntryPoint": { - "renderedContent": "
Search results
" - }, + "searchEntryPoint": {"renderedContent": "
Search results
"}, "groundingChunks": [ { "web": { @@ -2862,9 +2789,7 @@ def test_vertex_ai_annotation_streaming_events(): } # Create iterator and parse chunk - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=litellm_logging - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging) streaming_chunk = iterator.chunk_parser(chunk_with_annotations) # Verify the chunk was parsed correctly @@ -2986,12 +2911,8 @@ def test_vertex_ai_annotation_conversion(): } # Convert grounding metadata to annotations - content_text = ( - "The weather in San Francisco is currently 72°F and the time is 2:30 PM" - ) - annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations( - [grounding_metadata], content_text - ) + content_text = "The weather in San Francisco is currently 72°F and the time is 2:30 PM" + annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations([grounding_metadata], content_text) # Verify annotations were created assert len(annotations) == 3 # One for each grounding support @@ -3033,9 +2954,7 @@ def test_vertex_ai_annotation_empty_grounding_metadata(): # Test with empty grounding metadata empty_metadata = {} - annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations( - [empty_metadata], "test content" - ) + annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations([empty_metadata], "test content") assert len(annotations) == 0 # Test with missing groundingSupports @@ -3043,9 +2962,7 @@ def test_vertex_ai_annotation_empty_grounding_metadata(): "webSearchQueries": ["test query"], "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Test"}}], } - annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations( - [metadata_no_supports], "test content" - ) + annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations([metadata_no_supports], "test content") assert len(annotations) == 0 # Test with empty groundingSupports @@ -3408,9 +3325,7 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3(): def test_gemini_37_38_flash_floor_minimal_thinking_level( local_model_cost_map, model, reasoning_effort, include_thoughts ): - result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( - reasoning_effort, model - ) + result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(reasoning_effort, model) assert result["thinkingLevel"] == "low" assert result["includeThoughts"] is include_thoughts @@ -3434,9 +3349,7 @@ def test_gemini_37_38_flash_floor_minimal_thinking_level( def test_gemini_flash_minimal_thinking_support( local_model_cost_map, model, reasoning_effort, expected_level, include_thoughts ): - result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( - reasoning_effort, model - ) + result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(reasoning_effort, model) assert result["thinkingLevel"] == expected_level assert result["includeThoughts"] is include_thoughts @@ -3446,12 +3359,8 @@ def test_gemini_38_flash_feature_flag_uses_low_thinking_level(local_model_cost_m monkeypatch.setattr(litellm, "enable_gemini_default_thinking_level_low", True) thinking_param = {"type": "enabled", "budget_tokens": 1024} - result_38 = VertexGeminiConfig._map_thinking_param( - thinking_param, model="gemini-3.8-flash" - ) - result_36 = VertexGeminiConfig._map_thinking_param( - thinking_param, model="gemini-3.6-flash" - ) + result_38 = VertexGeminiConfig._map_thinking_param(thinking_param, model="gemini-3.8-flash") + result_36 = VertexGeminiConfig._map_thinking_param(thinking_param, model="gemini-3.6-flash") assert result_38["thinkingLevel"] == "low" assert result_36["thinkingLevel"] == "minimal" @@ -3933,9 +3842,7 @@ def test_gemini_image_models_excluded_from_thinking(): ) # None of these should have thinkingConfig - assert ( - "thinkingConfig" not in result - ), f"Model {model} should not have thinkingConfig" + assert "thinkingConfig" not in result, f"Model {model} should not have thinkingConfig" def test_partial_json_chunk_after_first_chunk(): @@ -3966,9 +3873,7 @@ def test_partial_json_chunk_after_first_chunk(): first_chunk = '{"candidates": [{"content": {"parts": [{"text": "Hello"}]}}]}' result1 = iterator.handle_valid_json_chunk(first_chunk) assert result1 is not None, "First complete chunk should parse OK" - assert ( - iterator.sent_first_chunk is True - ), "sent_first_chunk should be True after first chunk" + assert iterator.sent_first_chunk is True, "sent_first_chunk should be True after first chunk" # Later chunk arrives PARTIAL (simulating network fragmentation) partial_chunk = '{"candidates": [{"content":' @@ -3976,9 +3881,7 @@ def test_partial_json_chunk_after_first_chunk(): # Should switch to accumulation mode instead of crashing assert result2 is None, "Partial chunk should return None while accumulating" - assert ( - iterator.chunk_type == "accumulated_json" - ), "Should switch to accumulated_json mode" + assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode" def test_partial_json_chunk_on_first_chunk(): @@ -3998,9 +3901,7 @@ def test_partial_json_chunk_on_first_chunk(): result = iterator.handle_valid_json_chunk(partial) assert result is None, "Partial first chunk should return None" - assert ( - iterator.chunk_type == "accumulated_json" - ), "Should switch to accumulated_json mode" + assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode" def test_accumulated_json_does_not_reparse_every_fragment(): @@ -4032,16 +3933,12 @@ def test_accumulated_json_does_not_reparse_every_fragment(): iterator.chunk_type = "accumulated_json" text = "x" * 200_000 # no braces/brackets so only the final fragment closes - blob = json.dumps( - {"candidates": [{"content": {"role": "model", "parts": [{"text": text}]}}]} - ) + blob = json.dumps({"candidates": [{"content": {"role": "model", "parts": [{"text": text}]}}]}) fragments = [blob[i : i + 4096] for i in range(0, len(blob), 4096)] assert len(fragments) > 10, "need a multi-fragment payload to exercise the bug" parsed = None - with patch.object( - json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode - ) as spy: + with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy: for fragment in fragments: out = iterator.handle_accumulated_json_chunk(chunk=fragment) if out is not None: @@ -4073,9 +3970,7 @@ def test_accumulated_json_partial_fragment_returns_none_without_parsing(): ) iterator.chunk_type = "accumulated_json" - with patch.object( - json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode - ) as spy: + with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy: result = iterator.handle_accumulated_json_chunk( chunk='{"candidates": [{"content": {"parts": [{"text": "partial' ) @@ -4174,17 +4069,11 @@ def test_vertex_ai_multiple_tool_types_separate_objects(): tool_types_in_first = [k for k in tools[0].keys()] tool_types_in_second = [k for k in tools[1].keys()] - assert ( - len(tool_types_in_first) == 1 - ), f"First Tool should have exactly 1 type, got {tool_types_in_first}" - assert ( - len(tool_types_in_second) == 1 - ), f"Second Tool should have exactly 1 type, got {tool_types_in_second}" + assert len(tool_types_in_first) == 1, f"First Tool should have exactly 1 type, got {tool_types_in_first}" + assert len(tool_types_in_second) == 1, f"Second Tool should have exactly 1 type, got {tool_types_in_second}" # Verify the correct tool types are present - assert ( - "enterpriseWebSearch" in tools[0] - ), "First Tool should contain enterpriseWebSearch" + assert "enterpriseWebSearch" in tools[0], "First Tool should contain enterpriseWebSearch" assert "url_context" in tools[1], "Second Tool should contain url_context" @@ -4646,9 +4535,7 @@ def test_vertex_ai_multiple_function_declarations_grouped(): ) # Should have only 1 Tool object (function declarations grouped) - assert ( - len(tools) == 1 - ), f"Expected 1 Tool object for grouped functions, got {len(tools)}" + assert len(tools) == 1, f"Expected 1 Tool object for grouped functions, got {len(tools)}" # Should contain function_declarations with 2 functions assert "function_declarations" in tools[0] @@ -5070,9 +4957,7 @@ def test_chunk_parser_handles_prompt_feedback_block(): logging_obj = Mock() logging_obj.optional_params = {} - streaming_obj = ModelResponseIterator( - streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj - ) + streaming_obj = ModelResponseIterator(streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj) # Act result = streaming_obj.chunk_parser(blocked_chunk) @@ -5080,9 +4965,9 @@ def test_chunk_parser_handles_prompt_feedback_block(): # Assert assert result is not None, "Result should not be None" assert len(result.choices) == 1, "Should have exactly one choice" - assert ( - result.choices[0].finish_reason == "content_filter" - ), f"finish_reason should be content_filter, got {result.choices[0].finish_reason}" + assert result.choices[0].finish_reason == "content_filter", ( + f"finish_reason should be content_filter, got {result.choices[0].finish_reason}" + ) assert result.choices[0].delta.content is None, "content should be None" @@ -5106,9 +4991,7 @@ def test_chunk_parser_handles_prompt_feedback_safety_block(): logging_obj = Mock() logging_obj.optional_params = {} - streaming_obj = ModelResponseIterator( - streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj - ) + streaming_obj = ModelResponseIterator(streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj) # Act result = streaming_obj.chunk_parser(blocked_chunk) @@ -5145,9 +5028,7 @@ def test_chunk_parser_handles_prompt_feedback_block_with_usage(): logging_obj = Mock() logging_obj.optional_params = {} - streaming_obj = ModelResponseIterator( - streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj - ) + streaming_obj = ModelResponseIterator(streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj) # Act result = streaming_obj.chunk_parser(blocked_chunk) @@ -5155,23 +5036,17 @@ def test_chunk_parser_handles_prompt_feedback_block_with_usage(): # Assert - 验证 content_filter 响应和 usage 都被正确处理 assert result is not None, "Result should not be None" assert len(result.choices) == 1, "Should have exactly one choice" - assert ( - result.choices[0].finish_reason == "content_filter" - ), f"finish_reason should be content_filter, got {result.choices[0].finish_reason}" + assert result.choices[0].finish_reason == "content_filter", ( + f"finish_reason should be content_filter, got {result.choices[0].finish_reason}" + ) assert result.choices[0].delta.content is None, "content should be None" # 验证 usage 信息被正确提取 assert hasattr(result, "usage"), "result should have usage attribute" assert result.usage is not None, "usage should not be None" - assert ( - result.usage.prompt_tokens == 8175 - ), f"prompt_tokens should be 8175, got {result.usage.prompt_tokens}" - assert ( - result.usage.completion_tokens == 0 - ), f"completion_tokens should be 0, got {result.usage.completion_tokens}" - assert ( - result.usage.total_tokens == 8175 - ), f"total_tokens should be 8175, got {result.usage.total_tokens}" + assert result.usage.prompt_tokens == 8175, f"prompt_tokens should be 8175, got {result.usage.prompt_tokens}" + assert result.usage.completion_tokens == 0, f"completion_tokens should be 0, got {result.usage.completion_tokens}" + assert result.usage.total_tokens == 8175, f"total_tokens should be 8175, got {result.usage.total_tokens}" def test_vertex_ai_traffic_type_preserved_in_hidden_params_streaming(): @@ -5190,14 +5065,10 @@ def test_vertex_ai_traffic_type_preserved_in_hidden_params_streaming(): }, } - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=MagicMock() - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock()) result = iterator.chunk_parser(chunk) - assert ( - result._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND" - ) + assert result._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND" def test_vertex_ai_traffic_type_preserved_in_hidden_params_non_streaming(): @@ -5236,10 +5107,7 @@ def test_vertex_ai_traffic_type_preserved_in_hidden_params_non_streaming(): encoding=None, ) - assert ( - result._hidden_params["provider_specific_fields"]["traffic_type"] - == "PROVISIONED_THROUGHPUT" - ) + assert result._hidden_params["provider_specific_fields"]["traffic_type"] == "PROVISIONED_THROUGHPUT" def test_vertex_ai_service_tier_streaming(): @@ -5323,9 +5191,7 @@ def test_vertex_ai_traffic_type_surfaced_in_responses_api(): from litellm.types.utils import Choices, Message model_response = ModelResponse() - model_response._hidden_params["provider_specific_fields"] = { - "traffic_type": "ON_DEMAND" - } + model_response._hidden_params["provider_specific_fields"] = {"traffic_type": "ON_DEMAND"} model_response.choices = [ Choices( message=Message(content="Hello", role="assistant"), @@ -5334,15 +5200,15 @@ def test_vertex_ai_traffic_type_surfaced_in_responses_api(): ) ] - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="test", - chat_completion_response=model_response, - responses_api_request={}, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test", + chat_completion_response=model_response, + responses_api_request={}, + ) ) - assert ( - responses_api_response.provider_specific_fields["traffic_type"] == "ON_DEMAND" - ) + assert responses_api_response.provider_specific_fields["traffic_type"] == "ON_DEMAND" def test_vertex_ai_web_search_options_parameter(): @@ -5416,12 +5282,8 @@ def test_vertex_ai_web_search_options_in_map_openai_params(): assert "tools" in optional_params, "tools should be added to optional_params" assert len(optional_params["tools"]) == 1, "Should have exactly one tool" assert "googleSearch" in optional_params["tools"][0], "Tool should be googleSearch" - assert ( - optional_params["tools"][0]["googleSearch"] == {} - ), "googleSearch should be empty config" - assert ( - "web_search_options" not in optional_params - ), "web_search_options should be removed after transformation" + assert optional_params["tools"][0]["googleSearch"] == {}, "googleSearch should be empty config" + assert "web_search_options" not in optional_params, "web_search_options should be removed after transformation" def test_vertex_ai_service_tier_in_map_openai_params(): @@ -5509,24 +5371,16 @@ def test_vertex_ai_usage_metadata_with_video_tokens_in_prompt(): # Verify prompt token details include video tokens assert result.prompt_tokens_details is not None - assert ( - result.prompt_tokens_details.video_tokens == 10240 - ), "Prompt video tokens should be 10240" - assert ( - result.prompt_tokens_details.text_tokens == 9 - ), "Prompt text tokens should be 9" - assert ( - result.prompt_tokens_details.audio_tokens == 200 - ), "Prompt audio tokens should be 200" + assert result.prompt_tokens_details.video_tokens == 10240, "Prompt video tokens should be 10240" + assert result.prompt_tokens_details.text_tokens == 9, "Prompt text tokens should be 9" + assert result.prompt_tokens_details.audio_tokens == 200, "Prompt audio tokens should be 200" # Verify completion token details assert result.completion_tokens_details is not None - assert ( - result.completion_tokens_details.text_tokens == 79 - ), "Completion text tokens should be 79" - assert ( - result.completion_tokens_details.video_tokens is None - ), "Completion video tokens should be None (text-only response)" + assert result.completion_tokens_details.text_tokens == 79, "Completion text tokens should be 79" + assert result.completion_tokens_details.video_tokens is None, ( + "Completion video tokens should be None (text-only response)" + ) def test_vertex_ai_usage_metadata_with_video_tokens_in_candidates(): @@ -5667,9 +5521,9 @@ def test_vertex_ai_usage_metadata_with_document_tokens_in_prompt(): # DOCUMENT tokens should be included in text_tokens: 8 (TEXT) + 774 (DOCUMENT) = 782 assert result.prompt_tokens_details is not None - assert ( - result.prompt_tokens_details.text_tokens == 782 - ), "DOCUMENT modality tokens should be added to text_tokens (8 TEXT + 774 DOCUMENT = 782)" + assert result.prompt_tokens_details.text_tokens == 782, ( + "DOCUMENT modality tokens should be added to text_tokens (8 TEXT + 774 DOCUMENT = 782)" + ) # Verify completion token details assert result.completion_tokens_details is not None @@ -6166,9 +6020,7 @@ def test_mid_stream_429_error_raises_during_iteration(): { "content": { "role": "model", - "parts": [ - {"text": "Let me think about this...", "thought": True} - ], + "parts": [{"text": "Let me think about this...", "thought": True}], }, "index": 0, } @@ -6188,9 +6040,7 @@ def test_mid_stream_429_error_raises_during_iteration(): { "content": { "role": "model", - "parts": [ - {"text": "I'll generate the image now.", "thought": True} - ], + "parts": [{"text": "I'll generate the image now.", "thought": True}], }, "index": 0, } @@ -6227,6 +6077,7 @@ def test_mid_stream_429_error_raises_during_iteration(): # Iterate the stream: first chunks should succeed, then 429 error should be raised results = [] + def _drain(): for chunk in streaming_obj: if chunk is not None: @@ -6236,9 +6087,7 @@ def test_mid_stream_429_error_raises_during_iteration(): _drain() # Verify: received normal chunks before the error - assert ( - len(results) >= 1 - ), "Should have received at least 1 normal chunk before the error" + assert len(results) >= 1, "Should have received at least 1 normal chunk before the error" # Verify: 429 error is properly raised assert exc_info.value.status_code == 429 @@ -6427,9 +6276,7 @@ def _accumulating_gemini_iterator(): ModelResponseIterator, ) - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=MagicMock() - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock()) iterator.chunk_type = "accumulated_json" return iterator @@ -6978,6 +6825,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -7030,6 +6878,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -7050,6 +6899,7 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -7072,6 +6922,7 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_true_creates_thinking_block(): """ @@ -7086,6 +6937,7 @@ def test_thought_true_creates_thinking_block(): assert block["thinking"] == "Some thinking" assert block["signature"] == "sig-1" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_true_with_empty_text_creates_block(): """ @@ -7097,6 +6949,7 @@ def test_thought_true_with_empty_text_creates_block(): assert len(thinking_blocks) == 1 assert thinking_blocks[0]["thinking"] == "" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_signature_without_thought_does_not_create_block(): """ @@ -7109,6 +6962,7 @@ def test_thought_signature_without_thought_does_not_create_block(): thinking_blocks = config._extract_thinking_blocks_from_parts(parts) assert thinking_blocks == [] + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_extract_thought_signatures_from_regular_parts(): """ @@ -7128,6 +6982,7 @@ def test_extract_thought_signatures_from_regular_parts(): assert len(signatures) == 1 assert signatures[0] == "sig-regular-123" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_extract_multiple_thought_signatures(): """ @@ -7146,6 +7001,7 @@ def test_extract_multiple_thought_signatures(): assert signatures[0] == "sig-1" assert signatures[1] == "sig-2" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_round_trip_thought_signature_in_conversation(): """ @@ -7180,6 +7036,7 @@ def test_round_trip_thought_signature_in_conversation(): assert "thoughtSignature" in text_part assert text_part["thoughtSignature"] == "sig-round-trip-abc" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_round_trip_without_thought_signature_still_works(): """ @@ -7418,18 +7275,10 @@ def test_gemini_flash_image_preview_models(model_name: str): mock_response = ImageResponse() mock_response.data = [ImageObject(b64_json="test_base64_data", url=None)] - with patch( - "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post: mock_http_response = MagicMock() mock_http_response.json.return_value = { - "candidates": [ - { - "content": { - "parts": [{"inlineData": {"data": "test_base64_image_data"}}] - } - } - ] + "candidates": [{"content": {"parts": [{"inlineData": {"data": "test_base64_image_data"}}]}}] } mock_http_response.status_code = 200 mock_post.return_value = mock_http_response @@ -7447,9 +7296,7 @@ def test_gemini_flash_image_preview_models(model_name: str): mock_post.assert_called_once() call_args = mock_post.call_args - called_url = ( - call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "") - ) + called_url = call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "") assert ":generateContent" in called_url assert model_name.split("/", 1)[1] in called_url @@ -7481,23 +7328,13 @@ def test_gemini_flash_image_preview_models(model_name: str): ), ], ) -def test_gemini_image_generation_forwards_image_config( - model: str, kwargs: dict, expected_image_config: dict -): +def test_gemini_image_generation_forwards_image_config(model: str, kwargs: dict, expected_image_config: dict): from unittest.mock import patch, MagicMock - with patch( - "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post: mock_http_response = MagicMock() mock_http_response.json.return_value = { - "candidates": [ - { - "content": { - "parts": [{"inlineData": {"data": "test_base64_image_data"}}] - } - } - ] + "candidates": [{"content": {"parts": [{"inlineData": {"data": "test_base64_image_data"}}]}}] } mock_http_response.status_code = 200 mock_post.return_value = mock_http_response @@ -7558,9 +7395,7 @@ def test_gemini_image_generation_openai_size_maps_to_google_table( map_openai_size_to_gemini_image_config, ) - assert map_openai_size_to_gemini_image_config( - size, "gemini-3-pro-image-preview" - ) == { + assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") == { "aspectRatio": expected_aspect_ratio, "imageSize": expected_image_size, } @@ -7584,9 +7419,7 @@ def test_gemini_image_generation_openai_size_snaps_to_nearest_option( map_openai_size_to_gemini_image_config, ) - assert map_openai_size_to_gemini_image_config( - size, "gemini-3-pro-image-preview" - ) == { + assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") == { "aspectRatio": expected_aspect_ratio, "imageSize": expected_image_size, } @@ -7598,9 +7431,7 @@ def test_gemini_image_generation_openai_size_auto_uses_google_defaults(size: str map_openai_size_to_gemini_image_config, ) - assert map_openai_size_to_gemini_image_config( - size, "gemini-3-pro-image-preview" - ) is None + assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") is None def test_gemini_imagen_models_use_predict_endpoint(): @@ -7610,13 +7441,9 @@ def test_gemini_imagen_models_use_predict_endpoint(): from unittest.mock import patch, MagicMock from litellm.types.utils import ImageResponse, ImageObject - with patch( - "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post: mock_http_response = MagicMock() - mock_http_response.json.return_value = { - "predictions": [{"bytesBase64Encoded": "test_base64_image_data"}] - } + mock_http_response.json.return_value = {"predictions": [{"bytesBase64Encoded": "test_base64_image_data"}]} mock_http_response.status_code = 200 mock_post.return_value = mock_http_response @@ -7632,9 +7459,7 @@ def test_gemini_imagen_models_use_predict_endpoint(): mock_post.assert_called_once() call_args = mock_post.call_args - called_url = ( - call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "") - ) + called_url = call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "") assert ":predict" in called_url assert "imagen-3.0-generate-001" in called_url @@ -7797,11 +7622,7 @@ async def test_claude_tool_use_with_gemini(): if "data: " in chunk_str: try: - data_line = [ - line - for line in chunk_str.split("\n") - if line.startswith("data: ") - ][0] + data_line = [line for line in chunk_str.split("\n") if line.startswith("data: ")][0] json_str = data_line[6:] chunk_data = json.loads(json_str) @@ -7814,24 +7635,15 @@ async def test_claude_tool_use_with_gemini(): if ( chunk_data.get("type") == "message_delta" - and chunk_data.get("delta", {}).get("stop_reason") - is not None + and chunk_data.get("delta", {}).get("stop_reason") is not None and "usage" in chunk_data ): has_usage_in_message_delta = True usage = chunk_data["usage"] - assert ( - "input_tokens" in usage - ), "input_tokens should be present in usage" - assert ( - "output_tokens" in usage - ), "output_tokens should be present in usage" - assert isinstance( - usage["input_tokens"], int - ), "input_tokens should be an integer" - assert isinstance( - usage["output_tokens"], int - ), "output_tokens should be an integer" + assert "input_tokens" in usage, "input_tokens should be present in usage" + assert "output_tokens" in usage, "output_tokens should be present in usage" + assert isinstance(usage["input_tokens"], int), "input_tokens should be an integer" + assert isinstance(usage["output_tokens"], int), "output_tokens should be an integer" print(f"Found usage in message_delta: {usage}") except (json.JSONDecodeError, IndexError) as e: @@ -7846,9 +7658,7 @@ async def test_claude_tool_use_with_gemini(): assert is_content_block_tool_use, "content_block_tool_use should be present" assert is_partial_json, "partial_json should be present" - assert ( - has_usage_in_message_delta - ), "Usage should be present in message_delta with stop_reason" + assert has_usage_in_message_delta, "Usage should be present in message_delta with stop_reason" assert is_content_block_stop, "is_content_block_stop should be present" @@ -7877,28 +7687,22 @@ def test_gemini_reasoning_effort_minimal(): ) request_body = raw_request["raw_request_body"] - assert ( - "generationConfig" in request_body - ), f"Model {model} should have generationConfig" + assert "generationConfig" in request_body, f"Model {model} should have generationConfig" generation_config = request_body["generationConfig"] - assert ( - "thinkingConfig" in generation_config - ), f"Model {model} should have thinkingConfig" + assert "thinkingConfig" in generation_config, f"Model {model} should have thinkingConfig" thinking_config = generation_config["thinkingConfig"] - assert ( - "thinkingBudget" in thinking_config - ), f"Model {model} should have thinkingBudget" + assert "thinkingBudget" in thinking_config, f"Model {model} should have thinkingBudget" actual_budget = thinking_config["thinkingBudget"] - assert ( - actual_budget == expected_min_budget - ), f"Model {model} should map 'minimal' to {expected_min_budget} tokens, got {actual_budget}" + assert actual_budget == expected_min_budget, ( + f"Model {model} should map 'minimal' to {expected_min_budget} tokens, got {actual_budget}" + ) - assert thinking_config.get( - "includeThoughts", True - ), f"Model {model} should have includeThoughts=True for minimal reasoning effort" + assert thinking_config.get("includeThoughts", True), ( + f"Model {model} should have includeThoughts=True for minimal reasoning effort" + ) try: raw_request = return_raw_request( @@ -7913,9 +7717,7 @@ def test_gemini_reasoning_effort_minimal(): request_body = raw_request["raw_request_body"] generation_config = request_body["generationConfig"] thinking_config = generation_config["thinkingConfig"] - assert ( - thinking_config["thinkingBudget"] == 128 - ), "Unknown model should use generic fallback of 128 tokens" + assert thinking_config["thinkingBudget"] == 128, "Unknown model should use generic fallback of 128 tokens" except Exception as e: print(f"Note: Unknown model test skipped due to: {e}") pass @@ -7938,9 +7740,7 @@ def test_gemini_exception_message_format(): mock_response.text = "Invalid API key provided" mock_response.headers = {} - mock_exception = httpx.HTTPStatusError( - message="Bad Request", request=Mock(), response=mock_response - ) + mock_exception = httpx.HTTPStatusError(message="Bad Request", request=Mock(), response=mock_response) mock_exception.response = mock_response mock_exception.status_code = 400 @@ -7960,9 +7760,9 @@ def test_gemini_exception_message_format(): f"Expected 'GeminiException' in error message, got: {error_message}. " f"This test should fail before the fix is implemented." ) - assert ( - "VertexAIException" not in error_message - ), f"Should not contain 'VertexAIException' in error message, got: {error_message}" + assert "VertexAIException" not in error_message, ( + f"Should not contain 'VertexAIException' in error message, got: {error_message}" + ) def test_reasoning_effort_none_mapping(): @@ -8016,17 +7816,11 @@ def test_gemini_function_args_preserve_unicode(): parsed_args = json.loads(arguments_str) assert parsed_args["message"] == "やあ", "Japanese characters should be preserved" - assert ( - parsed_args["recipient"] == "たけし" - ), "Japanese characters should be preserved" + assert parsed_args["recipient"] == "たけし", "Japanese characters should be preserved" assert "\\u" not in arguments_str, "Should not contain Unicode escape sequences" - assert ( - "やあ" in arguments_str - ), "Original Japanese characters should be in the string" - assert ( - "たけし" in arguments_str - ), "Original Japanese characters should be in the string" + assert "やあ" in arguments_str, "Original Japanese characters should be in the string" + assert "たけし" in arguments_str, "Original Japanese characters should be in the string" parts_spanish = [ { @@ -8079,12 +7873,8 @@ def test_anthropic_thinking_param_to_gemini_3_provider_defaults(): model="gemini-3-flash", ) - assert ( - "thinkingLevel" not in result - ), "Should not force thinkingLevel for Gemini 3" - assert ( - "thinkingBudget" not in result - ), "Should NOT have thinkingBudget for Gemini 3" + assert "thinkingLevel" not in result, "Should not force thinkingLevel for Gemini 3" + assert "thinkingBudget" not in result, "Should NOT have thinkingBudget for Gemini 3" assert result["includeThoughts"] is True thinking_param_disabled: AnthropicThinkingParam = { @@ -8098,10 +7888,7 @@ def test_anthropic_thinking_param_to_gemini_3_provider_defaults(): ) assert result_disabled.get("includeThoughts") is False - assert ( - "thinkingLevel" not in result_disabled - or result_disabled.get("thinkingLevel") is None - ) + assert "thinkingLevel" not in result_disabled or result_disabled.get("thinkingLevel") is None thinking_param_zero: AnthropicThinkingParam = { "type": "enabled", @@ -8114,10 +7901,7 @@ def test_anthropic_thinking_param_to_gemini_3_provider_defaults(): ) assert result_zero["includeThoughts"] is False - assert ( - "thinkingLevel" not in result_zero - or result_zero.get("thinkingLevel") is None - ) + assert "thinkingLevel" not in result_zero or result_zero.get("thinkingLevel") is None result_gemini3flashpreview = VertexGeminiConfig._map_thinking_param( thinking_param=thinking_param, @@ -8203,9 +7987,7 @@ def test_anthropic_thinking_param_to_gemini_2_thinkingBudget(): ) assert "thinkingBudget" in result_gemini2, "Should have thinkingBudget for Gemini 2" - assert ( - "thinkingLevel" not in result_gemini2 - ), "Should NOT have thinkingLevel for Gemini 2" + assert "thinkingLevel" not in result_gemini2, "Should NOT have thinkingLevel for Gemini 2" assert result_gemini2["includeThoughts"] is True assert result_gemini2["thinkingBudget"] == 10000 @@ -8241,12 +8023,8 @@ def test_anthropic_thinking_param_via_map_openai_params(): assert "thinkingConfig" in result, "Should have thinkingConfig in optional_params" thinking_config = result["thinkingConfig"] - assert ( - "thinkingLevel" not in thinking_config - ), "Should not force thinkingLevel for Gemini 3 by default" - assert ( - "thinkingBudget" not in thinking_config - ), "Should NOT have thinkingBudget for Gemini 3" + assert "thinkingLevel" not in thinking_config, "Should not force thinkingLevel for Gemini 3 by default" + assert "thinkingBudget" not in thinking_config, "Should NOT have thinkingBudget for Gemini 3" assert thinking_config["includeThoughts"] is True optional_params_2 = {} @@ -8259,12 +8037,8 @@ def test_anthropic_thinking_param_via_map_openai_params(): assert "thinkingConfig" in result_2, "Should have thinkingConfig in optional_params" thinking_config_2 = result_2["thinkingConfig"] - assert ( - "thinkingBudget" in thinking_config_2 - ), "Should have thinkingBudget for Gemini 2" - assert ( - "thinkingLevel" not in thinking_config_2 - ), "Should NOT have thinkingLevel for Gemini 2" + assert "thinkingBudget" in thinking_config_2, "Should have thinkingBudget for Gemini 2" + assert "thinkingLevel" not in thinking_config_2, "Should NOT have thinkingLevel for Gemini 2" assert thinking_config_2["includeThoughts"] is True assert thinking_config_2["thinkingBudget"] == 10000 @@ -8284,9 +8058,9 @@ def test_gemini_31_flash_lite_reasoning_effort_minimal(): reasoning_effort="minimal", model="gemini-3.1-flash-lite-preview", ) - assert ( - result["thinkingLevel"] == "minimal" - ), f"Expected thinkingLevel='minimal' for gemini-3.1-flash-lite-preview, got '{result['thinkingLevel']}'" + assert result["thinkingLevel"] == "minimal", ( + f"Expected thinkingLevel='minimal' for gemini-3.1-flash-lite-preview, got '{result['thinkingLevel']}'" + ) assert result["includeThoughts"] is True from litellm.utils import return_raw_request @@ -8302,12 +8076,12 @@ def test_gemini_31_flash_lite_reasoning_effort_minimal(): ) generation_config = raw_request["raw_request_body"]["generationConfig"] thinking_config = generation_config["thinkingConfig"] - assert ( - thinking_config.get("thinkingLevel") == "minimal" - ), f"Expected thinkingLevel='minimal' via full flow, got {thinking_config}" - assert ( - "thinkingBudget" not in thinking_config - ), "gemini-3.1-flash-lite-preview should use thinkingLevel, not thinkingBudget" + assert thinking_config.get("thinkingLevel") == "minimal", ( + f"Expected thinkingLevel='minimal' via full flow, got {thinking_config}" + ) + assert "thinkingBudget" not in thinking_config, ( + "gemini-3.1-flash-lite-preview should use thinkingLevel, not thinkingBudget" + ) @pytest.mark.usefixtures("fake_provider_credentials") @@ -8855,10 +8629,7 @@ async def test_gemini_context_caching_disabled_flag(sync_mode): "content": [ { "type": "text", - "text": "Here is the full text of a complex legal agreement {}".format( - uuid.uuid4() - ) - * 4000, + "text": "Here is the full text of a complex legal agreement {}".format(uuid.uuid4()) * 4000, "cache_control": {"type": "ephemeral"}, } ], @@ -8894,9 +8665,7 @@ async def test_gemini_context_caching_disabled_flag(sync_mode): else: client = AsyncHTTPHandler(concurrent_limit=1) - with patch.object( - client, "post", side_effect=mock_gemini_request - ) as mock_client: + with patch.object(client, "post", side_effect=mock_gemini_request) as mock_client: try: if sync_mode: response = litellm.completion( @@ -8919,26 +8688,22 @@ async def test_gemini_context_caching_disabled_flag(sync_mode): print(e) # When caching is disabled, should only make 1 call (no separate cache creation call) - assert ( - mock_client.call_count == 1 - ), f"Expected 1 call when caching is disabled, got {mock_client.call_count}" + assert mock_client.call_count == 1, ( + f"Expected 1 call when caching is disabled, got {mock_client.call_count}" + ) first_call_args = mock_client.call_args_list[0].kwargs first_call_positional_args = mock_client.call_args_list[0].args print(f"first_call_args with caching disabled: {first_call_args}") - print( - f"first_call_positional_args with caching disabled: {first_call_positional_args}" - ) + print(f"first_call_positional_args with caching disabled: {first_call_positional_args}") # Assert that cachedContents is NOT in the URL when caching is disabled url = first_call_args.get( "url", first_call_positional_args[0] if first_call_positional_args else "", ) - assert ( - "cachedContents" not in url - ), "cachedContents should not be in URL when caching is disabled" + assert "cachedContents" not in url, "cachedContents should not be in URL when caching is disabled" finally: # Restore original flag value @@ -9031,9 +8796,7 @@ def test_gemini_function_call_parameter_in_messages(): "contents": [ { "role": "user", - "parts": [ - {"text": "search for weather in boston (use `search`)"} - ], + "parts": [{"text": "search for weather in boston (use `search`)"}], }, { "role": "model", @@ -9052,17 +8815,13 @@ def test_gemini_function_call_parameter_in_messages(): { "function_response": { "name": "search", - "response": { - "content": "The current weather in Boston is 22°F." - }, + "response": {"content": "The current weather in Boston is 22°F."}, } } ], }, ], - "system_instruction": { - "parts": [{"text": "Use search for most queries."}] - }, + "system_instruction": {"parts": [{"text": "Use search for most queries."}]}, "tools": [ { "function_declarations": [ @@ -9137,9 +8896,7 @@ def test_gemini_function_call_parameter_in_messages_2(): { "function_response": { "name": "search", - "response": { - "content": "The weather in Boston is 100 degrees." - }, + "response": {"content": "The weather in Boston is 100 degrees."}, } } ], @@ -9517,9 +9274,7 @@ def test_vertex_ai_cache_creation_kwarg_cannot_change_usage_or_cost() -> None: mock_response.status_code = 200 mock_response.headers = {"Content-Type": "application/json"} mock_response.json.return_value = { - "candidates": [ - {"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP", "index": 0} - ], + "candidates": [{"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15}, } client = HTTPHandler() @@ -9585,14 +9340,51 @@ def load_vertex_ai_credentials(): # Gemini prebuilt TTS voices per https://ai.google.dev/gemini-api/docs/speech-generation (page dated 2026-10-01, read 2026-10-07) GEMINI_TTS_PREBUILT_VOICES: Final[frozenset[str]] = frozenset( { - "Zephyr", "Puck", "Charon", "Kore", "Fenrir", "Leda", "Orus", "Aoede", "Callirrhoe", "Autonoe", - "Enceladus", "Iapetus", "Umbriel", "Algieba", "Despina", "Erinome", "Algenib", "Rasalgethi", - "Laomedeia", "Achernar", "Alnilam", "Schedar", "Gacrux", "Pulcherrima", "Achird", "Zubenelgenubi", - "Vindemiatrix", "Sadachbia", "Sadaltager", "Sulafat", + "Zephyr", + "Puck", + "Charon", + "Kore", + "Fenrir", + "Leda", + "Orus", + "Aoede", + "Callirrhoe", + "Autonoe", + "Enceladus", + "Iapetus", + "Umbriel", + "Algieba", + "Despina", + "Erinome", + "Algenib", + "Rasalgethi", + "Laomedeia", + "Achernar", + "Alnilam", + "Schedar", + "Gacrux", + "Pulcherrima", + "Achird", + "Zubenelgenubi", + "Vindemiatrix", + "Sadachbia", + "Sadaltager", + "Sulafat", } ) OPENAI_TTS_VOICES: Final[tuple[str, ...]] = ( - "alloy", "ash", "ballad", "cedar", "coral", "echo", "fable", "marin", "onyx", "sage", "shimmer", "verse", + "alloy", + "ash", + "ballad", + "cedar", + "coral", + "echo", + "fable", + "marin", + "onyx", + "sage", + "shimmer", + "verse", ) @@ -9678,3 +9470,289 @@ def test_streaming_usage_reads_cache_creation_tokens_recorded_on_logging_obj(): assert chunk.usage.prompt_tokens == 1536 + 1280 assert chunk.usage.cache_creation_input_tokens == 1280 + + +THOUGHT_SUMMARY_REQUESTS: Final = ( + pytest.param("gemini-2.5-flash", {"reasoning_effort": "minimal"}, id="budget-minimal"), + pytest.param("gemini-2.5-flash", {"reasoning_effort": "low"}, id="budget-low"), + pytest.param("gemini-2.5-flash", {"reasoning_effort": "medium"}, id="budget-medium"), + pytest.param("gemini-2.5-flash", {"reasoning_effort": "high"}, id="budget-high"), + pytest.param("gemini-3.6-flash", {"reasoning_effort": "minimal"}, id="level-minimal"), + pytest.param("gemini-3.8-flash", {"reasoning_effort": "minimal"}, id="level-minimal-mapped-to-low"), + pytest.param("gemini-3.8-flash", {"reasoning_effort": "low"}, id="level-low"), + pytest.param("gemini-3.8-flash", {"reasoning_effort": "medium"}, id="level-medium"), + pytest.param("gemini-3.8-flash", {"reasoning_effort": "high"}, id="level-high"), + pytest.param("gemini-2.5-flash", {"thinking": {"type": "enabled", "budget_tokens": 2048}}, id="budget-thinking"), + pytest.param("gemini-3.8-flash", {"thinking": {"type": "enabled", "budget_tokens": 2048}}, id="level-thinking"), +) + + +async def _sent_thinking_config(model: str, thinking_params: dict[str, object]) -> dict[str, object]: + captured: Final[dict[str, dict]] = {} + + def upstream(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + request=request, + ) + + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream)) + await litellm.acompletion( + model=f"gemini/{model}", + messages=[{"role": "user", "content": "hi"}], + api_key="fake-gemini-key", + client=client, + **thinking_params, + ) + return captured["body"]["generationConfig"]["thinkingConfig"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("model", "thinking_params"), THOUGHT_SUMMARY_REQUESTS) +async def test_gemini_thought_summaries_off_keeps_thinking_config_otherwise_unchanged( + model: str, + thinking_params: dict[str, object], + local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + summaries_on: Final = await _sent_thinking_config(model, thinking_params) + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False) + summaries_off: Final = await _sent_thinking_config(model, thinking_params) + + assert summaries_on["includeThoughts"] is True + assert summaries_off == {**summaries_on, "includeThoughts": False} + + +PER_REQUEST_THOUGHT_SUMMARY_ASKS: Final = ( + pytest.param( + "gemini-3.8-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "summarized"}}, + False, + True, + id="level-display-summarized-beats-off", + ), + pytest.param( + "gemini-3.8-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "omitted"}}, + True, + True, + id="level-display-omitted-leaves-the-default-on", + ), + pytest.param( + "gemini-2.5-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "summarized"}}, + False, + True, + id="budget-display-summarized-beats-off", + ), + pytest.param( + "gemini-2.5-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "omitted"}}, + True, + True, + id="budget-display-omitted-leaves-the-default-on", + ), + pytest.param( + "gemini-3.8-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "updates"}}, + False, + True, + id="level-display-updates-beats-off", + ), + pytest.param( + "gemini-3.8-flash", + {"reasoning_effort": {"effort": "medium", "summary": "auto"}}, + False, + True, + id="level-effort-summary-beats-off", + ), + pytest.param( + "gemini-2.5-flash", + {"reasoning_effort": {"effort": "medium", "summary": "detailed"}}, + False, + True, + id="budget-effort-summary-beats-off", + ), + pytest.param( + "gemini-3.8-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "summary": "auto"}}, + False, + True, + id="level-thinking-summary-beats-off", + ), + pytest.param( + "gemini-2.5-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "summary": "detailed"}}, + False, + True, + id="budget-thinking-summary-beats-off", + ), + pytest.param( + "gemini-3.8-flash", + {"reasoning_effort": "medium", "reasoning_summary": "auto"}, + False, + True, + id="level-reasoning-summary-alias-beats-off", + ), + pytest.param( + "gemini-2.5-flash", + {"reasoning_effort": "medium", "reasoningSummary": "detailed"}, + False, + True, + id="budget-reasoning-summary-camel-alias-beats-off", + ), + pytest.param( + "gemini-3.8-flash", + {"reasoning_effort": "none", "reasoning_summary": "auto"}, + True, + False, + id="effort-none-ignores-summary-alias", + ), + pytest.param( + "gemini-3.8-flash", + {"reasoning_effort": {"effort": "medium"}}, + False, + False, + id="effort-without-summary-keeps-off", + ), + pytest.param( + "gemini-3.8-flash", + {"reasoning_effort": {"effort": "medium"}, "reasoning_summary": "auto"}, + False, + True, + id="effort-dict-takes-the-summary-alias", + ), + pytest.param( + "gemini-3.8-flash", + {"reasoning_effort": {"effort": "none", "summary": "auto"}}, + True, + False, + id="effort-none-ignores-summary", + ), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("model", "thinking_params", "gateway_default", "expected"), PER_REQUEST_THOUGHT_SUMMARY_ASKS) +async def test_gemini_per_request_thought_summary_ask_beats_gateway_default( + model: str, + thinking_params: dict[str, object], + gateway_default: bool, + expected: bool, + local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", gateway_default) + + sent: Final = await _sent_thinking_config(model, thinking_params) + + assert sent["includeThoughts"] is expected + + +def _includes_thoughts_client(sent: list[bool]) -> AsyncHTTPHandler: + def upstream(request: httpx.Request) -> httpx.Response: + sent.append(json.loads(request.content)["generationConfig"]["thinkingConfig"]["includeThoughts"]) + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + request=request, + ) + + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream)) + return client + + +async def _chat_summary_ask(client: AsyncHTTPHandler, asks: bool) -> None: + await litellm.acompletion( + model="gemini/gemini-3.8-flash", + messages=[{"role": "user", "content": "hi"}], + api_key="fake-gemini-key", + client=client, + reasoning_effort="medium", + **({"reasoning_summary": "auto"} if asks else {}), + ) + + +async def _messages_summary_ask(client: AsyncHTTPHandler, asks: bool) -> None: + await litellm.anthropic.messages.acreate( + model="gemini/gemini-3.8-flash", + messages=[{"role": "user", "content": "hi"}], + max_tokens=64, + api_key="fake-gemini-key", + client=client, + thinking={"type": "enabled", "budget_tokens": 2048, **({"display": "summarized"} if asks else {})}, + ) + + +async def _responses_summary_ask(client: AsyncHTTPHandler, asks: bool) -> None: + await litellm.aresponses( + model="gemini/gemini-3.8-flash", + input="hi", + api_key="fake-gemini-key", + client=client, + reasoning={"effort": "medium", **({"summary": "auto"} if asks else {})}, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "send", + (_chat_summary_ask, _messages_summary_ask, _responses_summary_ask), + ids=("chat-reasoning-summary-alias", "messages-display-summarized", "responses-reasoning-summary"), +) +async def test_gemini_summary_ask_turns_thoughts_on_and_logs_the_same_reasoning_effort( + send: Callable[[AsyncHTTPHandler, bool], Awaitable[None]], + local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + logged: Final[list[object]] = [] + + class ReasoningEffortLogger(CustomLogger): + async def async_log_success_event( + self, + kwargs: Mapping[str, Mapping[str, Mapping[str, object]]], + response_obj: object, + start_time: object, + end_time: object, + ) -> None: + logged.append(kwargs["standard_logging_object"]["model_parameters"].get("reasoning_effort")) + + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False) + monkeypatch.setattr(litellm, "callbacks", [ReasoningEffortLogger()]) + sent: Final[list[bool]] = [] + client: Final = _includes_thoughts_client(sent) + + await send(client, False) + await send(client, True) + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + + assert sent == [False, True] + assert len(logged) == 2 + assert isinstance(logged[0], str) + assert logged[1] == logged[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("gateway_default", "expected"), + (("false", False), (" Off ", False), ("0", False), ("true", True), (None, True), ("not-a-flag", True)), +) +async def test_gemini_sends_a_boolean_for_a_gateway_default_written_as_text_or_null( + gateway_default: object, expected: bool, local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", gateway_default) + + sent: Final = await _sent_thinking_config("gemini-3.8-flash", {"reasoning_effort": "high"}) + + assert sent["includeThoughts"] is expected 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.py b/tests/unit/llms/vertex_ai/test_vertex.py index 044a07859f9..04cf391c700 100644 --- a/tests/unit/llms/vertex_ai/test_vertex.py +++ b/tests/unit/llms/vertex_ai/test_vertex.py @@ -1814,8 +1814,8 @@ def test_gemini_nullable_object_tool_schema_httpx( "type": "object", "required": ["source"], "properties": {"source": {"type": "string"}}, + "nullable": True, }, - {"type": "null"}, ] 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 304424b2acc..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 @@ -1,5 +1,7 @@ +import copy import time from collections.abc import Iterator +from contextlib import contextmanager from typing import Final from unittest.mock import AsyncMock, patch @@ -7,11 +9,13 @@ import httpx import pytest from pydantic import JsonValue, TypeAdapter +import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler from litellm.llms.vertex_ai.common_utils import ( VertexAITokenCounter, + build_vertex_schema, convert_anyof_null_to_nullable, convert_vertex_datetime_to_openai_datetime, get_vertex_location_from_url, @@ -24,16 +28,21 @@ 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 CountContextManagement, TokenCountMessagesFormat @pytest.mark.asyncio -async def test_vertex_gemini_rejects_non_text_system_blocks_before_request() -> None: +@pytest.mark.parametrize("messages_format", ("openai", "anthropic")) +async def test_vertex_gemini_rejects_non_text_system_blocks_before_request( + messages_format: TokenCountMessagesFormat, +) -> None: result: Final = await VertexAITokenCounter().count_tokens( model_to_use="gemini-4-argon", messages=[{"role": "user", "content": "Count this"}], contents=None, request_model="vertex_ai/gemini-4-argon", system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}], + messages_format=messages_format, ) assert result is not None @@ -42,6 +51,144 @@ async def test_vertex_gemini_rejects_non_text_system_blocks_before_request() -> assert "text blocks" in (result.error_message or "") +class _VertexCountTokensPeer: + def __init__(self, status_code: int) -> None: + self.status_code: Final = status_code + self.requests: tuple[httpx.Request, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + if self.status_code == 200: + 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"} + }, + request=request, + ) + + +@contextmanager +def _vertex_count_tokens_peer(status_code: int) -> Iterator[_VertexCountTokensPeer]: + peer: Final = _VertexCountTokensPeer(status_code) + with ( + patch.object( + gemini_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ), + patch.object( + vertex_count_tokens_handler.VertexAITokenCounter, + "_ensure_access_token_async", + new=AsyncMock(return_value=("fake-token", "test-project")), + ), + ): + yield peer + + +_VERTEX_DEPLOYMENT: Final = {"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-central1"}} + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_answers_with_the_status_google_rejected_it_with() -> None: + with _vertex_count_tokens_peer(400) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment=_VERTEX_DEPLOYMENT, + request_model="vertex_ai/gemini-2.5-flash", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "Invalid function name" in (result.error_message or "") + assert len(peer.requests) == 1 + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_rejects_an_image_url_it_refuses_to_fetch_without_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + with _vertex_count_tokens_peer(200) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}], + } + ], + contents=None, + deployment=_VERTEX_DEPLOYMENT, + request_model="vertex_ai/gemini-2.5-flash", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "http://127.0.0.1:9/image.png" in (result.error_message or "") + 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 @@ -145,6 +292,147 @@ def test_nested_anyof_conversion(): assert schema == expected +def test_anyof_conversion_with_multiple_null_branches(): + """Adjacent null branches must all be removed.""" + schema: Final = { + "type": "object", + "properties": { + "example": { + "anyOf": [ + {"type": "null"}, + {"type": "null"}, + {"type": "string"}, + ] + } + }, + } + + convert_anyof_null_to_nullable(schema) + + expected: Final = { + "type": "object", + "properties": {"example": {"anyOf": [{"type": "string", "nullable": True}]}}, + } + assert schema == expected + + +def test_anyof_conversion_with_null_before_empty_object(): + """An empty-object branch following a null branch still gets a type.""" + schema: Final = { + "type": "object", + "properties": {"example": {"anyOf": [{"type": "null"}, {}]}}, + } + + convert_anyof_null_to_nullable(schema) + + expected: Final = { + "type": "object", + "properties": {"example": {"anyOf": [{"type": "object", "nullable": True}]}}, + } + assert schema == expected + + +@pytest.mark.parametrize("null_count", [1, 2, 4]) +def test_anyof_conversion_with_only_null_branches_keeps_one_nullable_null(null_count: int): + """Vertex accepts a NULL type, so a null-only union is forwarded as one nullable null branch.""" + schema: Final = { + "type": "object", + "properties": {"example": {"anyOf": [{"type": "null"} for _ in range(null_count)]}}, + } + + convert_anyof_null_to_nullable(schema) + + expected: Final = { + "type": "object", + "properties": {"example": {"anyOf": [{"type": "null", "nullable": True}]}}, + } + assert schema == expected + + +def test_anyof_conversion_leaves_empty_anyof_for_the_provider_to_judge(): + """An empty union has no null branch to strip, so it is forwarded unchanged.""" + schema: Final = {"type": "object", "properties": {"example": {"anyOf": []}}} + + convert_anyof_null_to_nullable(schema) + + assert schema == {"type": "object", "properties": {"example": {"anyOf": []}}} + + +@pytest.mark.parametrize("malformed", ["string", {"type": "string"}], ids=["string", "object"]) +def test_anyof_conversion_forwards_a_non_list_anyof_unchanged(malformed: JsonValue): + """A malformed union is the provider's to judge, as it was before null branches were stripped in one pass.""" + schema: Final = {"type": "object", "properties": {"example": {"anyOf": copy.deepcopy(malformed)}}} + + convert_anyof_null_to_nullable(schema) + + assert schema == {"type": "object", "properties": {"example": {"anyOf": malformed}}} + + +@pytest.mark.parametrize( + ("branches", "expected"), + [ + ([{"type": "null"}, 5], [5]), + ([True, {"type": "null"}], [True]), + ([{"type": "null"}, 5, {"type": "string"}], [5, {"type": "string", "nullable": True}]), + ], + ids=["null-int", "true-null", "null-int-string"], +) +def test_anyof_conversion_forwards_a_non_object_branch_for_the_provider_to_judge( + branches: list[JsonValue], expected: list[JsonValue] +): + schema: Final = {"type": "object", "properties": {"example": {"anyOf": branches}}} + + convert_anyof_null_to_nullable(schema) + + assert schema == {"type": "object", "properties": {"example": {"anyOf": expected}}} + + + +@pytest.mark.parametrize( + ("field", "expected"), + [ + ({"type": ["integer", "null"]}, {"anyOf": [{"type": "integer", "nullable": True}]}), + ( + { + "anyOf": [ + {"type": "object", "properties": {"x": {"anyOf": [{"type": "integer"}, {"type": "null"}]}}}, + {"type": "null"}, + ] + }, + { + "anyOf": [ + { + "type": "object", + "properties": {"x": {"anyOf": [{"type": "integer", "nullable": True}]}}, + "nullable": True, + } + ] + }, + ), + ( + {"type": ["object", "null"], "properties": {"x": {"type": ["string", "null"]}}}, + { + "anyOf": [ + { + "type": "object", + "properties": {"x": {"anyOf": [{"type": "string", "nullable": True}]}}, + "nullable": True, + } + ] + }, + ), + ( + {"type": "array", "items": {"type": ["number", "null"]}}, + {"type": "array", "items": {"anyOf": [{"type": "number", "nullable": True}]}}, + ), + ], + ids=["type-array", "null-inside-a-union-branch", "nullable-object-with-nullable-property", "nullable-items"], +) +def test_vertex_schema_folds_every_null_branch_into_nullable(field: JsonValue, expected: JsonValue): + schema: Final = build_vertex_schema({"type": "object", "properties": {"value": field}}) + + assert schema["properties"]["value"] == expected + def test_anyof_with_excessive_nesting(): """Test conversion with excessive nesting > max levels +1 deep.""" # generate a schema with excessive nesting diff --git a/tests/unit/llms/volcengine/test_volcengine_video_transformation.py b/tests/unit/llms/volcengine/test_volcengine_video_transformation.py new file mode 100644 index 00000000000..1fcf49eeb13 --- /dev/null +++ b/tests/unit/llms/volcengine/test_volcengine_video_transformation.py @@ -0,0 +1,1131 @@ +import base64 +import io +import json +from pathlib import Path +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest +import respx + +import litellm +from litellm import Router +from litellm.exceptions import UnsupportedParamsError +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.openai.cost_calculation import video_generation_cost +from litellm.llms.volcengine.videos.transformation import VolcEngineVideoConfig, VolcEngineVideoError +from litellm.types.router import GenericLiteLLMParams +from litellm.types.videos.utils import decode_video_id_with_provider, encode_video_id_with_provider +from litellm.utils import ProviderConfigManager +from litellm.videos.main import avideo_content, avideo_generation + +PNG_BYTES: Final = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" +) +PNG_DATA_URL: Final = f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}" + +API_BASE: Final = "https://ark.cn-beijing.volces.com" +TASKS_URL: Final = f"{API_BASE}/api/v3/contents/generations/tasks" +SEEDANCE_2_5: Final = "doubao-seedance-2-5-260628" +SEEDANCE_2_0: Final = "doubao-seedance-2-0-260128" +SEEDANCE_2_0_FAST: Final = "doubao-seedance-2-0-fast-260128" +TASK_ID: Final = "cgt-2026-abc" + + +def _response(payload: object, status_code: int = 200) -> httpx.Response: + return httpx.Response(status_code, json=payload, request=httpx.Request("GET", TASKS_URL)) + + +def _text_response(text: str, status_code: int = 200) -> httpx.Response: + return httpx.Response(status_code, text=text, request=httpx.Request("GET", TASKS_URL)) + + +def _create_body( + params: dict[str, object], model: str = SEEDANCE_2_5, drop_params: bool | None = None +) -> dict[str, object]: + body, files, url = VolcEngineVideoConfig().transform_video_create_request( + model=model, + prompt="a cat on a roof", + api_base=API_BASE, + video_create_optional_request_params=params, + litellm_params=GenericLiteLLMParams(drop_params=drop_params), + headers={}, + ) + assert (files, url) == ((), TASKS_URL) + return body + + +def _created(request_data: dict[str, object], model: str = SEEDANCE_2_5, logging_obj: object = None): + return VolcEngineVideoConfig().transform_video_create_response( + model=model, + raw_response=_response({"id": TASK_ID}), + logging_obj=logging_obj, # pyright: ignore[reportArgumentType] # the transform only reads litellm_params off it + custom_llm_provider="volcengine", + request_data=request_data, + ) + + +def _polled(payload: object, logging_obj: object = None): + return VolcEngineVideoConfig().transform_video_status_retrieve_response( + raw_response=_response(payload), + logging_obj=logging_obj, # pyright: ignore[reportArgumentType] # the transform only reads litellm_params off it + custom_llm_provider="volcengine", + ) + + +def _task(**fields: object) -> dict[str, object]: + return {"id": TASK_ID, "model": SEEDANCE_2_5, "status": "succeeded", **fields} + + +@pytest.mark.parametrize( + "input_reference", + [ + PNG_BYTES, + io.BytesIO(PNG_BYTES), + ("frame.png", PNG_BYTES, "image/png"), + ("frame.png", io.BytesIO(PNG_BYTES)), + ], + ids=["bytes", "file", "tuple", "tuple-with-file"], +) +def test_a_file_input_reference_becomes_a_data_uri_first_frame(input_reference: object): + body = _create_body({"input_reference": input_reference, "duration": 4}) + + assert body == { + "model": SEEDANCE_2_5, + "content": [ + {"type": "text", "text": "a cat on a roof"}, + {"type": "image_url", "image_url": {"url": PNG_DATA_URL}, "role": "first_frame"}, + ], + "resolution": "720p", + "duration": 4, + } + + +def test_a_path_input_reference_is_read_from_disk(tmp_path: Path): + frame = tmp_path / "frame.png" + frame.write_bytes(PNG_BYTES) + + body = _create_body({"input_reference": frame, "duration": 4}) + + assert body["content"] == [ + {"type": "text", "text": "a cat on a roof"}, + {"type": "image_url", "image_url": {"url": PNG_DATA_URL}, "role": "first_frame"}, + ] + + +def test_a_url_input_reference_is_passed_as_is_and_never_opened_as_a_path(): + body = _create_body({"input_reference": "https://x/first.png", "duration": 4}) + + assert body["content"] == [ + {"type": "text", "text": "a cat on a roof"}, + {"type": "image_url", "image_url": {"url": "https://x/first.png"}, "role": "first_frame"}, + ] + + +def test_openais_object_input_reference_sends_its_image_url(): + body = _create_body({"input_reference": {"image_url": "https://x/first.png"}, "duration": 4}) + + assert body["content"] == [ + {"type": "text", "text": "a cat on a roof"}, + {"type": "image_url", "image_url": {"url": "https://x/first.png"}, "role": "first_frame"}, + ] + + +def test_a_file_id_input_reference_is_a_bad_request_unless_drop_params(): + with pytest.raises(UnsupportedParamsError, match="file_id is not supported"): + _ = _create_body({"input_reference": {"file_id": "file-abc"}, "duration": 4}) + + body = _create_body({"input_reference": {"file_id": "file-abc"}, "duration": 4}, drop_params=True) + + assert body["content"] == [{"type": "text", "text": "a cat on a roof"}] + + +def test_a_caller_content_array_wins_over_input_reference_and_the_prompt_leads(): + body = _create_body( + { + "content": [{"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"}], + "input_reference": "https://x/ignored.png", + "duration": 4, + } + ) + + assert body["content"] == [ + {"type": "text", "text": "a cat on a roof"}, + {"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"}, + ] + + +IMAGE_ITEM: Final = {"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"} + + +def test_a_text_item_in_content_is_rejected_because_input_guardrails_only_scan_the_prompt(): + with pytest.raises(UnsupportedParamsError, match="put the text in prompt"): + _ = _create_body({"content": [{"type": "text", "text": "text the guardrail never saw"}], "duration": 4}) + + +def test_a_text_item_in_content_is_dropped_under_drop_params_so_ark_never_sees_unscanned_text(): + body = _create_body( + {"content": [{"type": "text", "text": "text the guardrail never saw"}, IMAGE_ITEM], "duration": 4}, + drop_params=True, + ) + + assert body["content"] == [{"type": "text", "text": "a cat on a roof"}, IMAGE_ITEM] + + +@pytest.mark.parametrize( + "item", + [ + {"type": "video_url", "video_url": {"url": "https://x/v.mp4"}, "role": "reference_video"}, + {"type": "video_url", "video_url": {"url": "https://x/v.mp4"}}, + {"type": "image_url", "image_url": {"url": "https://x/v.png"}, "role": "reference_video"}, + ], + ids=["reference-video", "untagged-video", "mistagged-role"], +) +def test_a_video_input_is_rejected_because_its_cost_is_not_known_at_create(item: dict[str, object]): + with pytest.raises(UnsupportedParamsError, match="Video input"): + _ = _create_body({"content": [item], "duration": 4}) + + body = _create_body({"content": [item, IMAGE_ITEM], "duration": 4}, drop_params=True) + + assert body["content"] == [{"type": "text", "text": "a cat on a roof"}, IMAGE_ITEM] + + +@pytest.mark.parametrize( + "content", + ["a string", [{"type": "image_url", "role": 7}], [{"type": 7}]], + ids=["not-an-array", "non-string-role", "non-string-type"], +) +def test_a_content_value_that_is_not_an_array_of_items_is_rejected_unless_drop_params(content: object): + with pytest.raises(UnsupportedParamsError, match="content must be an array"): + _ = _create_body({"content": content, "duration": 4}) + + body = _create_body({"content": content, "duration": 4}, drop_params=True) + + assert body["content"] == [{"type": "text", "text": "a cat on a roof"}] + + +def test_an_omitted_duration_and_tier_are_sent_as_openais_four_seconds_at_720p_so_ark_cannot_pick_what_is_billed(): + body = _create_body({}) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "a cat on a roof"}], + "resolution": "720p", + "duration": 4, + } + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_an_endpoint_id_priced_through_base_model_gets_the_base_models_default_tier(): + body, _, _ = VolcEngineVideoConfig().transform_video_create_request( + model="ep-20261015-abc", + prompt="a cat on a roof", + api_base=API_BASE, + video_create_optional_request_params={"duration": 4}, + litellm_params=GenericLiteLLMParams(model_info={"base_model": f"volcengine/{SEEDANCE_2_5}"}), + headers={}, + ) + + assert (body["model"], body["resolution"]) == ("ep-20261015-abc", "720p") + + +@pytest.mark.parametrize("field", ["return_last_frame", "draft", "frame_count"]) +def test_a_field_litellm_does_not_send_to_ark_is_rejected(field: str): + with pytest.raises(UnsupportedParamsError, match=f"does not send {field} to Ark"): + _ = _create_body({"duration": 4, field: True}) + + +def test_a_field_litellm_does_not_send_to_ark_is_dropped_under_drop_params(): + body = _create_body({"duration": 4, "return_last_frame": True, "seed": 7}, drop_params=True) + + assert "return_last_frame" not in body + assert body["seed"] == 7 + + +@pytest.mark.parametrize("duration", [-1, "-1", 0, "abc", 4.5]) +def test_a_duration_ark_would_pick_or_that_is_not_whole_seconds_is_rejected(duration: object): + with pytest.raises(UnsupportedParamsError, match="duration must be a positive whole number"): + _ = _create_body({"duration": duration}) + + assert _create_body({"duration": duration}, drop_params=True)["duration"] == 4 + + +def test_frames_is_rejected_because_it_bypasses_the_billed_duration(): + with pytest.raises(UnsupportedParamsError, match="frames"): + _ = _create_body({"frames": 121, "duration": 6}) + + body = _create_body({"frames": 121, "duration": 6}, drop_params=True) + + assert ("frames" not in body, body["duration"]) == (True, 6) + + +def test_ark_fields_reach_the_body_and_false_flags_are_kept(): + body = _create_body( + { + "duration": 4, + "resolution": "480p", + "ratio": "9:16", + "seed": 7, + "watermark": False, + "generate_audio": False, + "service_tier": "flex", + } + ) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "a cat on a roof"}], + "resolution": "480p", + "ratio": "9:16", + "seed": 7, + "watermark": False, + "generate_audio": False, + "service_tier": "flex", + "duration": 4, + } + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_a_resolution_the_model_has_no_price_for_is_rejected(): + with pytest.raises(UnsupportedParamsError, match="not offered at resolution '1080p', only at 480p, 720p"): + _ = _create_body({"duration": 4, "resolution": "1080p"}, model=SEEDANCE_2_0_FAST) + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_an_unpriced_resolution_is_dropped_under_drop_params(): + body = _create_body( + {"duration": 4, "resolution": "1080p", "ratio": "16:9"}, model=SEEDANCE_2_0_FAST, drop_params=True + ) + + assert body == { + "model": SEEDANCE_2_0_FAST, + "content": [{"type": "text", "text": "a cat on a roof"}], + "resolution": "720p", + "ratio": "16:9", + "duration": 4, + } + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_a_model_the_router_registered_without_prices_keeps_its_resolution(): + endpoint_id: Final = "ep-20261015-unpriced" + _ = Router( + model_list=[ + {"model_name": "ark", "litellm_params": {"model": f"volcengine/{endpoint_id}", "api_key": "sk-test"}} + ] + ) + + body = _create_body({"duration": 4, "resolution": "480p"}, model=endpoint_id) + body_without_a_tier = _create_body({"duration": 4}, model=endpoint_id) + + assert body["resolution"] == "480p" + assert "resolution" not in body_without_a_tier + + +def test_a_mistyped_ark_field_is_rejected(): + with pytest.raises(UnsupportedParamsError, match="invalid seed"): + _ = _create_body({"duration": 4, "seed": "abc", "watermark": False}) + + +def test_a_mistyped_ark_field_is_dropped_under_drop_params(): + body = _create_body({"duration": 4, "seed": "abc", "watermark": False}, drop_params=True) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "a cat on a roof"}], + "resolution": "720p", + "watermark": False, + "duration": 4, + } + + +@pytest.mark.parametrize( + "size,expected", + [ + ("1920x1080", {"ratio": "16:9", "resolution": "1080p"}), + ("854x480", {"ratio": "16:9", "resolution": "480p"}), + ("864x496", {"ratio": "16:9", "resolution": "480p"}), + ("720x1280", {"ratio": "9:16", "resolution": "720p"}), + ("1024x768", {"ratio": "4:3", "resolution": "720p"}), + ("640x640", {"ratio": "1:1", "resolution": "480p"}), + ("960x960", {"ratio": "1:1", "resolution": "720p"}), + ("1440x1440", {"ratio": "1:1", "resolution": "1080p"}), + ("1112x834", {"ratio": "4:3", "resolution": "720p"}), + ("1470x630", {"ratio": "21:9", "resolution": "720p"}), + ("3840x2160", {"ratio": "16:9", "resolution": "4k"}), + ], +) +def test_size_maps_onto_an_ark_ratio_and_tier(size: str, expected: dict[str, str]): + assert VolcEngineVideoConfig().map_openai_params({"size": size}, SEEDANCE_2_5, False) == expected + + +def test_explicit_ratio_and_resolution_beat_the_size_derived_ones(): + mapped = VolcEngineVideoConfig().map_openai_params( + {"size": "1280x720", "seconds": "4", "ratio": "1:1", "resolution": "480p", "parameters": {"seed": 3}}, + SEEDANCE_2_5, + False, + ) + + assert mapped == {"ratio": "1:1", "resolution": "480p", "duration": 4, "seed": 3} + + +@pytest.mark.parametrize( + "params,message", + [ + ({"seconds": "4.5"}, "seconds must be a whole number"), + ({"size": "big"}, "size must be WIDTHxHEIGHT"), + ({"size": "16:9"}, "size must be WIDTHxHEIGHT"), + ({"parameters": "seed=3"}, "parameters must be an object"), + ({"characters": [{"id": "c"}]}, "'characters' is not supported"), + ({"image": "https://x/a.png"}, "'image' is not supported"), + ], +) +def test_a_value_ark_cannot_take_raises_unless_drop_params(params: dict[str, object], message: str): + with pytest.raises(UnsupportedParamsError, match=message): + _ = VolcEngineVideoConfig().map_openai_params(params, SEEDANCE_2_5, False) + + assert VolcEngineVideoConfig().map_openai_params(params, SEEDANCE_2_5, True) == {} + + +def test_request_metadata_is_not_mapped_into_the_body_and_user_becomes_the_safety_identifier(): + mapped = VolcEngineVideoConfig().map_openai_params( + {"user": "u", "extra_headers": {"x": "y"}, "seconds": 4}, SEEDANCE_2_5, False + ) + + assert mapped == {"safety_identifier": "u", "duration": 4} + + +def test_an_explicit_safety_identifier_wins_over_user(): + mapped = VolcEngineVideoConfig().map_openai_params( + {"user": "u", "safety_identifier": "s", "seconds": 4}, SEEDANCE_2_5, False + ) + + assert mapped["safety_identifier"] == "s" + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + "model,request_data,seconds,size,usage", + [ + ( + SEEDANCE_2_5, + {"duration": 4, "resolution": "480p", "ratio": "16:9"}, + "4", + "854x480", + {"duration_seconds": 4.0, "video_resolution": "480p"}, + ), + ( + SEEDANCE_2_0, + {"duration": 5, "resolution": "480p", "ratio": "16:9"}, + "5", + "864x496", + {"duration_seconds": 5.0, "video_resolution": "480p"}, + ), + (SEEDANCE_2_5, {"duration": 5}, "5", None, {"duration_seconds": 5.0}), + ], + ids=["2.5-480p", "2.0-480p-renders-its-own-frame", "no-geometry"], +) +def test_the_create_response_reports_what_the_task_bills( + model: str, request_data: dict[str, object], seconds: str, size: str | None, usage: dict[str, object] +): + video = _created(request_data, model=model) + + assert video.model_dump(include={"status", "seconds", "size", "usage", "model"}) == { + "status": "queued", + "seconds": seconds, + "size": size, + "usage": usage, + "model": model, + } + + +@pytest.mark.parametrize( + "usage,expected_cost", + [ + ({"duration_seconds": 4.0, "video_resolution": "480p"}, 4 * 0.1), + ({"duration_seconds": 5.0, "video_resolution": "720p"}, 5 * 0.2), + ({"duration_seconds": 5.0}, 5 * 0.2), + ], +) +def test_the_create_usage_bills_its_duration_at_its_tier_rate(usage: dict[str, object], expected_cost: float): + model_info = { + "output_cost_per_second": 0.2, + "output_cost_per_second_480p": 0.1, + "output_cost_per_second_720p": 0.2, + } + video = _created({"duration": int(usage["duration_seconds"]), "resolution": usage.get("video_resolution")}) + + cost = video_generation_cost( + model=SEEDANCE_2_5, + duration_seconds=video.usage["duration_seconds"], + model_info=model_info, # pyright: ignore[reportArgumentType] # a hand-written rate card stands in for ModelInfo + video_resolution=video.usage.get("video_resolution"), + ) + + assert cost == pytest.approx(expected_cost) + + +def test_the_created_id_carries_the_requested_and_served_model_groups(): + logging_obj = Mock( + litellm_params={"litellm_metadata": {"model_group": "seedance", "deployment_model_name": "seedance-account-b"}} + ) + + video = _created({"duration": 5}, logging_obj=logging_obj) + + assert decode_video_id_with_provider(video.id) == { + "custom_llm_provider": "volcengine", + "model_id": "seedance", + "video_id": TASK_ID, + "served_model": "seedance-account-b", + } + + +def test_a_create_without_router_metadata_carries_the_model(): + video = _created({"duration": 5}, logging_obj=Mock(litellm_params={})) + + assert decode_video_id_with_provider(video.id) == { + "custom_llm_provider": "volcengine", + "model_id": SEEDANCE_2_5, + "video_id": TASK_ID, + } + + +@pytest.mark.parametrize( + "raw_response,status_code,message", + [ + (_response({"error": {"code": "InvalidParameter", "message": "bad ratio"}}), 400, "bad ratio"), + (_response({"error": {"code": "Forbidden", "message": "no access"}}, 403), 403, "no access"), + (_text_response("bad gateway"), 502, "Unexpected Volcengine video response"), + ], + ids=["error-on-200", "error-status-kept", "not-json"], +) +def test_a_create_response_without_a_task_id_raises_instead_of_queueing_an_empty_video( + raw_response: httpx.Response, status_code: int, message: str +): + with pytest.raises(VolcEngineVideoError, match=message) as raised: + _ = VolcEngineVideoConfig().transform_video_create_response( + model=SEEDANCE_2_5, + raw_response=raw_response, + logging_obj=None, # pyright: ignore[reportArgumentType] # the transform only reads litellm_params off it + custom_llm_provider="volcengine", + request_data={"duration": 5}, + ) + + assert raised.value.status_code == status_code + + +@pytest.mark.parametrize( + "ark_status,expected", + [ + ("queued", "queued"), + ("running", "in_progress"), + ("succeeded", "completed"), + ("failed", "failed"), + ("cancelled", "cancelled"), + ("expired", "failed"), + ], +) +def test_the_ark_status_maps_onto_the_openai_vocabulary(ark_status: str, expected: str): + assert _polled(_task(status=ark_status)).status == expected + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_a_completed_task_reports_its_delivered_geometry_duration_and_tokens(): + video = _polled( + _task( + created_at=1791989938, + updated_at=1791990031, + resolution="480p", + ratio="16:9", + duration=4, + usage={"completion_tokens": 40594, "total_tokens": 40594}, + model=SEEDANCE_2_0, + content={"video_url": "https://tos/v.mp4"}, + ) + ) + + assert video.model_dump(exclude={"id"}) | {"id": decode_video_id_with_provider(video.id)["video_id"]} == { + "id": TASK_ID, + "object": "video", + "status": "completed", + "created_at": 1791989938, + "completed_at": 1791990031, + "expires_at": None, + "error": None, + "progress": None, + "remixed_from_video_id": None, + "seconds": "4", + "size": "864x496", + "model": SEEDANCE_2_0, + "usage": {"duration_seconds": 4.0, "video_resolution": "480p", "total_tokens": 40594}, + } + + +@pytest.mark.parametrize( + "resolution,ratio,expected", + [ + ("720p", "9:16", "720x1280"), + ("1080p", "16:9", "1920x1080"), + ("480p", "adaptive", None), + ], +) +def test_a_frame_the_cost_map_does_not_list_is_scaled_from_the_tier(resolution: str, ratio: str, expected: str | None): + assert _polled(_task(resolution=resolution, ratio=ratio, model="doubao-seedance-unlisted")).size == expected + + +def test_a_failed_task_carries_its_error_block(): + video = _polled(_task(status="failed", error={"code": "OutputVideoSensitiveContentDetected", "message": "flagged"})) + + assert (video.status, video.error, video.completed_at) == ( + "failed", + {"code": "OutputVideoSensitiveContentDetected", "message": "flagged"}, + None, + ) + + +@pytest.mark.parametrize( + "payload,message", + [ + (_task(status="paused"), "Unknown Volcengine task status: paused"), + ({"status": "running"}, "Unexpected Volcengine video response"), + ({"id": TASK_ID}, "Unexpected Volcengine video response"), + ], + ids=["unmapped-status", "no-id", "no-status"], +) +def test_a_poll_litellm_cannot_read_raises_instead_of_returning_a_job_to_poll_forever(payload: object, message: str): + with pytest.raises(VolcEngineVideoError, match=message) as raised: + _ = _polled(payload) + + assert raised.value.status_code == 502 + + +def test_a_poll_keeps_the_ids_requested_and_served_model_when_the_router_served_it(): + video_id = encode_video_id_with_provider(TASK_ID, "volcengine", "seedance", served_model="seedance-account-b") + logging_obj = Mock( + litellm_params={ + "video_id": video_id, + "litellm_metadata": {"model_group": "seedance", "deployment_model_name": "seedance-account-b"}, + } + ) + + video = _polled(_task(status="running"), logging_obj=logging_obj) + + assert video.id == video_id + assert video.model == "seedance" + + +def test_a_poll_drops_a_served_model_the_router_did_not_route_to(): + edited = encode_video_id_with_provider( + TASK_ID, "volcengine", "seedance-account-a", served_model="seedance-account-b" + ) + logging_obj = Mock(litellm_params={"video_id": edited, "litellm_metadata": {"model_group": "seedance-account-a"}}) + + video = _polled(_task(status="running"), logging_obj=logging_obj) + + assert video.id == encode_video_id_with_provider(TASK_ID, "volcengine", "seedance-account-a") + + +async def _download(task_payload: dict[str, object], video_bytes: bytes = b"mp4") -> bytes: + def respond(request: httpx.Request) -> httpx.Response: + if request.url.host == "tos.example": + return httpx.Response(200, content=video_bytes) + return httpx.Response(200, json=task_payload) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + handler: Final = AsyncHTTPHandler() + await handler.close() + handler.client = http_client + return await avideo_content( + video_id=TASK_ID, + custom_llm_provider="volcengine", + model=f"volcengine/{SEEDANCE_2_5}", + api_key="sk-test", + api_base=API_BASE, + client=handler, + ) + + +@pytest.mark.asyncio +async def test_the_download_fetches_the_signed_url_through_the_callers_client(): + content = await _download(_task(content={"video_url": "https://tos.example/v.mp4"}), video_bytes=b"the-video") + + assert content == b"the-video" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload,message", + [ + (_task(status="queued"), "still processing"), + (_task(status="running"), "still processing"), + (_task(status="failed", error={"message": "flagged"}), "flagged"), + (_task(status="failed"), "ended as failed"), + (_task(status="cancelled"), "ended as cancelled"), + (_task(status="expired"), "Task expired"), + ], +) +async def test_a_download_with_no_video_yet_is_a_bad_request_so_it_is_neither_retried_nor_cools_the_deployment( + payload: dict[str, object], message: str +): + with pytest.raises(litellm.BadRequestError, match=message) as raised: + _ = await _download(payload) + + assert raised.value.status_code == 400 + + +def test_the_list_maps_ark_paging_and_filters(): + url, params = VolcEngineVideoConfig().transform_video_list_request( + api_base=API_BASE, + litellm_params=GenericLiteLLMParams(), + headers={}, + limit=20, + extra_query={"filter.status": "succeeded"}, + ) + + assert (url, params) == (TASKS_URL, {"page_size": "20", "filter.status": "succeeded"}) + + +@pytest.mark.parametrize("paging", [{"after": "video_abc"}, {"order": "asc"}]) +def test_a_list_cursor_or_order_is_rejected_unless_drop_params(paging: dict[str, str]): + config = VolcEngineVideoConfig() + + with pytest.raises(UnsupportedParamsError, match="pages by page_num"): + _ = config.transform_video_list_request( + api_base=API_BASE, litellm_params=GenericLiteLLMParams(), headers={}, **paging + ) + _, params = config.transform_video_list_request( + api_base=API_BASE, litellm_params=GenericLiteLLMParams(drop_params=True), headers={}, limit=5, **paging + ) + + assert params == {"page_size": "5"} + + +def test_the_list_response_is_an_openai_list_body(): + response = VolcEngineVideoConfig().transform_video_list_response( + raw_response=_response({"total": 1, "items": [_task(status="running")]}), + logging_obj=Mock(), # pyright: ignore[reportArgumentType] # unused by the transform + custom_llm_provider="volcengine", + ) + + encoded = encode_video_id_with_provider(TASK_ID, "volcengine", SEEDANCE_2_5) + assert {key: value for key, value in response.items() if key != "data"} == { + "object": "list", + "total": 1, + "has_more": False, + "first_id": encoded, + "last_id": encoded, + } + assert [(item["id"], item["status"]) for item in response["data"]] == [(encoded, "in_progress")] + + +def test_an_empty_list_still_carries_the_object_type(): + response = VolcEngineVideoConfig().transform_video_list_response( + raw_response=_response({"total": 0, "items": []}), + logging_obj=Mock(), # pyright: ignore[reportArgumentType] # unused by the transform + ) + + assert response == {"object": "list", "data": [], "total": 0, "has_more": False} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call", ["avideo_edit", "avideo_extension", "avideo_remix"]) +async def test_edit_extension_and_remix_are_a_bad_request_because_a_video_input_cannot_be_billed_at_create(call: str): + with pytest.raises(litellm.BadRequestError, match="Ark bills a task with a video input") as raised: + _ = await getattr(litellm, call)( + video_id="https://x/source.mp4", + prompt="make it brighter", + seconds="5", + model=f"volcengine/{SEEDANCE_2_5}", + custom_llm_provider="volcengine", + api_key="sk-test", + api_base=API_BASE, + ) + + assert raised.value.status_code == 400 + + +@pytest.mark.parametrize( + "api_base", + [ + None, + API_BASE, + f"{API_BASE}/", + f"{API_BASE}/api/v3", + f"{API_BASE}/api/compatible", + f"{API_BASE}/api/v3/responses", + f"{API_BASE}/v1", + ], +) +def test_any_ark_surface_on_the_shared_credential_normalizes_to_the_host_root( + api_base: str | None, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.delenv("VOLCENGINE_API_BASE", raising=False) + monkeypatch.delenv("ARK_API_BASE", raising=False) + + assert VolcEngineVideoConfig().get_complete_url("m", api_base, {}) == API_BASE + + +@pytest.mark.parametrize("env_var", ["VOLCENGINE_API_BASE", "ARK_API_BASE"]) +def test_an_api_base_env_var_routes_a_deployment_without_api_base(env_var: str, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("VOLCENGINE_API_BASE", raising=False) + monkeypatch.delenv("ARK_API_BASE", raising=False) + monkeypatch.setenv(env_var, "https://ark.ap-southeast.bytepluses.com/api/v3") + + assert VolcEngineVideoConfig().get_complete_url("m", None, {}) == "https://ark.ap-southeast.bytepluses.com" + assert VolcEngineVideoConfig().get_complete_url("m", f"{API_BASE}/api/v3", {}) == API_BASE + + +@pytest.mark.parametrize( + "api_base,expected", + [ + ("https://gw.internal/ark/api/v3", "https://gw.internal/ark"), + ("https://gw.internal/api/ark/api/v3", "https://gw.internal/api/ark"), + ], +) +def test_a_path_prefix_before_the_api_marker_survives(api_base: str, expected: str): + assert VolcEngineVideoConfig().get_complete_url("m", api_base, {}) == expected + + +def test_a_missing_api_key_raises(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("ARK_API_KEY", raising=False) + monkeypatch.delenv("VOLCENGINE_API_KEY", raising=False) + + with pytest.raises(ValueError, match="Volcengine API key is required"): + _ = VolcEngineVideoConfig().validate_environment({}, SEEDANCE_2_5) + + +def test_volcengine_resolves_to_the_video_config(): + config = ProviderConfigManager.get_provider_video_config( + model=SEEDANCE_2_5, provider=litellm.LlmProviders.VOLCENGINE + ) + + assert type(config) is VolcEngineVideoConfig + + +ARK_VIDEO_TIERS: Final = { + SEEDANCE_2_5: {"480p": ((854, 480), 70), "720p": ((1280, 720), 70), "1080p": ((1920, 1080), 77)}, + SEEDANCE_2_0: {"480p": ((864, 496), 46), "720p": ((1280, 720), 46), "1080p": ((1920, 1080), 51), "4k": ((3840, 2160), 26)}, + SEEDANCE_2_0_FAST: {"480p": ((864, 496), 37), "720p": ((1280, 720), 37)}, + "doubao-seedance-2-0-mini-260615": {"480p": ((864, 496), 23), "720p": ((1280, 720), 23)}, +} +CNY_PER_USD: Final = 7.0 + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("model", sorted(ARK_VIDEO_TIERS)) +def test_each_tier_rate_is_its_rendered_pixels_at_arks_price_for_that_tier(model: str): + """Ark prices 1080p and 4k per token above and below the 480p/720p rate, so a tier is not the base + rate scaled by pixels. Prices and frames are Ark's published per-resolution table, read 2026-10-16 + from https://www.volcengine.com/docs/82379/1544106, whose own 2.0 examples agree: 1080p 16:9 5s bills + 2.48 CNY/s and 4k bills 5.05 CNY/s.""" + info = litellm.get_model_info(model=model, custom_llm_provider="volcengine") + frames = info["provider_specific_entry"]["volcengine_video_frames"] + + assert {tier: info[f"output_cost_per_second_{tier}"] for tier in ARK_VIDEO_TIERS[model]} == { + tier: pytest.approx(width * height * 24 / 1024 * cny_per_million / 1e6 / CNY_PER_USD, rel=1e-5) + for tier, ((width, height), cny_per_million) in ARK_VIDEO_TIERS[model].items() + } + + assert { + tier: tuple(int(side) for side in frames[tier]["16:9"].split("x")) + for tier in ARK_VIDEO_TIERS[model] + if tier in frames + } == {tier: frame for tier, (frame, _) in ARK_VIDEO_TIERS[model].items() if tier in frames} + + +async def _wire_body(**kwargs: object) -> dict[str, object]: + sent: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={"id": TASK_ID}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + handler: Final = AsyncHTTPHandler() + await handler.close() + handler.client = http_client + _ = await avideo_generation(api_key="sk-test", api_base=API_BASE, client=handler, **kwargs) + + (request,) = sent + assert (str(request.url), request.headers["authorization"]) == (TASKS_URL, "Bearer sk-test") + return json.loads(request.content) + + +@pytest.mark.asyncio +async def test_seconds_size_and_extra_body_content_reach_the_wire(): + body = await _wire_body( + model=f"volcengine/{SEEDANCE_2_5}", + prompt="use these", + seconds="4", + size="854x480", + extra_body={ + "content": [{"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"}] + }, + ) + + assert body == { + "model": SEEDANCE_2_5, + "content": [ + {"type": "text", "text": "use these"}, + {"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"}, + ], + "ratio": "16:9", + "resolution": "480p", + "duration": 4, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fields", + [{"parameters": {"seed": 3}}, {"extra_body": {"parameters": {"seed": 3}}}], + ids=["top_level", "extra_body"], +) +async def test_parameters_reach_the_wire_unpacked_whether_sent_top_level_or_in_extra_body(fields: dict[str, object]): + body = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", seconds="4", **fields) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "p"}], + "resolution": "720p", + "duration": 4, + "seed": 3, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "extra_body,expected", + [ + ({"size": "1920x1080"}, {"resolution": "1080p"}), + ({"seconds": "8"}, {"duration": 8}), + ({"user": "someone"}, {"safety_identifier": "someone"}), + ], + ids=["size", "seconds", "user"], +) +async def test_an_extra_body_field_litellm_maps_itself_is_honoured_and_its_leftover_copy_ignored( + extra_body: dict[str, object], expected: dict[str, object] +): + """The shared request builder copies extra_body onto the top-level params, where this adapter maps it, and then + merges the raw extra_body back over the mapped params, returning the same key a second time as a field of its + own. The second copy has to be ignored rather than rejected as a field litellm never sends.""" + body = await _wire_body( + model=f"volcengine/{SEEDANCE_2_5}", + prompt="p", + seconds="4", + size="1280x720", + extra_body=extra_body, + ) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "p"}], + "ratio": "16:9", + "resolution": "720p", + "duration": 4, + **expected, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fields,misplaced", + [ + ({"parameters": {"seconds": "8"}}, "seconds"), + ({"extra_body": {"parameters": {"seconds": "8"}}}, "seconds"), + ({"extra_body": {"parameters": {"size": "1920x1080"}}}, "size"), + ({"extra_body": {"parameters": {"user": "someone"}}}, "user"), + ({"extra_body": {"parameters": {"prompt": "a dog"}}}, "prompt"), + ({"extra_body": {"parameters": {"parameters": {"seed": 3}}}}, "parameters"), + ], + ids=["top_level_seconds", "seconds", "size", "user", "prompt", "nested_parameters"], +) +async def test_a_request_field_litellm_maps_itself_is_rejected_inside_parameters( + fields: dict[str, object], misplaced: str +): + with pytest.raises(UnsupportedParamsError, match=f"send {misplaced} as top-level request fields"): + _ = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", **fields) + + +@pytest.mark.asyncio +async def test_a_request_field_inside_parameters_is_dropped_under_drop_params_and_the_rest_is_sent(): + body = await _wire_body( + model=f"volcengine/{SEEDANCE_2_5}", + prompt="p", + drop_params=True, + extra_body={"parameters": {"seconds": "8", "seed": 3}}, + ) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "p"}], + "resolution": "720p", + "duration": 4, + "seed": 3, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "extra_body,field", + [ + ({"prompt": "a dog"}, "prompt"), + ({"model": SEEDANCE_2_0_FAST}, "model"), + ({"extra_headers": {"x-a": "1"}}, "extra_headers"), + ], + ids=["prompt", "model", "extra_headers"], +) +async def test_an_extra_body_copy_of_a_request_field_litellm_does_not_map_is_rejected( + extra_body: dict[str, object], field: str +): + with pytest.raises(UnsupportedParamsError, match=f"does not send {field} to Ark"): + _ = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", extra_body=extra_body) + + +@pytest.mark.asyncio +async def test_a_user_that_is_not_a_string_is_rejected_rather_than_ignored(): + with pytest.raises(UnsupportedParamsError, match="user must be a string, got int"): + _ = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", extra_body={"user": 5}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "control", + [{"drop_params": False}, {"organization": "org-1"}], + ids=["drop_params", "organization"], +) +async def test_a_litellm_control_field_is_not_mistaken_for_an_ark_field(control: dict[str, object]): + body = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", seconds="4", **control) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "p"}], + "resolution": "720p", + "duration": 4, + } + + +@pytest.mark.asyncio +async def test_additional_drop_params_removes_the_named_ark_fields_from_the_wire(): + body = await _wire_body( + model=f"volcengine/{SEEDANCE_2_5}", + prompt="p", + seconds="4", + additional_drop_params=["seed"], + extra_body={"seed": 3, "watermark": False}, + ) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "p"}], + "resolution": "720p", + "duration": 4, + "watermark": False, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("dropped", "request_fields"), + [ + pytest.param("user", {"user": "u-1"}, id="user"), + pytest.param("seconds", {"seconds": "8"}, id="seconds"), + pytest.param("size", {"size": "1920x1080"}, id="size"), + pytest.param("parameters", {"parameters": {"seed": 3}}, id="parameters"), + pytest.param("parameters", {"extra_body": {"parameters": {"seed": 3}}}, id="extra_body_parameters"), + pytest.param("image", {"image": {"gcsUri": "gs://b/a.png"}}, id="unsupported_image"), + ], +) +async def test_additional_drop_params_naming_a_request_field_sends_what_omitting_it_sends( + dropped: str, request_fields: dict[str, object] +): + request: Final = {"model": f"volcengine/{SEEDANCE_2_5}", "prompt": "p"} + + with_drop = await _wire_body(**request, **request_fields, additional_drop_params=[dropped]) + + assert with_drop == await _wire_body(**request) + + +@pytest.mark.asyncio +async def test_an_extra_body_duration_of_minus_one_is_rejected_before_anything_is_sent(): + with pytest.raises(UnsupportedParamsError, match="Ark-picked length"): + _ = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", extra_body={"duration": -1}) + + +ACCOUNT_A_HOST: Final = "https://ark-a.example" +ACCOUNT_B_HOST: Final = "https://ark-b.example" + + +def _two_account_router(**router_kwargs: object) -> Router: + return Router( + model_list=[ + { + "model_name": "seedance-a", + "litellm_params": { + "model": f"volcengine/{SEEDANCE_2_5}", + "api_key": "key-a", + "api_base": ACCOUNT_A_HOST, + }, + }, + { + "model_name": "seedance-b", + "litellm_params": { + "model": f"volcengine/{SEEDANCE_2_5}", + "api_key": "key-b", + "api_base": ACCOUNT_B_HOST, + }, + }, + ], + **router_kwargs, # pyright: ignore[reportArgumentType] # forwarded Router options + ) + + +def _proxy_poll_model(router: Router, video_id: str) -> str | None: + decoded: Final = decode_video_id_with_provider(video_id) + requested: Final = router.resolve_model_name_from_model_id( + decoded.get("model_id"), decoded.get("custom_llm_provider") + ) + served: Final = decoded.get("served_model") + return router.resolve_served_model_name(requested, served) if requested and served else requested + + +@pytest.mark.parametrize("caller_metadata", [None, {"tags": ["proxy-request"]}]) +@respx.mock +def test_a_video_is_polled_on_the_ark_account_that_created_it(caller_metadata: dict[str, object] | None): + router = _two_account_router() + _ = respx.post(f"{ACCOUNT_B_HOST}/api/v3/contents/generations/tasks").mock( + return_value=httpx.Response(200, json={"id": TASK_ID}) + ) + task_get = respx.get(f"{ACCOUNT_B_HOST}/api/v3/contents/generations/tasks/{TASK_ID}").mock( + return_value=httpx.Response(200, json=_task(status="running")) + ) + + created = router.video_generation(model="seedance-b", prompt="a cat", seconds="4", metadata=caller_metadata) + _ = router.video_status(model=_proxy_poll_model(router, created.id), video_id=created.id) + + assert [call.request.headers["authorization"] for call in task_get.calls] == ["Bearer key-b"] + + +@respx.mock +def test_a_video_created_through_a_routing_group_is_polled_on_the_member_that_created_it(): + router = _two_account_router( + routing_groups=[ + {"group_name": "seedance", "models": ["seedance-a", "seedance-b"], "routing_strategy": "simple-shuffle"} + ] + ) + for host in (ACCOUNT_A_HOST, ACCOUNT_B_HOST): + _ = respx.post(f"{host}/api/v3/contents/generations/tasks").mock( + return_value=httpx.Response(200, json={"id": TASK_ID}) + ) + _ = respx.get(f"{host}/api/v3/contents/generations/tasks/{TASK_ID}").mock( + return_value=httpx.Response(200, json=_task(status="running")) + ) + + for _ in range(8): + created = router.video_generation(model="seedance", prompt="a cat", seconds="4") + polled = router.video_status(model=_proxy_poll_model(router, created.id), video_id=created.id) + assert polled.id == created.id + + hosts = [call.request.url.host for call in respx.calls] + assert hosts[1::2] == hosts[0::2] 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..bb228544e14 --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py @@ -0,0 +1,327 @@ +"""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_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" + + +@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..af62594e3b8 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,58 @@ 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_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 69d182667ec..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): @@ -440,8 +443,6 @@ class TestErrorLogCarriesCallId: @pytest.mark.asyncio async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture): - from fastapi import HTTPException - import litellm.proxy.anthropic_endpoints.endpoints as ep import litellm.proxy.proxy_server as proxy_server from litellm.proxy._types import UserAPIKeyAuth @@ -460,16 +461,160 @@ class TestErrorLogCarriesCallId: proxy_server, "count_request_tokens", new=AsyncMock(side_effect=RuntimeError("tokenizer down")) ), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), - pytest.raises(HTTPException) as raised, ): - await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + response = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) - assert raised.value.status_code == 500 + assert response.status_code == 500 record = self._error_record(caplog) assert record.litellm_call_id == call_id 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/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index c17aed49014..70f7774e691 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -435,7 +435,7 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client): def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): - """End-user spend is zeroed and the tier's window advances, in one batch.""" + """End-user spend is zeroed in its own committed page, then the tier's window advances.""" now = datetime.now(timezone.utc) test_budget = _budget_row(budget_id="test-budget-1", budget_duration="1d", budget_reset_at=now) @@ -446,11 +446,12 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): "spend": 20.0, "litellm_budget_table": test_budget, "user_id": "test-enduser-1", + "budget_id": "test-budget-1", }, ) mock_prisma_client.data["budget"] = [test_budget] - mock_prisma_client.data["enduser"] = [test_enduser] + mock_prisma_client.db.litellm_endusertable.set_find_many_results([test_enduser]) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) @@ -458,10 +459,13 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): { "table": "enduser", "op": "update_many", - "where": {"budget_id": {"in": ["test-budget-1"]}, "spend": {"gt": 0}}, + "where": {"user_id": {"in": ["test-enduser-1"]}, "spend": {"gt": 0}}, "data": {"spend": 0}, } ] + assert [batcher.committed for batcher in mock_prisma_client.db.batchers] == [True, True] + assert [call["table"] for call in mock_prisma_client.db.batchers[0].calls] == ["enduser"] + assert "enduser" not in {call["table"] for call in mock_prisma_client.db.batchers[1].calls} budget_writes = _batch_writes(mock_prisma_client, "budget") assert len(budget_writes) == 1 @@ -470,6 +474,112 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): assert set(budget_writes[0]["data"].keys()) == {"budget_reset_at"} +def test_enduser_reset_commits_in_pages_before_the_window_advances(reset_budget_job, mock_prisma_client): + """Regression for #46404: one statement over every end user a tier gates + reads the whole table and outlives the prisma client's 30s request timeout + once the table holds millions of rows, so the batch rolls back and + budget_reset_at never moves. The reset has to land one committed page at a + time, keyed by user id, with the window advance after the last page.""" + mock_prisma_client.data["budget"] = [_budget_row(budget_id="shared-tier", budget_duration="1d")] + population: Final = RESET_BUDGET_JOB_BATCH_SIZE * 2 + 3 + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [ + types.SimpleNamespace(spend=4.0, user_id=f"cust-{index:06d}", budget_id="shared-tier") + for index in range(population) + ] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + batchers: Final = mock_prisma_client.db.batchers + assert [batcher.committed for batcher in batchers] == [True, True, True, True] + pages: Final = batchers[:3] + assert [[call["table"] for call in page.calls] for page in pages] == [["enduser"]] * 3 + assert [len(page.calls[0]["where"]["user_id"]["in"]) for page in pages] == [ + RESET_BUDGET_JOB_BATCH_SIZE, + RESET_BUDGET_JOB_BATCH_SIZE, + 3, + ] + assert [page.calls[0]["where"]["spend"] for page in pages] == [{"gt": 0}] * 3 + assert [page.calls[0]["data"] for page in pages] == [{"spend": 0}] * 3 + assert pages[0].calls[0]["where"]["user_id"]["in"][0] == "cust-000000" + assert pages[2].calls[0]["where"]["user_id"]["in"] == ["cust-001000", "cust-001001", "cust-001002"] + cascade: Final = batchers[3] + assert "enduser" not in {call["table"] for call in cascade.calls} + assert [call["where"] for call in cascade.calls if call["table"] == "budget"] == [{"budget_id": "shared-tier"}] + + +def test_enduser_page_failure_leaves_the_tier_due(reset_budget_job, mock_prisma_client): + """A page that fails after earlier pages committed must not advance the + window: the committed pages are at spend 0 and drop out of the next tick's + ``spend > 0`` read, so the next tick resumes where this one stopped.""" + mock_prisma_client.data["budget"] = [_budget_row(budget_id="shared-tier", budget_duration="1d")] + endusers: Final = mock_prisma_client.db.litellm_endusertable + endusers.set_find_many_results( + [ + types.SimpleNamespace(spend=4.0, user_id=f"cust-{index:06d}", budget_id="shared-tier") + for index in range(RESET_BUDGET_JOB_BATCH_SIZE + 3) + ] + ) + endusers.set_find_many_error(1, RuntimeError("connection reset while paging end users")) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + assert _batch_writes(mock_prisma_client, "budget") == [] + enduser_writes: Final = _batch_writes(mock_prisma_client, "enduser") + assert len(enduser_writes) == 1 + assert len(enduser_writes[0]["where"]["user_id"]["in"]) == RESET_BUDGET_JOB_BATCH_SIZE + + +def test_enduser_page_keeps_each_tier_s_rollover_statements_to_its_own_users( + rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch +): + """One page can hold end users of several expiring tiers. Each tier's + zero-then-decrement pair must name only that tier's user ids and its own + cap, or a 15.0 spend on a 20.0 tier would be carried as 5.0.""" + _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [ + _budget_row(budget_id="tier-10", budget_duration="1d", max_budget=10.0), + _budget_row(budget_id="tier-20", budget_duration="1d", max_budget=20.0), + ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [ + types.SimpleNamespace(spend=15.0, user_id="cust-a", budget_id="tier-10"), + types.SimpleNamespace(spend=15.0, user_id="cust-b", budget_id="tier-20"), + types.SimpleNamespace(spend=25.0, user_id="cust-c", budget_id="tier-10"), + ] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + assert _batch_writes(mock_prisma_client, "enduser") == [ + { + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["cust-a", "cust-c"]}, "spend": {"gt": 0, "lte": 10.0}}, + "data": {"spend": 0}, + }, + { + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["cust-a", "cust-c"]}, "spend": {"gt": 10.0}}, + "data": {"spend": {"decrement": 10.0}}, + }, + { + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["cust-b"]}, "spend": {"gt": 0, "lte": 20.0}}, + "data": {"spend": 0}, + }, + { + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["cust-b"]}, "spend": {"gt": 20.0}}, + "data": {"spend": {"decrement": 20.0}}, + }, + ] + + def test_reset_budget_all(reset_budget_job, mock_prisma_client): # Setup test data with timezone-aware datetime now = datetime.now(timezone.utc) @@ -529,6 +639,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): "spend": 20.0, "litellm_budget_table": test_budget, "user_id": "test-enduser-1", + "budget_id": "test-budget-1", }, ) @@ -536,7 +647,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): mock_prisma_client.data["user"] = [test_user] mock_prisma_client.data["team"] = [test_team] mock_prisma_client.data["budget"] = [test_budget] - mock_prisma_client.data["enduser"] = [test_enduser] + mock_prisma_client.db.litellm_endusertable.set_find_many_results([test_enduser]) # Run the test asyncio.run(reset_budget_job.reset_budget()) @@ -559,7 +670,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): { "table": "enduser", "op": "update_many", - "where": {"budget_id": {"in": ["test-budget-1"]}, "spend": {"gt": 0}}, + "where": {"user_id": {"in": ["test-enduser-1"]}, "spend": {"gt": 0}}, "data": {"spend": 0}, } ] @@ -579,7 +690,6 @@ _LINKED_TABLE_CASES = [ ("org", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ("tag", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ("model_access_group", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), - ("enduser", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ] @@ -620,34 +730,33 @@ def _bind_count(where: Dict[str, Any]) -> int: @pytest.mark.parametrize("population", [3, 40_000], ids=["small", "over-pg-bind-ceiling"]) def test_enduser_reset_bind_count_does_not_scale_with_population(reset_budget_job, mock_prisma_client, population): - """Regression for #40564. + """Regression for #40564 and #46404. - Enumerating every dependent user id put one bind variable per customer into - a single prepared statement. Past PostgreSQL's ceiling the statement could - not be parsed at all, so the whole atomic cascade rolled back, - budget_reset_at never advanced, and the tier stayed due on every later tick - forever. Matching on the budget link keeps the statement the same size no - matter how many customers share a tier. + Enumerating every dependent user id in one statement put one bind variable + per end user into a single prepared statement, which past PostgreSQL's + ceiling could not be parsed at all. Matching on the budget link instead made + one statement read and update the whole end-user table, which past a few + million rows outlives the prisma client's 30s request timeout. Paging by + user id keeps every statement at one page of binds however many end users + share a tier, and the statements have to cover the whole population. """ budget = _budget_row(budget_id="shared-tier", budget_duration="1d") mock_prisma_client.data["budget"] = [budget] - mock_prisma_client.data["enduser"] = [ - types.SimpleNamespace( - spend=1.0, - litellm_budget_table=budget, - user_id=f"cust-{index:08d}", - budget_id="shared-tier", - ) - for index in range(population) - ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [ + types.SimpleNamespace(spend=1.0, user_id=f"cust-{index:08d}", budget_id="shared-tier") + for index in range(population) + ] + ) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) writes = _batch_writes(mock_prisma_client, "enduser") - assert [_bind_count(write["where"]) for write in writes] == [2], ( - f"the cascade must not enumerate {population} user ids: past " - f"{_POSTGRES_MAX_BIND_VARIABLES} binds PostgreSQL refuses the statement, got {writes[:1]}" - ) + assert max(_bind_count(write["where"]) for write in writes) <= RESET_BUDGET_JOB_BATCH_SIZE + 1 + assert RESET_BUDGET_JOB_BATCH_SIZE + 1 < _POSTGRES_MAX_BIND_VARIABLES + assert sorted(user_id for write in writes for user_id in write["where"]["user_id"]["in"]) == [ + f"cust-{index:08d}" for index in range(population) + ] assert _batch_writes(mock_prisma_client, "budget")[0]["data"]["budget_reset_at"] is not None @@ -778,6 +887,7 @@ def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock "spend": 30.0, "litellm_budget_table": test_budget, "user_id": "enduser-explicit", + "budget_id": default_budget_id, }, ) @@ -811,40 +921,34 @@ def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock ) mock_prisma_client.data["budget"] = [test_budget] - mock_prisma_client.data["enduser"] = [enduser_with_budget] - - # Set up the DB mock for NULL-budget-id end users - mock_prisma_client.db.litellm_endusertable.set_find_many_results([enduser_no_budget_row]) + mock_prisma_client.db.litellm_endusertable.set_find_many_results([enduser_with_budget, enduser_no_budget_row]) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - # Both end users are zeroed: the linked rows on the tier's budget_id, the - # implicit ones on the NULL branch that stands in for the default tier. + # Both end users are zeroed in one page: the linked row on the tier's + # budget_id and the implicit one whose NULL budget_id stands in for the + # default tier. assert _batch_writes(mock_prisma_client, "enduser") == [ { "table": "enduser", "op": "update_many", - "where": {"budget_id": {"in": [default_budget_id]}, "spend": {"gt": 0}}, - "data": {"spend": 0}, - }, - { - "table": "enduser", - "op": "update_many", - "where": {"budget_id": None, "spend": {"gt": 0}}, + "where": {"user_id": {"in": ["enduser-explicit", "enduser-implicit"]}, "spend": {"gt": 0}}, "data": {"spend": 0}, }, ] - # The post-commit invalidation walk covers both branches, so implicitly - # created customers on the default tier get their cached spend dropped too, - # and it is paged rather than reading the whole customer population. + # The paged reset and the post-commit invalidation walk both cover both + # branches, so implicitly created customers on the default tier are reset + # and get their cached spend dropped too. find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls - assert len(find_many_calls) == 1 - assert find_many_calls[0]["where"]["OR"] == [ - {"budget_id": {"in": [default_budget_id]}}, - {"budget_id": None}, - ] - assert find_many_calls[0]["take"] == RESET_BUDGET_JOB_BATCH_SIZE + assert len(find_many_calls) == 2 + assert find_many_calls[0]["where"]["spend"] == {"gt": 0} + for call in find_many_calls: + assert call["where"]["OR"] == [ + {"budget_id": {"in": [default_budget_id]}}, + {"budget_id": None}, + ] + assert call["take"] == RESET_BUDGET_JOB_BATCH_SIZE litellm.max_end_user_budget_id = None @@ -879,7 +983,8 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_configured( # ride a default tier that is not expiring, so their spend stays put. find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls assert [call["where"] for call in find_many_calls] == [ - {"budget_id": {"in": ["some-budget"]}, "user_id": {"gt": ""}} + {"budget_id": {"in": ["some-budget"]}, "spend": {"gt": 0}, "user_id": {"gt": ""}}, + {"budget_id": {"in": ["some-budget"]}, "user_id": {"gt": ""}}, ] litellm.max_end_user_budget_id = None @@ -919,7 +1024,8 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_in_reset_li # ride a default tier that is not expiring, so their spend stays put. find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls assert [call["where"] for call in find_many_calls] == [ - {"budget_id": {"in": ["other-budget"]}, "user_id": {"gt": ""}} + {"budget_id": {"in": ["other-budget"]}, "spend": {"gt": 0}, "user_id": {"gt": ""}}, + {"budget_id": {"in": ["other-budget"]}, "user_id": {"gt": ""}}, ] litellm.max_end_user_budget_id = None @@ -1671,7 +1777,8 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma on a large enough deployment rather than a slow tick. Awaiting one cache call per customer makes the last customer wait out every customer ahead of it. Both regress silently, so pin the page size, the strictly advancing cursor, - and one batched call per page. + and one batched call per page. The paged reset ahead of the commit + invalidates its pages too, so the cache calls count both walks. """ counter_cache: Final = _make_counter_invalidation_job(monkeypatch) mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")] @@ -1685,7 +1792,9 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - reads: Final = mock_prisma_client.db.litellm_endusertable.find_many_calls + all_reads: Final = mock_prisma_client.db.litellm_endusertable.find_many_calls + reads: Final = [read for read in all_reads if "spend" not in read["where"]] + assert len(all_reads) == 6 assert [read["take"] for read in reads] == [RESET_BUDGET_JOB_BATCH_SIZE] * 3 assert [read["where"]["user_id"]["gt"] for read in reads] == [ "", @@ -1693,8 +1802,8 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma f"cust-{RESET_BUDGET_JOB_BATCH_SIZE * 2 - 1:06d}", ] - assert counter_cache.async_delete_cache_keys.await_count == 3 - assert counter_cache.user_api_key_cache.async_delete_cache_keys.await_count == 3 + assert counter_cache.async_delete_cache_keys.await_count == 6 + assert counter_cache.user_api_key_cache.async_delete_cache_keys.await_count == 6 counter_cache.async_delete_cache.assert_not_called() invalidated: Final = { @@ -1727,7 +1836,7 @@ def test_enduser_invalidation_reports_a_page_read_failure_instead_of_a_clean_fin for i in range(RESET_BUDGET_JOB_BATCH_SIZE + 3) ] ) - endusers.set_find_many_error(1, RuntimeError("connection reset while paging customers")) + endusers.set_find_many_error(3, RuntimeError("connection reset while paging customers")) logging_obj: Final = RecordingProxyLogging() job: Final = ResetBudgetJob(proxy_logging_obj=logging_obj, prisma_client=mock_prisma_client) @@ -1985,19 +2094,22 @@ def test_budget_reset_at_is_not_advanced_when_the_cascade_fails(db_factory, monk def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mock_prisma_client, monkeypatch): - """Dependent spend and the budget_reset_at advance ride one batch.""" + """Dependent spend and the budget_reset_at advance ride one batch, after + the end-user pages have committed on their own.""" _make_counter_invalidation_job(monkeypatch) now = datetime.now(timezone.utc) budget = _budget_row(budget_id="budget-1", budget_duration="7d") mock_prisma_client.data["budget"] = [budget] - mock_prisma_client.data["enduser"] = [ - type("EndUser", (), {"spend": 5.0, "litellm_budget_table": budget, "user_id": "enduser-1", "budget_id": "budget-1"}) - ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [type("EndUser", (), {"spend": 5.0, "user_id": "enduser-1", "budget_id": "budget-1"})] + ) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - assert len(mock_prisma_client.db.batchers) == 1, "the cascade must not be split across transactions" - batcher = mock_prisma_client.db.batchers[0] + assert len(mock_prisma_client.db.batchers) == 2, "one end-user page, then the cascade" + enduser_page, batcher = mock_prisma_client.db.batchers + assert enduser_page.committed is True + assert {(call["table"], call["op"]) for call in enduser_page.calls} == {("enduser", "update_many")} assert batcher.committed is True assert {(call["table"], call["op"]) for call in batcher.calls} == { ("team_membership", "update_many"), @@ -2006,7 +2118,6 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo ("tag", "update_many"), ("model_access_group", "update_many"), ("project", "update_many"), - ("enduser", "update_many"), ("budget", "update_many"), } budget_write = next(call for call in batcher.calls if call["table"] == "budget") @@ -3288,13 +3399,9 @@ def test_budget_cascade_carries_enduser_overage_when_rollover_enabled( _make_counter_invalidation_job(monkeypatch) budget = _budget_row(budget_id="budget-roll", budget_duration="1d", max_budget=10.0) mock_prisma_client.data["budget"] = [budget] - mock_prisma_client.data["enduser"] = [ - type( - "EndUser", - (), - {"spend": 15.0, "litellm_budget_table": budget, "user_id": "enduser-roll", "budget_id": "budget-roll"}, - ) - ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [type("EndUser", (), {"spend": 15.0, "user_id": "enduser-roll", "budget_id": "budget-roll"})] + ) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) @@ -3302,13 +3409,13 @@ def test_budget_cascade_carries_enduser_overage_when_rollover_enabled( assert { "table": "enduser", "op": "update_many", - "where": {"budget_id": "budget-roll", "spend": {"gt": 10.0}}, + "where": {"user_id": {"in": ["enduser-roll"]}, "spend": {"gt": 10.0}}, "data": {"spend": {"decrement": 10.0}}, } in enduser_writes assert { "table": "enduser", "op": "update_many", - "where": {"budget_id": "budget-roll", "spend": {"gt": 0, "lte": 10.0}}, + "where": {"user_id": {"in": ["enduser-roll"]}, "spend": {"gt": 0, "lte": 10.0}}, "data": {"spend": 0}, } in enduser_writes @@ -3377,13 +3484,9 @@ def test_cascade_rollover_writes_survive_sequential_execution( {"user_id": "member-1", "team_id": "team-1", "spend": 15.0, "budget_id": "budget-roll"}, ) mock_prisma_client.db.litellm_teammembership.set_find_many_results([membership]) - mock_prisma_client.data["enduser"] = [ - type( - "EndUser", - (), - {"spend": 15.0, "litellm_budget_table": budget, "user_id": "enduser-roll", "budget_id": "budget-roll"}, - ) - ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [type("EndUser", (), {"spend": 15.0, "user_id": "enduser-roll", "budget_id": "budget-roll"})] + ) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) @@ -3702,8 +3805,9 @@ def _wire_cascade_reads_for_test(prisma_client, endusers=()): tables an awaitable find_many so the reads resolve instead of falling into the job's warn-and-continue path. - End users are read by the post-commit invalidation walk rather than by - ``get_data``, so callers that care about customers pass them here. + End users are read by the paged reset and the post-commit invalidation + walk rather than by ``get_data``, so callers that care about them pass + them here. """ for table in ( "litellm_teammembership", @@ -4010,6 +4114,7 @@ async def test_reset_budget_endusers_are_zeroed_with_the_budget_window_advance(f prisma_client.get_data.side_effect = get_data_mock prisma_client.update_data = AsyncMock() batch_calls: Final = _wire_batcher_for_test(prisma_client) + _wire_cascade_reads_for_test(prisma_client, endusers=endusers) proxy_logging_obj: Final = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -4020,11 +4125,14 @@ async def test_reset_budget_endusers_are_zeroed_with_the_budget_window_advance(f await _await_tasks_spawned_by(job.reset_budget_for_litellm_budget_table()) - assert prisma_client.db.batch_.call_count == 1, "the cascade must be one transaction" + assert prisma_client.db.batch_.call_count == 2, "one end-user page, then the cascade in one transaction" enduser_writes: Final = [c for c in batch_calls if c["table"] == "enduser"] assert len(enduser_writes) == 1 - assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}} + assert enduser_writes[0]["where"] == { + "user_id": {"in": ["user1", "user2", "user3", "user4", "user5", "user6"]}, + "spend": {"gt": 0}, + } assert enduser_writes[0]["data"] == {"spend": 0} budget_writes: Final = [c for c in batch_calls if c["table"] == "budget"] @@ -4098,7 +4206,7 @@ async def test_reset_budget_continues_other_categories_on_failure(frozen_reset_c assert len([c for c in batch_calls if c["table"] == "team_membership"]) == 1 enduser_writes: Final = [c for c in batch_calls if c["table"] == "enduser"] assert len(enduser_writes) == 1 - assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}} + assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1"]}, "spend": {"gt": 0}} assert enduser_writes[0]["data"] == {"spend": 0} key_writes: Final = [c for c in batch_calls if c["table"] == "key" and c["op"] == "update"] @@ -4576,7 +4684,7 @@ async def test_service_logger_endusers_success(): enduser_writes = [c for c in batch_calls if c["table"] == "enduser"] assert len(enduser_writes) == 1 - assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}} + assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1", "user2"]}, "spend": {"gt": 0}} proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_called_once() ( 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/google_endpoints/test_google_api_endpoints.py b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py index 451792109de..114bbb184c6 100644 --- a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py @@ -182,6 +182,29 @@ def test_google_count_tokens_unchanged(): assert body["totalTokens"] == 7 +def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope(): + try: + client = _build_test_client() + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + with patch( + "litellm.proxy.proxy_server.token_counter", + new_callable=AsyncMock, + side_effect=RuntimeError("tokenizer unavailable"), + ): + response = client.post( + "/v1beta/models/test-model:countTokens", + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + ) + + assert response.status_code == 500 + error = response.json()["error"] + assert error["code"] == 500 + assert error["status"] == "INTERNAL" + assert "tokenizer unavailable" not in error["message"] + + def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs(): """A provider that counts in Anthropic shape (Bedrock, Anthropic) carries no totalTokens key, so the normalized total must win over the raw response's missing field.""" try: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py b/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py index 5d706a65a79..abcbb714186 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py @@ -9,7 +9,11 @@ Tests 40 different sentences to validate the conditional matching logic: """ import importlib +import json import os +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final import pytest from fastapi import HTTPException @@ -19,9 +23,11 @@ from litellm.proxy.guardrails.content_filter_data import POLICY_TEMPLATES_DIR from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( ContentFilterGuardrail, ) +from litellm.types.guardrails import BlockedWord, ContentFilterAction, ContentFilterPattern, GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( ContentFilterCategoryConfig, ) +from litellm.types.utils import Choices, Message, ModelResponse from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome # Test cases: (sentence, expected_result, reason) @@ -1793,3 +1799,237 @@ class TestSGPDPAPerformance: if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) + + +AWS_KEY: Final = "AKIA" + "SYNTHETIC0000000" +GITHUB_TOKEN: Final = "ghp_" + "a1" * 18 +FORBIDDEN_WORD: Final = "synthetic-forbidden-word" +SECRET_PATTERNS: Final = ( + ContentFilterPattern( + pattern_type="regex", pattern=r"AKIA[0-9A-Z]{16}", name="aws_access_key", action=ContentFilterAction.BLOCK + ), + ContentFilterPattern( + pattern_type="regex", pattern=r"ghp_[A-Za-z0-9]{36}", name="github_pat", action=ContentFilterAction.BLOCK + ), +) + + +def _competitor_intent_config(action_hint: str) -> Mapping[str, object]: + return MappingProxyType( + { + "brand_self": ["emirates"], + "competitors": ["qatar airways"], + "domain_words": ["airline", "flight"], + "policy": {"competitor_comparison": action_hint}, + "threshold_high": 0.5, + } + ) + + +def _observing_filter( + patterns: tuple[ContentFilterPattern, ...] = (), + categories: tuple[ContentFilterCategoryConfig, ...] = (), + competitor_intent_config: Mapping[str, object] | None = None, + mask_patterns: bool = False, + image_describer: object = None, +) -> ContentFilterGuardrail: + return ContentFilterGuardrail( + guardrail_name="cf-observe", + event_hook=GuardrailEventHooks.logging_only, + llm_router=image_describer, + image_model="vision" if image_describer else None, + patterns=[ + pattern.model_copy(update={"action": ContentFilterAction.MASK}) if mask_patterns else pattern + for pattern in patterns + ], + blocked_words=[BlockedWord(keyword=FORBIDDEN_WORD, action=ContentFilterAction.BLOCK)], + categories=list(categories), + competitor_intent_config=dict(competitor_intent_config) if competitor_intent_config else None, + ) + + +def _policy_category(template: str, category: str) -> ContentFilterCategoryConfig: + return ContentFilterCategoryConfig( + category=category, + category_file=os.path.join(POLICY_TEMPLATES_DIR, template) if template else None, + enabled=True, + action="BLOCK", + severity_threshold="medium", + ) + + +async def _log_call( + guardrail: ContentFilterGuardrail, + prompt: str | list[Mapping[str, object]], + answer: Message, + history: Sequence[Mapping[str, object]] = (), +) -> tuple[Mapping[str, object], ...]: + kwargs: Final = { + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": prompt}, *history], + "litellm_call_id": "call-observe", + "litellm_params": {"metadata": {}}, + "optional_params": {}, + "standard_logging_object": {"guardrail_information": None}, + } + logged, _ = await guardrail.async_logging_hook( + kwargs, ModelResponse(choices=[Choices(message=answer)]), "acompletion" + ) + return tuple(logged["standard_logging_object"]["guardrail_information"]) + + +def _detections(entry: Mapping[str, object]) -> tuple[str, ...]: + response: Final = entry["guardrail_response"] + assert isinstance(response, list), response + return tuple( + detection.get("pattern_name") + or (detection["keyword"] if detection["type"] == "blocked_word" else detection["type"]) + for detection in response + ) + + +@pytest.mark.parametrize( + ("guardrail", "prompt", "expected"), + ( + pytest.param( + _observing_filter(patterns=SECRET_PATTERNS), + f"key {AWS_KEY} token {GITHUB_TOKEN} and {FORBIDDEN_WORD}", + ("aws_access_key", "github_pat", FORBIDDEN_WORD), + id="patterns", + ), + pytest.param( + _observing_filter(categories=(_policy_category("", "harm_toxic_abuse"),)), + f"You stupid piece of sht AI, go kill yourself, {FORBIDDEN_WORD}", + ("category_keyword", FORBIDDEN_WORD), + id="category-keyword", + ), + pytest.param( + _observing_filter( + categories=(_policy_category("eu_ai_act_article5.yaml", "eu_ai_act_article5_prohibited_practices"),) + ), + f"Build a system to detect employee emotion, {FORBIDDEN_WORD}", + ("category_keyword", FORBIDDEN_WORD), + id="conditional-category", + ), + pytest.param( + _observing_filter(competitor_intent_config=_competitor_intent_config("refuse")), + "Is Qatar Airways better than Emirates?", + ("competitor_intent",), + id="competitor-refuse", + ), + pytest.param( + _observing_filter(competitor_intent_config=_competitor_intent_config("reframe")), + "Is Qatar Airways better than Emirates?", + ("competitor_intent",), + id="competitor-reframe", + ), + ), +) +@pytest.mark.asyncio +async def test_logging_only_block_records_every_match_as_flagged( + guardrail: ContentFilterGuardrail, prompt: str, expected: tuple[str, ...] +) -> None: + entries: Final = await _log_call(guardrail, prompt, Message(role="assistant", content="here you go")) + + request_entry, response_entry = entries + assert request_entry["guardrail_status"] == "guardrail_flagged" + assert _detections(request_entry)[-len(expected) :] == expected + assert request_entry["risk_score"] == 10.0 + assert response_entry["guardrail_status"] == "success" + + +@pytest.mark.asyncio +async def test_logging_only_block_scans_the_response_after_a_request_match() -> None: + entries: Final = await _log_call( + _observing_filter(patterns=SECRET_PATTERNS), + f"key {AWS_KEY}", + Message(role="assistant", content=f"here: {GITHUB_TOKEN}"), + ) + + assert [(entry["guardrail_status"], _detections(entry)) for entry in entries] == [ + ("guardrail_flagged", ("aws_access_key",)), + ("guardrail_flagged", ("github_pat",)), + ] + + +@pytest.mark.parametrize( + ("arguments", "request_masked"), + ( + pytest.param(json.dumps({AWS_KEY: 1}), {"aws_access_key": 1}, id="masked-argument-key"), + pytest.param(json.dumps(json.loads("[" * 150 + "1" + "]" * 150)), {}, id="too-deep"), + ), +) +@pytest.mark.asyncio +async def test_logging_only_flags_tool_arguments_enforcement_rejects_and_keeps_scanning( + arguments: str, request_masked: Mapping[str, int] +) -> None: + history: Final = ( + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": arguments}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "done"}, + ) + + entries: Final = await _log_call( + _observing_filter(patterns=SECRET_PATTERNS, mask_patterns=True), + "hi", + Message(role="assistant", content=f"here: {GITHUB_TOKEN}"), + history, + ) + + assert [(entry["guardrail_status"], entry["masked_entity_count"]) for entry in entries] == [ + ("guardrail_flagged", request_masked), + ("success", {"github_pat": 1}), + ] + + +@pytest.mark.asyncio +async def test_logging_only_scan_leaves_enforcement_of_the_same_guardrail_blocking() -> None: + guardrail: Final = _observing_filter(patterns=SECRET_PATTERNS) + await _log_call(guardrail, f"key {AWS_KEY}", Message(role="assistant", content="ok")) + request_data: Final = {"metadata": {}} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": [f"key {AWS_KEY} token {GITHUB_TOKEN}"]}, request_data=request_data, input_type="request" + ) + + assert exc_info.value.status_code == 400 + assert [ + entry["guardrail_status"] for entry in request_data["metadata"]["standard_logging_guardrail_information"] + ] == ["guardrail_intervened"] + + +class _ImageDescriberRunningAnotherFilter: + def __init__(self, enforcer: ContentFilterGuardrail) -> None: + self.enforcer = enforcer + self.enforcer_blocked: list[bool] = [] + + async def acompletion(self, **_: object) -> ModelResponse: + try: + await self.enforcer.apply_guardrail( + inputs={"texts": [f"key {AWS_KEY}"]}, request_data={"metadata": {}}, input_type="request" + ) + self.enforcer_blocked.append(False) + except HTTPException: + self.enforcer_blocked.append(True) + return ModelResponse(choices=[Choices(message=Message(role="assistant", content="a plain desk"))]) + + +@pytest.mark.asyncio +async def test_logging_only_scan_does_not_relax_another_filter_called_during_it() -> None: + describer: Final = _ImageDescriberRunningAnotherFilter( + ContentFilterGuardrail(guardrail_name="cf-enforce", patterns=list(SECRET_PATTERNS)) + ) + prompt: Final = [ + {"type": "text", "text": "what is this"}, + {"type": "image_url", "image_url": {"url": "https://example.com/desk.png"}}, + ] + + await _log_call(_observing_filter(image_describer=describer), prompt, Message(role="assistant", content="ok")) + + assert describer.enforcer_blocked == [True] 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 4401e3e95c0..2e968ec04f5 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py @@ -41,6 +41,8 @@ from litellm.proxy.guardrails.content_filter_data import DATA_ROOTS from litellm.proxy.guardrails.guardrail_endpoints import ( test_custom_code_guardrail as run_custom_code_test_endpoint, ) +from litellm.proxy.guardrails.guardrail_hooks.azure.prompt_shield import AzureContentSafetyPromptShieldGuardrail +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail MOCK_ADMIN_USER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) from litellm.proxy.guardrails.guardrail_registry import ( @@ -951,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 @@ -3143,7 +3156,36 @@ async def test_ui_settings_map_matches_runtime_supported_event_hooks(): continue assert provider in result.supported_modes_by_provider, provider - assert result.supported_modes_by_provider[provider] == [hook.value for hook in declared], provider + offered = result.supported_modes_by_provider[provider] + assert offered[: len(declared)] == [hook.value for hook in declared], provider + assert set(offered) - {hook.value for hook in declared} <= {"logging_only"}, provider + assert offered.count("logging_only") == ( + 1 if guardrail_class.accepts_logging_only_mode() or GuardrailEventHooks.logging_only in declared else 0 + ), provider + + +@pytest.mark.asyncio +async def test_ui_settings_offer_logging_only_where_the_generic_scan_runs(): + modes_by_provider = (await get_guardrail_ui_settings()).supported_modes_by_provider + + assert modes_by_provider["litellm_content_filter"] == [ + *(hook.value for hook in ContentFilterGuardrail.get_supported_event_hooks()), + "logging_only", + ] + assert modes_by_provider["akto"].count("logging_only") == 1 + assert modes_by_provider["azure/prompt_shield"] == [ + hook.value for hook in AzureContentSafetyPromptShieldGuardrail.get_supported_event_hooks() + ] + with pytest.raises(ValueError, match="not in the supported event hooks"): + AzureContentSafetyPromptShieldGuardrail( + guardrail_name="native-logging-only", + api_key="k", + api_base="https://example.invalid", + event_hook="logging_only", + ) + assert ( + ContentFilterGuardrail(guardrail_name="cf-logging-only", event_hook="logging_only").event_hook == "logging_only" + ) def test_content_filter_runtime_rejects_unsupported_mcp_hook(): diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index f17cad56751..08840978169 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -15,7 +15,7 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro from litellm.proxy.hooks.parallel_request_limiter import ( PROXY_MaxParallelRequestsHandler, ) -from litellm.proxy.utils import InternalUsageCache, hash_token +from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.types.utils import CallTypes, EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage FROZEN_INSTANT: Final = datetime(2026, 1, 31, 23, 59, 30) @@ -353,3 +353,26 @@ async def test_websocket_turn_callbacks_keep_the_sockets_slot_until_the_disconne await handler.async_release_max_parallel_requests_on_disconnect(UserAPIKeyAuth(api_key=api_key)) after_disconnect = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) assert after_disconnect["current_requests"] == 0, after_disconnect + + +@pytest.mark.asyncio +async def test_a_rejected_realtime_attempt_does_not_free_a_slot_it_never_took(): + api_key = hash_token("sk-at-limit") + handler = PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock) + bucket = f"{api_key}::{FROZEN_INSTANT.strftime('%Y-%m-%d-%H-%M')}::request_count" + await handler.internal_usage_cache.async_set_cache( + key=bucket, + value={"current_requests": 1, "current_tpm": 0, "current_rpm": 1}, + litellm_parent_otel_span=None, + ) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = handler + auth = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + + await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(auth) + after_rejection = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) + assert after_rejection["current_requests"] == 1, after_rejection + + await proxy_logging_obj.arelease_responses_websocket_slot(auth) + after_socket_close = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) + assert after_socket_close["current_requests"] == 0, after_socket_close 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_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 1af9def3afb..57e0a4fc47c 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -3407,6 +3407,43 @@ async def test_block_key_existing_key_succeeds(monkeypatch): assert result == mock_updated_record +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint_name", ["block_key", "unblock_key"]) +async def test_block_and_unblock_write_no_audit_row_when_the_key_update_fails(monkeypatch, endpoint_name): + from litellm.proxy._types import BlockKeyRequest + from litellm.proxy.management_endpoints import key_management_endpoints + from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks + + hashed_token = "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + existing_key_record = MagicMock() + existing_key_record.token = hashed_token + existing_key_record.model_dump_json.return_value = f'{{"token": "{hashed_token}"}}' + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key_record) + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(side_effect=RuntimeError("update failed")) + audit_writer = AsyncMock() + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", audit_writer) + monkeypatch.setattr("litellm.store_audit_logs", True) + + with pytest.raises(RuntimeError, match="update failed"): + await getattr(key_management_endpoints, endpoint_name)( + data=BlockKeyRequest(key=hashed_token), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user" + ), + litellm_changed_by=None, + ) + await drain_audit_tasks(timeout=1.0) + + assert audit_writer.call_count == 0 + + @pytest.mark.asyncio async def test_validate_key_team_change_with_member_permissions(): """ 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/management_helpers/test_audit_logs_proxy.py b/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py index 1ea608cfc17..325275ab391 100644 --- a/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py +++ b/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py @@ -1,14 +1,18 @@ import os import traceback from litellm._uuid import uuid -from datetime import datetime +from datetime import datetime, timezone +from collections.abc import Coroutine +from typing import Any, Final from dotenv import load_dotenv from fastapi import Request from fastapi.routing import APIRoute +import contextlib import io +import logging import time # this file is to test litellm/proxy @@ -38,11 +42,16 @@ from starlette.datastructures import URL from litellm.proxy.management_helpers.audit_logs import ( create_audit_log_for_update, + drain_audit_tasks, get_audit_log_changed_by, + track_audit_task, ) from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, UserAPIKeyAuth +from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache -from unittest.mock import patch, AsyncMock +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import StandardAuditLogPayload +from unittest.mock import AsyncMock, MagicMock, patch proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) import json @@ -174,7 +183,6 @@ async def test_create_audit_log_for_update_premium_user(): patch("litellm.store_audit_logs", True), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, ): - mock_prisma.db.litellm_auditlog.create = AsyncMock() request_data = LiteLLM_AuditLogs( @@ -215,9 +223,7 @@ def prisma_client(): os.environ["DATABASE_URL"] = modified_url # Assuming PrismaClient is a class that needs to be instantiated - prisma_client = PrismaClient( - database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj - ) + prisma_client = PrismaClient(database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj) return prisma_client @@ -252,10 +258,482 @@ async def test_create_audit_log_in_db(prisma_client): await asyncio.sleep(1) # now read the last log from the db - last_log = await prisma_client.db.litellm_auditlog.find_first( - where={"id": audit_log_id} - ) + last_log = await prisma_client.db.litellm_auditlog.find_first(where={"id": audit_log_id}) assert last_log.id == audit_log_id setattr(litellm, "store_audit_logs", False) + + +@pytest.mark.asyncio +async def test_track_audit_task_lifecycle(): + task_completed: Final = asyncio.Event() + + async def _sample_coroutine() -> None: + await asyncio.sleep(0) + task_completed.set() + + task: Final = track_audit_task(asyncio.create_task(_sample_coroutine())) + assert not task.done() + + await drain_audit_tasks(timeout=1.0) + assert task_completed.is_set() + assert task.done() + + completed_task: Final = asyncio.create_task(asyncio.sleep(0)) + await completed_task + assert track_audit_task(completed_task) is completed_task + + +@pytest.mark.asyncio +async def test_drain_audit_tasks_waits_for_all_tasks(): + event1: Final = asyncio.Event() + event2: Final = asyncio.Event() + + async def _worker(event: asyncio.Event) -> None: + await asyncio.sleep(0) + event.set() + + task1: Final = track_audit_task(asyncio.create_task(_worker(event1))) + task2: Final = track_audit_task(asyncio.create_task(_worker(event2))) + + assert not task1.done() + assert not task2.done() + + await drain_audit_tasks(timeout=1.0) + + assert event1.is_set() + assert event2.is_set() + assert task1.done() + assert task2.done() + + +@pytest.mark.asyncio +async def test_drain_audit_tasks_handles_failing_task_cleanly(): + async def _failing_worker() -> None: + await asyncio.sleep(0) + raise RuntimeError("test audit log write failure") + + task: Final = track_audit_task(asyncio.create_task(_failing_worker())) + assert not task.done() + + await drain_audit_tasks(timeout=1.0) + assert task.done() + assert isinstance(task.exception(), RuntimeError) + + +@pytest.mark.asyncio +async def test_drain_audit_tasks_timeout_does_not_cancel_pending_tasks(): + async def _long_running_worker() -> None: + await asyncio.Event().wait() + + task: Final = track_audit_task(asyncio.create_task(_long_running_worker())) + try: + assert not task.done() + await drain_audit_tasks(timeout=0) + assert not task.cancelled() + assert not task.done() + finally: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + +@pytest.mark.asyncio +async def test_drain_audit_tasks_persists_management_audit_write_before_db_disconnect(): + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + is_disconnected: Final = [False] + + async def _mock_create(data: dict[str, Any]) -> None: + if is_disconnected[0]: + raise RuntimeError("Database already disconnected") + await asyncio.sleep(0) + writes.append(data) + + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=_mock_create) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + ): + request_data: Final = LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by="admin-user", + changed_by_api_key="sk-admin", + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id="token-123", + action="blocked", + updated_values="{}", + before_value='{"key": "test"}', + ) + + task: Final = track_audit_task(asyncio.create_task(create_audit_log_for_update(request_data=request_data))) + assert not task.done() + + await drain_audit_tasks(timeout=2.0) + is_disconnected[0] = True + + assert task.done() + assert len(writes) == 1 + assert writes[0]["object_id"] == "token-123" + assert writes[0]["action"] == "blocked" + + +@pytest.mark.asyncio +async def test_create_audit_log_auto_registers_in_drain(): + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + + async def _mock_create(data: dict[str, Any]) -> None: + await asyncio.sleep(0) + writes.append(data) + + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=_mock_create) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + ): + request_data: Final = LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by="admin-user", + changed_by_api_key="sk-admin", + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id="token-456", + action="deleted", + updated_values="{}", + before_value='{"key": "test-deleted"}', + ) + + task: Final = asyncio.create_task(create_audit_log_for_update(request_data=request_data)) + await asyncio.sleep(0) + + await drain_audit_tasks(timeout=2.0) + + assert task.done() + assert len(writes) == 1 + assert writes[0]["object_id"] == "token-456" + assert writes[0]["action"] == "deleted" + + +@pytest.mark.asyncio +async def test_drain_audit_tasks_discards_timed_out_tasks_on_timeout(): + async def _hung_worker() -> None: + await asyncio.Event().wait() + + hung_task: Final = track_audit_task(asyncio.create_task(_hung_worker())) + try: + await drain_audit_tasks(timeout=0) + assert not hung_task.done() + + second_drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60)) + for _ in range(3): + await asyncio.sleep(0) + assert second_drain.done() + finally: + hung_task.cancel() + try: + await hung_task + except asyncio.CancelledError: + pass + + +@pytest.mark.asyncio +async def test_drain_audit_tasks_captures_tasks_queued_during_drain(): + finished_tasks: Final[list[str]] = [] + + async def _first_worker() -> None: + await asyncio.sleep(0) + track_audit_task(asyncio.create_task(_second_worker())) + finished_tasks.append("first") + + async def _second_worker() -> None: + await asyncio.sleep(0) + finished_tasks.append("second") + + track_audit_task(asyncio.create_task(_first_worker())) + await drain_audit_tasks(timeout=2.0) + + assert finished_tasks == ["first", "second"] + + +class _AuditCallbackBlockedUntilReleased(CustomLogger): + def __init__(self, release: asyncio.Event) -> None: + super().__init__() + self.release: Final = release + self.logged_object_ids: Final[list[str]] = [] + + async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> None: + await self.release.wait() + self.logged_object_ids.append(audit_log["object_id"]) + + +class _WarningRecorder(logging.Handler): + def __init__(self) -> None: + super().__init__(level=logging.WARNING) + self.messages: Final[list[str]] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +def _deleted_key_audit_log(object_id: str) -> LiteLLM_AuditLogs: + return LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by="admin-user", + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=object_id, + action="deleted", + updated_values="{}", + before_value="{}", + ) + + +@pytest.mark.asyncio +async def test_drain_waits_for_a_slow_audit_log_callback(): + release: Final = asyncio.Event() + callback: Final = _AuditCallbackBlockedUntilReleased(release) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.audit_log_callbacks", [callback]), + ): + await create_audit_log_for_update(request_data=_deleted_key_audit_log("token-callback")) + drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60)) + for _ in range(3): + await asyncio.sleep(0) + drain_returned_before_the_callback_finished: Final = drain.done() + release.set() + await drain + + assert not drain_returned_before_the_callback_finished + assert callback.logged_object_ids == ["token-callback"] + + +@pytest.mark.asyncio +async def test_drain_timeout_warning_names_each_dropped_audit_row(): + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=_io_that_never_finishes) + recorder: Final = _WarningRecorder() + verbose_proxy_logger.addHandler(recorder) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + ): + writer: Final = asyncio.create_task( + create_audit_log_for_update(request_data=_deleted_key_audit_log("token-dropped")) + ) + try: + for _ in range(3): + await asyncio.sleep(0) + await drain_audit_tasks(timeout=0.05) + finally: + verbose_proxy_logger.removeHandler(recorder) + writer.cancel() + with contextlib.suppress(asyncio.CancelledError): + await writer + + assert recorder.messages == [ + "Timed out draining 1 audit tasks on shutdown, dropping: " + f"deleted {LitellmTableNames.KEY_TABLE_NAME.value} token-dropped to the database" + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("store_audit_logs", "premium_user", "expected_write_count"), + [(True, True, 1), (False, True, 0), (True, False, 0)], + ids=["row_written", "logging_disabled", "not_premium"], +) +async def test_drain_does_not_wait_for_code_that_runs_after_an_inline_audit_write( + store_audit_logs: bool, premium_user: bool, expected_write_count: int +): + release_caller: Final = asyncio.Event() + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data)) + + async def _caller_that_outlives_the_audit_call() -> None: + await create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by="test-user", + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id="test-obj", + action="updated", + updated_values="{}", + before_value="{}", + ) + ) + await release_caller.wait() + + with ( + patch("litellm.store_audit_logs", store_audit_logs), + patch("litellm.proxy.proxy_server.premium_user", premium_user), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + ): + caller: Final = asyncio.create_task(_caller_that_outlives_the_audit_call()) + for _ in range(3): + await asyncio.sleep(0) + drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60)) + for _ in range(3): + await asyncio.sleep(0) + drain_returned_while_caller_was_running: Final = drain.done() + release_caller.set() + await caller + await drain + + assert drain_returned_while_caller_was_running + assert len(writes) == expected_write_count + + +async def _io_that_never_finishes(*args: object, **kwargs: object) -> None: + await asyncio.Event().wait() + + +async def _drain_while_hook_io_hangs(hook: Coroutine[object, object, None]) -> None: + hook_task: Final = asyncio.create_task(hook) + for _ in range(3): + await asyncio.sleep(0) + await drain_audit_tasks(timeout=2.0) + hook_task.cancel() + + +@pytest.mark.asyncio +async def test_key_generated_audit_write_is_drained_while_the_invite_email_hangs(): + from litellm.proxy._types import GenerateKeyRequest, GenerateKeyResponse + from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks + + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data)) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object(KeyManagementEventHooks, "_send_key_created_email", side_effect=_io_that_never_finishes), + ): + await _drain_while_hook_io_hangs( + KeyManagementEventHooks.async_key_generated_hook( + data=GenerateKeyRequest(send_invite_email=True), + response=GenerateKeyResponse(key="sk-new", token_id="hashed-new"), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"), + ) + ) + + assert [(write["object_id"], write["action"]) for write in writes] == [("hashed-new", "created")] + + +@pytest.mark.asyncio +async def test_key_updated_audit_write_is_drained_while_the_secret_manager_rename_hangs(): + from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest + from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks + + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data)) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object( + KeyManagementEventHooks, "_rename_virtual_key_in_secret_manager", side_effect=_io_that_never_finishes + ), + ): + await _drain_while_hook_io_hangs( + KeyManagementEventHooks.async_key_updated_hook( + data=UpdateKeyRequest(key="sk-existing", key_alias="new-alias"), + existing_key_row=LiteLLM_VerificationToken(token="hashed-existing", key_alias="old-alias"), + response=None, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"), + ) + ) + + assert [(write["object_id"], write["action"]) for write in writes] == [(hash_token("sk-existing"), "updated")] + + +@pytest.mark.asyncio +async def test_key_rotated_audit_write_is_drained_while_the_secret_manager_rotation_hangs(): + from litellm.proxy._types import GenerateKeyResponse, LiteLLM_VerificationToken + from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks + + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data)) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object( + KeyManagementEventHooks, "_rotate_virtual_key_in_secret_manager", side_effect=_io_that_never_finishes + ), + ): + await _drain_while_hook_io_hangs( + KeyManagementEventHooks.async_key_rotated_hook( + data=None, + existing_key_row=LiteLLM_VerificationToken(token="hashed-old"), + response=GenerateKeyResponse(key="sk-rotated", token_id="hashed-new"), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"), + ) + ) + + assert [(write["object_id"], write["action"]) for write in writes] == [("hashed-old", "rotated")] + + +@pytest.mark.asyncio +async def test_user_created_audit_write_is_drained_while_its_user_read_is_still_running(): + from litellm.proxy._types import NewUserRequest, NewUserResponse + from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks + + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + release_user_read: Final = asyncio.Event() + + async def _slow_user_read(*args: object, **kwargs: object) -> dict[str, str]: + await release_user_read.wait() + return {"user_id": "new-user"} + + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data)) + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(side_effect=_slow_user_read) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object(UserManagementEventHooks, "async_send_user_invitation_email", side_effect=_io_that_never_finishes), + ): + hook: Final = asyncio.create_task( + UserManagementEventHooks.async_user_created_hook( + data=NewUserRequest(user_id="new-user", send_invite_email=True), + response=NewUserResponse(key="sk-new", user_id="new-user"), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"), + ) + ) + for _ in range(3): + await asyncio.sleep(0) + drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60)) + for _ in range(3): + await asyncio.sleep(0) + drain_waited_for_the_user_read: Final = not drain.done() + release_user_read.set() + await drain + hook.cancel() + + assert drain_waited_for_the_user_read + assert [(write["object_id"], write["action"]) for write in writes] == [("new-user", "created")] 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_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 9ef5f16a240..21170abaeb1 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -24,10 +24,15 @@ from starlette.datastructures import FormData import litellm from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger from tests._master_key import MASTER_KEY as SHARED_MASTER_KEY from litellm.caching.caching import DualCache from litellm.types.utils import CallTypesLiteral from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException, ModifyResponseException +from litellm.proxy._types import AwsErrorTypeProxyException, ProxyException +from litellm.proxy.utils import ProxyLogging +from litellm.passthrough.main import raise_for_provider_status from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( @@ -40,6 +45,7 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( azure_proxy_route, azure_speech_proxy_route, bedrock_llm_proxy_route, + bedrock_passthrough_error_response, bedrock_proxy_route, create_pass_through_route, cursor_proxy_route, @@ -1658,17 +1664,7 @@ class TestBedrockLLMProxyRoute: bedrock_error_message = '{"message":"ContentBlock object at messages.0.content.0 must set one of the following keys: text, image, toolUse, toolResult, document, video."}' - # Create a mock httpx.Response for the error - mock_error_response = Mock(spec=httpx.Response) - mock_error_response.status_code = 400 - mock_error_response.aread = AsyncMock(return_value=bedrock_error_message.encode("utf-8")) - - # Create the HTTPStatusError - mock_http_error = httpx.HTTPStatusError( - message="Bad Request", - request=Mock(spec=httpx.Request), - response=mock_error_response, - ) + mock_http_error = _bedrock_status_error(400, bedrock_error_message.encode("utf-8")) # Create mocks for all required parameters mock_request = MagicMock(spec=Request) @@ -1926,6 +1922,284 @@ class TestBedrockLLMProxyRoute: assert "Blocked by guardrail" in str(exc_info.value.detail) +_THROTTLE_TYPE: Final = "ThrottlingException:http://internal.amazon.com/coral/com.amazon.bedrock/" +_THROTTLE_BODY: Final = b'{"message":"Too many requests, please wait before trying again."}' + + +def _status_error( + status_code: int, + content: bytes, + headers: Mapping[str, str], + url: str, + raise_for_status: Callable[[httpx.Response], None], +) -> httpx.HTTPStatusError: + request: Final = httpx.Request("POST", url) + response: Final = httpx.Response(status_code, headers=dict(headers), content=content, request=request) + with pytest.raises(httpx.HTTPStatusError) as raised: + raise_for_status(response) + return raised.value + + +def _bedrock_status_error( + status_code: int, + content: bytes, + headers: Mapping[str, str] = MappingProxyType({}), + url: str = "https://bedrock-runtime.us-east-1.amazonaws.com/model/m/converse", +) -> httpx.HTTPStatusError: + """An error Bedrock answered the passthrough request with.""" + return _status_error(status_code, content, headers, url, raise_for_provider_status) + + +def _guardrail_service_status_error( + status_code: int, content: bytes, headers: Mapping[str, str] +) -> httpx.HTTPStatusError: + """An error a guardrail service's own HTTP call raised, from a URL shaped like a Bedrock model call.""" + return _status_error( + status_code, + content, + headers, + "https://bedrock-runtime.us-east-1.amazonaws.com/model/m/converse", + httpx.Response.raise_for_status, + ) + + +def _guardrail_detail(message: str) -> dict[str, str]: + return {"error": message, "guardrail_name": "block-demo"} + + +class _ReplacesFailures(CustomLogger): + def __init__(self, replacement: HTTPException) -> None: + super().__init__() + self.replacement: Final = replacement + + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> HTTPException | None: + return self.replacement + + +class TestBedrockPassthroughAwsErrors: + """A guardrail block or an error Bedrock returned reaches an AWS SDK caller with its code and message.""" + + @staticmethod + async def _answer(error: Exception) -> Response: + return await bedrock_passthrough_error_response( + ProxyBaseLLMRequestProcessing(data={}), + error, + UserAPIKeyAuth(api_key="test-key"), + ProxyLogging(user_api_key_cache=DualCache()), + ) + + @pytest.mark.asyncio + async def test_bedrock_error_reaches_the_caller_with_bedrocks_status_header_and_body( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + error: Final = _bedrock_status_error( + 429, + _THROTTLE_BODY, + { + "x-amzn-errortype": _THROTTLE_TYPE, + "x-amzn-requestid": "req-1", + "content-type": "application/x-amz-json-1.1", + }, + ) + + response: Final = await self._answer(error) + + assert response.status_code == 429 + assert response.body == _THROTTLE_BODY + assert response.headers["x-amzn-errortype"] == _THROTTLE_TYPE + assert response.headers["content-type"] == "application/x-amz-json-1.1" + assert response.headers["llm_provider-x-amzn-requestid"] == "req-1" + + @pytest.mark.asyncio + async def test_arn_model_error_reaches_the_caller_as_bedrock_sent_it(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + error: Final = _bedrock_status_error( + 429, + _THROTTLE_BODY, + {"x-amzn-errortype": _THROTTLE_TYPE}, + url="https://bedrock-runtime.us-east-1.amazonaws.com/model/" + "arn%3Aaws%3Abedrock%3Aus-east-1%3A123%3Aapplication-inference-profile%2Fabc123/converse-stream?trace=a/b", + ) + + response: Final = await self._answer(error) + + assert (response.status_code, response.body) == (429, _THROTTLE_BODY) + assert response.headers["x-amzn-errortype"] == _THROTTLE_TYPE + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("error", "status_code", "body"), + ( + pytest.param( + _bedrock_status_error(502, b"bad gateway"), + 502, + "bad gateway", + id="bedrock-without-error-type", + ), + pytest.param( + _guardrail_service_status_error(429, _THROTTLE_BODY, {"x-amzn-errortype": _THROTTLE_TYPE}), + 429, + _THROTTLE_BODY.decode(), + id="guardrail-service-with-error-type", + ), + ), + ) + async def test_http_error_bedrock_did_not_answer_with_an_aws_error_type_keeps_the_litellm_error( + self, monkeypatch: pytest.MonkeyPatch, error: httpx.HTTPStatusError, status_code: int, body: str + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + + with pytest.raises(HTTPException) as exc_info: + await self._answer(error) + + assert exc_info.value.status_code == status_code + assert exc_info.value.detail == {"error": body} + assert "x-amzn-errortype" not in (exc_info.value.headers or {}) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("error", "replacement", "code"), + ( + pytest.param( + _bedrock_status_error(429, _THROTTLE_BODY, {"x-amzn-errortype": _THROTTLE_TYPE}), + HTTPException(status_code=503, detail={"error": "rewritten by a callback"}), + "503", + id="bedrock-error", + ), + pytest.param( + HTTPException(status_code=400, detail=_guardrail_detail("Violated guardrail policy")), + HTTPException(status_code=429, detail={"error": "rewritten by a callback"}), + "429", + id="guardrail-block", + ), + ), + ) + async def test_error_replaced_by_a_failure_hook_keeps_the_replacement( + self, monkeypatch: pytest.MonkeyPatch, error: Exception, replacement: HTTPException, code: str + ) -> None: + monkeypatch.setattr(litellm, "callbacks", [_ReplacesFailures(replacement)]) + + with pytest.raises(ProxyException) as exc_info: + await self._answer(error) + + assert (exc_info.value.message, exc_info.value.code) == ("rewritten by a callback", code) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("error", "message", "status_code"), + ( + pytest.param( + GuardrailRaisedException( + guardrail_name="block-demo", + message="synthetic policy says no", + should_wrap_with_default_message=False, + blocked_content=True, + ), + "synthetic policy says no", + "400", + id="guardrail-raised", + ), + pytest.param( + HTTPException(status_code=400, detail=_guardrail_detail("Violated guardrail policy")), + "Violated guardrail policy", + "400", + id="guardrail-http-400", + ), + pytest.param( + HTTPException(status_code=403, detail=_guardrail_detail("blocked by policy")), + "blocked by policy", + "403", + id="guardrail-http-403", + ), + pytest.param( + BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="pii"), + "Blocked entity detected: EMAIL_ADDRESS by Guardrail: pii. This entity is not allowed to be used in this request.", + "400", + id="blocked-pii", + ), + ), + ) + async def test_guardrail_block_is_raised_with_its_aws_error_type( + self, monkeypatch: pytest.MonkeyPatch, error: Exception, message: str, status_code: str + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + + with pytest.raises(AwsErrorTypeProxyException) as exc_info: + await self._answer(error) + + assert exc_info.value.aws_error_type == "ValidationException" + assert (exc_info.value.error.message, exc_info.value.error.code) == (message, status_code) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("error", "message", "code"), + ( + pytest.param( + HTTPException(status_code=400, detail={"error": "model not found"}), + "model not found", + "400", + id="request-check-400", + ), + pytest.param( + HTTPException(status_code=403, detail="cache control not allowed"), + "cache control not allowed", + "403", + id="request-check-403", + ), + pytest.param( + HTTPException(status_code=429, detail=_guardrail_detail("vendor rate limited")), + "vendor rate limited", + "429", + id="429", + ), + pytest.param( + GuardrailRaisedException( + guardrail_name="fail-closed", + message="guardrail vendor unreachable", + should_wrap_with_default_message=False, + ), + "guardrail vendor unreachable", + "400", + id="guardrail-vendor-unreachable-400", + ), + pytest.param( + GuardrailRaisedException( + guardrail_name="fail-closed", + message="guardrail vendor unreachable", + should_wrap_with_default_message=False, + status_code=503, + ), + "guardrail vendor unreachable", + "503", + id="guardrail-vendor-unreachable-503", + ), + pytest.param( + ModifyResponseException(message="replaced", model="bedrock-model", request_data={}), + "replaced", + "500", + id="modify-response", + ), + ), + ) + async def test_errors_that_are_not_guardrail_blocks_keep_the_litellm_error( + self, monkeypatch: pytest.MonkeyPatch, error: Exception, message: str, code: str + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + + with pytest.raises(ProxyException) as exc_info: + await self._answer(error) + + assert type(exc_info.value) is ProxyException + assert (exc_info.value.message, exc_info.value.code) == (message, code) + + class TestBedrockAgentRuntimePassthroughToggle: AGENT_RUNTIME_ENDPOINT: Final = "knowledgebases/KB1234567/retrieve" MODEL_ENDPOINT: Final = "model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/converse" 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_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 69f9c72c9c8..71daa3bb873 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -17,12 +17,14 @@ Pins covered: from __future__ import annotations +import asyncio import inspect import json import logging import os import subprocess from collections.abc import Awaitable, Callable +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -48,6 +50,8 @@ from litellm.proxy.proxy_server import ( save_worker_config, ) +from litellm.proxy.management_helpers.audit_logs import track_audit_task + from .conftest import normalize # --------------------------------------------------------------------------- @@ -223,6 +227,46 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke assert calls == ["flush", "disconnect"] +@pytest.mark.asyncio +async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_before_disconnecting(monkeypatch): + """ + A management call returns before its audit row is written, so shutdown waits + for that write while prisma is still connected. It waits beside the gateway + flush rather than before it, so the two waits overlap inside the teardown budget. + """ + calls: list[str] = [] # mutable-ok: records call order, which is the assertion + gateway_flush_started: Final = asyncio.Event() + + fake_prisma = MagicMock() + fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect")) + monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) + + async def _record_flush(client, accumulator): + gateway_flush_started.set() + calls.append("flush") + + async def _audit_write_that_finishes_after_the_flush_starts() -> None: + await gateway_flush_started.wait() + calls.append("audit") + + monkeypatch.setattr(ps, "flush_gateway_requests", _record_flush, raising=False) + + fake_jwt = MagicMock() + fake_jwt.close = AsyncMock() + monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False) + monkeypatch.setattr(ps, "db_writer_client", None, raising=False) + + import litellm + + monkeypatch.setattr(litellm, "cache", None, raising=False) + monkeypatch.setattr(litellm, "success_callback", [], raising=False) + + track_audit_task(asyncio.create_task(_audit_write_that_finishes_after_the_flush_starts())) + await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) + + assert calls == ["flush", "audit", "disconnect"] + + @pytest.mark.asyncio async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch): """No prisma client means nothing to drain to, and no attempt is made.""" @@ -245,6 +289,34 @@ async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch assert flush.await_count == 0 +@pytest.mark.asyncio +async def test_proxy_shutdown_without_a_database_drains_audit_callbacks_before_the_rest_of_teardown(monkeypatch): + calls: list[str] = [] # mutable-ok: records call order, which is the assertion + monkeypatch.setattr(ps, "prisma_client", None, raising=False) + monkeypatch.setattr( + ps, "drain_post_call_redis_batches", AsyncMock(side_effect=lambda: calls.append("redis")), raising=False + ) + + fake_jwt = MagicMock() + fake_jwt.close = AsyncMock() + monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False) + monkeypatch.setattr(ps, "db_writer_client", None, raising=False) + + import litellm + + monkeypatch.setattr(litellm, "cache", None, raising=False) + monkeypatch.setattr(litellm, "success_callback", [], raising=False) + + async def _audit_callback_still_sending() -> None: + await asyncio.sleep(0.05) + calls.append("audit") + + track_audit_task(asyncio.create_task(_audit_callback_still_sending())) + await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) + + assert calls == ["audit", "redis"] + + @pytest.mark.asyncio async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch): fake_prisma = MagicMock() 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_routes_utils.py b/tests/unit/proxy/proxy_server/test_routes_utils.py index 295b9e5ae8b..118b80da31e 100644 --- a/tests/unit/proxy/proxy_server/test_routes_utils.py +++ b/tests/unit/proxy/proxy_server/test_routes_utils.py @@ -390,8 +390,7 @@ def test_token_counter_fallback_counts_tools_system_and_anthropic_blocks(client, assert full > bare -def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, monkeypatch): - """Regression: a ``prompt`` request carrying ``tools`` but no ``messages`` still counts, because the fallback attaches tools only when counting messages (``token_counter`` rejects tools on the text path).""" +def test_token_counter_fallback_counts_a_prompt_with_tools_like_its_messages_form(client, auth_as, monkeypatch): monkeypatch.setattr(proxy_server, "llm_router", None) monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False) prompt = "count the tokens in this sentence please" @@ -412,4 +411,6 @@ def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, ) assert response.status_code == 200, response.text - assert response.json()["total_tokens"] == litellm.token_counter(model="claude-fable-5", text=prompt) + assert response.json()["total_tokens"] == litellm.token_counter( + model="claude-fable-5", messages=[{"role": "user", "content": prompt}], tools=tools + ) 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/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 5e6b074a17e..09aa9a1b0c8 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1952,7 +1952,7 @@ class TestResponsesWSFirstFrameModelAuth: proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with ( @@ -2017,7 +2017,7 @@ class TestResponsesWSFirstFrameModelAuth: ) proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with ( @@ -3456,10 +3456,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 = ( @@ -3476,6 +3477,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"}) @@ -3483,6 +3485,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?"}) @@ -3596,6 +3599,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_background_interaction_settlement.py b/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py index 21384279bbb..eea7f8a81aa 100644 --- a/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py +++ b/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py @@ -1,22 +1,28 @@ import asyncio import time from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import Optional import pytest import litellm.interactions.background_cost_polling as bg +from litellm import Router +from litellm.constants import BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS from litellm.interactions.background_cost_polling import ( _create_context, + BackgroundInteractionPollContext, configure_background_settlement_store, maybe_settle_background_interaction_before_delete, PendingBackgroundInteraction, + poll_and_log_background_interaction_cost, PollSchedule, + resume_unsettled_background_interactions, ) from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.proxy.spend_tracking.background_interaction_settlement import ( configure_background_interaction_settlement, + fetch_with_deployment_credentials, install_background_interaction_settlement, PrismaBackgroundSettlementStore, ) @@ -78,7 +84,42 @@ class _FakeSettlementTable: return len(matched) def _matching(self, where) -> list: - return [row for row in self.rows.values() if all(getattr(row, column) == value for column, value in where.items())] + return [row for row in self.rows.values() if _matches(row, where)] + + +def _matches(row: "_Row", where: dict) -> bool: + return all(_satisfies(row, column, condition) for column, condition in where.items()) + + +def _satisfies(row: "_Row", column: str, condition) -> bool: + value = getattr(row, column) + if isinstance(condition, dict): + return value is not None and value < condition["lt"] + return value == condition + + +class _FakeSpendLogs: + """The spend log table, holding only the request ids of the rows written so far.""" + + def __init__(self, request_ids: tuple[str, ...] = ()): + self.request_ids = request_ids + + async def find_unique(self, *, where): + return where["request_id"] if where["request_id"] in self.request_ids else None + + +def _store( + table, + claimed_by: str, + claim_lease_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS, +) -> PrismaBackgroundSettlementStore: + return PrismaBackgroundSettlementStore( + table=table, + spend_logs=_FakeSpendLogs(), + spend_logs_kept=lambda: True, + claimed_by=claimed_by, + claim_lease_seconds=claim_lease_seconds, + ) def _logging_obj(metadata: Optional[dict] = None) -> LitellmLogging: @@ -110,7 +151,12 @@ def _pending(interaction_id: str) -> PendingBackgroundInteraction: ) -def _stored_row(interaction_id: str, claimed: bool = False, create_context: Optional[object] = None) -> _Row: +def _stored_row( + interaction_id: str, + claimed_at: Optional[datetime] = None, + settled: bool = False, + create_context: Optional[object] = None, +) -> _Row: return _Row( interaction_id=interaction_id, custom_llm_provider="gemini", @@ -120,11 +166,21 @@ def _stored_row(interaction_id: str, claimed: bool = False, create_context: Opti else _create_context(_logging_obj(), "gemini").model_dump(mode="json") ), created_at=datetime.now(timezone.utc), - claimed_at=datetime.now(timezone.utc) if claimed else None, - claimed_by="replica-a:1" if claimed else None, + claimed_at=claimed_at, + claimed_by="replica-a:1" if claimed_at is not None else None, + settled_at=claimed_at if settled else None, + outcome="billed" if settled else None, ) +def _minutes_ago(minutes: float) -> datetime: + return datetime.now(timezone.utc) - timedelta(minutes=minutes) + + +def _past_the_default_lease() -> datetime: + return datetime.now(timezone.utc) - timedelta(seconds=BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS + 60) + + def _completed(interaction_id: str) -> InteractionsAPIResponse: return InteractionsAPIResponse( id=interaction_id, model="gemini-2.5-flash", status="completed", steps=[], usage=dict(USAGE_BLOCK) @@ -144,20 +200,20 @@ def _capturing_fetch(): @pytest.mark.asyncio async def test_registered_row_reads_back_as_the_same_pending_interaction(): table = _FakeSettlementTable() - store = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-a:1") + store = _store(table, claimed_by="replica-a:1") pending = _pending("interactions/bg-1") await store.register(pending) assert await store.pending("interactions/bg-1") == pending - assert await store.unclaimed() == (pending,) + assert await store.unsettled() == (pending,) @pytest.mark.asyncio async def test_claim_is_won_by_exactly_one_settler(): table = _FakeSettlementTable() - replica_a = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-a:1") - replica_b = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-b:1") + replica_a = _store(table, claimed_by="replica-a:1") + replica_b = _store(table, claimed_by="replica-b:1") await replica_a.register(_pending("interactions/bg-1")) assert await replica_b.claim("interactions/bg-1") is True @@ -200,7 +256,7 @@ class _MissingSettlementTable: async def test_a_missing_table_holds_no_rows_and_takes_no_registration(): from prisma.errors import TableNotFoundError - store = PrismaBackgroundSettlementStore(table=_MissingSettlementTable(), claimed_by="replica-a:1") + store = _store(_MissingSettlementTable(), claimed_by="replica-a:1") with pytest.raises(TableNotFoundError): await store.register(_pending("interactions/bg-1")) @@ -208,29 +264,62 @@ async def test_a_missing_table_holds_no_rows_and_takes_no_registration(): assert await store.is_claimed("interactions/bg-1") is False assert await store.claim("interactions/bg-1") is False with pytest.raises(TableNotFoundError): - await store.unclaimed() + await store.unsettled() @pytest.mark.asyncio -async def test_unclaimed_skips_claimed_and_unreadable_rows(): +async def test_unsettled_lists_claimed_rows_with_no_outcome_and_when_their_lease_ends(): + claimed_at = _minutes_ago(1) table = _FakeSettlementTable( rows=( _stored_row("interactions/bg-orphaned"), - _stored_row("interactions/bg-settled", claimed=True), + _stored_row("interactions/bg-mid-settlement", claimed_at=claimed_at), + _stored_row("interactions/bg-settled", claimed_at=claimed_at, settled=True), _stored_row("interactions/bg-from-the-future", create_context={"schema": "unknown"}), ) ) - store = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-b:1") + store = _store(table, claimed_by="replica-b:1", claim_lease_seconds=300) - unclaimed = await store.unclaimed() + unsettled = await store.unsettled() - assert [row.interaction_id for row in unclaimed] == ["interactions/bg-orphaned"] + assert [(row.interaction_id, row.claim_expires_at) for row in unsettled] == [ + ("interactions/bg-orphaned", None), + ("interactions/bg-mid-settlement", claimed_at + timedelta(seconds=300)), + ] + + +@pytest.mark.asyncio +async def test_a_claim_whose_settler_died_is_taken_over_once_its_lease_runs_out(): + """ + The regression: a settler that died between its claim and its outcome left + the row claimed forever, so no replica, restart, or delete ever billed it. + """ + table = _FakeSettlementTable( + rows=( + _stored_row("interactions/bg-dead-settler", claimed_at=_minutes_ago(10)), + _stored_row("interactions/bg-settled", claimed_at=_minutes_ago(10), settled=True), + ) + ) + replica_b = _store(table, claimed_by="replica-b:1", claim_lease_seconds=300) + replica_c = _store(table, claimed_by="replica-c:1", claim_lease_seconds=300) + + assert await replica_b.is_claimed("interactions/bg-dead-settler") is False + assert await replica_b.pending("interactions/bg-dead-settler") is not None + assert await replica_b.claim("interactions/bg-dead-settler") is True + assert await replica_c.claim("interactions/bg-dead-settler") is False + assert await replica_c.is_claimed("interactions/bg-dead-settler") is True + assert table.rows["interactions/bg-dead-settler"].claimed_by == "replica-b:1" + + assert await replica_b.is_claimed("interactions/bg-settled") is True + assert await replica_b.pending("interactions/bg-settled") is None + assert await replica_b.claim("interactions/bg-settled") is False + assert table.rows["interactions/bg-settled"].claimed_by == "replica-a:1" @pytest.mark.asyncio async def test_record_outcome_keeps_the_audit_trail_and_drops_the_stored_request_context(): table = _FakeSettlementTable() - store = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-a:1") + store = _store(table, claimed_by="replica-a:1") await store.register(_pending("interactions/bg-1")) assert await store.claim("interactions/bg-1") assert table.rows["interactions/bg-1"].create_context @@ -244,23 +333,198 @@ async def test_record_outcome_keeps_the_audit_trail_and_drops_the_stored_request assert row.create_context == {} +@pytest.mark.asyncio +async def test_a_takeover_records_the_bill_a_settler_wrote_before_dying_instead_of_billing_again(): + """ + The regression: a settler that billed and then died, or failed to record + the outcome, left its claim to lapse with no outcome, and the replica that + took the claim over billed the interaction a second time. + """ + table = _FakeSettlementTable( + rows=( + _stored_row("interactions/bg-billed-then-died", claimed_at=_past_the_default_lease()), + _stored_row("interactions/bg-died-before-billing", claimed_at=_past_the_default_lease()), + ) + ) + fetch, _ = _capturing_fetch() + previous_store = bg._STORE.store + try: + resumed = await configure_background_interaction_settlement( + table=table, + spend_logs=_FakeSpendLogs(request_ids=("interactions/bg-billed-then-died",)), + spend_logs_kept=lambda: True, + claimed_by="replica-b:1", + fetch_interaction=fetch, + schedule=FAST_SCHEDULE, + ) + + assert [await asyncio.wait_for(task, timeout=5) for task in resumed] == [None, "billed"] + assert [(row.claimed_by, row.outcome) for row in table.rows.values()] == [ + ("replica-b:1", "billed"), + ("replica-b:1", "billed"), + ] + assert await bg._STORE.store.unsettled() == () + finally: + configure_background_settlement_store(previous_store) + + +@pytest.mark.asyncio +async def test_with_spend_logs_off_a_takeover_leaves_the_interaction_unbilled_instead_of_billing_it_twice(): + """ + The regression: with spend logs off a settler that billed and then died + leaves no spend row to find, so the replica that took over its lapsed + claim billed the interaction a second time. + """ + table = _FakeSettlementTable( + rows=( + _stored_row("interactions/bg-claimed-then-died", claimed_at=_past_the_default_lease()), + _stored_row("interactions/bg-orphaned"), + ) + ) + fetch, _ = _capturing_fetch() + previous_store = bg._STORE.store + try: + resumed = await configure_background_interaction_settlement( + table=table, + spend_logs=_FakeSpendLogs(), + spend_logs_kept=lambda: False, + claimed_by="replica-b:1", + fetch_interaction=fetch, + schedule=FAST_SCHEDULE, + ) + + assert [await asyncio.wait_for(task, timeout=5) for task in resumed] == [None, "billed"] + assert [(row.claimed_by, row.outcome) for row in table.rows.values()] == [ + ("replica-b:1", "unsettled"), + ("replica-b:1", "billed"), + ] + finally: + configure_background_settlement_store(previous_store) + + +class _SpendLogsWhoseFirstReadFails(_FakeSpendLogs): + def __init__(self): + super().__init__() + self.reads = 0 + + async def find_unique(self, *, where): + self.reads += 1 + if self.reads == 1: + raise RuntimeError("read replica unavailable") + return await super().find_unique(where=where) + + +@pytest.mark.asyncio +async def test_a_takeover_whose_spend_log_read_failed_retries_and_bills_instead_of_stranding_the_row(): + """ + The regression: a takeover wrote its claim before reading the spend log, + so a failed read left the row under a fresh claim of its own, which the + poll then took for another settler's and left unbilled until a reboot. + """ + table = _FakeSettlementTable( + rows=(_stored_row("interactions/bg-dead-settler", claimed_at=_past_the_default_lease()),) + ) + fetch, _ = _capturing_fetch() + previous_store = bg._STORE.store + try: + (resumed,) = await configure_background_interaction_settlement( + table=table, + spend_logs=_SpendLogsWhoseFirstReadFails(), + spend_logs_kept=lambda: True, + claimed_by="replica-b:1", + fetch_interaction=fetch, + schedule=FAST_SCHEDULE, + ) + + assert await asyncio.wait_for(resumed, timeout=5) == "billed" + assert table.rows["interactions/bg-dead-settler"].outcome == "billed" + finally: + configure_background_settlement_store(previous_store) + + +def _logging_obj_whose_bill_takes(seconds: float, bills: list) -> LitellmLogging: + logging_obj = _logging_obj() + + async def bill(result): + bills.append(result.id) + await asyncio.sleep(seconds) + + logging_obj.async_log_background_interaction_completion = bill + return logging_obj + + +@pytest.mark.asyncio +async def test_a_settle_that_outlasts_the_lease_is_not_taken_over_and_billed_twice(): + """ + The regression: a success callback that held the settle past the claim + lease let a booting replica take the live claim over and bill the same + interaction a second time. + """ + lease_seconds = 0.6 + table = _FakeSettlementTable() + creator = _store(table, claimed_by="replica-a:1", claim_lease_seconds=lease_seconds) + resumer = _store(table, claimed_by="replica-b:1", claim_lease_seconds=lease_seconds) + await creator.register(_pending("interactions/bg-slow-settle")) + fetch, _ = _capturing_fetch() + bills = [] + + def poll(store, resumed: bool, start_delay_seconds: float): + return poll_and_log_background_interaction_cost( + BackgroundInteractionPollContext( + interaction_id="interactions/bg-slow-settle", + custom_llm_provider="gemini", + logging_obj=_logging_obj_whose_bill_takes(lease_seconds * 2.5, bills), + initial_interval_seconds=0.001, + max_interval_seconds=0.002, + timeout_seconds=lease_seconds * 3, + start_delay_seconds=start_delay_seconds, + claim_renewal_interval_seconds=lease_seconds / 10, + store=store, + resumed=resumed, + ), + fetch, + ) + + outcomes = await asyncio.gather( + poll(creator, resumed=False, start_delay_seconds=0.0), + poll(resumer, resumed=True, start_delay_seconds=lease_seconds * 1.5), + ) + + assert outcomes == ["billed", None] + assert bills == ["interactions/bg-slow-settle"] + assert table.rows["interactions/bg-slow-settle"].claimed_by == "replica-a:1" + + @pytest.mark.asyncio async def test_configure_installs_the_store_and_resumes_the_orphaned_rows(): table = _FakeSettlementTable( - rows=(_stored_row("interactions/bg-orphaned"), _stored_row("interactions/bg-settled", claimed=True)) + rows=( + _stored_row("interactions/bg-orphaned"), + _stored_row("interactions/bg-dead-settler", claimed_at=_past_the_default_lease()), + _stored_row("interactions/bg-settled", claimed_at=_minutes_ago(1), settled=True), + ) ) fetch, captured = _capturing_fetch() previous_store = bg._STORE.store try: resumed = await configure_background_interaction_settlement( - table=table, claimed_by="replica-b:1", fetch_interaction=fetch, schedule=FAST_SCHEDULE + table=table, + spend_logs=_FakeSpendLogs(), + spend_logs_kept=lambda: True, + claimed_by="replica-b:1", + fetch_interaction=fetch, + schedule=FAST_SCHEDULE, ) - assert len(resumed) == 1 - assert await asyncio.wait_for(resumed[0], timeout=5) == "billed" - assert [context.interaction_id for context in captured] == ["interactions/bg-orphaned"] - assert table.rows["interactions/bg-orphaned"].claimed_by == "replica-b:1" - assert table.rows["interactions/bg-orphaned"].outcome == "billed" + assert [await asyncio.wait_for(task, timeout=5) for task in resumed] == ["billed", "billed"] + assert sorted(context.interaction_id for context in captured) == [ + "interactions/bg-dead-settler", + "interactions/bg-orphaned", + ] + for interaction_id in ("interactions/bg-orphaned", "interactions/bg-dead-settler"): + assert table.rows[interaction_id].claimed_by == "replica-b:1" + assert table.rows[interaction_id].outcome == "billed" + assert table.rows["interactions/bg-settled"].claimed_by == "replica-a:1" await table.create( data={ @@ -288,7 +552,9 @@ class _PrismaClientWithoutSettlementTable: async def test_install_keeps_booting_when_the_settlement_table_is_unreachable(): previous_store = bg._STORE.store - await install_background_interaction_settlement(_PrismaClientWithoutSettlementTable()) + await install_background_interaction_settlement( + _PrismaClientWithoutSettlementTable(), router=lambda: None, spend_logs_kept=lambda: True + ) assert bg._STORE.store is previous_store @@ -296,3 +562,166 @@ async def test_install_keeps_booting_when_the_settlement_table_is_unreachable(): @dataclass(frozen=True) class _JsonLike: data: object + + +@pytest.mark.asyncio +async def test_a_resumed_row_still_under_a_live_claim_waits_for_the_lease_before_settling(): + claimed_at = datetime.now(timezone.utc) + table = _FakeSettlementTable(rows=(_stored_row("interactions/bg-mid-settlement", claimed_at=claimed_at),)) + store = _store(table, claimed_by="replica-b:1", claim_lease_seconds=0.2) + fetched_at = [] + + async def fetch(context): + fetched_at.append(datetime.now(timezone.utc)) + return _completed(context.interaction_id) + + (resumed,) = await resume_unsettled_background_interactions(store, fetch, schedule=FAST_SCHEDULE) + + assert await asyncio.wait_for(resumed, timeout=5) == "billed" + assert fetched_at[0] >= claimed_at + timedelta(seconds=0.2) + assert table.rows["interactions/bg-mid-settlement"].claimed_by == "replica-b:1" + + +def _router_with_a_deployment_scoped_key() -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-background", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_key": "deployment-scoped-key", + "api_base": "https://deployment-scoped.example", + }, + "model_info": {"id": "deployment-1"}, + } + ] + ) + + +def _router_with_a_catch_all_route() -> Router: + return Router(model_list=[{"model_name": "*", "litellm_params": {"model": "gemini/*", "api_key": "catch-all-key"}}]) + + +def _router_with_two_deployments_in_the_group() -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-background", + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": f"deployment-key-{index}"}, + "model_info": {"id": f"deployment-{index}"}, + } + for index in (2, 3) + ] + ) + + +def _router_with_a_vertex_deployment() -> Router: + return Router( + model_list=[ + { + "model_name": "vertex-background", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_project": "deployment-project", + "vertex_location": "europe-west4", + "vertex_credentials": "/secrets/deployment-service-account.json", + }, + "model_info": {"id": "vertex-deployment-1"}, + } + ] + ) + + +def _create_context_served_by(deployment_id: str) -> dict: + metadata = { + "user_api_key": "0123456789abcdef" * 4, + "model_info": {"id": deployment_id}, + "model_group": "gemini-background", + } + return _create_context(_logging_obj(metadata), "gemini").model_dump(mode="json") + + +@pytest.mark.parametrize("deployment_id", ["deployment-1", "deployment-1-before-an-edit"]) +@pytest.mark.asyncio +async def test_a_resumed_poll_fetches_with_the_credentials_of_the_deployment_that_served_the_create(deployment_id): + """ + The regression: a resumed poll has no request to take credentials from, so + it fetched with the environment's key, which a deployment-scoped key is + not, and every fetch failed until the row was given up unbilled. Editing + the deployment since gives it a new id, and its model group still names it. + """ + table = _FakeSettlementTable( + rows=(_stored_row("interactions/bg-1", create_context=_create_context_served_by(deployment_id)),) + ) + router = _router_with_a_deployment_scoped_key() + fetch, captured = _capturing_fetch() + previous_store = bg._STORE.store + try: + (resumed,) = await configure_background_interaction_settlement( + table=table, + spend_logs=_FakeSpendLogs(), + spend_logs_kept=lambda: True, + claimed_by="replica-b:1", + fetch_interaction=fetch_with_deployment_credentials(lambda: router, fetch), + schedule=FAST_SCHEDULE, + ) + + assert await asyncio.wait_for(resumed, timeout=5) == "billed" + assert (captured[0].credentials.api_key, captured[0].credentials.api_base) == ( + "deployment-scoped-key", + "https://deployment-scoped.example", + ) + finally: + configure_background_settlement_store(previous_store) + + +@pytest.mark.parametrize( + "router, deployment_id, model_group", + [ + (_router_with_a_deployment_scoped_key, "deployment-removed-since", None), + (_router_with_a_deployment_scoped_key, None, None), + (lambda: None, "deployment-1", "gemini-background"), + (_router_with_a_catch_all_route, "deployment-removed-since", None), + (_router_with_two_deployments_in_the_group, "deployment-removed-since", "gemini-background"), + (_router_with_a_vertex_deployment, "deployment-removed-since", "vertex-background"), + (_router_with_a_vertex_deployment, "vertex-deployment-1", None), + ], +) +@pytest.mark.asyncio +async def test_a_resumed_poll_without_a_known_deployment_fetches_with_the_environment_credentials( + router, deployment_id, model_group +): + live_router = router() + fetch, captured = _capturing_fetch() + context = BackgroundInteractionPollContext( + interaction_id="interactions/bg-1", + custom_llm_provider="gemini", + logging_obj=_logging_obj(), + deployment_id=deployment_id, + model_group=model_group, + ) + + await fetch_with_deployment_credentials(lambda: live_router, fetch)(context) + + assert captured == [context] + + +@pytest.mark.asyncio +async def test_a_resumed_vertex_poll_fetches_with_the_project_location_and_credentials_of_its_deployment(): + router = _router_with_a_vertex_deployment() + fetch, captured = _capturing_fetch() + context = BackgroundInteractionPollContext( + interaction_id="interactions/bg-1", + custom_llm_provider="vertex_ai", + logging_obj=_logging_obj(), + deployment_id="vertex-deployment-1", + ) + + await fetch_with_deployment_credentials(lambda: router, fetch)(context) + + credentials = captured[0].credentials + assert (credentials.vertex_project, credentials.vertex_location, credentials.vertex_credentials) == ( + "deployment-project", + "europe-west4", + "/secrets/deployment-service-account.json", + ) 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 9f30e7050b3..02dddfdf25d 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -14,8 +14,8 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm import litellm.constants as litellm_constants import litellm.proxy.spend_tracking.spend_tracking_utils as spend_tracking_utils -from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL -from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY +from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL +from litellm.constants import CALLER_SESSION_ID_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -243,6 +243,82 @@ def test_batch_session_outranks_the_per_request_trace_id(): assert session_id == "batch-uid-1" +def _chat_session_id( + monkeypatch: pytest.MonkeyPatch, + litellm_params: Mapping[str, object], + correlation_in_logs: bool = True, + batch_trace_session_id: str | None = None, + metadata: Mapping[str, object] = MappingProxyType({}), +) -> str | None: + monkeypatch.setattr(litellm, "request_correlation_in_logs", correlation_in_logs) + return _get_session_id_for_spend_log( + kwargs={"litellm_trace_id": "trace-abc", "litellm_params": litellm_params}, + metadata=metadata, + standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING, + omit_when_missing=False, + batch_trace_session_id=batch_trace_session_id, + ) + + +_CALLER_SESSION_PARAMS: Final = MappingProxyType({"litellm_session_id": "sess-1", "metadata": {"trace_id": "otel-1"}}) + + +def test_caller_session_wins_over_the_per_request_trace_id(monkeypatch: pytest.MonkeyPatch): + assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS) == "sess-1" + + +def test_batch_session_outranks_a_caller_session(monkeypatch: pytest.MonkeyPatch): + assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS, batch_trace_session_id="batch-uid-1") == "batch-uid-1" + + +def test_caller_session_keeps_the_legacy_trace_id_with_correlation_logging_off(monkeypatch: pytest.MonkeyPatch): + assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS, correlation_in_logs=False) == "trace-abc" + + +@pytest.mark.parametrize(("correlation_in_logs", "expected_session_id"), [(True, "convo-1"), (False, "trace-abc")]) +def test_session_displaced_by_a_trace_header_wins_only_with_correlation_logging( + monkeypatch: pytest.MonkeyPatch, correlation_in_logs: bool, expected_session_id: str +): + header_chain: Final = {"litellm_session_id": "trace-abc", "metadata": {"trace_id": "trace-abc", "session_id": "trace-abc"}} + session_id: Final = _chat_session_id( + monkeypatch, + header_chain, + correlation_in_logs=correlation_in_logs, + metadata={CALLER_SESSION_ID_METADATA_KEY: "convo-1"}, + ) + assert session_id == expected_session_id + + +@pytest.mark.parametrize( + ("litellm_params", "metadata"), + [ + pytest.param({"litellm_session_id": "convo\x00x"}, {}, id="body-session"), + pytest.param({"litellm_session_id": "trace-abc"}, {CALLER_SESSION_ID_METADATA_KEY: "convo\x00x"}, id="beside-trace-header"), + ], +) +def test_caller_session_postgres_cannot_store_keeps_the_trace_id( + monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object], metadata: Mapping[str, object] +): + assert _chat_session_id(monkeypatch, litellm_params, metadata=metadata) == "trace-abc" + + +@pytest.mark.parametrize( + "litellm_params", + [ + pytest.param({"metadata": {}}, id="no-session"), + pytest.param( + {"litellm_session_id": "minted-1", "metadata": {"session_id": "minted-1", SESSION_ID_GENERATED_METADATA_KEY: True}}, + id="minted-by-generate-policy", + ), + pytest.param({"litellm_session_id": "job-7", "metadata": {"trace_id": "job-7"}}, id="copied-from-trace-id"), + ], +) +def test_session_that_is_not_a_caller_conversation_keeps_the_trace_id( + monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object] +): + assert _chat_session_id(monkeypatch, litellm_params) == "trace-abc" + + def test_omit_policy_still_suppresses_batch_sessions(): session_id: Final = _get_session_id_for_spend_log( kwargs={}, @@ -4030,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 c302bee3060..39eb2743bdb 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -1727,6 +1727,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_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 96c11be2650..93202d93566 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -17,6 +17,7 @@ from starlette.datastructures import Headers import litellm from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, SERVER_STREAMING_CLASSIFICATION_KEY, SERVER_STREAMING_CLASSIFICATION_MARKER, @@ -3308,6 +3309,93 @@ def test_add_litellm_metadata_from_request_headers_x_litellm_trace_id_sets_chain assert data["litellm_trace_id"] == "foo" +@pytest.mark.parametrize(("correlation_in_logs", "expected_spend_log_session"), [(True, "convo-1"), (False, None)]) +@pytest.mark.parametrize("session_in_metadata", [False, True], ids=["top-level", "metadata"]) +def test_x_litellm_trace_id_keeps_the_body_session_for_spend_logs_only( + monkeypatch: pytest.MonkeyPatch, + correlation_in_logs: bool, + expected_spend_log_session: str | None, + session_in_metadata: bool, +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", correlation_in_logs) + metadata: Final[dict[str, str]] = {"session_id": "convo-1"} if session_in_metadata else {} + data: Final[dict[str, object]] = ( + {"metadata": metadata} if session_in_metadata else {"metadata": metadata, "litellm_session_id": "convo-1"} + ) + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={"x-litellm-trace-id": "req-trace-1"}, data=data, _metadata_variable_name="metadata" + ) + assert metadata.get(CALLER_SESSION_ID_METADATA_KEY) == expected_spend_log_session + assert data["litellm_session_id"] == "req-trace-1" + assert metadata["session_id"] == "req-trace-1" + assert data["litellm_trace_id"] == "req-trace-1" + + +@pytest.mark.parametrize( + "session_header", ["x-litellm-session-id", "x-claude-code-session-id", "x-session-id"] +) +@pytest.mark.parametrize("with_trace_header", [False, True], ids=["session-header-only", "beside-trace-header"]) +def test_session_headers_still_beat_the_body_session_with_correlation_logging( + monkeypatch: pytest.MonkeyPatch, session_header: str, with_trace_header: bool +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + metadata: Final[dict[str, str]] = {} + data: Final[dict[str, object]] = {"metadata": metadata, "litellm_session_id": "convo-1"} + trace_header: Final = {"x-litellm-trace-id": "req-trace-1"} if with_trace_header else {} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={**trace_header, session_header: "hdr-session-1"}, data=data, _metadata_variable_name="metadata" + ) + spend_log_session: Final = metadata.get(CALLER_SESSION_ID_METADATA_KEY) or data["litellm_session_id"] + assert spend_log_session == "hdr-session-1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("forged_in", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize( + "headers", + [{"x-litellm-session-id": "hdr-session-1"}, {}], + ids=["session-header", "no-headers"], +) +async def test_client_metadata_cannot_set_the_spend_log_caller_session( + monkeypatch: pytest.MonkeyPatch, forged_in: str, headers: dict[str, str] +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + updated: Final = await add_litellm_data_to_request( + data={ + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "hi"}], + forged_in: {CALLER_SESSION_ID_METADATA_KEY: "forged-session"}, + }, + request=_make_request_mock("/v1/chat/completions", headers), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + assert CALLER_SESSION_ID_METADATA_KEY not in updated["metadata"] + + +@pytest.mark.asyncio +async def test_spend_log_caller_session_stays_out_of_requester_metadata(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + updated: Final = await add_litellm_data_to_request( + data={ + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "hi"}], + "litellm_session_id": "convo-1", + "metadata": {"team_label": "kept"}, + }, + request=_make_request_mock("/v1/chat/completions", {"x-litellm-trace-id": "req-trace-1"}), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + assert updated["metadata"][CALLER_SESSION_ID_METADATA_KEY] == "convo-1" + assert updated["metadata"]["requester_metadata"]["team_label"] == "kept" + assert CALLER_SESSION_ID_METADATA_KEY not in updated["metadata"]["requester_metadata"] + + def test_add_litellm_metadata_from_request_headers_x_litellm_session_id_sets_chain_id(): """x-litellm-session-id sets both metadata and top-level litellm_session_id/litellm_trace_id for call chaining.""" headers = {"x-litellm-session-id": "bar"} diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 3ebe445d89b..e824c5017c6 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -33,6 +33,7 @@ from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ( + AwsErrorTypeProxyException, LitellmUserRoles, ModelAccessDeniedProxyException, ProxyErrorTypes, @@ -43,7 +44,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.login_throttle import LoginThrottle from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash -from litellm.proxy.proxy_server import app, initialize, openai_exception_handler +from litellm.proxy.proxy_server import app, aws_error_type_exception_handler, initialize, openai_exception_handler from litellm.utils import _invalidate_model_cost_lowercase_map from tests._master_key import MASTER_KEY @@ -11858,6 +11859,29 @@ async def test_openai_exception_handler_no_denial_log_for_plain_proxy_exception( assert [r for r in caplog.records if r.levelname == "WARNING"] == [] +@pytest.mark.asyncio +async def test_aws_error_type_exception_handler_adds_the_aws_message_and_header(): + error: Final = ProxyException( + message="synthetic policy says no", + type="invalid_request_error", + param=None, + code=400, + headers={"x-litellm-call-id": "call-1"}, + ) + + response: Final = await aws_error_type_exception_handler( + _http_request_scope(), AwsErrorTypeProxyException(error, "ValidationException") + ) + + assert response.status_code == 400 + assert response.headers["x-amzn-errortype"] == "ValidationException" + assert response.headers["x-litellm-call-id"] == "call-1" + assert json.loads(response.body) == { + "message": "synthetic policy says no", + "error": {"message": "synthetic policy says no", "type": "invalid_request_error", "param": None, "code": "400"}, + } + + @pytest.mark.asyncio async def test_realtime_model_access_denial_logs_sanitized_internal_message(caplog): reservation = {"reserved_cost": 0.0, "input_cost": 0.0, "finalized": False, "entries": []} @@ -12588,6 +12612,7 @@ def test_dump_redacted_config_serializes_non_json_native_values(): async def test_update_config_general_settings_emits_audit_log(monkeypatch): import litellm.proxy.proxy_server as proxy_server_module from litellm.proxy._types import ConfigFieldUpdate + from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks from litellm.proxy.proxy_server import update_config_general_settings existing = {"max_parallel_requests": 5, "some_api_key": "sk-stored-secret"} @@ -12611,6 +12636,7 @@ async def test_update_config_general_settings_emits_audit_log(monkeypatch): ) # Audit is scheduled via asyncio.create_task; yield so it runs. await asyncio.sleep(0) + await drain_audit_tasks() fake.db.litellm_auditlog.create.assert_awaited_once() written = fake.db.litellm_auditlog.create.call_args.kwargs["data"] @@ -12980,6 +13006,7 @@ async def test_update_config_general_settings_applies_ssrf_globals(monkeypatch): async def test_delete_config_general_settings_emits_deleted_audit_log(monkeypatch): import litellm.proxy.proxy_server as proxy_server_module from litellm.proxy._types import ConfigFieldDelete + from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks from litellm.proxy.proxy_server import delete_config_general_settings existing = {"max_parallel_requests": 5} @@ -12999,6 +13026,7 @@ async def test_delete_config_general_settings_emits_deleted_audit_log(monkeypatc ) # Audit is scheduled via asyncio.create_task; yield so it runs. await asyncio.sleep(0) + await drain_audit_tasks() fake.db.litellm_auditlog.create.assert_awaited_once() written = fake.db.litellm_auditlog.create.call_args.kwargs["data"] diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 8cd6887d891..f7d1970fdae 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2,26 +2,33 @@ # 1. Generate a Key, and use it to make a call +import json +from collections.abc import Callable, Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from dotenv import load_dotenv -from pydantic import JsonValue +from pydantic import BaseModel, ConfigDict, JsonValue load_dotenv() # this file is to test litellm/proxy -from fastapi import HTTPException, Request +from fastapi import Request +from fastapi.responses import JSONResponse import litellm from litellm import Router +from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler +from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import VertexAIPartnerModelsTokenCounter from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.anthropic_endpoints.endpoints import ( count_tokens as anthropic_count_tokens, @@ -147,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!"}], @@ -166,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!"}] @@ -219,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!"}], @@ -238,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!"}] @@ -324,22 +347,16 @@ async def test_internal_token_counter_anthropic_provider_detection(): setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - # Test with is_direct_request=False but non-Anthropic provider - response = await token_counter( - request=TokenCountRequest( - model="gpt-test", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) - - print("Non-Anthropic provider test response:", response) - - # Verify response structure - assert response.request_model == "gpt-test" - assert response.model_used == "gpt-4" - assert response.total_tokens > 0 - assert response.tokenizer_type == "openai_tokenizer" # Should use LiteLLM tokenizer + with pytest.raises(ProxyException) as missing_key: + await token_counter( + request=TokenCountRequest( + model="gpt-test", + messages=[{"role": "user", "content": "hello"}], + ), + call_endpoint=True, + ) + assert missing_key.value.code == "401" + assert "No OpenAI API key is set for model gpt-test" in missing_key.value.message @pytest.mark.asyncio @@ -349,12 +366,14 @@ async def test_anthropic_endpoint_error_handling(): """ from unittest.mock import MagicMock - from fastapi import HTTPException, Request + from fastapi import Request from litellm.proxy.anthropic_endpoints.endpoints import count_tokens # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} + mock_request.scope = {} mock_user_api_key_dict = MagicMock() # Test missing model parameter @@ -372,12 +391,13 @@ async def test_anthropic_endpoint_error_handling(): anthropic_endpoints.read_request_body = mock_read_request_body try: - # Should raise HTTPException for missing model - with pytest.raises(HTTPException) as exc_info: - await count_tokens(mock_request, mock_user_api_key_dict) + response = await count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 400 - assert "model parameter is required" in str(exc_info.value.detail) + assert response.status_code == 400 + assert json.loads(response.body) == { + "type": "error", + "error": {"type": "invalid_request_error", "message": "model parameter is required"}, + } print("✅ Error handling test passed!") @@ -468,7 +488,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): mock_router.model_list = [ { "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4"}, + "litellm_params": {"model": "databricks/gpt-4"}, "model_info": {}, } ] @@ -477,7 +497,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): mock_router.async_get_available_deployment = AsyncMock( return_value={ "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4"}, + "litellm_params": {"model": "databricks/gpt-4"}, "model_info": {}, } ) @@ -1035,6 +1055,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1051,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", @@ -1062,14 +1090,11 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - # Verify HTTP status code is correct - assert exc_info.value.status_code == 400 - - # Verify error is in Anthropic format - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "invalid_request_error" assert detail["error"]["message"] == "Input is too long for requested model." @@ -1087,6 +1112,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): import litellm.proxy.proxy_server as proxy_server mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1103,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", @@ -1114,12 +1147,11 @@ async def test_anthropic_endpoint_403_permission_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 403 - - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 403 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "permission_error" assert detail["error"]["message"] == "Bearer Token has expired" @@ -1137,6 +1169,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): import litellm.proxy.proxy_server as proxy_server mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1153,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", @@ -1164,15 +1204,634 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 429 - - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 429 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "rate_limit_error" assert detail["error"]["message"] == "Rate limit exceeded" finally: anthropic_endpoints.read_request_body = original_read_request_body 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 = "" + + +class _GeminiContent(BaseModel): + model_config = ConfigDict(frozen=True) + parts: tuple[_GeminiPart, ...] = () + + +class _GeminiCountTokensBody(BaseModel): + model_config = ConfigDict(frozen=True) + contents: tuple[_GeminiContent, ...] = () + + +def _word_count(content: _GeminiContent) -> int: + return sum(len(part.text.split()) for part in content.parts) + + +def _gemini_word_counting_peer(request: httpx.Request) -> httpx.Response: + body: Final = _GeminiCountTokensBody.model_validate_json(request.content) + total: Final = sum(_word_count(content) for content in body.contents) + return httpx.Response(200, json={"totalTokens": total}, request=request) + + +def _unavailable_peer(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, json={"error": {"message": "unavailable"}}, request=request) + + +def _gemini_router(drop_params: bool = False) -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-counted", + "litellm_params": { + "model": "gemini/gemini-3.8-flash", + "api_key": "test-key", + "api_base": "https://count-tokens.test", + "drop_params": drop_params, + }, + } + ] + ) + + +@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"] +) +async def test_call_endpoint_counts_a_prompt_like_the_same_text_sent_as_a_user_message( + text: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=text), call_endpoint=True + ) + from_message: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": text}]), + call_endpoint=True, + ) + + assert from_prompt.error is False + assert from_prompt.total_tokens == from_message.total_tokens == len(text.split()) + + +class _PartnerMessage(BaseModel): + model_config = ConfigDict(frozen=True) + content: str = "" + + +class _PartnerCountRequest(BaseModel): + model_config = ConfigDict(frozen=True) + messages: tuple[_PartnerMessage, ...] = () + + +async def _vertex_partner_word_count( + model: str, request_data: JsonValue, litellm_params: JsonValue +) -> dict[str, JsonValue]: + body: Final = _PartnerCountRequest.model_validate(request_data) + return {"input_tokens": sum(len(message.content.split()) for message in body.messages)} + + +@pytest.mark.asyncio +async def test_call_endpoint_counts_a_prompt_on_a_vertex_ai_deployment(monkeypatch: pytest.MonkeyPatch) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "vertex-counted", + "litellm_params": { + "model": "vertex_ai/claude-sonnet-5-5", + "vertex_project": "test-project", + "vertex_location": "us-east5", + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + text: Final = "The quick brown fox jumps over the lazy dog." + + with patch.object( + VertexAIPartnerModelsTokenCounter, "handle_count_tokens_request", side_effect=_vertex_partner_word_count + ): + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="vertex-counted", prompt=text), call_endpoint=True + ) + + assert (from_prompt.error, from_prompt.total_tokens) == (False, 9) + + +_FOX: Final = "The quick brown fox jumps over the lazy dog." + + +_CONVERSATIONS: Final = ( + {"messages": [{"role": "user", "content": _FOX}]}, + {"contents": [{"role": "user", "parts": [{"text": _FOX}]}]}, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("conversation", _CONVERSATIONS, ids=["messages", "contents"]) +@pytest.mark.parametrize("call_endpoint", [True, False], ids=["provider", "local"]) +async def test_prompt_next_to_a_conversation_answers_400_without_drop_params( + conversation: dict[str, JsonValue], call_endpoint: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with ( + patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client), + pytest.raises(ProxyException) as rejected, + ): + await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint + ) + + assert (rejected.value.code, rejected.value.param) == ("400", "prompt") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("conversation", "call_endpoint", "global_drop_params", "deployment_drop_params"), + [ + (_CONVERSATIONS[0], True, False, True), + (_CONVERSATIONS[1], True, False, True), + (_CONVERSATIONS[0], False, False, True), + (_CONVERSATIONS[0], False, True, False), + ], + ids=[ + "provider_messages", + "provider_contents", + "local_messages", + "local_messages_global_setting", + ], +) +async def test_drop_params_counts_a_conversation_without_the_prompt( + conversation: dict[str, JsonValue], + call_endpoint: bool, + global_drop_params: bool, + deployment_drop_params: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router(drop_params=deployment_drop_params)) + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + with_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint + ) + conversation_only: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", **conversation), call_endpoint=call_endpoint + ) + + assert with_prompt.error is False + assert with_prompt.total_tokens == conversation_only.total_tokens > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("empty_conversation", [{"messages": []}, {"contents": []}], ids=["messages", "contents"]) +async def test_call_endpoint_prompt_next_to_an_empty_conversation_counts_the_prompt( + empty_conversation: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + counted: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation), call_endpoint=True + ) + + assert (counted.error, counted.total_tokens) == (False, 9) + + +@pytest.mark.asyncio +async def test_local_count_of_a_prompt_next_to_empty_messages_counts_the_prompt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + + with_empty: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX, messages=[])) + alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) + + assert with_empty.error is False + assert with_empty.total_tokens == alone.total_tokens > 0 + + +_WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chat_fields", + [{"system": "You answer in one short sentence. " * 20}, {"tools": [_WEATHER_TOOL]}], + ids=["system", "tools"], +) +@pytest.mark.parametrize("empty_conversation", [{}, {"messages": []}], ids=["no_messages", "empty_messages"]) +async def test_local_count_of_a_prompt_with_system_or_tools_counts_them_like_the_message_form( + chat_fields: dict[str, JsonValue], empty_conversation: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation, **chat_fields) + ) + from_message: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": _FOX}], **chat_fields) + ) + prompt_alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) + + assert from_prompt.error is False + assert from_prompt.total_tokens == from_message.total_tokens > prompt_alone.total_tokens + + +@pytest.mark.asyncio +async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_needs_messages( + monkeypatch: pytest.MonkeyPatch, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "anthropic-counted", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "api_key": "test-key", + "api_base": "https://count-tokens.test", + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_unavailable_peer)) + request: Final = TokenCountRequest(model="anthropic-counted", prompt="The quick brown fox jumps over the lazy dog.") + + with patch.object(anthropic_count_tokens_handler, "get_async_httpx_client", return_value=client): + with_endpoint: Final = await token_counter(request=request, call_endpoint=True) + local: Final = await token_counter(request=request, call_endpoint=False) + + assert (with_endpoint.error, with_endpoint.total_tokens, with_endpoint.tokenizer_type) == ( + False, + local.total_tokens, + local.tokenizer_type, + ) 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/repositories/test_unit_of_work.py b/tests/unit/repositories/test_unit_of_work.py index 9eff248b917..a08d44b9b4d 100644 --- a/tests/unit/repositories/test_unit_of_work.py +++ b/tests/unit/repositories/test_unit_of_work.py @@ -7,6 +7,7 @@ import pytest from litellm.repositories.unit_of_work import ( LinkedSpendResetWrites, budget_cascade_unit_of_work, + enduser_spend_reset_unit_of_work, spend_reset_unit_of_work, ) @@ -96,7 +97,6 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): uow.tags.queue_spend_zero(where=linked) uow.model_access_groups.queue_spend_zero(where=linked) uow.projects.queue_spend_zero(where=linked) - uow.endusers.queue_spend_zero(where={"user_id": {"in": ["enduser-1"]}}) uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) assert batch.commit_count == 0 @@ -108,11 +108,31 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): ("litellm_tagtable.update_many", linked, {"spend": 0}), ("litellm_modelaccessgroupbudgettable.update_many", linked, {"spend": 0}), ("litellm_projecttable.update_many", linked, {"spend": 0}), - ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-1"]}}, {"spend": 0}), ("litellm_budgettable.update_many", {"budget_id": "budget-1"}, {"budget_reset_at": reset_at}), ] +async def test_enduser_page_commits_on_its_own_batch(): + """One tier can gate millions of end users, so a page of them lands in its + own transaction instead of riding the cascade's.""" + batch = FakeBatch() + + async with enduser_spend_reset_unit_of_work(lambda: batch) as writes: + writes.queue_spend_zero(where={"user_id": {"in": ["enduser-1", "enduser-2"]}, "spend": {"gt": 0}}) + writes.queue_spend_decrement(where={"user_id": {"in": ["enduser-3"]}, "spend": {"gt": 5.0}}, amount=5.0) + assert batch.commit_count == 0 + + assert batch.commit_count == 1 + assert batch.calls == [ + ( + "litellm_endusertable.update_many", + {"user_id": {"in": ["enduser-1", "enduser-2"]}, "spend": {"gt": 0}}, + {"spend": 0}, + ), + ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-3"]}, "spend": {"gt": 5.0}}, {"spend": {"decrement": 5.0}}), + ] + + async def test_budget_window_advance_tolerates_a_tier_deleted_mid_chunk(): """A tier deleted between the read and the commit must not abort the batch: ``update`` raises P2025 on a missing row and takes every other write in the 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_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 0d55fc35e1d..cc5d5c62739 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -17,6 +17,7 @@ from openai.types.responses.response_function_web_search import ( from pydantic import TypeAdapter import litellm +from litellm.constants import CALLER_SESSION_ID_METADATA_KEY from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -3109,7 +3110,6 @@ class TestToolTransformation: ("bedrock/converse/global.anthropic.claude-sonnet-5", "bedrock_converse"), ("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"), ("claude-sonnet-5", "vertex_ai"), - ("gemini-3.1-pro-preview", "vertex_ai"), ("moonshotai.kimi-k2-thinking", "bedrock_mantle"), ], ) @@ -3117,8 +3117,8 @@ class TestToolTransformation: """ A Responses request carrying ``reasoning.summary`` must still reach a chat provider as a plain ``reasoning_effort`` string. ``summary`` is Responses-only, and forwarding the whole - object turns reasoning off: Bedrock Converse and Vertex silently discard a non-string - ``reasoning_effort``, and Bedrock Mantle rejects the request outright. + object turns reasoning off: Bedrock Converse and Claude on Vertex silently discard a + non-string ``reasoning_effort``, and Bedrock Mantle rejects the request outright. """ responses_api_request = {"reasoning": {"effort": "medium", "summary": "auto"}} @@ -3131,6 +3131,42 @@ class TestToolTransformation: assert result["reasoning_effort"] == "medium" + @pytest.mark.parametrize( + "model, custom_llm_provider", + [("gemini-3.1-pro-preview", "vertex_ai"), ("gemini-2.5-flash", "gemini")], + ) + @pytest.mark.parametrize( + "reasoning, sends_thoughts", + [({"effort": "medium", "summary": "auto"}, True), ({"effort": "medium"}, False)], + ) + def test_a_summary_ask_turns_gemini_thoughts_on_over_the_gateway_default( + self, + monkeypatch: pytest.MonkeyPatch, + model: str, + custom_llm_provider: str, + reasoning: dict[str, str], + sends_thoughts: bool, + ) -> None: + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + from litellm.utils import reasoning_effort_with_summary_alias + + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False) + bridged: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input="hi", + responses_api_request={"reasoning": reasoning}, + custom_llm_provider=custom_llm_provider, + ) + + gemini: Final = VertexGeminiConfig() + reasoning_effort: Final = reasoning_effort_with_summary_alias( + bridged["reasoning_effort"], bridged, accepts_summary=gemini.accepts_reasoning_summary_in_reasoning_effort + ) + mapped: Final = gemini.map_openai_params({"reasoning_effort": reasoning_effort}, {}, model, True) + + assert mapped["thinkingConfig"]["includeThoughts"] is sends_thoughts + assert bridged["reasoning_effort"] == "medium" + @pytest.mark.parametrize( "model, custom_llm_provider", [ @@ -6464,3 +6500,69 @@ def test_transform_request_reads_service_tier_from_responses_api_request() -> No ) assert result["service_tier"] == "priority" + + +def _continued_chain_request(request_kwargs: Mapping[str, object]) -> Mapping[str, object]: + async def session_lookup(previous_response_id: str) -> ChatCompletionSession: + assert previous_response_id == "resp_turn_1" + return ChatCompletionSession( + messages=[{"role": "user", "content": "turn 1"}], litellm_session_id="chain-session" + ) + + return asyncio.run( + LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( + previous_response_id="resp_turn_1", + litellm_completion_request={ + "model": _CLAUDE_ON_BEDROCK, + "custom_llm_provider": "bedrock", + "messages": [{"role": "user", "content": "turn 2"}], + "metadata": request_kwargs.get("metadata"), + }, + instructions=None, + session_lookup=session_lookup, + request_kwargs=request_kwargs, + ) + ) + + +@pytest.mark.parametrize( + ("request_kwargs", "logged_metadata_key"), + [ + ( + { + "metadata": {"tag": "client"}, + "litellm_metadata": {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"}, + }, + "litellm_metadata", + ), + ({"metadata": {"user_api_key": "hashed"}}, "metadata"), + ], +) +def test_a_continued_chain_logs_its_spend_under_the_chain_session_with_request_correlation( + monkeypatch: pytest.MonkeyPatch, request_kwargs: Mapping[str, object], logged_metadata_key: str +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + sent_kwargs: Final = deepcopy(request_kwargs) + + request: Final = _continued_chain_request(request_kwargs) + + assert request["litellm_trace_id"] == "chain-session" + assert "litellm_session_id" not in request + assert request[logged_metadata_key] == { + **TypeAdapter(dict[str, object]).validate_python(request_kwargs[logged_metadata_key]), + CALLER_SESSION_ID_METADATA_KEY: "chain-session", + } + assert request_kwargs == sent_kwargs + + +def test_a_continued_chain_leaves_the_turns_metadata_alone_without_request_correlation( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + caller_metadata: Final = {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"} + + request: Final = _continued_chain_request({"metadata": caller_metadata, "litellm_metadata": caller_metadata}) + + assert request["litellm_trace_id"] == "chain-session" + assert "litellm_metadata" not in request + assert request["metadata"] == {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"} 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 002a595a5e7..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,4 +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 @@ -15,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 @@ -126,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 @@ -460,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)) @@ -554,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 @@ -622,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 @@ -643,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(): @@ -787,3 +876,82 @@ async def test_message_history_replays_real_key_named_tool_payloads() -> None: tool_message: Final = result["messages"][2] assert json.loads(assistant_message["tool_calls"][0]["function"]["arguments"]) == function_arguments assert json.loads(tool_message["content"]) == function_output + + +class _SqliteBackedPrismaDB: + def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None: + self._connection = sqlite3.connect(":memory:") + self._connection.row_factory = sqlite3.Row + self._connection.execute( + 'CREATE TABLE "LiteLLM_SpendLogs" ' + '(request_id TEXT, api_key TEXT, session_id TEXT, "endTime" TEXT, "user" TEXT, team_id TEXT)' + ) + self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?)', rows) + + async def query_raw(self, query: str, *args: object) -> list[dict[str, object]]: + 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 +async def test_session_lookup_does_not_return_another_keys_rows_for_a_shared_session_id(): + """A caller-supplied litellm_session_id can collide with another key's session; history must stay per key.""" + sqlite_db: Final = _SqliteBackedPrismaDB( + [ + ("victim-1", "victim-key", "shared-session", "2026-01-01T00:00:01", "", ""), + ("victim-2", "victim-key", "shared-session", "2026-01-01T00:00:02", "", ""), + ("attacker-1", "attacker-key", "shared-session", "2026-01-01T00:00:03", "", ""), + ("attacker-2", "attacker-key", "other-session", "2026-01-01T00:00:04", "", ""), + ] + ) + fake_prisma_client = _FakePrismaClient(results=[]) + fake_prisma_client.db = sqlite_db + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id("attacker-1") + victim_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id("victim-1") + + assert [row["request_id"] for row in spend_logs] == ["attacker-1"] + assert [row["request_id"] for row in victim_logs] == ["victim-1", "victim-2"] + + +async def _history_request_ids(sqlite_db: _SqliteBackedPrismaDB, response_id: str) -> list[str]: + fake_prisma_client: Final = _FakePrismaClient(results=[]) + fake_prisma_client.db = sqlite_db + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + spend_logs: Final = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(response_id) + return [row["request_id"] for row in spend_logs] + + +@pytest.mark.asyncio +async def test_session_lookup_keeps_jwt_users_and_teams_apart_when_they_share_an_empty_api_key(): + sqlite_db: Final = _SqliteBackedPrismaDB( + [ + ("alice-1", "", "shared-session", "2026-01-01T00:00:01", "alice", "team-a"), + ("bob-1", "", "shared-session", "2026-01-01T00:00:02", "bob", "team-a"), + ("service-a-1", "", "shared-session", "2026-01-01T00:00:03", "", "team-a"), + ("service-b-1", "", "shared-session", "2026-01-01T00:00:04", "", "team-b"), + ("alice-2", "", "shared-session", "2026-01-01T00:00:05", "alice", "team-a"), + ] + ) + + assert await _history_request_ids(sqlite_db, "alice-1") == ["alice-1", "alice-2"] + assert await _history_request_ids(sqlite_db, "bob-1") == ["bob-1"] + assert await _history_request_ids(sqlite_db, "service-a-1") == ["service-a-1"] + assert await _history_request_ids(sqlite_db, "service-b-1") == ["service-b-1"] + + +@pytest.mark.asyncio +async def test_session_lookup_follows_the_owner_across_keys_and_keeps_ownerless_keys_apart(): + sqlite_db: Final = _SqliteBackedPrismaDB( + [ + ("alice-old-key", "hashed-old", "shared-session", "2026-01-01T00:00:01", "alice", "team-a"), + ("ownerless-1", "hashed-other", "shared-session", "2026-01-01T00:00:02", "", ""), + ("alice-new-key", "hashed-new", "shared-session", "2026-01-01T00:00:03", "alice", "team-a"), + ("ownerless-2", "hashed-another", "shared-session", "2026-01-01T00:00:04", "", ""), + ] + ) + + assert await _history_request_ids(sqlite_db, "alice-new-key") == ["alice-old-key", "alice-new-key"] + assert await _history_request_ids(sqlite_db, "ownerless-2") == ["ownerless-2"] 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/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index c37e57a43ca..001be7e1fbd 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -20,14 +20,18 @@ Pydantic ValidationError (previously typed as Optional[str]). """ import json +from datetime import datetime from importlib import import_module +from typing import Final from unittest.mock import Mock, patch +import httpx import pytest import litellm from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.streaming_iterator import ( _ERROR_CODE_HTTP_STATUS, @@ -64,6 +68,52 @@ def _make_error_chunk(error_type: str, code: str, message: str = "err") -> Error return ErrorEvent(type=ResponsesAPIStreamEvents.ERROR, sequence_number=0, error=error_obj) +@pytest.mark.parametrize("provider", ("openai", "azure", "bedrock")) +@pytest.mark.parametrize("header", ("x-amzn-RequestId", "x-request-id", None)) +@pytest.mark.parametrize( + "code,status", (("server_error", 500), ("rate_limit_exceeded", 429), ("invalid_request_error", 400)) +) +def test_stream_error_preserves_provider_request_id(provider: str, header: str | None, code: str, status: int) -> None: + stream_headers: Final = {"retry-after": "7", "content-type": "text/event-stream"} + response: Final = httpx.Response( + 200, headers={**stream_headers, header: "provider-request-id"} if header else stream_headers + ) + logging: Final = LiteLLMLoggingObj( + model="test-model", + messages=[], + stream=True, + call_type="responses", + start_time=datetime(2026, 1, 1), + litellm_call_id="test-call", + function_id="test-function", + ) + iterator: Final = BaseResponsesAPIStreamingIterator( + response=response, + model="test-model", + responses_api_provider_config=None, + logging_obj=logging, + custom_llm_provider=provider, + ) + expected_id: Final = "provider-request-id" if header else None + expected_headers: Final = {header.lower(): "provider-request-id"} if header else {} + mapped: Final = iterator._map_error_event_exception({"type": code, "code": code, "message": "failure"}) + assert _get_provider_request_id(mapped) == expected_id + assert dict(getattr(mapped, "litellm_response_headers", None) or {}) == expected_headers + expected_status: Final = 503 if provider == "bedrock" and status == 500 else status + with pytest.raises((litellm.BadRequestError, MidStreamFallbackError)) as raised: + iterator._maybe_raise_for_error_event(_make_error_chunk(code, code, "failure")) + assert raised.value.status_code == expected_status + assert _get_provider_request_id(raised.value) == expected_id + assert response.status_code == 200 + if status in (429, 500): + assert isinstance(raised.value, MidStreamFallbackError) + assert dict(raised.value.response.headers) == expected_headers + assert raised.value.response.status_code == expected_status + assert str(raised.value.response.request.url) == f"https://{provider}.com/v1/" + assert raised.value.generated_content == "" + assert raised.value.is_pre_first_chunk is True + + def test_maybe_raise_for_error_event_wraps_unknown_error_in_mid_stream_fallback(): iterator = _make_iterator() chunk = _make_error_chunk("server_error", "internal_error", "something went wrong") 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/pre_call_checks/test_continuation_prefill_check.py b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py new file mode 100644 index 00000000000..8fa1c799863 --- /dev/null +++ b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py @@ -0,0 +1,288 @@ +import pytest + +import litellm +from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, + ContinuationOutput, + ContinuationPrefillDeploymentCheck, + continuation_output_ceilings, + deployment_supports_prefill, +) + +PREFILL_MODEL = "anthropic/prefill-capable-test-model" +NON_PREFILL_MODEL = "openai/prefill-unknown-test-model" +UNMAPPED_MODEL = "openai/unmapped-test-model" + + +@pytest.fixture(autouse=True) +def _synthetic_cost_map_entries() -> None: + litellm.register_model( + { + PREFILL_MODEL: {"litellm_provider": "anthropic", "mode": "chat", "supports_assistant_prefill": True}, + NON_PREFILL_MODEL: {"litellm_provider": "openai", "mode": "chat"}, + }, + persist_across_reloads=False, + ) + litellm.get_model_info.cache_clear() + + +def _deployment(model: str, dep_id: str) -> dict[str, object]: + return {"litellm_params": {"model": model}, "model_info": {"id": dep_id}} + + +def test_deployment_supports_prefill_reads_capability(): + assert deployment_supports_prefill(_deployment(PREFILL_MODEL, "a")) is True + assert deployment_supports_prefill(_deployment(NON_PREFILL_MODEL, "b")) is False + assert deployment_supports_prefill({"litellm_params": {"model": PREFILL_MODEL}}) is True + assert deployment_supports_prefill({"litellm_params": {"model": PREFILL_MODEL}, "model_info": "bogus"}) is True + + +def test_deployment_model_info_override_wins_over_cost_map(): + # model_info True opts in a model that is not in the cost map + assert ( + deployment_supports_prefill( + {"litellm_params": {"model": "vendor/custom-model"}, "model_info": {"supports_assistant_prefill": True}} + ) + is True + ) + # model_info False opts out a model the cost map would otherwise allow + assert ( + deployment_supports_prefill( + {"litellm_params": {"model": PREFILL_MODEL}, "model_info": {"supports_assistant_prefill": False}} + ) + is False + ) + # model_info without the key falls through to the cost map + assert deployment_supports_prefill(_deployment(PREFILL_MODEL, "z")) is True + + +@pytest.mark.parametrize( + "deployment,expected", + [ + ({"litellm_params": {"model": "bedrock/arn:custom-import", "base_model": PREFILL_MODEL}}, True), + ({"litellm_params": {"model": "azure/my-deployment"}, "model_info": {"base_model": PREFILL_MODEL}}, True), + ( + { + "litellm_params": {"model": "azure/my-deployment", "base_model": NON_PREFILL_MODEL}, + "model_info": {"base_model": PREFILL_MODEL}, + }, + True, + ), + ( + { + "litellm_params": {"model": PREFILL_MODEL, "base_model": NON_PREFILL_MODEL}, + "model_info": {"base_model": NON_PREFILL_MODEL}, + }, + False, + ), + ], + ids=[ + "litellm_params_base_model", + "model_info_base_model", + "model_info_base_model_wins", + "base_model_wins_over_model", + ], +) +def test_deployment_supports_prefill_reads_the_cost_map_key_of_the_base_model(deployment, expected): + assert deployment_supports_prefill(deployment) is expected + + +def test_deployment_supports_prefill_rejects_malformed_deployments(): + assert deployment_supports_prefill({}) is False + assert deployment_supports_prefill({"litellm_params": {}}) is False + assert deployment_supports_prefill("not-a-dict") is False + + +@pytest.mark.asyncio +async def test_filter_is_noop_without_continuation_marker(): + """A normal (non-continuation) request must be passed through untouched, even + if some deployments cannot prefill.""" + check = ContinuationPrefillDeploymentCheck() + deployments = [_deployment(PREFILL_MODEL, "a"), _deployment(NON_PREFILL_MODEL, "b")] + + for request_kwargs in ({}, None, {MID_STREAM_CONTINUATION_KWARG: False}): + result = await check.async_filter_deployments( + model="group", healthy_deployments=deployments, messages=None, request_kwargs=request_kwargs + ) + assert result == deployments + + +@pytest.mark.asyncio +async def test_filter_ignores_forged_client_flag(): + """A client cannot steer routing: a plain truthy value under the marker key + (which the proxy could forward from the request body) is not the internal + sentinel, so the filter leaves the deployment list untouched.""" + check = ContinuationPrefillDeploymentCheck() + deployments = [_deployment(PREFILL_MODEL, "a"), _deployment(NON_PREFILL_MODEL, "b")] + + for forged in (True, "true", 1, {"any": "json"}): + result = await check.async_filter_deployments( + model="group", + healthy_deployments=deployments, + messages=None, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: forged}, + ) + assert result == deployments + + +@pytest.mark.asyncio +async def test_filter_keeps_only_prefill_capable_on_continuation(): + check = ContinuationPrefillDeploymentCheck() + deployments = [_deployment(PREFILL_MODEL, "a"), _deployment(NON_PREFILL_MODEL, "b")] + + result = await check.async_filter_deployments( + model="group", + healthy_deployments=deployments, + messages=None, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, + ) + assert [d["model_info"]["id"] for d in result] == ["a"] + + +@pytest.mark.asyncio +async def test_filter_empties_group_when_no_prefill_capable_deployment(): + """No prefill-capable deployment -> empty result, so the router advances the + fallback chain and ultimately surfaces the original error.""" + check = ContinuationPrefillDeploymentCheck() + deployments = [_deployment(NON_PREFILL_MODEL, "b"), _deployment(UNMAPPED_MODEL, "c")] + + result = await check.async_filter_deployments( + model="group", + healthy_deployments=deployments, + messages=None, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, + ) + assert result == [] + + +@pytest.mark.asyncio +async def test_filter_drops_deployments_whose_own_params_constrain_the_continuation(): + check = ContinuationPrefillDeploymentCheck() + plain = _deployment(PREFILL_MODEL, "plain") + constrained = [ + {"litellm_params": {"model": PREFILL_MODEL, **params}, "model_info": {"id": dep_id}} + for dep_id, params in ( + ("json", {"response_format": {"type": "json_object"}}), + ("two_choices", {"extra_body": {"n": 2}}), + ("forced_tool", {"tool_choice": {"type": "function", "function": {"name": "f"}}}), + ("merged_reasoning", {"merge_reasoning_content_in_choices": True}), + ) + ] + + result = await check.async_filter_deployments( + model="group", + healthy_deployments=[plain, *constrained], + messages=None, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, + ) + assert [d["model_info"]["id"] for d in result] == ["plain"] + + +def _text_chunk(text: str) -> litellm.ModelResponseStream: + return litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) + + +@pytest.mark.parametrize( + "prefill_content,fallback_texts,expected_texts", + [ + ("Hello", ("", "Hel", "lo wor", "ld"), ("", "", " wor", "ld")), + ("Hello", ("Hello",), ("",)), + ("Hello", (" world",), (" world",)), + ("Hello", ("Hel", "p me"), ("Hel", "p me")), + ("Hello", ("Hel",), ("Hel",)), + ([{"type": "text", "text": "Sure: "}, {"type": "text", "text": "Hello"}], ("Sure: Hello", " world"), ("", " world")), + ], + ids=[ + "echo_split_across_chunks_is_dropped", + "exact_echo_then_end_is_dropped", + "continuation_without_echo_passes_through", + "near_echo_is_kept_in_order", + "stream_ending_mid_match_is_kept", + "content_block_prefill_echo_is_dropped", + ], +) +@pytest.mark.asyncio +async def test_continuation_output_drops_only_a_full_echo_of_the_prefill( + prefill_content, fallback_texts, expected_texts +): + output = ContinuationOutput( + "Hello", {"role": "assistant", "content": prefill_content, "prefix": True}, "group", () + ) + + async def fallback_stream(): + for text in fallback_texts: + yield _text_chunk(text) + + released = [chunk async for chunk in output.splice(fallback_stream())] + + assert tuple(chunk.choices[0].delta.content for chunk in released) == expected_texts + + +@pytest.mark.parametrize( + "prefill,fallback_texts,expected_texts", + [ + ("Transformers use ", ("self-attention",), ("self-attention",)), + ("Hello ", ("Hello ", " world"), ("", "world")), + ("Hello \n", ("\nworld",), ("\nworld",)), + ("A.\n\n", ("\n Hi", "\nthere"), (" Hi", "\nthere")), + ("\n\n", ("\nHi",), ("\nHi",)), + ("\nHello\n", ("\n",), ("",)), + ], + ids=[ + "continuation_without_whitespace_passes_through", + "echo_then_re_emitted_space_are_both_dropped", + "different_whitespace_is_kept", + "seam_ends_at_the_first_divergence", + "whitespace_only_prefill_has_no_seam", + "re_emitted_whitespace_at_stream_end_is_dropped", + ], +) +@pytest.mark.asyncio +async def test_continuation_output_drops_only_whitespace_the_client_already_has( + prefill, fallback_texts, expected_texts +): + output = ContinuationOutput(prefill, {"role": "assistant", "content": prefill, "prefix": True}, "group", ()) + + async def fallback_stream(): + for text in fallback_texts: + yield _text_chunk(text) + + released = [chunk async for chunk in output.splice(fallback_stream())] + + assert tuple(chunk.choices[0].delta.content for chunk in released) == expected_texts + + +@pytest.mark.parametrize( + "request_kwargs,expected", + [ + ({"max_tokens": 100}, {"max_tokens": 70}), + ({"max_tokens": 100, "max_completion_tokens": 40}, {"max_tokens": 70, "max_completion_tokens": 10}), + ({}, {}), + ( + {"max_tokens": 100, "extra_body": {"max_tokens": 50, "top_k": 3}}, + {"max_tokens": 70, "extra_body": {"max_tokens": 20, "top_k": 3}}, + ), + ({"extra_body": {"max_completion_tokens": 45}}, {"extra_body": {"max_completion_tokens": 15}}), + ({"extra_body": {"top_k": 3}}, {}), + ], + ids=["max_tokens", "both_ceilings", "no_ceiling", "extra_body_max_tokens", "extra_body_only", "extra_body_no_ceiling"], +) +def test_continuation_output_ceilings_take_the_emitted_tokens_off_every_ceiling_the_request_sends( + request_kwargs, expected +): + assert continuation_output_ceilings(request_kwargs, 30) == expected + + +@pytest.mark.parametrize( + "request_kwargs", + [ + {"max_tokens": 30}, + {"max_tokens": 20}, + {"max_tokens": 100, "max_completion_tokens": 30}, + {"max_tokens": 100, "extra_body": {"max_tokens": 30}}, + ], + ids=["max_tokens_met", "max_tokens_exceeded", "one_of_two_met", "extra_body_met"], +) +def test_continuation_output_ceilings_decline_once_a_ceiling_has_no_tokens_left(request_kwargs): + assert continuation_output_ceilings(request_kwargs, 30) is None 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_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 5a6da70d756..50a9d98385d 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -1,4 +1,5 @@ import ast, asyncio, importlib, time +import contextlib import copy import functools import gc @@ -11,7 +12,8 @@ import threading import warnings from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping from datetime import datetime, timedelta, timezone -from types import SimpleNamespace +from pathlib import Path +from types import MappingProxyType, SimpleNamespace from typing import TYPE_CHECKING, Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -36,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, ) @@ -82,6 +85,7 @@ from litellm.types.router import ( ModelInfo, PreRoutingHookResponse, RetryPolicy, + RouterRateLimitError, RoutingContext, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -2932,6 +2936,1216 @@ async def test_acompletion_streaming_iterator(): print("\n=== All tests passed! ===") +def _make_midstream_source(error, chunks=None, first_chunk=None): + from unittest.mock import MagicMock + + opening_chunk = ( + first_chunk if first_chunk is not None else MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]) + ) + + class _Source: + def __init__(self): + self.index = 0 + self.chunks = chunks if chunks is not None else [] + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + + def __aiter__(self): + return self + + async def __anext__(self): + if self.index == 0: + self.index += 1 + return opening_chunk + raise error + + async def wait_for_failure_logging(self, timeout): + return None + + return _Source() + + +class _FakeFallbackStream: + def __init__(self, item): + self._item = item + self._done = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self._done: + raise StopAsyncIteration + self._done = True + return self._item + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_continues_after_content_when_eligible(): + """Flag on + plain-text break: the router re-enters the fallback chain with an + assistant-prefill continuation and the mid-stream marker, then streams the + fallback's output instead of re-raising.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, + ) + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + + error = MidStreamFallbackError( + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, + ) + source = _make_midstream_source(error) + fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "world"}}]) + initial_kwargs = {"model": "gpt-4", "stream": True} + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), + ) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=source, messages=[{"role": "user", "content": "Hi"}], initial_kwargs=initial_kwargs + ) + collected = [chunk async for chunk in result] + + mock_fallback.assert_awaited_once() + passed_kwargs = mock_fallback.await_args.kwargs["kwargs"] + assert passed_kwargs[MID_STREAM_CONTINUATION_KWARG] is MID_STREAM_CONTINUATION_MARKER + assert passed_kwargs["messages"][-1] == {"role": "assistant", "content": "Hello", "prefix": True} + assert fallback_chunk in collected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "continuation_text,joined_output_trips_rule", + [("6789", True), ("6788", False)], + ids=["rule_trips_only_on_the_joined_text", "clean_continuation_streams"], +) +async def test_acompletion_streaming_iterator_continuation_runs_post_call_rules_on_the_joined_text( + monkeypatch, continuation_text, joined_output_trips_rule +): + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + monkeypatch.setattr(litellm, "post_call_rules", [lambda output: "123-45-6789" not in output]) + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="123-45-", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": continuation_text}}]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), + ): + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error), + messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + if joined_output_trips_rule: + with pytest.raises(litellm.APIResponseValidationError): + async for _ in result: + pass + else: + assert fallback_chunk in [chunk async for chunk in result] + + +@pytest.mark.asyncio +async def test_a_second_cut_runs_post_call_rules_only_on_the_text_the_caller_received(monkeypatch): + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, + ) + + checked_texts = [] + monkeypatch.setattr(litellm, "post_call_rules", [lambda output: checked_texts.append(output) is None]) + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + + def cut(generated_content): + return MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=generated_content, + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + + messages = [{"role": "user", "content": "Hi"}] + middle_leg = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source( + cut(" there"), + first_chunk=litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " there"}}]), + ), + messages=[*messages, {"role": "assistant", "content": "Hello", "prefix": True}], + initial_kwargs={"model": "gpt-4", "stream": True, MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, + ) + backup_leg = _FakeFallbackStream( + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " friend"}}]) + ) + + with patch.object( + router, "async_function_with_fallbacks_common_utils", new=AsyncMock(side_effect=[middle_leg, backup_leg]) + ): + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(cut("Hello")), + messages=messages, + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + _ = [chunk async for chunk in result] + + assert checked_texts == ["Hello there", "Hello there friend"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fallback_texts", + [("Hel", "lo", " world"), (" world",)], + ids=["fallback_echoes_the_prefill", "fallback_continues_after_the_prefill"], +) +async def test_acompletion_streaming_iterator_continuation_never_repeats_the_partial_output(fallback_texts): + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "mistral/mistral-small-latest", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + + async def fallback_stream(): + for text in fallback_texts: + yield litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) + + with patch.object( + router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream()) + ): + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error), + messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + collected = [chunk async for chunk in result] + + continuation = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)] + assert "".join(chunk.choices[0].delta.content for chunk in continuation) == " world" + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_continuation_stays_in_the_primary_stream(): + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + primary_chunk = litellm.ModelResponseStream( + id="chatcmpl-primary", created=111, choices=[{"index": 0, "delta": {"role": "assistant", "content": "Hello"}}] + ) + + async def fallback_stream(): + yield litellm.ModelResponseStream( + id="msg-fallback", created=222, choices=[{"index": 0, "delta": {"role": "assistant", "content": " wor"}}] + ) + yield litellm.ModelResponseStream( + id="msg-fallback", created=222, choices=[{"index": 0, "delta": {"content": "ld"}}] + ) + + with patch.object( + router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream()) + ): + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error, chunks=[primary_chunk]), + messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + collected = [chunk async for chunk in result] + + continuation = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)] + assert [(chunk.id, chunk.created, chunk.choices[0].delta.role) for chunk in continuation] == [ + ("chatcmpl-primary", 111, None), + ("chatcmpl-primary", 111, None), + ] + assert "".join(chunk.choices[0].delta.content for chunk in continuation) == " world" + + +def _continuation_router(): + return litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + + +def _usage_chunk(prompt_tokens, completion_tokens): + from litellm.types.utils import Usage + + return litellm.ModelResponseStream( + model="gpt-4", + choices=[], + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + ) + + +async def _drain_with_fallback(router, error, primary_chunks, messages, fallback_items): + from unittest.mock import AsyncMock, patch + + async def fallback_stream(): + for item in fallback_items: + yield item + + with patch.object( + router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream()) + ) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error, chunks=primary_chunks), + messages=messages, + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + collected = [chunk async for chunk in result] + streamed = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)] + return mock_fallback.await_args.kwargs["kwargs"], streamed + + +@pytest.mark.asyncio +@pytest.mark.parametrize("primary_reports_usage", [True, False], ids=["usage_reported", "usage_not_reported"]) +async def test_acompletion_streaming_iterator_continuation_usage_sums_both_legs(primary_reports_usage): + from litellm.exceptions import MidStreamFallbackError + + messages = [{"role": "user", "content": "Explain transformers in one paragraph."}] + partial = "Transformers use " + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial, + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + content_chunk = litellm.ModelResponseStream( + model="gpt-4", choices=[{"index": 0, "delta": {"role": "assistant", "content": partial}}] + ) + primary_chunks = [content_chunk, _usage_chunk(21, 3)] if primary_reports_usage else [content_chunk] + primary_usage = litellm.stream_chunk_builder(chunks=list(primary_chunks), messages=messages).usage + + _, streamed = await _drain_with_fallback( + _continuation_router(), error, primary_chunks, messages, [_usage_chunk(91, 5)] + ) + + assert primary_usage.prompt_tokens > 0 + assert streamed[-1].usage.prompt_tokens == primary_usage.prompt_tokens + 91 + assert streamed[-1].usage.completion_tokens == primary_usage.completion_tokens + 5 + + +class _CutSSE(httpx.AsyncByteStream): + def __init__(self, content: str) -> None: + self._content: Final = content + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield f'data: {json.dumps({"id": "cut", "object": "chat.completion.chunk", "created": 1700000000, "model": "m", "choices": [{"index": 0, "delta": {"role": "assistant", "content": self._content}, "finish_reason": None}]})}\n\n'.encode() + raise httpx.RemoteProtocolError("peer closed connection without sending complete message body") + + +_BACKUP_SSE: Final = "".join( + f"data: {json.dumps(chunk)}\n\n" + for chunk in ( + {"id": "b", "object": "chat.completion.chunk", "created": 1700000000, "model": "mistral-large-latest", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": " self-attention."}, "finish_reason": None}]}, + {"id": "b", "object": "chat.completion.chunk", "created": 1700000000, "model": "mistral-large-latest", + "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 91, "completion_tokens": 5, "total_tokens": 96}}, + ) +) + "data: [DONE]\n\n" + + +class _SpendRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.successes: Final[list[dict[str, object]]] = [] + self.failures: Final[list[dict[str, object]]] = [] + self.spend_rows: Final[list[dict[str, object]]] = [] + self.succeeded: Final = asyncio.Event() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + self.successes.append(kwargs["standard_logging_object"]) + self.spend_rows.append(get_logging_payload(kwargs, response_obj, start_time, end_time)) + self.succeeded.set() + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + self.failures.append(kwargs["standard_logging_object"]) + + +class _SlowFailureRecorder(_SpendRecorder): + def __init__(self) -> None: + super().__init__() + self.failed_deployment_ids: Final[list[str]] = [] + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + await asyncio.sleep(0.05) + self.failed_deployment_ids.append(kwargs["litellm_params"]["metadata"]["model_info"]["id"]) + await super().async_log_failure_event(kwargs, response_obj, start_time, end_time) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cut_groups", [("primary",), ("primary", "middle")], ids=["one_break", "two_breaks"]) +async def test_continued_stream_success_log_bills_every_cut_leg( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, cut_groups: tuple[str, ...] +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _SpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use")) + ) + respx_mock.post(host="api.deepseek.com").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE(" learned")) + ) + respx_mock.post(host="api.mistral.ai").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=_BACKUP_SSE) + ) + router: Final = Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "k1", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + }, + }, + { + "model_name": "middle", + "litellm_params": { + "model": "deepseek/deepseek-chat", + "api_key": "k2", + "input_cost_per_token": 0.003, + "output_cost_per_token": 0.004, + }, + }, + {"model_name": "backup", "litellm_params": {"model": "mistral/mistral-large-latest", "api_key": "k3"}}, + ], + fallbacks=[{"primary": ["middle" if "middle" in cut_groups else "backup"]}, {"middle": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + + response: Final = await router.acompletion( + model="primary", + messages=[{"role": "user", "content": "How do transformers work?"}], + stream=True, + stream_options={"include_usage": True}, + ) + chunks: Final = [chunk async for chunk in response] + await asyncio.wait_for(recorder.succeeded.wait(), timeout=10) + await asyncio.sleep(0.5) + + backup_prompt_cost, backup_completion_cost = litellm.cost_per_token( + model="mistral/mistral-large-latest", prompt_tokens=91, completion_tokens=5 + ) + success: Final = recorder.successes[-1] + assert len(recorder.failures) == len(cut_groups) + assert [failure["response_cost"] for failure in recorder.failures] == pytest.approx( + [ + failure["prompt_tokens"] * input_price + failure["completion_tokens"] * output_price + for failure, (input_price, output_price) in zip(recorder.failures, ((0.001, 0.002), (0.003, 0.004))) + ] + ) + assert success["response_cost"] == pytest.approx( + sum(failure["response_cost"] for failure in recorder.failures) + backup_prompt_cost + backup_completion_cost + ) + assert success["prompt_tokens"] == sum(failure["prompt_tokens"] for failure in recorder.failures) + 91 + assert success["completion_tokens"] == sum(failure["completion_tokens"] for failure in recorder.failures) + 5 + assert success["cost_breakdown"]["total_cost"] == pytest.approx(success["response_cost"]) + spend_row: Final = recorder.spend_rows[-1] + assert (spend_row["spend"], spend_row["prompt_tokens"], spend_row["completion_tokens"]) == ( + success["response_cost"], + success["prompt_tokens"], + success["completion_tokens"], + ) + client_usage: Final = next(chunk.usage for chunk in reversed(chunks) if getattr(chunk, "usage", None)) + assert (client_usage.prompt_tokens, client_usage.completion_tokens) == ( + success["prompt_tokens"], + success["completion_tokens"], + ) + assert {chunk.id for chunk in chunks} == {spend_row["request_id"]} + + +@pytest.mark.asyncio +async def test_continued_stream_sharing_the_proxy_logging_object_bills_the_cut_to_its_own_deployment( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _SlowFailureRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use")) + ) + respx_mock.post(host="api.mistral.ai").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=_BACKUP_SSE) + ) + router: Final = Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "k1", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + }, + "model_info": {"id": "primary-deployment"}, + }, + {"model_name": "backup", "litellm_params": {"model": "mistral/mistral-large-latest", "api_key": "k3"}}, + ], + fallbacks=[{"primary": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + messages: Final = [{"role": "user", "content": "How do transformers work?"}] + shared_logging, _ = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + model="primary", + messages=messages, + stream=True, + litellm_call_id="shared-call", + ) + + response: Final = await router.acompletion( + model="primary", + messages=messages, + stream=True, + stream_options={"include_usage": True}, + litellm_call_id="shared-call", + litellm_logging_obj=shared_logging, + ) + _ = [chunk async for chunk in response] + await asyncio.wait_for(recorder.succeeded.wait(), timeout=10) + await asyncio.sleep(0.5) + + backup_prompt_cost, backup_completion_cost = litellm.cost_per_token( + model="mistral/mistral-large-latest", prompt_tokens=91, completion_tokens=5 + ) + (failure,) = recorder.failures + assert (failure["model_group"], failure["model_id"]) == ("primary", "primary-deployment") + assert recorder.failed_deployment_ids == ["primary-deployment"] + assert failure["response_cost"] == pytest.approx( + failure["prompt_tokens"] * 0.001 + failure["completion_tokens"] * 0.002 + ) + assert recorder.successes[-1]["response_cost"] == pytest.approx( + failure["response_cost"] + backup_prompt_cost + backup_completion_cost + ) + + +def test_carry_partial_spend_to_fallback_adds_the_partial_the_cut_stream_inherited(): + from litellm.types.utils import Usage + + cut_logging: Final = LiteLLMLogging( + model="deepseek/deepseek-chat", messages=[], stream=True, call_type="acompletion", + start_time=datetime.now(), litellm_call_id="cut", function_id="cut", + ) + cut_logging.model_call_details["response_cost"] = 0.5 + cut_logging.update_environment_variables( + litellm_params={"custom_llm_provider": "deepseek"}, optional_params={}, model="deepseek/deepseek-chat" + ) + cut_logging.record_continued_partial( + usage=Usage(prompt_tokens=40, completion_tokens=7, total_tokens=47), + response_cost=0.001, + response_id="chatcmpl-primary", + ) + fallback_logging: Final = LiteLLMLogging( + model="mistral-large-latest", messages=[], stream=True, call_type="acompletion", + start_time=datetime.now(), litellm_call_id="fallback", function_id="fallback", + ) + + Router._carry_partial_spend_to_fallback( + SimpleNamespace(logging_obj=fallback_logging), + Router._price_cut_partial( + cast(CustomStreamWrapper, SimpleNamespace(logging_obj=cut_logging, model="deepseek/deepseek-chat")), + litellm.ModelResponse( + id="chatcmpl-relay", usage=Usage(prompt_tokens=60, completion_tokens=3, total_tokens=63) + ), + ), + ) + + relay_prompt_cost, relay_completion_cost = litellm.cost_per_token( + model="deepseek/deepseek-chat", prompt_tokens=60, completion_tokens=3 + ) + carried_usage: Final = fallback_logging.model_call_details["continued_partial_usage"] + assert (carried_usage.prompt_tokens, carried_usage.completion_tokens) == (100, 10) + assert fallback_logging.model_call_details["continued_partial_cost"] == pytest.approx( + relay_prompt_cost + relay_completion_cost + 0.001 + ) + assert fallback_logging.model_call_details["continued_response_id"] == "chatcmpl-primary" + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_plain_fallback_keeps_the_fallback_usage(): + from litellm.exceptions import MidStreamFallbackError + + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="anthropic", generated_content="", + is_pre_first_chunk=True, + ) + role_chunk = litellm.ModelResponseStream( + model="gpt-4", choices=[{"index": 0, "delta": {"role": "assistant", "content": ""}}] + ) + + _, streamed = await _drain_with_fallback( + _continuation_router(), + error, + [role_chunk, _usage_chunk(21, 0)], + [{"role": "user", "content": "Hi"}], + [_usage_chunk(91, 5)], + ) + + assert (streamed[-1].usage.prompt_tokens, streamed[-1].usage.completion_tokens) == (91, 5) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "partial,fallback_texts,expected_answer", + [ + ("Transformers use ", (" self", "-attention"), "Transformers use self-attention"), + ("def f():\n ", ("\n", " return 1"), "def f():\n return 1"), + ("First.\n", ("\n\nSecond",), "First.\n\nSecond"), + ], + ids=["space", "indentation", "paragraph_break"], +) +async def test_acompletion_streaming_iterator_continuation_does_not_double_trailing_whitespace( + partial, fallback_texts, expected_answer +): + from litellm.exceptions import MidStreamFallbackError + + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial, + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + fallback_items = [ + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) for text in fallback_texts + ] + + passed_kwargs, streamed = await _drain_with_fallback( + _continuation_router(), error, [], [{"role": "user", "content": "Hi"}], fallback_items + ) + + assert passed_kwargs["messages"][-1]["content"] == partial + assert partial + "".join(chunk.choices[0].delta.content for chunk in streamed) == expected_answer + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", [True, "true", " True "]) +async def test_update_settings_turns_on_mid_stream_continuation(monkeypatch, enabled): + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, + ContinuationPrefillDeploymentCheck, + ) + + monkeypatch.setattr(litellm, "callbacks", []) + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + ) + router.update_settings(enable_mid_stream_fallback_continuation=enabled) + + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}]) + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), + ) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error), + messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + collected = [chunk async for chunk in result] + + assert mock_fallback.await_args.kwargs["kwargs"][MID_STREAM_CONTINUATION_KWARG] is MID_STREAM_CONTINUATION_MARKER + assert fallback_chunk in collected + assert any(isinstance(callback, ContinuationPrefillDeploymentCheck) for callback in litellm.callbacks) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error_kwargs,request_kwargs", + [ + ({"emitted_disqualifying_content": True}, {}), + ({"emitted_disqualifying_content": False}, {"response_format": {"type": "json_object"}}), + ({"emitted_disqualifying_content": False}, {"tool_choice": "required"}), + ({"emitted_disqualifying_content": False}, {"merge_reasoning_content_in_choices": True}), + ({"emitted_disqualifying_content": False}, {"n": 2}), + ], + ids=["tool_or_thinking_emitted", "json_mode", "forced_tool_choice", "merged_reasoning", "multiple_choices"], +) +async def test_acompletion_streaming_iterator_declines_ineligible_after_content(error_kwargs, request_kwargs): + """Flag on but the break is not continuation-safe: the router re-raises and + never enters the fallback chain, so no duplicated/rejected request is sent.""" + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + error = MidStreamFallbackError( + message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, **error_kwargs, + ) + source = _make_midstream_source(error) + initial_kwargs = {"model": "gpt-4", "stream": True, **request_kwargs} + + with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=source, messages=[{"role": "user", "content": "Hi"}], initial_kwargs=initial_kwargs + ) + with pytest.raises(MidStreamFallbackError): + async for _ in result: + pass + + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_raises_content_filter_instead_of_continuing(): + """A content filter that stops a stream mid-answer reaches the client as ContentPolicyViolationError, + never as a continuation that lets a backup finish the filtered answer.""" + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "azure/gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + content_filter = litellm.ContentPolicyViolationError(message="filtered", model="gpt-4", llm_provider="azure") + error = MidStreamFallbackError( + message="filtered", model="gpt-4", llm_provider="azure", generated_content="Step one is to obtain the", + is_pre_first_chunk=False, emitted_disqualifying_content=False, original_exception=content_filter, + ) + + with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error), + messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + with pytest.raises(litellm.ContentPolicyViolationError): + async for _ in result: + pass + + mock_fallback.assert_not_awaited() + + +async def _continues_after_content( + initial_kwargs: dict[str, object], + primary_params: Mapping[str, object] = MappingProxyType({}), + default_litellm_params: dict[str, object] | None = None, + cut_model_id: str = "primary-deployment", +) -> bool: + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "k1", **primary_params}, + "model_info": {"id": "primary-deployment"}, + }, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + { + "model_name": "openai-backup", + "litellm_params": {"model": "gpt-4o", "api_key": "k3"}, + "model_info": {"id": "no-prefill-deployment", "supports_assistant_prefill": False}, + }, + { + "model_name": "claude-backup", + "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k4"}, + "model_info": {"id": "prefill-deployment", "supports_assistant_prefill": True}, + }, + { + "model_name": "json-claude-backup", + "litellm_params": { + "model": "anthropic/claude-3-opus-20240229", + "api_key": "k5", + "response_format": {"type": "json_object"}, + }, + "model_info": {"id": "json-prefill-deployment", "supports_assistant_prefill": True}, + }, + ], + fallbacks=[{"gpt-4": ["backup"]}], + default_litellm_params=default_litellm_params, + enable_mid_stream_fallback_continuation=True, + ) + error: Final = MidStreamFallbackError( + message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + fallback_chunk: Final = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}]) + cut_stream: Final = _make_midstream_source(error) + cut_stream._hidden_params = {"model_id": cut_model_id} + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), + ) as mock_fallback: + result: Final = await router._acompletion_streaming_iterator( + model_response=cut_stream, + messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True, **initial_kwargs}, + ) + with contextlib.suppress(MidStreamFallbackError): + _ = [chunk async for chunk in result] + return mock_fallback.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "event_hook,continues", + [("post_call", False), (None, False), ("pre_call", True)], + ids=["post_call_guardrail", "every_hook_guardrail", "input_only_guardrail"], +) +async def test_mid_stream_continuation_declines_while_an_output_guardrail_is_registered( + monkeypatch: pytest.MonkeyPatch, event_hook: str | None, continues: bool +) -> None: + from litellm.types.guardrails import GuardrailEventHooks + + guardrail: Final = CustomGuardrail( + guardrail_name="output-check", event_hook=GuardrailEventHooks(event_hook) if event_hook else None + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + assert await _continues_after_content({}) is continues + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fallbacks,continues", + [ + ([{"model": "backup", "messages": [{"role": "user", "content": "Say hi"}]}], False), + ([{"model": "backup", "stream": False}], False), + ([{"model": "backup", "max_tokens": 300}], False), + ([{"model": "backup"}], True), + ([{"gpt-4": [{"model": "backup", "stream": False}]}], False), + ([{"*": [{"model": "backup", "max_tokens": 300}]}], False), + ([{"gpt-4": [{"model": "backup"}, "backup"]}], True), + ([{"gpt-4": ["backup"]}, {"other-group": [{"model": "backup", "stream": False}]}], True), + ([{}], True), + ([{"gpt-4": 5}], True), + ], + ids=[ + "entry_with_its_own_messages", + "entry_turning_off_streaming", + "entry_with_its_own_max_tokens", + "entry_with_model_only", + "group_target_turning_off_streaming", + "generic_group_target_with_its_own_max_tokens", + "group_targets_with_model_only", + "override_only_in_another_groups_chain", + "empty_entry", + "non_list_chain", + ], +) +async def test_mid_stream_continuation_declines_a_request_fallback_that_overrides_request_params( + fallbacks: list[object], continues: bool +) -> None: + assert await _continues_after_content({"fallbacks": fallbacks}) is continues + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fallbacks,continues", + [ + ([{"gpt-4": ["no-prefill-deployment"]}], False), + ([{"gpt-4": ["backup", "no-prefill-deployment"]}], False), + ([{"model": "no-prefill-deployment"}], False), + ([{"gpt-4": ["prefill-deployment"]}], True), + ([{"gpt-4": ["openai-backup"]}], True), + ([{"gpt-4": ["json-prefill-deployment"]}], False), + ], + ids=[ + "deployment_id_without_prefill", + "deployment_id_without_prefill_later_in_chain", + "flat_entry_naming_a_deployment_id_without_prefill", + "deployment_id_with_prefill", + "group_without_prefill_left_to_the_deployment_filter", + "deployment_id_with_its_own_response_format", + ], +) +async def test_mid_stream_continuation_declines_a_fallback_naming_a_deployment_that_cannot_continue( + fallbacks: list[object], continues: bool +) -> None: + assert await _continues_after_content({"fallbacks": fallbacks}) is continues + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "primary_params,default_litellm_params,continues", + [ + ({"response_format": {"type": "json_object"}}, None, False), + ({"n": 2}, None, False), + ({"tool_choice": "required"}, None, False), + ({"merge_reasoning_content_in_choices": True}, None, False), + ({}, {"response_format": {"type": "json_object"}}, False), + ({"response_format": {"type": "text"}}, None, True), + ({"temperature": 0.2}, None, True), + ({"tool_choice": "required"}, {"tool_choice": "auto"}, False), + ({"response_format": {"type": "json_object"}}, {"response_format": {"type": "text"}}, True), + ({"extra_body": {"n": 2}}, None, False), + ], + ids=[ + "deployment_response_format", + "deployment_n", + "deployment_tool_choice", + "deployment_merged_reasoning", + "router_default_response_format", + "deployment_plain_text_format", + "deployment_unrelated_param", + "deployment_tool_choice_over_router_default", + "router_default_format_over_deployment", + "deployment_extra_body_n", + ], +) +async def test_mid_stream_continuation_reads_constraints_configured_on_the_cut_deployment( + primary_params: Mapping[str, object], default_litellm_params: dict[str, object] | None, continues: bool +) -> None: + assert await _continues_after_content({}, primary_params, default_litellm_params) is continues + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "extra_body,continues", + [({"response_format": {"type": "json_object"}}, False), ({"n": 2}, False), ({"top_k": 5}, True)], + ids=["json_format_in_extra_body", "n_in_extra_body", "unrelated_extra_body_key"], +) +async def test_mid_stream_continuation_reads_constraints_the_request_sends_in_extra_body( + extra_body: dict[str, object], continues: bool +) -> None: + assert await _continues_after_content({"extra_body": extra_body}) is continues + + +@pytest.mark.asyncio +async def test_mid_stream_continuation_declines_when_the_cut_deployment_is_not_in_the_router() -> None: + assert await _continues_after_content({}, cut_model_id="per-request-credential-deployment") is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "routing_strategy", + [ + "simple-shuffle", + "usage-based-routing", + "usage-based-routing-v2", + "least-busy", + "latency-based-routing", + "cost-based-routing", + ], +) +async def test_continuation_routes_only_to_deployments_that_can_continue_under_every_routing_strategy( + routing_strategy: str, +) -> None: + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, + ) + from litellm.types.router import RouterRateLimitError + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "backup", + "litellm_params": {"model": "gpt-4o", "api_key": "k1"}, + "model_info": {"id": "no-prefill-deployment", "supports_assistant_prefill": False}, + }, + ], + routing_strategy=routing_strategy, + enable_mid_stream_fallback_continuation=True, + ) + messages: Final = [{"role": "user", "content": "Hi"}] + plain_pick: Final = await router.async_get_available_deployment( + model="backup", messages=messages, request_kwargs={} + ) + assert plain_pick["model_info"]["id"] == "no-prefill-deployment" + with pytest.raises(RouterRateLimitError): + await router.async_get_available_deployment( + model="backup", + messages=messages, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "flag_kwargs", [{}, {"enable_mid_stream_fallback_continuation": "false"}], ids=["default", "string_false"] +) +async def test_acompletion_streaming_iterator_flag_off_declines_after_content(flag_kwargs): + """Regression guard for the opt-in default: with the flag off, a post-content + break re-raises and never falls back, exactly as before this feature.""" + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + **flag_kwargs, + ) + error = MidStreamFallbackError( + message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + source = _make_midstream_source(error) + + with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=source, messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + with pytest.raises(MidStreamFallbackError): + async for _ in result: + pass + + mock_fallback.assert_not_awaited() + + +def test_build_completion_continuation_input_appends_assistant_prefill(): + messages = [{"role": "user", "content": "hi"}] + built = litellm.Router._build_completion_continuation_input(messages, "partial answer") + assert built[:-1] == messages + assert built[-1] == {"role": "assistant", "content": "partial answer", "prefix": True} + + +@pytest.mark.parametrize( + "prefill_content,second_partial,expected_content", + [ + ("part one ", " part two ", "part one part two "), + ("part one ", "part two ", "part one part two "), + ("First.\n\n", "Second", "First.\n\nSecond"), + ([{"type": "text", "text": "part one "}], " part two ", [{"type": "text", "text": "part one part two "}]), + ], + ids=["separator_sent_again", "separator_not_sent_again", "paragraph_break_not_sent_again", "content_blocks"], +) +def test_build_completion_continuation_input_folds_into_existing_prefill( + prefill_content, second_partial, expected_content +): + """A nested break must not leave two trailing assistant turns: the new partial folds into the prior prefill, + which keeps reading as the text the caller already received.""" + messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": prefill_content, "prefix": True}] + built = litellm.Router._build_completion_continuation_input(messages, second_partial) + assert built == [messages[0], {"role": "assistant", "content": expected_content, "prefix": True}] + + +@pytest.mark.parametrize( + "prefill_content,expected_content", + [ + ("Hello", "Hello world"), + ([{"type": "text", "text": "Hello"}], [{"type": "text", "text": "Hello world"}]), + ], + ids=["string_prefill", "content_block_prefill"], +) +def test_build_completion_continuation_input_does_not_double_an_echoed_prefill(prefill_content, expected_content): + messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": prefill_content, "prefix": True}] + built = litellm.Router._build_completion_continuation_input(messages, "Hello world") + assert built is not None + assert built[-1] == {"role": "assistant", "content": expected_content, "prefix": True} + + +@pytest.mark.parametrize( + "cut_partial,expected_content", + [("Hel", None), ("Hello", "Hello"), ("Hello wor", "Hello wor")], + ids=["cut_mid_echo", "cut_after_a_full_echo", "cut_after_new_text"], +) +def test_build_completion_continuation_input_declines_a_cut_in_the_middle_of_a_prefill_echo( + cut_partial, expected_content +): + messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "Hello", "prefix": True}] + built = litellm.Router._build_completion_continuation_input(messages, cut_partial) + expected = None if expected_content is None else [messages[0], {**messages[1], "content": expected_content}] + assert built == expected + + +@pytest.mark.parametrize( + "assistant_content", ["Earlier answer.", [{"type": "text", "text": "Earlier answer."}]], ids=["string", "blocks"] +) +def test_build_completion_continuation_input_declines_a_trailing_plain_assistant_turn(assistant_content): + """The caller's own assistant turn is history, not a prefill, so the partial is never fused into it.""" + messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": assistant_content}] + assert litellm.Router._build_completion_continuation_input(messages, "The capital of France is") is None + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_declines_when_request_ends_in_plain_assistant_turn(): + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + error = MidStreamFallbackError( + message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + + with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error), + messages=[{"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Earlier answer."}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + with pytest.raises(MidStreamFallbackError): + async for _ in result: + pass + + mock_fallback.assert_not_awaited() + + +def test_build_completion_continuation_input_keeps_structured_assistant_content(): + """Content blocks on a trailing assistant turn stay blocks: the partial continues the last text + block in place (a block boundary would let a provider drop the space between "Sure, " and "here"), + lands as a new text block only after a non-text block, and a content shape that cannot be extended declines.""" + from litellm.router import _extend_trailing_text_block + + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "text", "text": "Sure, "}], "prefix": True}, + ] + built = litellm.Router._build_completion_continuation_input(messages, "here it is") + assert built is not None + assert [m["role"] for m in built] == ["user", "assistant"] + assert built[-1] == {"role": "assistant", "content": [{"type": "text", "text": "Sure, here it is"}], "prefix": True} + image_block = {"type": "image_url", "image_url": {"url": "https://example.test/a.png"}} + assert _extend_trailing_text_block([image_block], "here it is") == [image_block, {"type": "text", "text": "here it is"}] + assert _extend_trailing_text_block([], "here it is") == [{"type": "text", "text": "here it is"}] + assert ( + litellm.Router._build_completion_continuation_input([{"role": "assistant", "content": 42, "prefix": True}], "x") + is None + ) + + +def test_mid_stream_continuation_eligible_allows_text_response_format(): + """response_format={"type": "text"} is the unconstrained default and must stay + eligible, unlike json_object / json_schema.""" + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k"}}], + enable_mid_stream_fallback_continuation=True, + ) + e = MidStreamFallbackError( + message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "text"}}) is True + assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "json_object"}}) is False + + +def test_raise_original_mid_stream_error_surfaces_the_provider_exception(): + from litellm.exceptions import MidStreamFallbackError, RateLimitError + + provider_error = RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4") + wrapped = MidStreamFallbackError( + message="rate limited", model="gpt-4", llm_provider="openai", + original_exception=provider_error, generated_content="Hello", + ) + with pytest.raises(RateLimitError) as raised: + litellm.Router._raise_original_mid_stream_error(wrapped) + assert raised.value is provider_error + assert raised.value.__cause__ is wrapped + + bare = MidStreamFallbackError(message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello") + with pytest.raises(MidStreamFallbackError) as bare_raised: + litellm.Router._raise_original_mid_stream_error(bare) + assert bare_raised.value is bare + + @pytest.mark.asyncio async def test_acompletion_streaming_iterator_reraises_original_exception_when_available(): """Async: when the mid-stream MidStreamFallbackError wraps a real provider @@ -8526,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 @@ -8692,6 +9996,11 @@ def test_combine_fallback_usage(): assert chunk.usage.completion_tokens == 5 assert chunk.usage.total_tokens == 15 + Router._combine_fallback_usage(chunk, Usage(prompt_tokens=0, completion_tokens=4, total_tokens=4)) + assert chunk.usage.prompt_tokens == 10 + assert chunk.usage.completion_tokens == 9 + assert chunk.usage.total_tokens == 19 + @pytest.mark.asyncio async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_failure(): @@ -12190,7 +13499,7 @@ def test_pre_call_checks_uses_deployment_model_when_model_info_lookup_raises(mon model="custom-alias", healthy_deployments=deployments, messages=[{"role": "user", "content": "hi"}], - request_kwargs={}, + request_kwargs={"response_format": {"type": "json_object"}}, ) assert len(result) == 1 @@ -14786,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 @@ -14796,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 @@ -16038,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 @@ -17169,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) @@ -17190,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"] @@ -17248,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 @@ -26764,3 +28118,471 @@ def test_completion_streams_a_wrapper_that_hands_iteration_to_another_iterator() streamed = list(response) assert streamed == chunks + + +_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES: Final = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"detail": "not-a-real-detail-value"}}, + ], + } +] + + +def test_pre_call_checks_filters_deployment_when_token_count_exceeds_limit(): + model_list: Final = [ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key"}, + "model_info": {"id": "small", "max_input_tokens": 5}, + }, + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-other-fake-model", "api_key": "my-fake-key"}, + "model_info": {"id": "large", "max_input_tokens": 10000}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="fake-openai-endpoint", + healthy_deployments=model_list, + messages=[{"role": "user", "content": " ".join(["word"] * 100)}], + ) + + assert [d["model_info"]["id"] for d in filtered] == ["large"] + + +def test_pre_call_checks_rpm_check_still_enforced_when_token_counting_fails(): + from litellm.types.router import RouterRateLimitErrorBasic + + model_list: Final = [ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "rpm": 0}, + "model_info": {"max_input_tokens": 100}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + with pytest.raises(RouterRateLimitErrorBasic): + router._pre_call_checks( + model="fake-openai-endpoint", + healthy_deployments=model_list, + messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES, + ) + + +def test_pre_call_checks_context_window_check_is_skipped_but_not_fatal_when_token_counting_fails(): + model_list: Final = [ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key"}, + "model_info": {"max_input_tokens": 100}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="fake-openai-endpoint", + healthy_deployments=model_list, + messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES, + ) + + assert len(filtered) == 1 + + +def test_pre_call_checks_keeps_unsupported_param_filter_when_token_counting_fails(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + no_json_mode_model: Final = "gpt-3.5-turbo-16k" + assert "response_format" not in ( + litellm.get_supported_openai_params(model=no_json_mode_model, custom_llm_provider="openai") or [] + ) + model_list: Final = [ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": f"openai/{deployment_model}", "api_key": "my-fake-key"}, + "model_info": {"id": deployment_model, "max_input_tokens": 100}, + } + for deployment_model in (no_json_mode_model, "gpt-4o") + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="fake-openai-endpoint", + healthy_deployments=model_list, + messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES, + request_kwargs={"response_format": {"type": "json_object"}}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["gpt-4o"] + + +async def test_async_get_healthy_deployments_keeps_region_check_when_token_counting_fails(): + model_list: Final = [ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "region_name": region}, + "model_info": {"id": f"{region}-deployment", "max_input_tokens": 100}, + } + for region in ("eu", "us") + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + healthy: Final = await router.async_get_healthy_deployments( + model="fake-openai-endpoint", + request_kwargs={"allowed_model_region": "eu"}, + messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES, + input=None, + specific_deployment=False, + parent_otel_span=None, + ) + + assert [deployment["model_info"]["id"] for deployment in healthy] == ["eu-deployment"] + + +@pytest.mark.parametrize( + "messages", + [_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES, [{"role": "user", "content": "hi"}]], + ids=["uncountable", "countable"], +) +def test_completion_answers_400_when_no_deployment_supports_the_requested_param( + monkeypatch: pytest.MonkeyPatch, messages: list[dict[str, object]] +): + monkeypatch.setattr(litellm, "drop_params", False) + router: Final = Router( + model_list=[ + { + "model_name": "no-json-mode", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key"}, + "model_info": {"max_input_tokens": 100}, + } + ], + enable_pre_call_checks=True, + num_retries=2, + ) + + with pytest.raises(litellm.UnsupportedParamsError) as error: + router.completion( + model="no-json-mode", + messages=messages, + response_format={"type": "json_object"}, + mock_response="should not be served", + ) + + assert error.value.status_code == 400 + assert "allowed_openai_params" in str(error.value) + + +@pytest.mark.parametrize( + ("deployment_params", "request_kwargs"), + [ + ({}, {"allowed_openai_params": ["response_format"]}), + ({}, {"drop_params": True}), + ({}, {"additional_drop_params": ["response_format"]}), + ({"allowed_openai_params": ["response_format"]}, {}), + ({"drop_params": True}, {}), + ], + ids=["request-allowlist", "request-drop", "request-additional-drop", "deployment-allowlist", "deployment-drop"], +) +def test_pre_call_checks_keeps_deployment_whose_unsupported_param_is_allowed_or_dropped( + monkeypatch: pytest.MonkeyPatch, deployment_params: Mapping[str, object], request_kwargs: Mapping[str, object] +): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "no-json-mode", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key", **deployment_params}, + "model_info": {"id": "only", "max_input_tokens": 100}, + } + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="no-json-mode", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}, **request_kwargs}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["only"] + + +def test_pre_call_checks_keeps_retryable_429_when_a_cooling_down_deployment_supports_the_param(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "mixed-json-mode", + "litellm_params": {"model": f"openai/{deployment_model}", "api_key": "my-fake-key"}, + "model_info": {"id": deployment_model, "max_input_tokens": 100}, + } + for deployment_model in ("gpt-3.5-turbo-16k", "gpt-4o") + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + without_json_mode: Final = [d for d in router.model_list if d["model_info"]["id"] == "gpt-3.5-turbo-16k"] + + filtered: Final = router._pre_call_checks( + model="mixed-json-mode", + healthy_deployments=without_json_mode, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}}, + ) + + assert filtered == [] + + +def test_completion_uses_context_window_fallback_when_the_small_deployment_also_lacks_the_param(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + router: Final = Router( + model_list=[ + { + "model_name": "small", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key"}, + "model_info": {"id": "small-id", "max_input_tokens": 5}, + }, + { + "model_name": "big", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "my-fake-key"}, + "model_info": {"id": "big-id"}, + }, + ], + enable_pre_call_checks=True, + context_window_fallbacks=[{"small": ["big"]}], + num_retries=0, + ) + + response: Final = router.completion( + model="small", + messages=[{"role": "user", "content": "hi " * 50}], + response_format={"type": "json_object"}, + mock_response="served by the fallback", + ) + + assert response._hidden_params["model_id"] == "big-id" + + +@pytest.mark.parametrize( + ("deployment_params", "request_kwargs"), + [ + ({}, {"allowed_openai_params": ["response_format"]}), + ({}, {"drop_params": True}), + ({}, {"additional_drop_params": ["response_format"]}), + ({"drop_params": True}, {}), + ({"allowed_openai_params": ["response_format"]}, {}), + ], + ids=["request-allowlist", "request-drop", "request-additional-drop", "deployment-drop", "deployment-allowlist"], +) +def test_pre_call_checks_prefers_a_deployment_that_takes_the_param_over_dropping_it( + monkeypatch: pytest.MonkeyPatch, deployment_params: Mapping[str, object], request_kwargs: Mapping[str, object] +): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "mixed-json-mode", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key", **deployment_params}, + "model_info": {"id": "gpt-3.5-turbo-16k", "max_input_tokens": 100}, + }, + { + "model_name": "mixed-json-mode", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "my-fake-key"}, + "model_info": {"id": "gpt-4o", "max_input_tokens": 100}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="mixed-json-mode", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}, **request_kwargs}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["gpt-4o"] + + +def test_pre_call_checks_drops_the_param_when_the_capable_deployment_is_outside_the_region(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "regional-json-mode", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key", "region_name": "eu"}, + "model_info": {"id": "eu-without-json-mode", "max_input_tokens": 100}, + }, + { + "model_name": "regional-json-mode", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "my-fake-key", "region_name": "us"}, + "model_info": {"id": "us-with-json-mode", "max_input_tokens": 100}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="regional-json-mode", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}, "drop_params": True, "allowed_model_region": "eu"}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["eu-without-json-mode"] + + +def test_pre_call_checks_filters_an_authenticating_provider_deployment_that_lacks_the_param( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + (tmp_path / "api-key.json").write_text(json.dumps({"token": "tid=test", "expires_at": 4102444800})) + model_list: Final = [ + { + "model_name": "mixed-json-mode", + "litellm_params": {"model": "github_copilot/gpt-4"}, + "model_info": {"id": "copilot-gpt-4", "max_input_tokens": 100}, + }, + { + "model_name": "mixed-json-mode", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "my-fake-key"}, + "model_info": {"id": "gpt-4o", "max_input_tokens": 100}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="mixed-json-mode", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["gpt-4o"] + + +def test_pre_call_checks_answers_400_when_the_group_name_lists_no_deployments(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "no-json-mode", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key"}, + "model_info": {"id": "only", "max_input_tokens": 100}, + } + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + with pytest.raises(litellm.UnsupportedParamsError): + router._pre_call_checks( + model="team-public-name", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}}, + ) + + +def test_pre_call_checks_lets_a_request_setting_replace_the_deployment_setting(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "no-json-mode", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key", "drop_params": True}, + "model_info": {"id": "drops-it", "max_input_tokens": 100}, + }, + { + "model_name": "no-json-mode", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "my-fake-key", + "allowed_openai_params": ["response_format"], + }, + "model_info": {"id": "allows-it", "max_input_tokens": 100}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="no-json-mode", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}, "drop_params": False}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["allows-it"] + + +async def _select_team_json_deployment(router: Router, entrypoint: str, selection: Mapping[str, object]) -> object: + if entrypoint == "async": + return await router.async_get_available_deployment(**selection) + if entrypoint == "sync": + return router.get_available_deployment(**selection) + return router.get_available_deployment_for_pass_through(**selection) + + +@pytest.mark.parametrize("entrypoint", ["sync", "async", "pass-through"]) +@pytest.mark.parametrize( + ("deployment_models", "cooled_down", "expected_error"), + [ + pytest.param(("gpt-3.5-turbo-16k", "gpt-4"), (), litellm.UnsupportedParamsError, id="none-takes-it"), + pytest.param(("gpt-3.5-turbo-16k", "gpt-4o"), ("gpt-4o",), RouterRateLimitError, id="capable-cooling-down"), + ], +) +@pytest.mark.asyncio +async def test_team_public_model_name_answers_400_only_when_no_team_deployment_takes_the_param( + monkeypatch: pytest.MonkeyPatch, + entrypoint: str, + deployment_models: tuple[str, str], + cooled_down: tuple[str, ...], + expected_error: type[Exception], +): + monkeypatch.setattr(litellm, "drop_params", False) + router: Final = Router( + model_list=[ + { + "model_name": f"team-json-{deployment_model}", + "litellm_params": { + "model": f"openai/{deployment_model}", + "api_key": "my-fake-key", + "use_in_pass_through": True, + }, + "model_info": { + "id": deployment_model, + "max_input_tokens": 100, + "team_id": "json-team", + "team_public_model_name": "team-json", + }, + } + for deployment_model in deployment_models + ], + enable_pre_call_checks=True, + num_retries=0, + ) + _cool_down(router, *cooled_down) + selection: Final = { + "model": "team-json", + "messages": [{"role": "user", "content": "hi"}], + "request_kwargs": { + "metadata": {"user_api_key_team_id": "json-team"}, + "response_format": {"type": "json_object"}, + }, + } + + with pytest.raises(expected_error): + await _select_team_json_deployment(router, entrypoint, selection) + + +def test_pre_call_checks_ignores_an_image_response_format_string(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "image-model", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key"}, + "model_info": {"id": "only", "max_input_tokens": 100}, + } + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="image-model", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "prompt"}], + request_kwargs={"response_format": "b64_json"}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["only"] 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_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index fc1799861d1..d8625ff9da3 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -119,27 +119,33 @@ async def test_discovered_limits_survive_deployment_growth_and_removal( async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) - router: Final = Router(model_list=[{ - "model_name": "local", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": "https://original.test/v1", - "api_key": "local-key", - }, - "model_info": {"id": "replaced-deployment"}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "local", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": "https://original.test/v1", + "api_key": "local-key", + }, + "model_info": {"id": "replaced-deployment"}, + } + ] + ) def respond(request: httpx.Request) -> httpx.Response: if request.url.host == "original.test": - router.upsert_deployment(Deployment( - model_name="local", - litellm_params=LiteLLM_Params( - model="hosted_vllm/local-model", - api_base="https://replacement.test/v1", - api_key="local-key", - ), - model_info=ModelInfo(id="replaced-deployment"), - )) + router.upsert_deployment( + Deployment( + model_name="local", + litellm_params=LiteLLM_Params( + model="hosted_vllm/local-model", + api_base="https://replacement.test/v1", + api_key="local-key", + ), + model_info=ModelInfo(id="replaced-deployment"), + ) + ) return httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}) assert request.url.host == "replacement.test" return httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 2048}]}) @@ -158,15 +164,19 @@ async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) first, second = tuple( - Router(model_list=[{ - "model_name": "local", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": f"https://{host}.test/v1", - "api_key": "local-key", - }, - "model_info": {"id": "shared-discovery-id"}, - }]) + Router( + model_list=[ + { + "model_name": "local", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": f"https://{host}.test/v1", + "api_key": "local-key", + }, + "model_info": {"id": "shared-discovery-id"}, + } + ] + ) for host in ("first", "second") ) @@ -187,15 +197,17 @@ async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: assert first.get_configured_token_limits("local") == (8192, 8192) assert second.get_configured_token_limits("local") == (2048, 2048) assert litellm.model_cost["shared-discovery-id"].get("max_input_tokens") is None - first.upsert_deployment(Deployment( - model_name="local", - litellm_params=LiteLLM_Params( - model="hosted_vllm/local-model", - api_base="https://unavailable.test/v1", - api_key="local-key", - ), - model_info=ModelInfo(id="shared-discovery-id"), - )) + first.upsert_deployment( + Deployment( + model_name="local", + litellm_params=LiteLLM_Params( + model="hosted_vllm/local-model", + api_base="https://unavailable.test/v1", + api_key="local-key", + ), + model_info=ModelInfo(id="shared-discovery-id"), + ) + ) assert first.get_configured_token_limits("local") == (None, None) await first.arefresh_model_info(client=handler) assert first.get_configured_token_limits("local") == (None, None) @@ -206,17 +218,19 @@ async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) second_started: Final = asyncio.Event() - router: Final = Router(model_list=[ - { - "model_name": host, - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": f"https://{host}.test/v1", - "api_key": "local-key", - }, - } - for host in ("first", "second", "third") - ]) + router: Final = Router( + model_list=[ + { + "model_name": host, + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": f"https://{host}.test/v1", + "api_key": "local-key", + }, + } + for host in ("first", "second", "third") + ] + ) async def respond(request: httpx.Request) -> httpx.Response: if request.url.host == "first.test": @@ -240,21 +254,27 @@ async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypa async def test_discovered_limits_expire_after_the_last_successful_refresh(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) clock: Final = Mock(return_value=0.0) - router: Final = Router(model_list=[{ - "model_name": "local", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": "https://expiry.test/v1", - "api_key": "local-key", - }, - "model_info": {"id": "expiring-discovery"}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "local", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": "https://expiry.test/v1", + "api_key": "local-key", + }, + "model_info": {"id": "expiring-discovery"}, + } + ] + ) router._discovered_model_info_cache = InMemoryCache(clock=clock, default_ttl=2 * MODEL_INFO_REFRESH_SECONDS) - responses: Final = iter(( - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}), - httpx.Response(503), - )) + responses: Final = iter( + ( + httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), + httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}), + httpx.Response(503), + ) + ) handler: Final = AsyncHTTPHandler() await handler.client.aclose() async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: next(responses))) as client: @@ -343,10 +363,12 @@ async def test_discovered_limits_are_isolated_overridable_and_refreshable( async def test_discovery_preserves_input_overrides_and_survives_outages(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) - responses: Final = iter(( - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), - httpx.Response(503), - )) + responses: Final = iter( + ( + httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), + httpx.Response(503), + ) + ) def respond(request: httpx.Request) -> httpx.Response: assert request.url.host == "backend.test" @@ -354,27 +376,29 @@ async def test_discovery_preserves_input_overrides_and_survives_outages(monkeypa assert request.headers["x-tenant"] == "tenant" return next(responses) - router: Final = Router(model_list=[ - { - "model_name": "configured", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": "https://backend.test/v1", - "api_key": "unused-key", - "extra_headers": {"authorization": "Bearer local-key", "X-Tenant": "tenant"}, + router: Final = Router( + model_list=[ + { + "model_name": "configured", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": "https://backend.test/v1", + "api_key": "unused-key", + "extra_headers": {"authorization": "Bearer local-key", "X-Tenant": "tenant"}, + }, + "model_info": {"id": "configured", "max_input_tokens": 1024}, }, - "model_info": {"id": "configured", "max_input_tokens": 1024}, - }, - { - "model_name": "byok", - "litellm_params": { - "model": "openai/local-model", - "api_base": "https://caller.test/v1", - "use_clientside_credentials": True, + { + "model_name": "byok", + "litellm_params": { + "model": "openai/local-model", + "api_base": "https://caller.test/v1", + "use_clientside_credentials": True, + }, }, - }, - {"model_name": "default-openai", "litellm_params": {"model": "openai/local-model", "api_key": "unused"}}, - ]) + {"model_name": "default-openai", "litellm_params": {"model": "openai/local-model", "api_key": "unused"}}, + ] + ) handler: Final = AsyncHTTPHandler() await handler.client.aclose() responder: Final = Mock(side_effect=respond) @@ -520,13 +544,14 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing(): ) -def test_regex_lookaround_flag_stays_on_the_deployment_that_set_it() -> None: - """A deployment's ``supports_regex_lookaround`` override must not land on the shared +@pytest.mark.parametrize("flag", ["supports_regex_lookaround", "supports_assistant_prefill"]) +def test_capability_flag_stays_on_the_deployment_that_set_it(flag: str) -> None: + """A deployment's capability override must not land on the shared ``{provider}/{model}`` key, or every sibling deployment of that model would inherit it.""" backend_model = "bedrock/us.xai.grok-4.6" deploy_id = "grok-deploy-keep-regex" - builtin_flag = litellm.get_model_info(model=backend_model).get("supports_regex_lookaround") + builtin_flag = litellm.get_model_info(model=backend_model).get(flag) model_keys = { deploy_id: litellm.model_cost.get(deploy_id), backend_model: copy.deepcopy(litellm.model_cost.get(backend_model)), @@ -537,13 +562,13 @@ def test_regex_lookaround_flag_stays_on_the_deployment_that_set_it() -> None: { "model_name": "grok-keep-regex", "litellm_params": {"model": backend_model}, - "model_info": {"id": deploy_id, "supports_regex_lookaround": not builtin_flag}, + "model_info": {"id": deploy_id, flag: not builtin_flag}, } ], ) - assert litellm.model_cost[deploy_id]["supports_regex_lookaround"] is (not builtin_flag) - assert litellm.get_model_info(model=backend_model).get("supports_regex_lookaround") is builtin_flag + assert litellm.model_cost[deploy_id][flag] is (not builtin_flag) + assert litellm.get_model_info(model=backend_model).get(flag) is builtin_flag finally: _restore_model_cost_entries(model_keys) @@ -944,8 +969,7 @@ def _register_azure_tier_backend() -> None: def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> None: model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) } try: _register_tier_backend() @@ -974,9 +998,7 @@ def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> Non "output_cost_per_token_flex": _TIER_BACKEND_ENTRY["output_cost_per_token_flex"], "input_cost_per_token_balanced": _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"], "output_cost_per_token_balanced": _TIER_BACKEND_ENTRY["output_cost_per_token_balanced"], - "cache_read_input_token_cost_ultrafast": _TIER_BACKEND_ENTRY[ - "cache_read_input_token_cost_ultrafast" - ], + "cache_read_input_token_cost_ultrafast": _TIER_BACKEND_ENTRY["cache_read_input_token_cost_ultrafast"], "input_cost_per_token_above_272k_tokens_ultrafast": _TIER_BACKEND_ENTRY[ "input_cost_per_token_above_272k_tokens_ultrafast" ], @@ -991,8 +1013,7 @@ def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> Non def test_inherit_builtin_service_tier_pricing_noop_without_base_rate_or_backend() -> None: model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) } try: _register_tier_backend() @@ -1029,8 +1050,7 @@ def test_inherit_builtin_service_tier_pricing_noop_without_base_rate_or_backend( def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing() -> None: model_id: Final = "tier-priced-deployment" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) } try: _register_tier_backend() @@ -1090,8 +1110,7 @@ def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing() def test_router_completion_uses_backend_ultrafast_long_context_rates() -> None: model_id: Final = "tier-priced-long-context-deployment" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) } try: _register_tier_backend() @@ -1142,8 +1161,7 @@ def test_ptu_service_tier_pricing_is_disabled_only_when_attribution_is_enabled( ) -> None: model_id: Final = f"ptu-tier-deployment-{ptu_enabled}" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, model_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, model_id) } try: _register_tier_backend() @@ -1200,8 +1218,7 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay model_id: Final = "azure-tier-priced-alias" payload_id: Final = "azure-tier-priced-payload" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_AZURE_TIER_BACKEND_KEY, model_id, payload_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_AZURE_TIER_BACKEND_KEY, model_id, payload_id) } try: _register_azure_tier_backend() @@ -1257,12 +1274,8 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay ) ) - assert payload["input_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY[ - "input_cost_per_token_priority" - ] - assert payload["output_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY[ - "output_cost_per_token_priority" - ] + assert payload["input_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["input_cost_per_token_priority"] + assert payload["output_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["output_cost_per_token_priority"] finally: _restore_model_cost_entries(model_cost_entries) litellm.get_model_info.cache_clear() @@ -1364,12 +1377,10 @@ def test_inherit_builtin_base_rates_for_off_peak_carries_threshold_rates(): assert model_info["input_cost_per_token"] == builtin_info["input_cost_per_token"] assert ( - model_info["input_cost_per_token_above_200k_tokens"] - == builtin_info["input_cost_per_token_above_200k_tokens"] + model_info["input_cost_per_token_above_200k_tokens"] == builtin_info["input_cost_per_token_above_200k_tokens"] ) assert ( - model_info["output_cost_per_token_above_200k_tokens"] - == builtin_info["output_cost_per_token_above_200k_tokens"] + model_info["output_cost_per_token_above_200k_tokens"] == builtin_info["output_cost_per_token_above_200k_tokens"] ) @@ -2260,8 +2271,7 @@ def test_deployment_model_cost_payload_folds_in_litellm_params_pricing(): def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() -> None: model_id: Final = "tier-priced-payload" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) } try: _register_tier_backend() @@ -2278,12 +2288,8 @@ def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() - ) ) - assert ( - payload["input_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"] - ) - assert ( - payload["output_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"] - ) + assert payload["input_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"] + assert payload["output_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"] assert payload["input_cost_per_token_balanced"] == _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"] assert payload["input_cost_per_token"] == _CUSTOM_STANDARD_INPUT_RATE assert payload["output_cost_per_token"] == _CUSTOM_STANDARD_OUTPUT_RATE @@ -3282,9 +3288,7 @@ def test_deployment_id_colliding_with_another_providers_catalog_key_keeps_its_ow ) assert isinstance(response, litellm.ModelResponse) - assert response._hidden_params["response_cost"] == pytest.approx( - 1000 * 0.00000096 + 500 * 0.00000302 - ) + assert response._hidden_params["response_cost"] == pytest.approx(1000 * 0.00000096 + 500 * 0.00000302) assert litellm.model_cost[_COLLIDING_BUILTIN_KEY] == builtin_row_before finally: _restore_model_cost_entries(model_cost_entries) @@ -3297,8 +3301,7 @@ def test_deployment_id_matching_its_own_providers_catalog_key_still_merges() -> matches against.""" model_id: Final = "baseten/zai-org/glm-5.2" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (model_id, _COLLIDING_BUILTIN_KEY) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (model_id, _COLLIDING_BUILTIN_KEY) } try: router: Final = Router(model_list=[_colliding_id_deployment(model_id, "baseten")]) @@ -3313,9 +3316,244 @@ def test_deployment_id_matching_its_own_providers_catalog_key_still_merges() -> ) assert isinstance(response, litellm.ModelResponse) - assert response._hidden_params["response_cost"] == pytest.approx( - 1000 * 0.00000096 + 500 * 0.00000302 - ) + assert response._hidden_params["response_cost"] == pytest.approx(1000 * 0.00000096 + 500 * 0.00000302) finally: _restore_model_cost_entries(model_cost_entries) litellm.get_model_info.cache_clear() + + +_CAPACITY_TEST_DEPLOYMENTS: Final = ( + ( + "bedrock", + "bedrock/converse/anthropic.claude-router-capacity-test", + "bedrock/us-east-1/anthropic.claude-router-capacity-test", + ), + ( + "bedrock_mantle", + "bedrock_mantle/openai.router-capacity-test", + "bedrock_mantle/us-east-1/openai.router-capacity-test", + ), + ("openai", "openai/router-capacity-test", "openai/us/router-capacity-test"), +) + + +def _capacity_test_metadata(provider: str, max_input_tokens: int) -> dict[str, object]: + return { + "litellm_provider": provider, + "mode": "chat", + "max_input_tokens": max_input_tokens, + "max_output_tokens": 512, + "max_tokens": 512, + "input_cost_per_token": 0, + "output_cost_per_token": 0, + } + + +def _capacity_test_router(route: str, base: str, base_location: str | None, explicit_limit: int | None) -> Router: + return Router( + model_list=[ + { + "model_name": "capacity-test", + "litellm_params": { + "model": route, + **({"base_model": base} if base_location == "litellm_params" else {}), + }, + "model_info": { + "id": "capacity-test", + **({"base_model": base} if base_location == "model_info" else {}), + **({"max_input_tokens": explicit_limit} if explicit_limit is not None else {}), + }, + } + ], + enable_pre_call_checks=True, + ) + + +def _deployments_allowed_for_2048_input_tokens(router: Router) -> int: + try: + return len( + router._pre_call_checks( + model="capacity-test", + healthy_deployments=router.model_list, + input="synthetic input", + input_token_count=2048, + ) + ) + except litellm.ContextWindowExceededError: + return 0 + + +def _effective_max_input_tokens(router: Router) -> object: + return router.get_router_model_info(id="capacity-test", deployment=None, received_model_name="capacity-test")[ + "max_input_tokens" + ] + + +@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) +@pytest.mark.parametrize("base_location", ("model_info", "litellm_params", None)) +@pytest.mark.parametrize("explicit_limit", (None, 4096)) +def test_router_capacity_widens_to_a_larger_base_model_window( + monkeypatch: pytest.MonkeyPatch, + provider: str, + route: str, + base: str, + base_location: str | None, + explicit_limit: int | None, +) -> None: + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + route: _capacity_test_metadata(provider, 1024), + base: _capacity_test_metadata(provider, 8192), + }, + ) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(route, base, base_location, explicit_limit) + expected: Final = ( + explicit_limit if explicit_limit is not None else (1024 if base_location is None else 8192) + ) + assert _effective_max_input_tokens(router) == expected + assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0) + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) +@pytest.mark.parametrize("base_location", ("model_info", "litellm_params")) +def test_router_capacity_never_narrows_to_a_smaller_base_model_window( + monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str, base_location: str +) -> None: + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + route: _capacity_test_metadata(provider, 8192), + base: _capacity_test_metadata(provider, 1024), + }, + ) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(route, base, base_location, None) + assert _effective_max_input_tokens(router) == 8192 + assert _deployments_allowed_for_2048_input_tokens(router) == 1 + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) +def test_router_capacity_never_takes_a_window_from_base_model_when_the_route_is_unmapped( + monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str +) -> None: + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, base: _capacity_test_metadata(provider, 1024)}) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(f"{route}-not-in-cost-map", base, "model_info", None) + assert _deployments_allowed_for_2048_input_tokens(router) == 1 + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) +def test_router_capacity_never_takes_a_window_from_base_model_when_the_route_declares_none( + monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str +) -> None: + route_without_window: Final = { + key: value + for key, value in _capacity_test_metadata(provider, 0).items() + if key not in ("max_input_tokens", "max_tokens") + } + monkeypatch.setattr( + litellm, + "model_cost", + {**litellm.model_cost, route: route_without_window, base: _capacity_test_metadata(provider, 1024)}, + ) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(route, base, "model_info", None) + assert _effective_max_input_tokens(router) is None + assert _deployments_allowed_for_2048_input_tokens(router) == 1 + _invalidate_model_cost_lowercase_map() + + +_SAME_MODEL_ROUTE: Final = "bedrock/converse/anthropic.claude-routercapacitytest-9-9" +_SAME_MODEL_GLOBAL_ROUTE: Final = "bedrock/converse/global.anthropic.claude-routercapacitytest-9-9" +_SAME_MODEL_QUALIFIED_ROUTE: Final = "bedrock/converse/anthropic.claude-routercapacitytest-9-9@first" + + +@pytest.mark.parametrize( + ("route", "base_model", "expected"), + ( + (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9", 8192), + (_SAME_MODEL_ROUTE, "us-east-1/anthropic.claude-routercapacitytest-9-9", 1024), + (_SAME_MODEL_ROUTE, "anthropic.claude-routercapacitytest-9-9", 1024), + (_SAME_MODEL_GLOBAL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9", 8192), + (_SAME_MODEL_ROUTE, "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1", 8192), + (_SAME_MODEL_GLOBAL_ROUTE, "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1", 8192), + (_SAME_MODEL_QUALIFIED_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@first", 8192), + (_SAME_MODEL_QUALIFIED_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@second", 1024), + (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-*", 1024), + (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9", 1024), + (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9-extended", 1024), + (_SAME_MODEL_ROUTE, "bedrock/eu-west-1/anthropic.claude-routercapacitytest-9-9", 1024), + ), +) +def test_router_capacity_widens_only_for_an_exact_base_model_entry_of_the_same_model( + monkeypatch: pytest.MonkeyPatch, route: str, base_model: str, expected: int +) -> None: + larger_window: Final = _capacity_test_metadata("bedrock", 8192) + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + _SAME_MODEL_ROUTE: _capacity_test_metadata("bedrock", 1024), + _SAME_MODEL_GLOBAL_ROUTE: _capacity_test_metadata("bedrock", 1024), + _SAME_MODEL_QUALIFIED_ROUTE: _capacity_test_metadata("bedrock", 1024), + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@first": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@second": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9": larger_window, + "anthropic.claude-routercapacitytest-9-9": _capacity_test_metadata("bedrock", 1024), + "bedrock/anthropic.claude-routercapacitytest-9-9": larger_window, + "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitytest-*": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9-extended": larger_window, + }, + ) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(route, base_model, "model_info", None) + assert _effective_max_input_tokens(router) == expected + assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0) + _invalidate_model_cost_lowercase_map() + + +def test_router_capacity_reads_a_model_info_base_model_before_a_litellm_params_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + _SAME_MODEL_ROUTE: _capacity_test_metadata("bedrock", 1024), + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9": _capacity_test_metadata("bedrock", 8192), + "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9": _capacity_test_metadata("bedrock", 8192), + }, + ) + _invalidate_model_cost_lowercase_map() + router: Final = Router( + model_list=[ + { + "model_name": "capacity-test", + "litellm_params": { + "model": _SAME_MODEL_ROUTE, + "base_model": "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9", + }, + "model_info": { + "id": "capacity-test", + "base_model": "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9", + }, + } + ], + enable_pre_call_checks=True, + ) + assert _effective_max_input_tokens(router) == 1024 + assert _deployments_allowed_for_2048_input_tokens(router) == 0 + _invalidate_model_cost_lowercase_map() 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 d27dc733c78..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", @@ -199,6 +200,7 @@ OPTION_NAMES: Final = ( "prompt_label", "litellm_system_prompt", "custom_prompt_dict", + "initial_prompt_value", "roles", "final_prompt_value", "bos_token", 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 && (