From 1d407c2f26d7587bf184f7cf19dfcaede7e860d7 Mon Sep 17 00:00:00 2001 From: Kent Date: Fri, 26 Jun 2026 17:24:21 +0800 Subject: [PATCH 001/504] fix(bedrock): validate file-content retrieval against the configured output bucket Bedrock batch jobs write their results to s3_output_bucket_name when it differs from the input bucket, but the file-content retrieval path validated the file id only against the input bucket (s3_bucket_name). A deployment that configures a separate output bucket therefore could not retrieve its own batch outputs: the id validated against the input bucket and was rejected as a foreign bucket. Resolve the trusted output bucket alongside the input bucket from the immutable credential snapshot (or AWS_S3_OUTPUT_BUCKET_NAME), and try the file id against each configured bucket, returning the first that validates. The SSRF guard is preserved: only server-configured buckets are tried, never a request param, and an id outside both is still rejected. --- litellm/llms/bedrock/files/transformation.py | 67 +++++++++++-- .../test_bedrock_files_transformation.py | 94 +++++++++++++++++++ 2 files changed, 151 insertions(+), 10 deletions(-) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 6cfaa88275d..df06c333d1f 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -81,11 +81,12 @@ class _BedrockS3RequestParams(BaseModel): class _TrustedS3ModelCredentials(BaseModel): - """The S3 bucket the server trusts file ids against, from the deployment snapshot.""" + """The S3 buckets the server trusts file ids against, from the deployment snapshot.""" model_config = ConfigDict(extra="ignore") s3_bucket_name: str | None = None + s3_output_bucket_name: str | None = None def extract_s3_uri_from_file_id(file_id: str) -> str: @@ -135,6 +136,41 @@ def get_configured_s3_bucket_name(litellm_params: Mapping[str, object]) -> str: return bucket_name +def get_configured_s3_bucket_names( + litellm_params: Mapping[str, object], +) -> tuple[str, ...]: + """ + Resolve the server-configured S3 buckets a Bedrock file id may live in. + + Bedrock batch outputs land in ``s3_output_bucket_name`` when it differs from + the input bucket, so retrieval validates against both. Same trust rules as + ``get_configured_s3_bucket_name``: only the immutable credential snapshot or + the environment, never a request param. + """ + trusted_model_credentials = litellm_params.get( + "_litellm_internal_model_credentials" + ) + input_bucket: str | None = None + output_bucket: str | None = None + if isinstance(trusted_model_credentials, MappingProxyType): + snapshot: dict[str, object] = {} + snapshot.update(trusted_model_credentials) # any-ok: untyped snapshot + trusted = _TrustedS3ModelCredentials.model_validate(snapshot) + input_bucket = trusted.s3_bucket_name + output_bucket = trusted.s3_output_bucket_name + input_bucket = input_bucket or os.getenv("AWS_S3_BUCKET_NAME") + output_bucket = output_bucket or os.getenv("AWS_S3_OUTPUT_BUCKET_NAME") + + buckets = tuple( + dict.fromkeys(bucket for bucket in (input_bucket, output_bucket) if bucket) + ) + if not buckets: + raise ValueError( + "S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval." + ) + return buckets + + class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ Config for Bedrock Files - handles S3 uploads for Bedrock batch processing @@ -1042,15 +1078,26 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): raise ValueError("file_id is required for Bedrock file content retrieval") s3_uri = extract_s3_uri_from_file_id(file_id) - bucket_name, object_key = validate_managed_cloud_file_id( - file_id=s3_uri, - scheme="s3://", - configured_bucket_name=get_configured_s3_bucket_name(litellm_params), - allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES, - allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids( - litellm_params - ), - ) + allow_legacy = should_allow_legacy_cloud_file_ids(litellm_params) + last_error: ValueError | None = None + bucket_name: str | None = None + object_key: str | None = None + for configured_bucket in get_configured_s3_bucket_names(litellm_params): + try: + bucket_name, object_key = validate_managed_cloud_file_id( + file_id=s3_uri, + scheme="s3://", + configured_bucket_name=configured_bucket, + allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES, + allow_legacy_cloud_file_ids=allow_legacy, + ) + break + except ValueError as e: + last_error = e + if bucket_name is None or object_key is None: + raise last_error or ValueError( + "file_id must reference a LiteLLM-managed storage object" + ) # The shared file-content handler passes optional_params={}, so AWS # credentials/region arrive via litellm_params here (unlike the upload diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index c548fe53e15..dd111969555 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1308,6 +1308,100 @@ class TestBedrockFileContentTransformation: litellm_params=self._litellm_params(), ) + def _trusted(self, **creds) -> dict: + from types import MappingProxyType + + params = self._litellm_params() + params["_litellm_internal_model_credentials"] = MappingProxyType(dict(creds)) + return params + + def test_retrieves_from_distinct_output_bucket(self, monkeypatch): + """Batch outputs can land in a separate s3_output_bucket_name. Retrieval + must validate the file id against the output bucket too, not just the + input bucket, or the very outputs the feature serves are unreachable.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + + url, _ = BedrockFilesConfig().transform_file_content_request( + file_content_request={ + "file_id": "s3://out-bucket/litellm-batch-outputs/job/in.jsonl.out" + }, + optional_params={}, + litellm_params=self._trusted( + s3_bucket_name="in-bucket", s3_output_bucket_name="out-bucket" + ), + ) + + assert ( + url + == "https://s3.us-west-2.amazonaws.com/out-bucket/litellm-batch-outputs/job/in.jsonl.out" + ) + + def test_output_bucket_falls_back_to_env(self, monkeypatch): + """The output bucket resolves from AWS_S3_OUTPUT_BUCKET_NAME when not in + the trusted snapshot, mirroring the input-bucket env fallback.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "in-bucket") + monkeypatch.setenv("AWS_S3_OUTPUT_BUCKET_NAME", "env-out-bucket") + + url, _ = BedrockFilesConfig().transform_file_content_request( + file_content_request={ + "file_id": "s3://env-out-bucket/litellm-batch-outputs/job/in.jsonl.out" + }, + optional_params={}, + litellm_params=self._litellm_params(), + ) + + assert ( + url + == "https://s3.us-west-2.amazonaws.com/env-out-bucket/litellm-batch-outputs/job/in.jsonl.out" + ) + + def test_input_bucket_still_validates_when_output_bucket_set(self, monkeypatch): + """Adding output-bucket support must not break retrieval of input-bucket + objects when both buckets are configured.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + + url, _ = BedrockFilesConfig().transform_file_content_request( + file_content_request={ + "file_id": "s3://in-bucket/litellm-batch-outputs/job/in.jsonl.out" + }, + optional_params={}, + litellm_params=self._trusted( + s3_bucket_name="in-bucket", s3_output_bucket_name="out-bucket" + ), + ) + + assert ( + url + == "https://s3.us-west-2.amazonaws.com/in-bucket/litellm-batch-outputs/job/in.jsonl.out" + ) + + def test_rejects_bucket_outside_input_and_output(self, monkeypatch): + """A file id whose bucket is neither the input nor the output bucket is + still rejected (SSRF / bucket-confusion guard).""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + + with pytest.raises(ValueError, match="configured storage bucket"): + BedrockFilesConfig().transform_file_content_request( + file_content_request={ + "file_id": "s3://other-bucket/litellm-batch-outputs/job/x.jsonl.out" + }, + optional_params={}, + litellm_params=self._trusted( + s3_bucket_name="in-bucket", s3_output_bucket_name="out-bucket" + ), + ) + def test_sign_request_without_botocore_raises_helpful_error(self, monkeypatch): """A missing botocore must surface an actionable 'install boto3' error rather than a raw import failure.""" From a9a322d63f26e48a517d84ad428e77fcb4cbca03 Mon Sep 17 00:00:00 2001 From: Kent Date: Fri, 26 Jun 2026 19:03:10 +0800 Subject: [PATCH 002/504] fix(router): keep s3_output_bucket_name in the trusted credential snapshot The output-bucket retrieval fix only worked via the AWS_S3_OUTPUT_BUCKET_NAME env var, never via per-model s3_output_bucket_name config. The proxy builds the trusted snapshot that retrieval validates against by round-tripping a deployment's litellm_params through CredentialLiteLLMParams in get_deployment_credentials_with_provider, and that strict allowlist did not declare s3_output_bucket_name, so the field was silently dropped before retrieval saw it (same trap as azure_ad_token in #30235). The snapshot branch of get_configured_s3_bucket_names was therefore dead in the model-routing path and output-bucket file ids were rejected as foreign. Declaring s3_output_bucket_name on CredentialLiteLLMParams lets it survive into the snapshot, so the existing multi-bucket validation works for per-model output buckets too. The PR's tests injected the field straight into the MappingProxyType, bypassing this filter, so they passed despite the live gap. _trusted now builds the snapshot through CredentialLiteLLMParams the way the proxy does, and a router-level regression test pins that get_deployment_credentials_with_provider preserves the output bucket. Both fail without this change. --- litellm/types/router.py | 6 ++ .../test_bedrock_files_transformation.py | 19 +++++- ...st_azure_ad_token_credential_resolution.py | 68 +++++++++++++++++++ 3 files changed, 90 insertions(+), 3 deletions(-) diff --git a/litellm/types/router.py b/litellm/types/router.py index a1c571ed7f7..333ec6d9b05 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -188,6 +188,12 @@ class CredentialLiteLLMParams(BaseModel): aws_bedrock_runtime_endpoint: Optional[str] = None aws_bedrock_project_id: Optional[str] = None s3_bucket_name: Optional[str] = None + # Like the fields above, must be declared here or the strict dump in + # ``get_deployment_credentials_with_provider`` drops it from the trusted + # snapshot, so per-model output-bucket config never reaches Bedrock + # file-content retrieval and output-bucket file ids are wrongly rejected + # (#26335). + s3_output_bucket_name: Optional[str] = None ## IBM WATSONX ## watsonx_region_name: Optional[str] = None diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index dd111969555..448132f20b1 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1308,17 +1308,30 @@ class TestBedrockFileContentTransformation: litellm_params=self._litellm_params(), ) - def _trusted(self, **creds) -> dict: + def _trusted(self, **deployment_litellm_params) -> dict: + """Build the trusted snapshot the way the proxy does: deployment + litellm_params funneled through ``CredentialLiteLLMParams`` (the strict + allowlist ``get_deployment_credentials_with_provider`` applies) before + retrieval ever sees them. Injecting a raw ``MappingProxyType`` would + bypass that filter and hide whether a bucket field actually survives + into the snapshot in production.""" from types import MappingProxyType + from litellm.types.router import CredentialLiteLLMParams + + snapshot = CredentialLiteLLMParams(**deployment_litellm_params).model_dump( + exclude_none=True + ) params = self._litellm_params() - params["_litellm_internal_model_credentials"] = MappingProxyType(dict(creds)) + params["_litellm_internal_model_credentials"] = MappingProxyType(snapshot) return params def test_retrieves_from_distinct_output_bucket(self, monkeypatch): """Batch outputs can land in a separate s3_output_bucket_name. Retrieval must validate the file id against the output bucket too, not just the - input bucket, or the very outputs the feature serves are unreachable.""" + input bucket, or the very outputs the feature serves are unreachable. + The snapshot is built through the production credential filter, so this + fails if s3_output_bucket_name is dropped from that allowlist.""" from litellm.llms.bedrock.files.transformation import BedrockFilesConfig monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) diff --git a/tests/test_litellm/test_azure_ad_token_credential_resolution.py b/tests/test_litellm/test_azure_ad_token_credential_resolution.py index 958b236c9b3..f6f47b3f4c4 100644 --- a/tests/test_litellm/test_azure_ad_token_credential_resolution.py +++ b/tests/test_litellm/test_azure_ad_token_credential_resolution.py @@ -143,3 +143,71 @@ class TestRouterCredentialResolution: assert credentials is not None assert credentials.get("api_key") == "sk-static-key" assert "azure_ad_token" not in credentials + + +class TestRouterCredentialResolutionS3OutputBucket: + """Same strict-dump trap as azure_ad_token (#30235), for Bedrock batch + file retrieval (#26335). Bedrock batch outputs land in a per-model + ``s3_output_bucket_name`` when it differs from the input bucket. The + file-content retrieval path validates a file id against the buckets in the + trusted credential snapshot, and that snapshot is built by round-tripping + the deployment's ``litellm_params`` through ``CredentialLiteLLMParams``. If + the field is undeclared it is dropped, so the output bucket never reaches + retrieval and output-bucket file ids are rejected as foreign.""" + + def test_credentials_preserve_s3_output_bucket_name(self): + from litellm import Router + + deployment_id = "bedrock-batch-output-bucket-fixed-uuid" + router = Router( + model_list=[ + { + "model_name": "bedrock-batch", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + "s3_bucket_name": "in-bucket", + "s3_output_bucket_name": "out-bucket", + "aws_region_name": "us-west-2", + }, + "model_info": {"id": deployment_id}, + } + ] + ) + + credentials = router.get_deployment_credentials_with_provider( + model_id=deployment_id + ) + assert credentials is not None + assert credentials.get("s3_output_bucket_name") == "out-bucket", ( + "Router credential resolution dropped s3_output_bucket_name; " + "Bedrock batch file-content retrieval will reject output-bucket " + "file ids as foreign for model-routed deployments (#26335)" + ) + assert credentials.get("s3_bucket_name") == "in-bucket" + + def test_credentials_without_output_bucket_unaffected(self): + """A deployment that configures only the input bucket keeps it and does + not gain a phantom output bucket in the resolved credentials.""" + from litellm import Router + + deployment_id = "bedrock-batch-input-only-fixed-uuid" + router = Router( + model_list=[ + { + "model_name": "bedrock-batch-input-only", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + "s3_bucket_name": "in-bucket", + "aws_region_name": "us-west-2", + }, + "model_info": {"id": deployment_id}, + } + ] + ) + + credentials = router.get_deployment_credentials_with_provider( + model_id=deployment_id + ) + assert credentials is not None + assert credentials.get("s3_bucket_name") == "in-bucket" + assert "s3_output_bucket_name" not in credentials From 6e6a508e0e46da5af34b9ccaa89602a2acf2adc1 Mon Sep 17 00:00:00 2001 From: Kent Date: Fri, 26 Jun 2026 19:18:09 +0800 Subject: [PATCH 003/504] chore(ui): regenerate schema.d.ts for s3_output_bucket_name Adding s3_output_bucket_name to CredentialLiteLLMParams changes the proxy OpenAPI spec, so the generated dashboard types need regenerating to match (Check UI API Types Sync). --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b53acf930f2..f4c124659a4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25410,6 +25410,8 @@ export interface components { s3_bucket_name?: string | null; /** S3 Encryption Key Id */ s3_encryption_key_id?: string | null; + /** S3 Output Bucket Name */ + s3_output_bucket_name?: string | null; /** Search Context Cost Per Query */ search_context_cost_per_query?: { [key: string]: unknown; @@ -33118,6 +33120,8 @@ export interface components { s3_bucket_name?: string | null; /** S3 Encryption Key Id */ s3_encryption_key_id?: string | null; + /** S3 Output Bucket Name */ + s3_output_bucket_name?: string | null; /** Search Context Cost Per Query */ search_context_cost_per_query?: { [key: string]: unknown; From 6ed1c6b420e197b3327cf37e9d54f3a6cd9e17fd Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Fri, 26 Jun 2026 13:25:28 +0000 Subject: [PATCH 004/504] fix(deps): bump langgraph-checkpoint to 4.1.1 to resolve OSV vulnerability --- uv.lock | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/uv.lock b/uv.lock index cac1696bf34..8c9e20eeb21 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-06-20T23:16:25.061268Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P3D" [manifest] @@ -3160,15 +3160,15 @@ wheels = [ [[package]] name = "langgraph-checkpoint" -version = "4.1.0" +version = "4.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "ormsgpack" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/02/b4/6005c5dd88ad484fe6235d4c43a0d2cee7e91b08ad85a180985c2662df87/langgraph_checkpoint-4.1.0.tar.gz", hash = "sha256:e5bb304e30fc1363ac8fcb5f7dee5ca2185d77fe475b0d01de2c5f91324c2c21", size = 181942, upload-time = "2026-05-12T03:33:49.888Z" } +sdist = { url = "https://files.pythonhosted.org/packages/83/47/886af6f886f0bff2273164a45f008694e48a96ff3cd25ff0228f2aa9480e/langgraph_checkpoint-4.1.1.tar.gz", hash = "sha256:6c2bdb530c91f91d7d9c1bd100925d0fc4f498d418c17f3587d1526279482a25", size = 184020, upload-time = "2026-05-22T16:57:38.503Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/93/74/d3be2b41955e20ccd624dba5f6fe9d38dcee385ba470a6e13ed86732fc86/langgraph_checkpoint-4.1.0-py3-none-any.whl", hash = "sha256:8bc2a0466a20c38b865ce6671b42093fd5c041133f32351cae4222e0eeaf7fb5", size = 56047, upload-time = "2026-05-12T03:33:48.548Z" }, + { url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" }, ] [[package]] @@ -3281,6 +3281,7 @@ dependencies = [ { name = "importlib-metadata" }, { name = "jinja2" }, { name = "jsonschema" }, + { name = "langgraph-checkpoint" }, { name = "openai" }, { name = "pydantic" }, { name = "python-dotenv" }, @@ -3505,6 +3506,7 @@ requires-dist = [ { name = "jinja2", specifier = ">=3.1.6,<4.0" }, { name = "jsonschema", specifier = ">=4.0.0,<5.0" }, { name = "langfuse", marker = "extra == 'proxy-runtime'", specifier = ">=2.59.7,<3.0" }, + { name = "langgraph-checkpoint", specifier = "==4.1.1" }, { name = "litellm-enterprise", marker = "extra == 'proxy'", editable = "enterprise" }, { name = "litellm-proxy-extras", marker = "extra == 'proxy'", editable = "litellm-proxy-extras" }, { name = "llm-sandbox", marker = "extra == 'proxy-runtime'", specifier = ">=0.3.39,<1.0" }, From 3f5186f9afcced38bdcd8a6095c11e6a8e206627 Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Fri, 26 Jun 2026 13:39:29 +0000 Subject: [PATCH 005/504] fix(ocr): use defensive getattr in load_rust_ocr --- litellm/ocr/rust_bridge.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index 1e3312c1473..253b35cb689 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/ocr/rust_bridge.py @@ -97,7 +97,7 @@ def load_rust_ocr() -> RustOcr | None: import litellm_python_bridge except ImportError: return None - return cast(RustOcr, litellm_python_bridge.ocr) + return cast(RustOcr, getattr(litellm_python_bridge, "ocr", None)) def load_rust_aocr() -> RustAocr | None: From 6a940ef3f4c3a6c89698d5fb082bd3a99c4841bd Mon Sep 17 00:00:00 2001 From: Kent Date: Tue, 30 Jun 2026 02:10:44 +0800 Subject: [PATCH 006/504] chore(types): type the dict-shim helpers to offset the budget gate Adding s3_output_bucket_name to CredentialLiteLLMParams adds one reportUnknownArgumentType error at each untyped **kwargs construction site of GenericLiteLLMParams repo-wide (~114 sites), which pushed the basedpyright budget just over its ceiling. Typing the key parameter of the get/__getitem__/ __setitem__/__contains__ dict-shim helpers on ModelInfo, GenericLiteLLMParams, LiteLLM_Params, and Deployment removes the unknown-argument errors at the getattr/setattr/hasattr calls in those bodies, bringing the repo total back under the cap without raising any other rule. --- litellm/types/router.py | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/litellm/types/router.py b/litellm/types/router.py index 87e3364ac7e..75370d94895 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -154,19 +154,19 @@ class ModelInfo(BaseModel): model_config = ConfigDict(extra="allow") - def __contains__(self, key): + def __contains__(self, key: str) -> bool: # Define custom behavior for the 'in' operator return hasattr(self, key) - def get(self, key, default=None): + def get(self, key: str, default=None): # Custom .get() method to access attributes with a default value if the attribute doesn't exist return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key: str): # Allow dictionary-style access to attributes return getattr(self, key) - def __setitem__(self, key, value): + def __setitem__(self, key: str, value) -> None: # Allow dictionary-style assignment of attributes setattr(self, key, value) @@ -303,19 +303,19 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): return filtered return data - def __contains__(self, key): + def __contains__(self, key: str) -> bool: # Define custom behavior for the 'in' operator return hasattr(self, key) - def get(self, key, default=None): + def get(self, key: str, default=None): # Custom .get() method to access attributes with a default value if the attribute doesn't exist return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key: str): # Allow dictionary-style access to attributes return getattr(self, key) - def __setitem__(self, key, value): + def __setitem__(self, key: str, value) -> None: # Allow dictionary-style assignment of attributes setattr(self, key, value) @@ -328,19 +328,19 @@ class LiteLLM_Params(GenericLiteLLMParams): model: str model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True) - def __contains__(self, key): + def __contains__(self, key: str) -> bool: # Define custom behavior for the 'in' operator return hasattr(self, key) - def get(self, key, default=None): + def get(self, key: str, default=None): # Custom .get() method to access attributes with a default value if the attribute doesn't exist return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key: str): # Allow dictionary-style access to attributes return getattr(self, key) - def __setitem__(self, key, value): + def __setitem__(self, key: str, value) -> None: # Allow dictionary-style assignment of attributes setattr(self, key, value) @@ -473,19 +473,19 @@ class Deployment(BaseModel): # if using pydantic v1 return self.dict(**kwargs) - def __contains__(self, key): + def __contains__(self, key: str) -> bool: # Define custom behavior for the 'in' operator return hasattr(self, key) - def get(self, key, default=None): + def get(self, key: str, default=None): # Custom .get() method to access attributes with a default value if the attribute doesn't exist return getattr(self, key, default) - def __getitem__(self, key): + def __getitem__(self, key: str): # Allow dictionary-style access to attributes return getattr(self, key) - def __setitem__(self, key, value): + def __setitem__(self, key: str, value) -> None: # Allow dictionary-style assignment of attributes setattr(self, key, value) From 39bd5fb8b7923f94712bac1bd9f1f6a93adb50ad Mon Sep 17 00:00:00 2001 From: Sujith Date: Tue, 14 Jul 2026 15:18:16 +0530 Subject: [PATCH 007/504] fix(main): forward store and prompt_cache_key params on chat completions (#33184) --- litellm/main.py | 8 ++++ litellm/utils.py | 2 + tests/test_litellm/test_main.py | 84 +++++++++++++++++++++++++++++++++ 3 files changed, 94 insertions(+) diff --git a/litellm/main.py b/litellm/main.py index 7d457d9cdd1..81b02082abd 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -435,6 +435,8 @@ async def acompletion( verbosity: Optional[Literal["low", "medium", "high"]] = None, safety_identifier: Optional[str] = None, service_tier: Optional[str] = None, + store: Optional[bool] = None, + prompt_cache_key: Optional[str] = None, # set api_base, api_version, api_key base_url: Optional[str] = None, api_version: Optional[str] = None, @@ -585,6 +587,8 @@ async def acompletion( "verbosity": verbosity, "safety_identifier": safety_identifier, "service_tier": service_tier, + "store": store, + "prompt_cache_key": prompt_cache_key, "extra_headers": extra_headers, "acompletion": True, # assuming this is a required parameter "thinking": thinking, @@ -4828,6 +4832,8 @@ def completion( # type: ignore extra_headers: Optional[dict] = None, safety_identifier: Optional[str] = None, service_tier: Optional[str] = None, + store: Optional[bool] = None, + prompt_cache_key: Optional[str] = None, # soon to be deprecated params by OpenAI functions: Optional[List] = None, function_call: Optional[str] = None, @@ -5249,6 +5255,8 @@ def completion( # type: ignore ), "safety_identifier": safety_identifier, "service_tier": service_tier, + "store": store, + "prompt_cache_key": prompt_cache_key, "allowed_openai_params": kwargs.get("allowed_openai_params"), "base_model": base_model, } diff --git a/litellm/utils.py b/litellm/utils.py index 18b89ee0d13..b16ecdf88be 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3791,6 +3791,8 @@ def get_optional_params( thinking: Optional[AnthropicThinkingParam] = None, web_search_options: Optional[OpenAIWebSearchOptions] = None, safety_identifier: Optional[str] = None, + store: Optional[bool] = None, + prompt_cache_key: Optional[str] = None, base_model: Optional[str] = None, **kwargs, ): diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 28cf4fa0744..0624b590df7 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2081,3 +2081,87 @@ def test_stream_chunk_builder_text_completion_combines_text_and_usage(): assert response.usage.prompt_tokens > 0 assert response.usage.completion_tokens > 0 assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + + +def test_completion_forwards_store_and_prompt_cache_key_to_openai(): + """ + Regression test for https://github.com/BerriAI/litellm/issues/33184 + + store and prompt_cache_key are documented OpenAI chat completion params that + were accepted as supported but silently dropped before the provider request + was built, because they were not named parameters of completion() and + get_optional_params() the way safety_identifier is. + """ + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key") + + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + try: + litellm.completion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + store=False, + prompt_cache_key="test-cache-key", + client=client, + ) + except Exception as e: + print(e) + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + assert request_body["store"] is False + assert request_body["prompt_cache_key"] == "test-cache-key" + + +@pytest.mark.asyncio +async def test_acompletion_forwards_store_and_prompt_cache_key_to_openai(): + """ + Async variant of the store/prompt_cache_key forwarding regression test for + https://github.com/BerriAI/litellm/issues/33184 + """ + from openai import AsyncOpenAI + + client = AsyncOpenAI(api_key="fake-api-key") + + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + try: + await litellm.acompletion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + store=False, + prompt_cache_key="test-cache-key", + client=client, + ) + except Exception as e: + print(e) + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + assert request_body["store"] is False + assert request_body["prompt_cache_key"] == "test-cache-key" + + +def test_completion_omits_store_and_prompt_cache_key_when_not_passed(): + """ + When store and prompt_cache_key are not passed, they must not appear in the + outbound request body (guards against always forwarding None defaults). + """ + from openai import OpenAI + + client = OpenAI(api_key="fake-api-key") + + with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + try: + litellm.completion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + client=client, + ) + except Exception as e: + print(e) + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + assert "store" not in request_body + assert "prompt_cache_key" not in request_body From 4eaa70440a247cee2767bd16e7dc830da559105a Mon Sep 17 00:00:00 2001 From: Sujith Date: Tue, 14 Jul 2026 15:49:46 +0530 Subject: [PATCH 008/504] fix(main): use PEP 604 unions for new store and prompt_cache_key params --- litellm/main.py | 8 ++++---- litellm/utils.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 81b02082abd..3b13e04bbe3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -435,8 +435,8 @@ async def acompletion( verbosity: Optional[Literal["low", "medium", "high"]] = None, safety_identifier: Optional[str] = None, service_tier: Optional[str] = None, - store: Optional[bool] = None, - prompt_cache_key: Optional[str] = None, + store: bool | None = None, + prompt_cache_key: str | None = None, # set api_base, api_version, api_key base_url: Optional[str] = None, api_version: Optional[str] = None, @@ -4832,8 +4832,8 @@ def completion( # type: ignore extra_headers: Optional[dict] = None, safety_identifier: Optional[str] = None, service_tier: Optional[str] = None, - store: Optional[bool] = None, - prompt_cache_key: Optional[str] = None, + store: bool | None = None, + prompt_cache_key: str | None = None, # soon to be deprecated params by OpenAI functions: Optional[List] = None, function_call: Optional[str] = None, diff --git a/litellm/utils.py b/litellm/utils.py index b16ecdf88be..bc5b2447761 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3791,8 +3791,8 @@ def get_optional_params( thinking: Optional[AnthropicThinkingParam] = None, web_search_options: Optional[OpenAIWebSearchOptions] = None, safety_identifier: Optional[str] = None, - store: Optional[bool] = None, - prompt_cache_key: Optional[str] = None, + store: bool | None = None, + prompt_cache_key: str | None = None, base_model: Optional[str] = None, **kwargs, ): From 371fa670d6f79dfd579945e2d357f5b978d9af21 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 17 Jul 2026 19:33:28 +0000 Subject: [PATCH 009/504] fix(proxy): forward Bedrock event-stream content-type on unbuffered passthrough Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 14 +++-- .../proxy/test_common_request_processing.py | 55 +++++++++++++++++++ 2 files changed, 65 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c7c9397d850..24831a64410 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1766,6 +1766,7 @@ class ProxyBaseLLMRequestProcessing: return StreamingResponse( content=generator, # type: ignore[arg-type] status_code=status.HTTP_200_OK, + media_type=self._passthrough_event_stream_media_type(), headers=custom_headers, ) else: @@ -2216,10 +2217,15 @@ class ProxyBaseLLMRequestProcessing: def _passthrough_event_stream_media_type(self) -> Optional[str]: """ - Content-type for a buffered passthrough event-stream response, resolved - from the provider handler so the proxy stays provider-agnostic. Mirrors - the upstream content-type the non-streaming path forwards, since the - buffered streaming generator carries no headers of its own. + Content-type for a passthrough event-stream response, resolved from the + provider handler so the proxy stays provider-agnostic. Mirrors the + upstream content-type the non-streaming path forwards, since the + streaming generator carries no headers of its own. Used for both the + buffered (guardrail-rewritten) and the unbuffered relay paths so + clients that enforce the event-stream content-type (e.g. Claude Code on + Bedrock invoke-with-response-stream) see the correct header instead of + Starlette's application/octet-stream default. Returns None for providers + with no event-stream media type, leaving the response default unchanged. """ from litellm.llms.pass_through.guardrail_translation.handler import ( LlmPassthroughRouteHandler, diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index ebfbb46053d..f1a3745e85b 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -4137,6 +4137,61 @@ class TestAllmPassthroughStreamingProviderGate: assert streamed == chunks mock_handler.assert_not_awaited() + @pytest.mark.asyncio + async def test_bedrock_invoke_stream_sets_event_stream_content_type(self, monkeypatch): + """ + Regression for LIT-4561. The unbuffered Bedrock event-stream relay + (invoke-with-response-stream, no post-call guardrail rewriting) must set + content-type: application/vnd.amazon.eventstream instead of leaving it to + Starlette's application/octet-stream default, which trips Claude Code's + content-type guard added in 2.1.208 + """ + processing_obj = self._build_processing_obj( + "bedrock", "model/us.anthropic.claude-sonnet-4-20250514-v1:0/invoke-with-response-stream" + ) + chunks = [b"raw-1", b"raw-2"] + + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, + ): + result = await self._run(processing_obj, monkeypatch, chunks) + + assert isinstance(result, StreamingResponse) + assert result.media_type == "application/vnd.amazon.eventstream" + assert result.headers["content-type"] == "application/vnd.amazon.eventstream" + streamed = [chunk async for chunk in result.body_iterator] + assert streamed == chunks + + @pytest.mark.asyncio + async def test_non_bedrock_stream_keeps_default_content_type(self, monkeypatch): + """ + A provider with no registered event-stream media type must not have one + forced onto its unbuffered stream, so the response default is unchanged + """ + processing_obj = self._build_processing_obj("anthropic") + chunks = [b"chunk-1", b"chunk-2"] + + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, + ): + result = await self._run(processing_obj, monkeypatch, chunks) + + assert isinstance(result, StreamingResponse) + assert result.media_type is None + assert result.headers.get("content-type") != "application/vnd.amazon.eventstream" + class TestResponseCostHeaderForTypedDictResponses: """ From bde00952b6aa68372e739e1f377cc9c32a9063f1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 23:19:21 +0000 Subject: [PATCH 010/504] fix(proxy): requeue Redis spend buffer transactions when DB commit fails The Redis transaction buffer leader drains the spend buffers with a destructive lpop before committing to the database. When the DB commit failed after exhausting retries, the popped transactions were only logged and then lost, permanently undercounting key/user/team/org/end-user/ team-member/tag/agent and daily spend after a database outage. Track each popped category and re-push the ones that were not committed back to their Redis buffers so a later scheduler tick retries them. Categories that already committed are not re-queued, so their spend is not double-counted. The daily tag spend path gets the same treatment. --- litellm/proxy/db/db_spend_update_writer.py | 45 +++++- .../redis_update_buffer.py | 55 +++++++ .../test_redis_update_buffer.py | 46 ++++++ .../proxy/db/test_db_spend_update_writer.py | 151 +++++++++++++++++- 4 files changed, 289 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 54a4c2dad91..cc266019ff7 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -797,6 +797,12 @@ class DBSpendUpdateWriter: ): verbose_proxy_logger.debug("acquired lock for spend updates") + # Track everything popped from Redis. Each category is removed once it + # has been committed to the DB, so whatever is left after a failure can + # be re-queued for the next tick instead of being lost. Committed + # categories are never re-queued, so their spend is not double-counted. + uncommitted: dict[str, Any] = {} # mutable-ok: drives which popped categories still need re-queuing + try: ( db_spend_update_transactions, @@ -807,6 +813,15 @@ class DBSpendUpdateWriter: daily_agent_spend_update_transactions, ) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + uncommitted = { # mutable-ok: drives which popped categories still need re-queuing + "db_spend_update_transactions": db_spend_update_transactions, + "daily_spend_update_transactions": daily_spend_update_transactions, + "daily_team_spend_update_transactions": daily_team_spend_update_transactions, + "daily_org_spend_update_transactions": daily_org_spend_update_transactions, + "daily_end_user_spend_update_transactions": daily_end_user_spend_update_transactions, + "daily_agent_spend_update_transactions": daily_agent_spend_update_transactions, + } + if db_spend_update_transactions is not None: verbose_proxy_logger.info( "Spend tracking - committing spend updates from Redis to DB: " @@ -826,6 +841,7 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, db_spend_update_transactions=db_spend_update_transactions, ) + uncommitted.pop("db_spend_update_transactions", None) if daily_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_user_spend( @@ -834,6 +850,8 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_spend_update_transactions, ) + uncommitted.pop("daily_spend_update_transactions", None) + if daily_team_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_team_spend( n_retry_times=n_retry_times, @@ -841,6 +859,7 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_team_spend_update_transactions, ) + uncommitted.pop("daily_team_spend_update_transactions", None) if daily_org_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_org_spend( @@ -849,6 +868,7 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_org_spend_update_transactions, ) + uncommitted.pop("daily_org_spend_update_transactions", None) if daily_end_user_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_end_user_spend( @@ -857,6 +877,8 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_end_user_spend_update_transactions, ) + uncommitted.pop("daily_end_user_spend_update_transactions", None) + if daily_agent_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_agent_spend( n_retry_times=n_retry_times, @@ -864,14 +886,20 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_agent_spend_update_transactions, ) + uncommitted.pop("daily_agent_spend_update_transactions", None) except Exception as e: spend_log_error( "Spend tracking - failed to commit spend updates from Redis to DB. " - "Data already popped from Redis may be lost. Error: %s", + "Re-queuing uncommitted transactions to Redis for retry on next tick. Error: %s", str(e), exc=e, ) finally: + to_restore = { # mutable-ok: transient kwargs payload consumed immediately below + name: txns for name, txns in uncommitted.items() if txns is not None + } + if to_restore: + await self.redis_update_buffer.restore_transactions_to_redis(**to_restore) await self.pod_lock_manager.release_lock( cronjob_id=DB_SPEND_UPDATE_JOB_NAME, ) @@ -1020,11 +1048,11 @@ class DBSpendUpdateWriter: cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, ): verbose_proxy_logger.debug("acquired lock for daily tag spend updates") + daily_tag_spend_update_transactions = ( + await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() + ) + committed = False try: - daily_tag_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() - ) - if daily_tag_spend_update_transactions: await DBSpendUpdateWriter.update_daily_tag_spend( n_retry_times=n_retry_times, @@ -1032,14 +1060,19 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_tag_spend_update_transactions, ) + committed = True except Exception as e: spend_log_error( "Spend tracking - failed to commit daily tag spend updates from Redis to DB. " - "Data already popped from Redis may be lost. Error: %s", + "Re-queuing to Redis for retry on next tick. Error: %s", str(e), exc=e, ) finally: + if not committed and daily_tag_spend_update_transactions: + await self.redis_update_buffer.restore_transactions_to_redis( + daily_tag_spend_update_transactions=daily_tag_spend_update_transactions, + ) await self.pod_lock_manager.release_lock( cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, ) diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index c924448669d..b30fadd86ab 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -8,6 +8,8 @@ import asyncio import json from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast +from redis.exceptions import RedisError + from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache from litellm.constants import ( @@ -374,6 +376,59 @@ class RedisUpdateBuffer: if daily_txns: await daily_queue.update_queue.put(daily_txns) + async def restore_transactions_to_redis( + self, + db_spend_update_transactions: DBSpendUpdateTransactions | None = None, + daily_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + daily_team_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + daily_org_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + daily_end_user_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + daily_agent_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + daily_tag_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + ) -> None: + """ + Re-push transactions that were popped from Redis but not committed to the DB. + + The leader drains the buffers with a destructive ``lpop`` before committing to + the database. When a commit fails after its retries are exhausted, the popped + transactions must be pushed back so a later scheduler tick can retry them; + otherwise the aggregated spend is lost permanently. The re-pushed payloads use + the same JSON encoding as the store path, so the next drain parses them normally. + """ + if self.redis_cache is None: + return + + _configs = ( + (db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY), + (daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY), + (daily_team_spend_update_transactions, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY), + (daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY), + (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY), + (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY), + (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY), + ) + + rpush_list: list[RedisPipelineRpushOperation] = [ # mutable-ok: async_rpush_pipeline requires a list arg + RedisPipelineRpushOperation(key=redis_key, values=[safe_dumps(transactions)]) + for transactions, redis_key in _configs + if transactions + ] + if len(rpush_list) == 0: + return + + try: + await self.redis_cache.async_rpush_pipeline(rpush_list=rpush_list) + verbose_proxy_logger.info( + "Spend tracking - restored %d uncommitted transaction set(s) to Redis for retry on next tick.", + len(rpush_list), + ) + except RedisError as e: + verbose_proxy_logger.error( + "Spend tracking - failed to restore uncommitted transactions to Redis. " + "These spend updates are lost. Error: %s", + str(e), + ) + @staticmethod def _number_of_transactions_to_store_in_redis( db_spend_update_transactions: DBSpendUpdateTransactions, diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 33372e7794a..79909561683 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -270,6 +270,52 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): assert result == (None, None, None, None, None, None) +@pytest.mark.asyncio +async def test_restore_transactions_to_redis_pushes_only_provided( + redis_update_buffer, mock_redis_cache +): + """ + restore_transactions_to_redis re-pushes only the transaction sets it was + given, to their matching buffer keys, so uncommitted spend can be retried. + """ + from litellm.constants import ( + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + REDIS_UPDATE_BUFFER_KEY, + ) + + mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1, 1]) + + db_spend = {"key_list_transactions": {"key1": 1.0}} + daily_user = {"user_key1": {"spend": 1.0}} + + await redis_update_buffer.restore_transactions_to_redis( + db_spend_update_transactions=db_spend, + daily_spend_update_transactions=daily_user, + ) + + mock_redis_cache.async_rpush_pipeline.assert_called_once() + rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] + pushed_keys = {op["key"] for op in rpush_list} + assert pushed_keys == { + REDIS_UPDATE_BUFFER_KEY, + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + } + # Payloads round-trip through the same JSON encoding used on the store path + payloads = {op["key"]: json.loads(op["values"][0]) for op in rpush_list} + assert payloads[REDIS_UPDATE_BUFFER_KEY] == db_spend + assert payloads[REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY] == daily_user + + +@pytest.mark.asyncio +async def test_restore_transactions_to_redis_noop_when_empty( + redis_update_buffer, mock_redis_cache +): + """Nothing to restore -> no Redis call.""" + mock_redis_cache.async_rpush_pipeline = AsyncMock() + await redis_update_buffer.restore_transactions_to_redis() + mock_redis_cache.async_rpush_pipeline.assert_not_called() + + def test_validate_redis_transaction_buffer_raises_without_redis(): """ When use_redis_transaction_buffer=true but no Redis cache is configured, diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 4c17c5d3482..10544e82453 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1532,9 +1532,9 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer = AsyncMock() mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock() - # Return all-None tuple (no data to commit) + # Return all-None tuple (no data to commit); the pipeline yields 6 slots mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = ( - AsyncMock(return_value=(None, None, None, None, None, None, None)) + AsyncMock(return_value=(None, None, None, None, None, None)) ) db_writer.redis_update_buffer = mock_redis_update_buffer @@ -1565,6 +1565,153 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer.assert_not_called() +@pytest.mark.asyncio +async def test_commit_with_redis_requeues_all_on_db_failure(): + """ + Regression for #33872: if the DB commit fails after the leader has already + popped transactions from Redis, the popped transactions must be re-queued to + Redis so a later tick can retry them, instead of being silently lost. + """ + db_writer = DBSpendUpdateWriter() + + db_spend = { + "user_list_transactions": {"user1": 1.5}, + "end_user_list_transactions": {}, + "key_list_transactions": {"key1": 1.5}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + daily_user = {"user_key1": {"spend": 1.5, "api_requests": 1}} + + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(db_spend, daily_user, None, None, None, None) + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + + mock_pod_lock_manager = AsyncMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + db_writer.pod_lock_manager = mock_pod_lock_manager + + # Every DB write raises -> simulates a full database outage + db_writer._commit_spend_updates_to_db = AsyncMock(side_effect=Exception("db down")) + + with patch.object( + DBSpendUpdateWriter, + "update_daily_user_spend", + new=AsyncMock(side_effect=Exception("db down")), + ): + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=MagicMock(), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + # Both failed categories must be re-queued to Redis, nothing lost + mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once() + _, kwargs = mock_redis_update_buffer.restore_transactions_to_redis.call_args + assert kwargs["db_spend_update_transactions"] == db_spend + assert kwargs["daily_spend_update_transactions"] == daily_user + # The lock must still be released + mock_pod_lock_manager.release_lock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_commit_with_redis_only_requeues_failed_category(): + """ + A partial DB failure must not re-queue categories that already committed, + otherwise their spend would be double-counted on the next tick. + """ + db_writer = DBSpendUpdateWriter() + + db_spend = { + "user_list_transactions": {"user1": 1.5}, + "end_user_list_transactions": {}, + "key_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + daily_user = {"user_key1": {"spend": 1.5, "api_requests": 1}} + + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(db_spend, daily_user, None, None, None, None) + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + + mock_pod_lock_manager = AsyncMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + db_writer.pod_lock_manager = mock_pod_lock_manager + + # db_spend commits fine; only the daily user commit fails + db_writer._commit_spend_updates_to_db = AsyncMock() + + with patch.object( + DBSpendUpdateWriter, + "update_daily_user_spend", + new=AsyncMock(side_effect=Exception("db down")), + ): + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=MagicMock(), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once() + _, kwargs = mock_redis_update_buffer.restore_transactions_to_redis.call_args + # Only the failed daily category is requeued; the committed db_spend is not + assert kwargs == {"daily_spend_update_transactions": daily_user} + + +@pytest.mark.asyncio +async def test_commit_with_redis_no_requeue_on_success(): + """When all commits succeed, nothing should be re-queued to Redis.""" + db_writer = DBSpendUpdateWriter() + + db_spend = { + "user_list_transactions": {"user1": 1.5}, + "end_user_list_transactions": {}, + "key_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(db_spend, None, None, None, None, None) + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + + mock_pod_lock_manager = AsyncMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + db_writer.pod_lock_manager = mock_pod_lock_manager + + db_writer._commit_spend_updates_to_db = AsyncMock() + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=MagicMock(), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + mock_redis_update_buffer.restore_transactions_to_redis.assert_not_awaited() + + @pytest.mark.parametrize( "bucket_name,input_dict,table_attr,method_name,where_key,expected_order", [ From 118b47a8a39a79922f09c851364e43e3c73d0ce1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 23:45:23 +0000 Subject: [PATCH 011/504] fix(proxy): keep tag drain inside try and cover requeue paths with tests Move the destructive daily-tag Redis drain back inside the try so a Redis read failure still releases the pod lock via the finally block, and use a covariant Mapping for the restore signature. Add regression tests for the daily-tag requeue-on-failure/no-requeue-on-success paths and the RedisError swallow branch in restore_transactions_to_redis. --- litellm/proxy/db/db_spend_update_writer.py | 13 ++-- .../redis_update_buffer.py | 13 ++-- .../test_redis_update_buffer.py | 18 +++++ .../proxy/db/test_db_spend_update_writer.py | 72 +++++++++++++++++++ 4 files changed, 102 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index cc266019ff7..f13bf2e2105 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -797,11 +797,7 @@ class DBSpendUpdateWriter: ): verbose_proxy_logger.debug("acquired lock for spend updates") - # Track everything popped from Redis. Each category is removed once it - # has been committed to the DB, so whatever is left after a failure can - # be re-queued for the next tick instead of being lost. Committed - # categories are never re-queued, so their spend is not double-counted. - uncommitted: dict[str, Any] = {} # mutable-ok: drives which popped categories still need re-queuing + uncommitted: dict[str, Any] = {} # mutable-ok: tracks popped categories still needing commit try: ( @@ -1048,11 +1044,12 @@ class DBSpendUpdateWriter: cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, ): verbose_proxy_logger.debug("acquired lock for daily tag spend updates") - daily_tag_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() - ) + daily_tag_spend_update_transactions: dict[str, DailyTagSpendTransaction] | None = None committed = False try: + daily_tag_spend_update_transactions = ( + await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() + ) if daily_tag_spend_update_transactions: await DBSpendUpdateWriter.update_daily_tag_spend( n_retry_times=n_retry_times, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index b30fadd86ab..660fd514d99 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -6,6 +6,7 @@ This is to prevent deadlocks and improve reliability import asyncio import json +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from redis.exceptions import RedisError @@ -379,12 +380,12 @@ class RedisUpdateBuffer: async def restore_transactions_to_redis( self, db_spend_update_transactions: DBSpendUpdateTransactions | None = None, - daily_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, - daily_team_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, - daily_org_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, - daily_end_user_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, - daily_agent_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, - daily_tag_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None = None, + daily_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + daily_team_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + daily_org_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + daily_end_user_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + daily_agent_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + daily_tag_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, ) -> None: """ Re-push transactions that were popped from Redis but not committed to the DB. diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 79909561683..3325893c5f6 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -316,6 +316,24 @@ async def test_restore_transactions_to_redis_noop_when_empty( mock_redis_cache.async_rpush_pipeline.assert_not_called() +@pytest.mark.asyncio +async def test_restore_transactions_to_redis_swallows_redis_error( + redis_update_buffer, mock_redis_cache +): + """A Redis failure during restore must not propagate to the caller's finally block.""" + from redis.exceptions import RedisError + + mock_redis_cache.async_rpush_pipeline = AsyncMock( + side_effect=RedisError("redis down") + ) + + await redis_update_buffer.restore_transactions_to_redis( + db_spend_update_transactions={"key_list_transactions": {"key1": 1.0}}, + ) + + mock_redis_cache.async_rpush_pipeline.assert_called_once() + + def test_validate_redis_transaction_buffer_raises_without_redis(): """ When use_redis_transaction_buffer=true but no Redis cache is configured, diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 10544e82453..06d06b50234 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1712,6 +1712,78 @@ async def test_commit_with_redis_no_requeue_on_success(): mock_redis_update_buffer.restore_transactions_to_redis.assert_not_awaited() +@pytest.mark.asyncio +async def test_commit_daily_tag_spend_requeues_on_db_failure(): + """A failed daily tag commit must re-queue the popped tag transactions and release the lock.""" + db_writer = DBSpendUpdateWriter() + + daily_tag = {"tag_key1": {"spend": 1.5, "api_requests": 1}} + + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis = AsyncMock() + mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer = AsyncMock( + return_value=daily_tag + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + + mock_pod_lock_manager = AsyncMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + db_writer.pod_lock_manager = mock_pod_lock_manager + + with patch.object( + DBSpendUpdateWriter, + "update_daily_tag_spend", + new=AsyncMock(side_effect=Exception("db down")), + ): + await db_writer._commit_daily_tag_spend_to_db_with_redis( + prisma_client=MagicMock(), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once_with( + daily_tag_spend_update_transactions=daily_tag, + ) + mock_pod_lock_manager.release_lock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_commit_daily_tag_spend_no_requeue_on_success(): + """A successful daily tag commit must not re-queue anything.""" + db_writer = DBSpendUpdateWriter() + + daily_tag = {"tag_key1": {"spend": 1.5, "api_requests": 1}} + + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis = AsyncMock() + mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer = AsyncMock( + return_value=daily_tag + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + + mock_pod_lock_manager = AsyncMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + db_writer.pod_lock_manager = mock_pod_lock_manager + + with patch.object( + DBSpendUpdateWriter, + "update_daily_tag_spend", + new=AsyncMock(), + ): + await db_writer._commit_daily_tag_spend_to_db_with_redis( + prisma_client=MagicMock(), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + mock_redis_update_buffer.restore_transactions_to_redis.assert_not_awaited() + mock_pod_lock_manager.release_lock.assert_awaited_once() + + @pytest.mark.parametrize( "bucket_name,input_dict,table_attr,method_name,where_key,expected_order", [ From fc36825dfd68ab3e5b142a401810de84a452fe62 Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Tue, 21 Jul 2026 06:16:02 +0000 Subject: [PATCH 012/504] fix(batches): support AWS Bedrock batch cancellation via StopModelInvocationJob (#33986) --- litellm/batches/main.py | 9 +- litellm/llms/bedrock/batches/handler.py | 162 +++++++++--------------- tests/test_bedrock_cancel_batch.py | 55 ++++++++ 3 files changed, 124 insertions(+), 102 deletions(-) create mode 100644 tests/test_bedrock_cancel_batch.py diff --git a/litellm/batches/main.py b/litellm/batches/main.py index f124882b5a4..0c2cf15d385 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -1087,9 +1087,16 @@ def cancel_batch( timeout=timeout, max_retries=optional_params.max_retries, ) + elif custom_llm_provider == "bedrock": + from litellm.llms.bedrock.batches.handler import BedrockBatchesHandler + + response = BedrockBatchesHandler.cancel_batch( + batch_id=batch_id, + **kwargs, + ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.".format( + message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', 'vertex_ai', and 'bedrock' are supported.".format( custom_llm_provider ), model="n/a", diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index c071f331337..55236eae525 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -6,9 +6,7 @@ from openai.types.batch import Metadata as OpenAIBatchMetadata from litellm.types.utils import LiteLLMBatch -# AWS Bedrock model-invocation-job statuses → OpenAI Batch statuses. -# Mirrors the mapping used by `BedrockBatchesConfig.transform_create_batch_response` -# so create / retrieve return consistent statuses. +# AWS Bedrock model-invocation-job statuses -> OpenAI Batch statuses. _BEDROCK_MIJ_STATUS_TO_OPENAI = { "Submitted": "validating", "Validating": "validating", @@ -44,17 +42,6 @@ def _extract_job_id_from_arn(arn: str) -> Optional[str]: def _predict_output_file_uri( output_prefix: str, input_uri: str, job_id: Optional[str] ) -> Optional[str]: - """ - Compute the deterministic per-job result file URI Bedrock writes to. - - Bedrock lays results out as:: - - //.out - - We compute it client-side so OpenAI-style ``client.files.content(output_file_id)`` - works without an extra S3 ``ListObjectsV2`` round-trip. Returns ``None`` if we - don't have enough info; callers should fall back to the bare prefix. - """ if not output_prefix or not input_uri or not job_id: return None if not output_prefix.endswith("/"): @@ -76,40 +63,76 @@ def _to_epoch(value: Any) -> Optional[int]: class BedrockBatchesHandler: - """ - Handler for Bedrock Batches. + """Handler for Bedrock Batches.""" - Specific providers/models needed some special handling. + @staticmethod + def cancel_batch( + batch_id: str, + aws_region_name: Optional[str] = None, + logging_obj=None, + **kwargs, + ) -> "LiteLLMBatch": + """ + Cancel an AWS Bedrock batch model invocation job using StopModelInvocationJob. + """ + try: + import boto3 + from botocore.exceptions import ClientError + except ImportError as exc: + raise ImportError( + "Missing boto3/botocore to call bedrock. Run 'pip install boto3'." + ) from exc - E.g. Twelve Labs Embedding Async Invoke - """ + region = ( + aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" + ) + + from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig + + creds = BedrockBatchesConfig().get_credentials( + aws_access_key_id=kwargs.get("aws_access_key_id"), + aws_secret_access_key=kwargs.get("aws_secret_access_key"), + aws_session_token=kwargs.get("aws_session_token"), + aws_region_name=region, + aws_session_name=kwargs.get("aws_session_name"), + aws_profile_name=kwargs.get("aws_profile_name"), + aws_role_name=kwargs.get("aws_role_name"), + aws_web_identity_token=kwargs.get("aws_web_identity_token"), + aws_sts_endpoint=kwargs.get("aws_sts_endpoint"), + aws_external_id=kwargs.get("aws_external_id"), + ) + + client = boto3.client( + "bedrock", + region_name=region, + aws_access_key_id=creds.access_key, + aws_secret_access_key=creds.secret_key, + aws_session_token=creds.token, + ) + + try: + client.stop_model_invocation_job(jobIdentifier=batch_id) + except ClientError as e: + # Idempotency: if job is already Stopping/Stopped/Completed, swallow ValidationException + if e.response.get("Error", {}).get("Code") != "ValidationException": + raise e + + return BedrockBatchesHandler._handle_model_invocation_job_status( + batch_id=batch_id, + aws_region_name=region, + logging_obj=logging_obj, + **kwargs, + ) @staticmethod def _handle_async_invoke_status( batch_id: str, aws_region_name: str, logging_obj=None, **kwargs ) -> "LiteLLMBatch": - """ - Handle async invoke status check for AWS Bedrock. - - This is for Twelve Labs Embedding Async Invoke. - - Args: - batch_id: The async invoke ARN - aws_region_name: AWS region name - **kwargs: Additional parameters - - Returns: - dict: Status information including status, output_file_id (S3 URL), etc. - """ import asyncio - from litellm.llms.bedrock.embed.embedding import BedrockEmbedding async def _async_get_status(): - # Create embedding handler instance embedding_handler = BedrockEmbedding() - - # Get the status of the async invoke job status_response = await embedding_handler._get_async_invoke_status( invocation_arn=batch_id, aws_region_name=aws_region_name, @@ -117,18 +140,13 @@ class BedrockBatchesHandler: **kwargs, ) - # Transform response to a LiteLLMBatch object - from litellm.types.utils import LiteLLMBatch - openai_batch_metadata: OpenAIBatchMetadata = { - "output_file_id": status_response["outputDataConfig"][ - "s3OutputDataConfig" - ]["s3Uri"], + "output_file_id": status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"], "failure_message": status_response.get("failureMessage") or "", "model_arn": status_response["modelArn"], } - result = LiteLLMBatch( + return LiteLLMBatch( id=status_response["invocationArn"], object="batch", status=status_response["status"], @@ -151,10 +169,6 @@ class BedrockBatchesHandler: input_file_id="", ) - return result - - # Since this function is called from within an async context via run_in_executor, - # we need to create a new event loop in a thread to avoid conflicts import concurrent.futures def run_in_thread(): @@ -176,37 +190,6 @@ class BedrockBatchesHandler: logging_obj=None, **kwargs, ) -> "LiteLLMBatch": - """ - Handle ``GetModelInvocationJob`` status check for AWS Bedrock bulk batch - inference jobs (the ARN type returned by ``CreateModelInvocationJob``). - - ``CreateModelInvocationJob`` lives on the Bedrock **control plane** - (``bedrock..amazonaws.com``), distinct from the data-plane - ``bedrock-runtime`` endpoint that serves Twelve Labs async-invoke ARNs. - The two ARN families therefore can't share a handler — see - ``litellm/batches/main.py`` for the dispatch. - - Args: - batch_id: A ``arn:aws:bedrock:::model-invocation-job/`` - ARN (or just the trailing job id; both are accepted by - ``GetModelInvocationJob``). - aws_region_name: Region for the boto3 ``bedrock`` client. If omitted, - we fall back to parsing the region out of ``batch_id`` itself. - logging_obj: Optional litellm logging object. - **kwargs: Optional AWS credential overrides - (``aws_access_key_id``, ``aws_secret_access_key``, - ``aws_session_token``, ``aws_profile_name``, - ``aws_role_name``, ``aws_session_name``, - ``aws_web_identity_token``, ``aws_sts_endpoint``, - ``aws_external_id``). Unknown keys are ignored. - - Returns: - ``LiteLLMBatch`` shaped like an OpenAI Batch resource. Note that - ``request_counts`` is always ``(0, 0, 0)`` because - ``GetModelInvocationJob`` does not surface per-record counts; - callers that need accurate counts should parse - ``manifest.json.out`` from the output S3 prefix. - """ try: import boto3 except ImportError as exc: @@ -214,15 +197,10 @@ class BedrockBatchesHandler: "Missing boto3 to call bedrock. Run 'pip install boto3'." ) from exc - # Resolve region: explicit > parsed-from-ARN > us-east-1 (boto3 default). region = ( aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" ) - # Resolve credentials through the same path the rest of the bedrock - # provider uses, so model_list / env / role-assumption configs are - # honored. We instantiate BedrockBatchesConfig (which extends - # BaseAWSLLM) lazily to avoid a circular import at module load. from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig creds = BedrockBatchesConfig().get_credentials( @@ -247,10 +225,6 @@ class BedrockBatchesHandler: ) if logging_obj is not None: - # Use the bare job id in the logged URL so we don't double up the - # `model-invocation-job/` segment when `batch_id` is a full ARN. - # `GetModelInvocationJob` accepts either form, but only the bare id - # produces a sensible-looking URL in logs. url_path_id = _extract_job_id_from_arn(batch_id) or batch_id logging_obj.pre_call( input=batch_id, @@ -291,26 +265,12 @@ class BedrockBatchesHandler: .get("s3Uri", "") ) - # Bedrock returns the output *prefix* the user supplied at job creation. - # Actual results land at //.out — we - # surface that single-file URI as `output_file_id` so the OpenAI-style - # download flow works without an extra S3 listing call. We deliberately - # do NOT fall back to the bare prefix when prediction fails: a prefix - # is not a downloadable object, so handing it back as `output_file_id` - # would reproduce the very NoSuchKey bug this handler exists to fix. - # The bare prefix is preserved in metadata for callers that want the - # `manifest.json.out` or want to do their own listing. job_arn = response.get("jobArn", batch_id) job_id = _extract_job_id_from_arn(job_arn) output_file_uri = _predict_output_file_uri(output_prefix, input_uri, job_id) completed_at = _to_epoch(response.get("endTime")) - # Note: metadata uses "" (not None) for unknown URIs to satisfy the - # OpenAI Batch metadata schema, which is `dict[str, str]`. The - # `output_file_id` field on the LiteLLMBatch itself does carry None - # correctly (see below), so callers should branch on that, not on - # `metadata["output_file_uri"]`. openai_batch_metadata: OpenAIBatchMetadata = { "model_arn": response.get("modelId", ""), "job_arn": job_arn, diff --git a/tests/test_bedrock_cancel_batch.py b/tests/test_bedrock_cancel_batch.py new file mode 100644 index 00000000000..e52cc2c447e --- /dev/null +++ b/tests/test_bedrock_cancel_batch.py @@ -0,0 +1,55 @@ +from unittest.mock import MagicMock, patch +import pytest + +import litellm +from litellm.llms.bedrock.batches.handler import BedrockBatchesHandler + + +@patch("boto3.client") +def test_bedrock_cancel_batch_handler(mock_boto_client): + mock_client_instance = MagicMock() + mock_boto_client.return_value = mock_client_instance + + mock_client_instance.stop_model_invocation_job.return_value = {} + mock_client_instance.get_model_invocation_job.return_value = { + "jobArn": "arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id", + "status": "Stopping", + "submitTime": 1700000000, + "lastModifiedTime": 1700000100, + } + + res = BedrockBatchesHandler.cancel_batch( + batch_id="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id", + aws_region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + + mock_client_instance.stop_model_invocation_job.assert_called_once_with( + jobIdentifier="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id" + ) + assert res.status == "cancelling" + + +@patch("boto3.client") +def test_litellm_cancel_batch_bedrock_dispatcher(mock_boto_client): + mock_client_instance = MagicMock() + mock_boto_client.return_value = mock_client_instance + + mock_client_instance.stop_model_invocation_job.return_value = {} + mock_client_instance.get_model_invocation_job.return_value = { + "jobArn": "arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id", + "status": "Stopped", + "submitTime": 1700000000, + "lastModifiedTime": 1700000100, + } + + res = litellm.cancel_batch( + batch_id="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id", + custom_llm_provider="bedrock", + aws_region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + + assert res.status == "cancelled" From 163ab6e34b5e1b59675d850f86e0f96cdbce4d64 Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Tue, 21 Jul 2026 06:24:17 +0000 Subject: [PATCH 013/504] fix(batches): refine bedrock cancel_batch type hints and validation error handling --- litellm/batches/main.py | 2 +- litellm/llms/bedrock/batches/handler.py | 9 +++++++-- litellm/ocr/rust_bridge.py | 2 +- uv.lock | 10 ++++------ 4 files changed, 13 insertions(+), 10 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 0c2cf15d385..9d4cb29e926 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -944,7 +944,7 @@ async def acancel_batch( def cancel_batch( batch_id: str, model: Optional[str] = None, - custom_llm_provider: Union[Literal["openai", "azure", "vertex_ai"], str] = "openai", + custom_llm_provider: Union[Literal["openai", "azure", "vertex_ai", "bedrock"], str] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 55236eae525..c710bcfdafa 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -113,8 +113,13 @@ class BedrockBatchesHandler: try: client.stop_model_invocation_job(jobIdentifier=batch_id) except ClientError as e: - # Idempotency: if job is already Stopping/Stopped/Completed, swallow ValidationException - if e.response.get("Error", {}).get("Code") != "ValidationException": + error_code = e.response.get("Error", {}).get("Code") + error_msg = e.response.get("Error", {}).get("Message", "").lower() + if error_code == "ValidationException" and any( + term in error_msg for term in ["stop", "terminal", "completed", "already"] + ): + pass + else: raise e return BedrockBatchesHandler._handle_model_invocation_job_status( diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index 253b35cb689..1e3312c1473 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/ocr/rust_bridge.py @@ -97,7 +97,7 @@ def load_rust_ocr() -> RustOcr | None: import litellm_python_bridge except ImportError: return None - return cast(RustOcr, getattr(litellm_python_bridge, "ocr", None)) + return cast(RustOcr, litellm_python_bridge.ocr) def load_rust_aocr() -> RustAocr | None: diff --git a/uv.lock b/uv.lock index 8c9e20eeb21..cac1696bf34 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer = "2026-06-20T23:16:25.061268Z" exclude-newer-span = "P3D" [manifest] @@ -3160,15 +3160,15 @@ wheels = [ [[package]] name = "langgraph-checkpoint" -version = "4.1.1" +version = "4.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "ormsgpack" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/83/47/886af6f886f0bff2273164a45f008694e48a96ff3cd25ff0228f2aa9480e/langgraph_checkpoint-4.1.1.tar.gz", hash = "sha256:6c2bdb530c91f91d7d9c1bd100925d0fc4f498d418c17f3587d1526279482a25", size = 184020, upload-time = "2026-05-22T16:57:38.503Z" } +sdist = { url = "https://files.pythonhosted.org/packages/02/b4/6005c5dd88ad484fe6235d4c43a0d2cee7e91b08ad85a180985c2662df87/langgraph_checkpoint-4.1.0.tar.gz", hash = "sha256:e5bb304e30fc1363ac8fcb5f7dee5ca2185d77fe475b0d01de2c5f91324c2c21", size = 181942, upload-time = "2026-05-12T03:33:49.888Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" }, + { url = "https://files.pythonhosted.org/packages/93/74/d3be2b41955e20ccd624dba5f6fe9d38dcee385ba470a6e13ed86732fc86/langgraph_checkpoint-4.1.0-py3-none-any.whl", hash = "sha256:8bc2a0466a20c38b865ce6671b42093fd5c041133f32351cae4222e0eeaf7fb5", size = 56047, upload-time = "2026-05-12T03:33:48.548Z" }, ] [[package]] @@ -3281,7 +3281,6 @@ dependencies = [ { name = "importlib-metadata" }, { name = "jinja2" }, { name = "jsonschema" }, - { name = "langgraph-checkpoint" }, { name = "openai" }, { name = "pydantic" }, { name = "python-dotenv" }, @@ -3506,7 +3505,6 @@ requires-dist = [ { name = "jinja2", specifier = ">=3.1.6,<4.0" }, { name = "jsonschema", specifier = ">=4.0.0,<5.0" }, { name = "langfuse", marker = "extra == 'proxy-runtime'", specifier = ">=2.59.7,<3.0" }, - { name = "langgraph-checkpoint", specifier = "==4.1.1" }, { name = "litellm-enterprise", marker = "extra == 'proxy'", editable = "enterprise" }, { name = "litellm-proxy-extras", marker = "extra == 'proxy'", editable = "litellm-proxy-extras" }, { name = "llm-sandbox", marker = "extra == 'proxy-runtime'", specifier = ">=0.3.39,<1.0" }, From ab997e04eb4f0f50bc2c6ae738231455cdd97329 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 25 Jul 2026 00:09:21 +0000 Subject: [PATCH 014/504] fix(caching): cache anthropic /v1/messages responses, including streaming anthropic_messages was missing from the cache's supported call types, so every /v1/messages request went to the provider. Adding it alone is not enough: the cache key is built from the OpenAI-ish param set, which has no system, top_k or stop_sequences, so two requests differing only by system prompt shared an entry and the second got the first one's answer. The Anthropic Messages request shape now feeds the key set as well. Streaming responses return to the caller before async_set_cache runs, so they are teed on the way out and the SSE events are stored verbatim once the stream reaches message_stop without a provider error. A hit replays those bytes and logs the request as a cache hit with zero cost. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/caching.py | 45 +---- litellm/caching/caching_handler.py | 41 +++- .../litellm_core_utils/model_param_helper.py | 19 +- .../messages/response_cache.py | 163 ++++++++++++++++ .../anthropic_passthrough_logging_handler.py | 16 +- .../streaming_handler.py | 2 +- litellm/types/caching.py | 19 ++ litellm/utils.py | 5 +- tests/test_litellm/caching/test_caching.py | 21 ++ .../messages/test_response_cache.py | 179 ++++++++++++++++++ 10 files changed, 457 insertions(+), 53 deletions(-) create mode 100644 litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 34badaa3e8a..88a5e08604e 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -67,20 +67,7 @@ class Cache: default_in_memory_ttl: Optional[float] = None, default_in_redis_ttl: Optional[float] = None, similarity_threshold: Optional[float] = None, - supported_call_types: Optional[List[CachingSupportedCallTypes]] = [ - "completion", - "acompletion", - "embedding", - "aembedding", - "atranscription", - "transcription", - "atext_completion", - "text_completion", - "arerank", - "rerank", - "responses", - "aresponses", - ], + supported_call_types: list[CachingSupportedCallTypes] | None = list(DEFAULT_CACHING_SUPPORTED_CALL_TYPES), # s3 Bucket, boto3 configuration azure_account_url: Optional[str] = None, azure_blob_container: Optional[str] = None, @@ -930,20 +917,7 @@ def enable_cache( host: Optional[str] = None, port: Optional[str] = None, password: Optional[str] = None, - supported_call_types: Optional[List[CachingSupportedCallTypes]] = [ - "completion", - "acompletion", - "embedding", - "aembedding", - "atranscription", - "transcription", - "atext_completion", - "text_completion", - "arerank", - "rerank", - "responses", - "aresponses", - ], + supported_call_types: list[CachingSupportedCallTypes] | None = list(DEFAULT_CACHING_SUPPORTED_CALL_TYPES), **kwargs, ): """ @@ -990,20 +964,7 @@ def update_cache( host: Optional[str] = None, port: Optional[str] = None, password: Optional[str] = None, - supported_call_types: Optional[List[CachingSupportedCallTypes]] = [ - "completion", - "acompletion", - "embedding", - "aembedding", - "atranscription", - "transcription", - "atext_completion", - "text_completion", - "arerank", - "rerank", - "responses", - "aresponses", - ], + supported_call_types: list[CachingSupportedCallTypes] | None = list(DEFAULT_CACHING_SUPPORTED_CALL_TYPES), **kwargs, ): """ diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index b17e055c7ea..8b2d033f24a 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -116,7 +116,8 @@ def _should_defer_streaming_cache_hit_callbacks(*, kwargs: Dict[str, Any]) -> bo When stream=True, do not run success callbacks at cache-hit time. Cached chat/text completion replay uses CustomStreamWrapper; cached Responses - replay uses CachedResponsesAPIStreamingIterator. Both invoke logging success + replay uses CachedResponsesAPIStreamingIterator; cached Anthropic Messages + replay uses CachedAnthropicMessagesStreamIterator. All invoke logging success handlers when the stream finishes; firing them here too would double-count spend and callback records. """ @@ -848,6 +849,18 @@ class LLMCachingHandler: response_type="audio_transcription", hidden_params=hidden_params, ) + elif ( + call_type == CallTypes.anthropic_messages.value or call_type == CallTypes.aanthropic_messages.value + ) and isinstance(cached_result, dict): + from litellm.llms.anthropic.experimental_pass_through.messages.response_cache import ( + convert_cached_anthropic_messages_result, + ) + + cached_result = convert_cached_anthropic_messages_result( + cached_result=cached_result, + logging_obj=logging_obj, + kwargs=kwargs, + ) elif (call_type == "aresponses" or call_type == "responses") and isinstance(cached_result, dict): use_chat_completion_cache = _is_chat_completion_cached_dict(cached_result) if use_chat_completion_cache: @@ -1044,6 +1057,32 @@ class LLMCachingHandler: and (kwargs.get("cache", {}).get("no-store", False) is not True) ) + def wrap_streaming_result_for_cache(self, result: Any, call_type: str) -> Any: + """ + Tee a streaming result so it still reaches the cache. + + Streaming responses are returned to the caller before ``async_set_cache`` + runs. Chat/text completion streams are teed inside ``CustomStreamWrapper`` + and Responses API streams inside their own iterator; Anthropic Messages + streams have no such hook, so they are wrapped here. + """ + if call_type not in ( + CallTypes.anthropic_messages.value, + CallTypes.aanthropic_messages.value, + ): + return result + if litellm.cache is None or not self._should_store_result_in_cache( + original_function=self.original_function, kwargs=self.request_kwargs + ): + return result + if not hasattr(result, "__anext__"): + return result + from litellm.llms.anthropic.experimental_pass_through.messages.response_cache import ( + AnthropicMessagesStreamCacheWriter, + ) + + return AnthropicMessagesStreamCacheWriter(stream=result, caching_handler=self) + def _is_call_type_supported_by_cache( self, original_function: Callable, diff --git a/litellm/litellm_core_utils/model_param_helper.py b/litellm/litellm_core_utils/model_param_helper.py index 39b3f0d5376..cf4eba933b8 100644 --- a/litellm/litellm_core_utils/model_param_helper.py +++ b/litellm/litellm_core_utils/model_param_helper.py @@ -18,6 +18,7 @@ from openai.types.responses.response_create_params import ( ) from litellm._logging import verbose_logger +from litellm.types.llms.anthropic import AnthropicMessagesRequest from litellm.types.rerank import RerankRequest @@ -40,7 +41,7 @@ class ModelParamHelper: @staticmethod def get_exclude_params_for_model_parameters() -> Set[str]: - return set(["messages", "prompt", "input"]) + return set(["messages", "prompt", "input", "system"]) @staticmethod def _get_relevant_args_to_use_for_logging() -> Set[str]: @@ -73,6 +74,7 @@ class ModelParamHelper: transcription_kwargs = ModelParamHelper._get_litellm_supported_transcription_kwargs() rerank_kwargs = ModelParamHelper._get_litellm_supported_rerank_kwargs() responses_api_kwargs = ModelParamHelper._get_litellm_supported_responses_api_kwargs() + anthropic_messages_kwargs = ModelParamHelper._get_litellm_supported_anthropic_messages_kwargs() exclude_kwargs = ModelParamHelper._get_exclude_kwargs() combined_kwargs = chat_completion_kwargs.union( @@ -81,6 +83,7 @@ class ModelParamHelper: transcription_kwargs, rerank_kwargs, responses_api_kwargs, + anthropic_messages_kwargs, ) combined_kwargs = combined_kwargs.difference(exclude_kwargs) return combined_kwargs @@ -167,12 +170,24 @@ class ModelParamHelper: streaming_params: Set[str] = set(getattr(ResponseCreateParamsStreaming, "__annotations__", {}).keys()) return non_streaming_params.union(streaming_params) + @staticmethod + def _get_litellm_supported_anthropic_messages_kwargs() -> set[str]: + """ + Get the litellm supported Anthropic /v1/messages kwargs + + This follows the Anthropic Messages API spec. `system`, `top_k` and + `stop_sequences` have no OpenAI equivalent, so without them the cache key + for a /v1/messages request ignores them and collides across requests that + differ only by system prompt. + """ + return set(getattr(AnthropicMessagesRequest, "__annotations__", {}).keys()) + @staticmethod def _get_exclude_kwargs() -> Set[str]: """ Get the kwargs to exclude from the cache key """ - return set(["metadata"]) + return set(["metadata", "litellm_metadata"]) ModelParamHelper._relevant_logging_args = frozenset(ModelParamHelper._get_relevant_args_to_use_for_logging()) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py new file mode 100644 index 00000000000..e94d8f6bbaf --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py @@ -0,0 +1,163 @@ +""" +Response caching for Anthropic Messages (`/v1/messages`) requests. + +Non-streaming responses are plain dicts and are stored by the generic caching +handler. Streaming responses are returned to the caller before +``LLMCachingHandler.async_set_cache`` runs, so they are teed here instead: the +SSE events are buffered while they are forwarded and persisted verbatim once the +stream completes, and a hit replays exactly what the provider sent. +""" + +from collections.abc import AsyncIterator +from typing import TYPE_CHECKING, Any, cast + +import litellm +from litellm._logging import verbose_logger +from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + BaseAnthropicMessagesStreamingIterator, + _is_message_stop_chunk, + _is_provider_error_chunk, + aclose_if_supported, +) +from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, +) + +if TYPE_CHECKING: + from litellm.caching.caching_handler import LLMCachingHandler + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +else: + LLMCachingHandler = Any + LiteLLMLoggingObj = Any + +CACHED_STREAM_EVENTS_KEY = "litellm_cached_anthropic_sse_events" + + +def _decode(chunk: bytes | str) -> str: + return chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk + + +class AnthropicMessagesStreamCacheWriter: + """ + Forwards a `/v1/messages` SSE stream unchanged while buffering it, then + writes the collected events to the response cache on normal completion. + + Only a stream that ran to a ``message_stop`` without a provider ``error`` + event is written, so partial or failed responses cannot be replayed. + """ + + def __init__( + self, + stream: AsyncIterator[bytes | str], + caching_handler: "LLMCachingHandler", + ) -> None: + self.stream = stream + self.caching_handler = caching_handler + self.collected_events: list[str] = [] + self.saw_message_stop = False + self.saw_provider_error = False + self.persisted = False + self._hidden_params: dict[str, Any] = getattr(stream, "_hidden_params", {}) or {} + + def __aiter__(self) -> "AnthropicMessagesStreamCacheWriter": + return self + + async def __anext__(self) -> bytes | str: + try: + chunk = await self.stream.__anext__() + except StopAsyncIteration: + await self._persist() + raise + chunk_bytes = chunk.encode("utf-8") if isinstance(chunk, str) else chunk + self.saw_message_stop = self.saw_message_stop or _is_message_stop_chunk(chunk_bytes) + self.saw_provider_error = self.saw_provider_error or _is_provider_error_chunk(chunk_bytes) + self.collected_events.append(_decode(chunk)) + return chunk + + async def aclose(self) -> None: + await aclose_if_supported(self.stream) + + async def _persist(self) -> None: + if self.persisted or litellm.cache is None: + return + if not self.saw_message_stop or self.saw_provider_error: + return + self.persisted = True + + request_kwargs = dict(self.caching_handler.request_kwargs) + if not self.caching_handler._should_store_result_in_cache( + original_function=self.caching_handler.original_function, + kwargs=request_kwargs, + ): + return + preset_cache_key = self.caching_handler.preset_cache_key + if preset_cache_key is not None: + request_kwargs["cache_key"] = preset_cache_key + + try: + await litellm.cache.async_add_cache( + {CACHED_STREAM_EVENTS_KEY: self.collected_events}, + dynamic_cache_object=self.caching_handler.dual_cache, + **request_kwargs, + ) + except Exception as e: + verbose_logger.exception("Anthropic Messages stream cache write failed: %s", e) + + +class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterator): + """ + Replays cached `/v1/messages` SSE events and logs the request as a cache hit + once the replay finishes, mirroring what the live stream logs at end of stream. + """ + + def __init__( + self, + events: list[str], + litellm_logging_obj: LiteLLMLoggingObj, + request_body: dict[str, Any], + ) -> None: + super().__init__(litellm_logging_obj=litellm_logging_obj, request_body=request_body) + self.chunks: list[bytes] = [event.encode("utf-8") for event in events] + self.current_index = 0 + self._hidden_params: dict[str, Any] = {"cache_hit": True} + litellm_logging_obj.model_call_details["cache_hit"] = True + + def __aiter__(self) -> "CachedAnthropicMessagesStreamIterator": + return self + + async def __anext__(self) -> bytes: + if self.current_index >= len(self.chunks): + await self._handle_streaming_logging(self.chunks) + raise StopAsyncIteration + chunk = self.chunks[self.current_index] + self.current_index += 1 + return chunk + + +def get_cached_stream_events(cached_result: dict[str, Any]) -> list[str] | None: + events = cached_result.get(CACHED_STREAM_EVENTS_KEY) + if isinstance(events, list): + return [_decode(event) for event in events] + return None + + +def convert_cached_anthropic_messages_result( + cached_result: dict[str, Any], + logging_obj: LiteLLMLoggingObj, + kwargs: dict[str, Any], +) -> AnthropicMessagesResponse | CachedAnthropicMessagesStreamIterator: + """ + Turn a cached `/v1/messages` entry back into what the caller expects: an + SSE replay iterator for a streamed entry, otherwise the response itself + (``AnthropicMessagesResponse`` is a TypedDict, i.e. a dict at runtime). + """ + events = get_cached_stream_events(cached_result) + if events is not None: + return CachedAnthropicMessagesStreamIterator( + events=events, + litellm_logging_obj=logging_obj, + request_body=kwargs, + ) + return cast( # cast-ok: AnthropicMessagesResponse is a TypedDict; validating would drop provider fields we must replay verbatim + AnthropicMessagesResponse, cached_result + ) 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 50e90699194..51813983876 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 @@ -255,12 +255,16 @@ class AnthropicPassthroughLoggingHandler: litellm_params=(logging_obj.litellm_params if hasattr(logging_obj, "litellm_params") else None) ) - response_cost = litellm.completion_cost( - completion_response=litellm_model_response, - model=model_for_cost, - custom_llm_provider=custom_llm_provider, - custom_pricing=custom_pricing, - router_model_id=router_model_id, + response_cost = ( + 0.0 + if logging_obj.model_call_details.get("cache_hit") is True + else litellm.completion_cost( + completion_response=litellm_model_response, + model=model_for_cost, + custom_llm_provider=custom_llm_provider, + custom_pricing=custom_pricing, + router_model_id=router_model_id, + ) ) kwargs["response_cost"] = response_cost diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 4dc1e0e70dd..24e5f1d16d5 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -161,7 +161,7 @@ class PassThroughStreamingHandler: result=standard_logging_response_object, start_time=start_time, end_time=end_time, - cache_hit=False, + cache_hit=litellm_logging_obj.model_call_details.get("cache_hit") is True, prefer_async_handlers=True, **kwargs, ) diff --git a/litellm/types/caching.py b/litellm/types/caching.py index eaa80c2f525..4255a8bd7fc 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -30,8 +30,27 @@ CachingSupportedCallTypes = Literal[ "rerank", "responses", "aresponses", + "anthropic_messages", + "aanthropic_messages", ] +DEFAULT_CACHING_SUPPORTED_CALL_TYPES: tuple[CachingSupportedCallTypes, ...] = ( + "completion", + "acompletion", + "embedding", + "aembedding", + "atranscription", + "transcription", + "atext_completion", + "text_completion", + "arerank", + "rerank", + "responses", + "aresponses", + "anthropic_messages", + "aanthropic_messages", +) + class RedisPipelineIncrementOperation(TypedDict): """ diff --git a/litellm/utils.py b/litellm/utils.py index a11c5500503..f5c8330c284 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1708,7 +1708,10 @@ def client(original_function): start_time=start_time, end_time=end_time, ) - return result + return _llm_caching_handler.wrap_streaming_result_for_cache( + result=result, + call_type=call_type, + ) elif call_type == CallTypes.arealtime.value: return result ### POST-CALL RULES ### diff --git a/tests/test_litellm/caching/test_caching.py b/tests/test_litellm/caching/test_caching.py index eaee54bac5a..b65e8773c85 100644 --- a/tests/test_litellm/caching/test_caching.py +++ b/tests/test_litellm/caching/test_caching.py @@ -1,6 +1,8 @@ import logging import re +import pytest + from litellm.caching.caching import Cache from litellm.types.caching import LiteLLMCacheType from litellm.types.utils import Embedding, EmbeddingResponse, Usage @@ -146,3 +148,22 @@ def test_exact_cache_key_still_includes_prompt(): model="gpt-4o-mini", messages=[{"role": "user", "content": "b"}] ) assert key_a != key_b + + +@pytest.mark.parametrize( + "anthropic_param", + [ + {"system": "answer ALPHA"}, + {"top_k": 5}, + {"stop_sequences": ["STOP"]}, + ], +) +def test_exact_cache_key_includes_anthropic_messages_params(anthropic_param): + """Anthropic /v1/messages params with no OpenAI equivalent must still key the + cache; without them two requests that differ only by system prompt collide.""" + cache = Cache(type=LiteLLMCacheType.LOCAL) + messages = [{"role": "user", "content": "which greek letter?"}] + baseline = cache.get_cache_key(model="claude-sonnet-4-5", messages=messages) + assert baseline != cache.get_cache_key( + model="claude-sonnet-4-5", messages=messages, **anthropic_param + ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py new file mode 100644 index 00000000000..344152cd828 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py @@ -0,0 +1,179 @@ +import asyncio +import os +import sys +from typing import Any, AsyncIterator, Dict, List + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +import litellm +from litellm.caching.caching import Cache, LiteLLMCacheType +from litellm.llms.anthropic.experimental_pass_through.messages import handler + +STREAM_EVENTS: List[bytes] = [ + b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_stream_1", "type": "message", ' + b'"role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": null, ' + b'"usage": {"input_tokens": 10, "output_tokens": 0}}}\n\n', + b'event: content_block_start\ndata: {"type": "content_block_start", "index": 0, ' + b'"content_block": {"type": "text", "text": ""}}\n\n', + b'event: content_block_delta\ndata: {"type": "content_block_delta", "index": 0, ' + b'"delta": {"type": "text_delta", "text": "ALPHA"}}\n\n', + b'event: content_block_stop\ndata: {"type": "content_block_stop", "index": 0}\n\n', + b'event: message_delta\ndata: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, ' + b'"usage": {"output_tokens": 3}}\n\n', + b'event: message_stop\ndata: {"type": "message_stop"}\n\n', +] + + +def _anthropic_response(message_id: str, text: str) -> Dict[str, Any]: + return { + "id": message_id, + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": text}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 3}, + } + + +class _CountingHandler: + """Stands in for the provider dispatch so cache hits are observable as skipped calls.""" + + def __init__(self, results: List[Any]) -> None: + self.results = results + self.calls: List[Dict[str, Any]] = [] + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + self.calls.append(kwargs) + return self.results[min(len(self.calls) - 1, len(self.results) - 1)] + + +async def _byte_stream(chunks: List[bytes]) -> AsyncIterator[bytes]: + for chunk in chunks: + yield chunk + + +async def _collect(stream: AsyncIterator[bytes]) -> List[bytes]: + return [chunk async for chunk in stream] + + +@pytest.fixture +def local_cache(): + previous_cache = litellm.cache + litellm.cache = Cache(type=LiteLLMCacheType.LOCAL) + yield litellm.cache + litellm.cache = previous_cache + + +@pytest.fixture +def request_kwargs() -> Dict[str, Any]: + return { + "model": "anthropic/claude-sonnet-4-5", + "custom_llm_provider": "anthropic", + "api_key": "fake-key", + "max_tokens": 64, + "messages": [{"role": "user", "content": "which greek letter?"}], + } + + +@pytest.mark.asyncio +async def test_non_streaming_request_is_served_from_cache(local_cache, request_kwargs, monkeypatch): + fake_handler = _CountingHandler([_anthropic_response("msg_1", "ALPHA"), _anthropic_response("msg_2", "BETA")]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + first = await litellm.anthropic_messages(**request_kwargs) + await asyncio.sleep(0) + second = await litellm.anthropic_messages(**request_kwargs) + + assert len(fake_handler.calls) == 1 + assert first == second + assert second["content"][0]["text"] == "ALPHA" + + +@pytest.mark.asyncio +async def test_cache_key_separates_different_system_prompts(local_cache, request_kwargs, monkeypatch): + """`system` has no OpenAI equivalent; if it is dropped from the cache key the + second request is answered with the first system prompt's response.""" + fake_handler = _CountingHandler([_anthropic_response("msg_1", "ALPHA"), _anthropic_response("msg_2", "BETA")]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + first = await litellm.anthropic_messages(**request_kwargs, system="Always answer ALPHA") + await asyncio.sleep(0) + second = await litellm.anthropic_messages(**request_kwargs, system="Always answer BETA") + + assert len(fake_handler.calls) == 2 + assert first["content"][0]["text"] == "ALPHA" + assert second["content"][0]["text"] == "BETA" + + +@pytest.mark.parametrize("anthropic_param", [{"top_k": 5}, {"stop_sequences": ["STOP"]}]) +@pytest.mark.asyncio +async def test_cache_key_separates_anthropic_native_params(local_cache, request_kwargs, monkeypatch, anthropic_param): + fake_handler = _CountingHandler([_anthropic_response("msg_1", "ALPHA"), _anthropic_response("msg_2", "BETA")]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + await litellm.anthropic_messages(**request_kwargs) + await asyncio.sleep(0) + await litellm.anthropic_messages(**request_kwargs, **anthropic_param) + + assert len(fake_handler.calls) == 2 + + +@pytest.mark.asyncio +async def test_streaming_request_is_replayed_from_cache(local_cache, request_kwargs, monkeypatch): + fake_handler = _CountingHandler([_byte_stream(STREAM_EVENTS), _byte_stream([b"event: never_used\n\n"])]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + first = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + second_stream = await litellm.anthropic_messages(**request_kwargs, stream=True) + second = await _collect(second_stream) + + assert len(fake_handler.calls) == 1 + assert first == STREAM_EVENTS + assert second == STREAM_EVENTS + assert second_stream._hidden_params["cache_hit"] is True + + +@pytest.mark.asyncio +async def test_streaming_cache_is_not_shared_with_non_streaming(local_cache, request_kwargs, monkeypatch): + fake_handler = _CountingHandler([_byte_stream(STREAM_EVENTS), _anthropic_response("msg_2", "ALPHA")]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + non_streaming = await litellm.anthropic_messages(**request_kwargs) + + assert len(fake_handler.calls) == 2 + assert non_streaming["content"][0]["text"] == "ALPHA" + + +@pytest.mark.asyncio +async def test_failed_stream_is_not_cached(local_cache, request_kwargs, monkeypatch): + error_events = STREAM_EVENTS[:3] + [ + b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "overloaded"}}\n\n' + ] + fake_handler = _CountingHandler([_byte_stream(error_events), _byte_stream(STREAM_EVENTS)]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + failed = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + replayed = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + + assert failed == error_events + assert len(fake_handler.calls) == 2 + assert replayed == STREAM_EVENTS + + +@pytest.mark.asyncio +async def test_abandoned_stream_is_not_cached(local_cache, request_kwargs, monkeypatch): + fake_handler = _CountingHandler([_byte_stream(STREAM_EVENTS), _byte_stream(STREAM_EVENTS)]) + monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) + + partial_stream = await litellm.anthropic_messages(**request_kwargs, stream=True) + await partial_stream.__anext__() + await partial_stream.aclose() + + replayed = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + + assert len(fake_handler.calls) == 2 + assert replayed == STREAM_EVENTS From d2a5de2e04d10042060c1e68c157857bdacfb165 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 25 Jul 2026 00:15:55 +0000 Subject: [PATCH 015/504] refactor(caching): tighten anthropic messages cache types and drop comments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/caching_handler.py | 11 ++------ .../litellm_core_utils/model_param_helper.py | 7 +---- .../messages/response_cache.py | 28 ------------------- 3 files changed, 3 insertions(+), 43 deletions(-) diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 8b2d033f24a..70a2e3fd1b2 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -18,6 +18,7 @@ import asyncio import datetime import inspect import time +from collections.abc import AsyncIterator from typing import ( TYPE_CHECKING, Any, @@ -1058,14 +1059,6 @@ class LLMCachingHandler: ) def wrap_streaming_result_for_cache(self, result: Any, call_type: str) -> Any: - """ - Tee a streaming result so it still reaches the cache. - - Streaming responses are returned to the caller before ``async_set_cache`` - runs. Chat/text completion streams are teed inside ``CustomStreamWrapper`` - and Responses API streams inside their own iterator; Anthropic Messages - streams have no such hook, so they are wrapped here. - """ if call_type not in ( CallTypes.anthropic_messages.value, CallTypes.aanthropic_messages.value, @@ -1075,7 +1068,7 @@ class LLMCachingHandler: original_function=self.original_function, kwargs=self.request_kwargs ): return result - if not hasattr(result, "__anext__"): + if not isinstance(result, AsyncIterator): return result from litellm.llms.anthropic.experimental_pass_through.messages.response_cache import ( AnthropicMessagesStreamCacheWriter, diff --git a/litellm/litellm_core_utils/model_param_helper.py b/litellm/litellm_core_utils/model_param_helper.py index cf4eba933b8..7e99e5fc5b2 100644 --- a/litellm/litellm_core_utils/model_param_helper.py +++ b/litellm/litellm_core_utils/model_param_helper.py @@ -174,13 +174,8 @@ class ModelParamHelper: def _get_litellm_supported_anthropic_messages_kwargs() -> set[str]: """ Get the litellm supported Anthropic /v1/messages kwargs - - This follows the Anthropic Messages API spec. `system`, `top_k` and - `stop_sequences` have no OpenAI equivalent, so without them the cache key - for a /v1/messages request ignores them and collides across requests that - differ only by system prompt. """ - return set(getattr(AnthropicMessagesRequest, "__annotations__", {}).keys()) + return set(AnthropicMessagesRequest.__annotations__.keys()) @staticmethod def _get_exclude_kwargs() -> Set[str]: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py index e94d8f6bbaf..4873b1fdb96 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py @@ -1,13 +1,3 @@ -""" -Response caching for Anthropic Messages (`/v1/messages`) requests. - -Non-streaming responses are plain dicts and are stored by the generic caching -handler. Streaming responses are returned to the caller before -``LLMCachingHandler.async_set_cache`` runs, so they are teed here instead: the -SSE events are buffered while they are forwarded and persisted verbatim once the -stream completes, and a hit replays exactly what the provider sent. -""" - from collections.abc import AsyncIterator from typing import TYPE_CHECKING, Any, cast @@ -38,14 +28,6 @@ def _decode(chunk: bytes | str) -> str: class AnthropicMessagesStreamCacheWriter: - """ - Forwards a `/v1/messages` SSE stream unchanged while buffering it, then - writes the collected events to the response cache on normal completion. - - Only a stream that ran to a ``message_stop`` without a provider ``error`` - event is written, so partial or failed responses cannot be replayed. - """ - def __init__( self, stream: AsyncIterator[bytes | str], @@ -105,11 +87,6 @@ class AnthropicMessagesStreamCacheWriter: class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterator): - """ - Replays cached `/v1/messages` SSE events and logs the request as a cache hit - once the replay finishes, mirroring what the live stream logs at end of stream. - """ - def __init__( self, events: list[str], @@ -146,11 +123,6 @@ def convert_cached_anthropic_messages_result( logging_obj: LiteLLMLoggingObj, kwargs: dict[str, Any], ) -> AnthropicMessagesResponse | CachedAnthropicMessagesStreamIterator: - """ - Turn a cached `/v1/messages` entry back into what the caller expects: an - SSE replay iterator for a streamed entry, otherwise the response itself - (``AnthropicMessagesResponse`` is a TypedDict, i.e. a dict at runtime). - """ events = get_cached_stream_events(cached_result) if events is not None: return CachedAnthropicMessagesStreamIterator( From b6cf4066f4e907c03f11065f52f4da149e649128 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 25 Jul 2026 01:17:43 +0000 Subject: [PATCH 016/504] fix(caching): log cached anthropic stream replay only once Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../messages/response_cache.py | 5 ++- .../messages/test_response_cache.py | 32 +++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py index 4873b1fdb96..d5b68c99130 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/response_cache.py @@ -96,6 +96,7 @@ class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterat super().__init__(litellm_logging_obj=litellm_logging_obj, request_body=request_body) self.chunks: list[bytes] = [event.encode("utf-8") for event in events] self.current_index = 0 + self.logged = False self._hidden_params: dict[str, Any] = {"cache_hit": True} litellm_logging_obj.model_call_details["cache_hit"] = True @@ -104,7 +105,9 @@ class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterat async def __anext__(self) -> bytes: if self.current_index >= len(self.chunks): - await self._handle_streaming_logging(self.chunks) + if not self.logged: + self.logged = True + await self._handle_streaming_logging(self.chunks) raise StopAsyncIteration chunk = self.chunks[self.current_index] self.current_index += 1 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py index 344152cd828..071580347a6 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_response_cache.py @@ -177,3 +177,35 @@ async def test_abandoned_stream_is_not_cached(local_cache, request_kwargs, monke assert len(fake_handler.calls) == 2 assert replayed == STREAM_EVENTS + +@pytest.mark.asyncio +async def test_cached_stream_replay_logs_once_when_polled_after_exhaustion(): + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.llms.anthropic.experimental_pass_through.messages.response_cache import ( + CachedAnthropicMessagesStreamIterator, + ) + from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, + ) + + logging_obj = MagicMock() + logging_obj.model_call_details = {} + iterator = CachedAnthropicMessagesStreamIterator( + events=[event.decode("utf-8") for event in STREAM_EVENTS], + litellm_logging_obj=logging_obj, + request_body={"model": "claude-sonnet-4-5"}, + ) + + with patch.object( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + new=AsyncMock(), + ) as mock_route: + assert await _collect(iterator) == STREAM_EVENTS + for _ in range(2): + with pytest.raises(StopAsyncIteration): + await iterator.__anext__() + await asyncio.sleep(0) + + mock_route.assert_called_once() From b2d2b29e2fc6fc8dc86a665adfb07c5ea5f69a3d Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 28 Jul 2026 23:05:55 +0000 Subject: [PATCH 017/504] fix(streaming): keep provider usage-only chunks for cost tracking without include_usage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/streaming_handler.py | 12 ++++ .../test_streaming_handler.py | 68 +++++++++++++++++++ 2 files changed, 80 insertions(+) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 60dbf7c644a..c2b0ca5e9db 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1465,6 +1465,7 @@ class CustomStreamWrapper: if self.stream_options is not None and self.stream_options["include_usage"] is True: model_response.choices = [] return model_response + self._record_usage_only_chunk(model_response=model_response) return ## CHECK FOR TOOL USE @@ -1691,6 +1692,17 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = "tool_calls" return model_response + def _record_usage_only_chunk(self, model_response: "ModelResponseStream") -> None: + """ + Keep provider usage-only chunks (e.g. OpenRouter's post-finish chunk, which carries a + provider-reported cost) available to cost tracking. They are never returned to the + caller; ``stream_options.include_usage`` only controls what the caller sees. + """ + if getattr(model_response, "usage", None) is None: + return + model_response.choices = [] + self.chunks.append(model_response) + @staticmethod def _propagate_usage_cost_to_hidden_params( response: "ModelResponse", diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 514714136fd..22daaf64dfd 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1507,6 +1507,74 @@ async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Loggin assert usage_chunks[-1].usage.cost == 0.00025 +@pytest.mark.asyncio +async def test_openrouter_streaming_usage_only_chunk_without_stream_options( + logging_obj: Logging, +): + """ + Regression: OpenRouter's post-finish chunk has `choices: []`. When the caller did not + pass stream_options.include_usage it was dropped before cost tracking, so the + provider-reported cost never reached the assembled response. + """ + from litellm.cost_calculator import get_response_cost_from_hidden_params + from litellm.utils import ModelResponseListIterator + + chunk1 = ModelResponseStream( + id="chatcmpl-or", + created=1742056047, + model="openrouter/claude", + choices=[ + StreamingChoices( + finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant") + ) + ], + usage=None, + ) + chunk2 = ModelResponseStream( + id="chatcmpl-or", + created=1742056048, + model="openrouter/claude", + choices=[ + StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) + ], + usage=None, + ) + usage_only_chunk = ModelResponseStream( + id="chatcmpl-or", + created=1742056049, + model="openrouter/claude", + choices=[], + usage=Usage( + completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025 + ), + ) + + response = CustomStreamWrapper( + completion_stream=ModelResponseListIterator( + model_responses=[chunk1, chunk2, usage_only_chunk] + ), + model="openrouter/claude", + custom_llm_provider="openrouter", + logging_obj=logging_obj, + ) + + collected_chunks = [chunk async for chunk in response] + + assert all(getattr(chunk, "usage", None) is None for chunk in collected_chunks) + + complete_response = litellm.stream_chunk_builder( + chunks=response.chunks, + messages=[{"role": "user", "content": "Hey"}], + ) + assert complete_response is not None + assert complete_response.usage.cost == 0.00025 + + CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response) + assert ( + get_response_cost_from_hidden_params(complete_response._hidden_params) == 0.00025 + ) + + def test_openrouter_streaming_cost_propagates_to_hidden_params(): """ Verify that provider-reported cost from usage.cost flows into From ef26590f72877a65947d2cb16df0bbbe3cd61892 Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 28 Jul 2026 23:12:06 +0000 Subject: [PATCH 018/504] test(streaming): assert logged cost from success callback for usage-only chunk Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_streaming_handler.py | 52 +++++++++++++------ 1 file changed, 35 insertions(+), 17 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 22daaf64dfd..f96a19a26c0 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1508,15 +1508,15 @@ async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Loggin @pytest.mark.asyncio -async def test_openrouter_streaming_usage_only_chunk_without_stream_options( - logging_obj: Logging, -): +async def test_openrouter_streaming_usage_only_chunk_without_stream_options(): """ Regression: OpenRouter's post-finish chunk has `choices: []`. When the caller did not pass stream_options.include_usage it was dropped before cost tracking, so the provider-reported cost never reached the assembled response. """ - from litellm.cost_calculator import get_response_cost_from_hidden_params + import time + + from litellm.integrations.custom_logger import CustomLogger from litellm.utils import ModelResponseListIterator chunk1 = ModelResponseStream( @@ -1549,30 +1549,48 @@ async def test_openrouter_streaming_usage_only_chunk_without_stream_options( ), ) + class MockCallback(CustomLogger): + pass + + mock_callback = MockCallback() + litellm.success_callback = [mock_callback] + litellm._async_success_callback = [mock_callback] + + stream_logging_obj = Logging( + model="openrouter/claude", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="12345", + function_id="1245", + ) + stream_logging_obj.update_environment_variables( + model="openrouter/claude", + optional_params={}, + litellm_params={}, + custom_llm_provider="openrouter", + ) + response = CustomStreamWrapper( completion_stream=ModelResponseListIterator( model_responses=[chunk1, chunk2, usage_only_chunk] ), model="openrouter/claude", custom_llm_provider="openrouter", - logging_obj=logging_obj, + logging_obj=stream_logging_obj, ) - collected_chunks = [chunk async for chunk in response] + with patch.object(mock_callback, "async_log_success_event") as mock_success_event: + collected_chunks = [chunk async for chunk in response] + await asyncio.sleep(1) assert all(getattr(chunk, "usage", None) is None for chunk in collected_chunks) - complete_response = litellm.stream_chunk_builder( - chunks=response.chunks, - messages=[{"role": "user", "content": "Hey"}], - ) - assert complete_response is not None - assert complete_response.usage.cost == 0.00025 - - CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response) - assert ( - get_response_cost_from_hidden_params(complete_response._hidden_params) == 0.00025 - ) + mock_success_event.assert_called_once() + logged_kwargs = mock_success_event.call_args.kwargs["kwargs"] + assert logged_kwargs["response_cost"] == 0.00025 + assert logged_kwargs["standard_logging_object"]["response_cost"] == 0.00025 def test_openrouter_streaming_cost_propagates_to_hidden_params(): From 09e1fb5ea2c71b11d717c86a5ae835e1a6ca7c9a Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 28 Jul 2026 23:16:37 +0000 Subject: [PATCH 019/504] test(streaming): restore success callbacks and await dispatch deterministically Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_streaming_handler.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index f96a19a26c0..99061decfb7 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1553,6 +1553,8 @@ async def test_openrouter_streaming_usage_only_chunk_without_stream_options(): pass mock_callback = MockCallback() + previous_success_callback = litellm.success_callback + previous_async_success_callback = litellm._async_success_callback litellm.success_callback = [mock_callback] litellm._async_success_callback = [mock_callback] @@ -1581,9 +1583,19 @@ async def test_openrouter_streaming_usage_only_chunk_without_stream_options(): logging_obj=stream_logging_obj, ) - with patch.object(mock_callback, "async_log_success_event") as mock_success_event: - collected_chunks = [chunk async for chunk in response] - await asyncio.sleep(1) + success_logged = asyncio.Event() + try: + with patch.object( + mock_callback, + "async_log_success_event", + new_callable=AsyncMock, + side_effect=lambda *args, **kwargs: success_logged.set(), + ) as mock_success_event: + collected_chunks = [chunk async for chunk in response] + await asyncio.wait_for(success_logged.wait(), timeout=30) + finally: + litellm.success_callback = previous_success_callback + litellm._async_success_callback = previous_async_success_callback assert all(getattr(chunk, "usage", None) is None for chunk in collected_chunks) From 6cfcb6cd839c1d23c7a59f247b1900346b9b2cab Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 14:08:18 +0000 Subject: [PATCH 020/504] fix(vertex_ai): translate /v1/embeddings batch rows to Gemini embedding shape Vertex batch files sent every jsonl line through the generateContent transform, so embeddings rows went out as {"request": {"contents": [...]}} and Vertex rejected each one with "no such field: 'contents'"; the OpenAI "input" was dropped along the way too. Route lines by their own url: embeddings lines now emit the EmbedContentRequest shape (singular content, embed_content_config sibling, custom_id round-tripping through the top-level key), and matching output rows come back as OpenAI embeddings responses. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 253 +++++++++++++--- .../test_vertex_ai_files_transformation.py | 277 ++++++++++++++++++ 2 files changed, 491 insertions(+), 39 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index dd877b52eb8..516a3ca7184 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -5,6 +5,7 @@ import json import os import re import time +from collections.abc import Mapping from typing import ( Any, Callable, @@ -16,6 +17,7 @@ from typing import ( Tuple, Union, ) +from urllib.parse import unquote import httpx from httpx import Headers, Response @@ -51,6 +53,9 @@ from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) +from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( + transform_openai_input_gemini_embed_content, +) from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, @@ -62,13 +67,26 @@ from litellm.types.llms.openai import ( ) from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.vertex_ai import GcsBucketResponse -from litellm.types.utils import LlmProviders, ModelResponse +from litellm.types.utils import ( + Embedding, + EmbeddingResponse, + LlmProviders, + ModelResponse, + Usage, +) from ..common_utils import VertexAIError from ..vertex_llm_base import VertexBase _GCP_LABEL_VALUE_MAX_LEN = 63 _CUSTOM_ID_RAW_LABEL_PREFIX = "b32_" +_VERTEX_BATCH_KEY_FIELD = "key" +_MANAGED_GCS_MODEL_PATH_PATTERN = re.compile(r"publishers/[^/]+/models/([^/?]+)") +_EMBED_CONTENT_CONFIG_FIELD_BY_GEMINI_PARAM = { + "outputDimensionality": "output_dimensionality", + "taskType": "task_type", + "title": "title", +} def _sanitize_gcp_label_value(value: str) -> str: @@ -131,6 +149,21 @@ def _set_litellm_batch_custom_id_labels(labels: Dict[str, str], custom_id: Any) labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk +def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, Any]) -> str: + """ + Resolve the OpenAI `custom_id` for a Vertex batch output row. + + Embedding rows carry it in the top-level `key` field that Vertex echoes back; + `generateContent` rows have no such field, so it is smuggled through request + labels instead (see `_set_litellm_batch_custom_id_labels`). + """ + key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) + if key is not None: + return str(key) + request_data = vertex_output_row.get("request") or {} + return _get_litellm_batch_custom_id_from_labels(request_data.get("labels") or {}) + + def _get_litellm_batch_custom_id_from_labels(labels: Dict[str, Any]) -> str: """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" raw = labels.get("litellm_custom_id_raw") @@ -149,10 +182,156 @@ def _get_litellm_batch_custom_id_from_labels(labels: Dict[str, Any]) -> str: return str(labels.get("litellm_custom_id", "unknown")) +def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, Any]) -> bool: + """ + Whether a Vertex batch output row came from an `EmbedContentRequest`. + + Successful rows hold the vector under `response.embedding.values`; failed rows only + carry `status`, so they are recognized from the singular `content` that the + embeddings request shape echoes back. + """ + if "request" not in vertex_output_row: + return False + response = vertex_output_row.get("response") + if isinstance(response, dict) and isinstance(response.get("embedding"), dict): + return True + request_data = vertex_output_row.get("request") + return bool(vertex_output_row.get("status")) and isinstance(request_data, dict) and "content" in request_data + + +def _openai_batch_output_row( + custom_id: str, + body: Mapping[str, Any] | None = None, + error: Mapping[str, str] | None = None, +) -> Mapping[str, Any]: + """ + One row of an OpenAI batch output file. Per the OpenAI Batch spec, failed rows set + `response` to null and populate `error` instead. + """ + return { + "id": f"batch_req_{uuid.uuid4()}", + "custom_id": custom_id, + "response": None + if body is None + else { + "status_code": 200, + "request_id": body.get("id", ""), + "body": body, + }, + "error": error, + } + + +def _transform_vertex_embeddings_batch_output_row_to_openai( + vertex_output_row: Mapping[str, Any], + model: str | None, +) -> Mapping[str, Any]: + """ + Transforms one Vertex Gemini Embedding batch output row into an OpenAI batch + output row holding an `/v1/embeddings` response body. + + Example Vertex jsonl + {"key": "id_1", "request": {...}, "response": {"tokenCount": "2", "embedding": {"values": [-0.015, 0.024]}}} + + `tokenCount` is serialized as a string by Vertex (int64 proto field), and the row + carries no `modelVersion`, so the model comes from the batch the row belongs to. + """ + custom_id = _get_litellm_batch_custom_id(vertex_output_row) + status = vertex_output_row.get("status", "") + if status: + return _openai_batch_output_row( + custom_id=custom_id, + error={"code": "vertex_ai_error", "message": status}, + ) + + vertex_response = vertex_output_row.get("response") or {} + token_count = int(vertex_response.get("tokenCount") or 0) + body = EmbeddingResponse( + model=model or "", + data=[ + Embedding( + embedding=vertex_response["embedding"]["values"], + index=0, + object="embedding", + ) + ], + usage=Usage(prompt_tokens=token_count, total_tokens=token_count), + ).model_dump() + return _openai_batch_output_row(custom_id=custom_id, body=body) + + +def _model_from_managed_gcs_url(url: str) -> str | None: + """ + Extracts the model from a LiteLLM-managed Vertex batch GCS url. + + Batch inputs and their sibling outputs are stored under + `.../publishers/google/models//...`, which is the only place the model of an + embeddings batch output row can be recovered from; unlike `generateContent` + responses, embedding rows carry no `modelVersion`. + """ + match = _MANAGED_GCS_MODEL_PATH_PATTERN.search(unquote(url)) + return match.group(1) if match else None + + +def _is_embeddings_batch_entry(openai_entry: Mapping[str, Any]) -> bool: + """ + Whether an OpenAI batch JSONL line targets the embeddings endpoint. + + OpenAI puts the target route on each line's `url` (e.g. `/v1/embeddings`); Vertex + has no equivalent per-line field, so the route decides which Vertex request shape + the line has to be translated into. + """ + url = openai_entry.get("url") + if not isinstance(url, str): + return False + path = url.split("?")[0].rstrip("/") + return path == "embeddings" or path.endswith("/embeddings") + + +def _openai_batch_jsonl_entry_to_vertex_embeddings_row( + openai_entry: Mapping[str, Any], +) -> Mapping[str, Any]: + """ + Transforms a single OpenAI `/v1/embeddings` batch entry into a Vertex Gemini + Embedding batch row. + + Example Vertex jsonl + {"key": "id_1", "request": {"content": {"parts": [{"text": "Hello World"}]}}, "embed_content_config": {"output_dimensionality": 768, "task_type": "RETRIEVAL_DOCUMENT"}} + + Note that `content` is singular (an `EmbedContentRequest`, not a + `GenerateContentRequest`), the per-row config is a sibling of `request` rather than + part of it, and the `custom_id` round-trips through the top-level `key`. + + API Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings + """ + openai_request_body = openai_entry.get("body") or {} + embedding_input = openai_request_body.get("input") + if embedding_input is None: + raise ValueError("`input` is required on /v1/embeddings batch requests, but was not provided") + + embed_content_request = transform_openai_input_gemini_embed_content( + input=embedding_input, + model=openai_request_body.get("model", ""), + optional_params=openai_request_body, + ) + embed_content_config = { + config_field: embed_content_request[gemini_param] + for gemini_param, config_field in _EMBED_CONTENT_CONFIG_FIELD_BY_GEMINI_PARAM.items() + if gemini_param in embed_content_request + } + + custom_id = openai_entry.get("custom_id") + return { + **({_VERTEX_BATCH_KEY_FIELD: str(custom_id)} if custom_id is not None else {}), + "request": {"content": embed_content_request["content"]}, + **({"embed_content_config": embed_content_config} if embed_content_config else {}), + } + + def _openai_batch_jsonl_entry_to_vertex_wrapped_request( openai_entry: Dict[str, Any], map_openai_to_vertex_params: Callable[[Dict[str, Any]], Dict[str, Any]], -) -> Dict[str, Any]: +) -> Mapping[str, Any]: """ Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. @@ -160,6 +339,9 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} """ + if _is_embeddings_batch_entry(openai_entry): + return _openai_batch_jsonl_entry_to_vertex_embeddings_row(openai_entry) + openai_request_body = openai_entry.get("body") or {} vertex_request_body = _transform_request_body( messages=openai_request_body.get("messages", []), @@ -629,6 +811,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): transformed_content = self._try_transform_vertex_batch_output_to_openai( content=content, logging_obj=logging_obj, + model=_model_from_managed_gcs_url(str(raw_response.request.url)), ) if transformed_content != content: # Create a new response with transformed content and updated Content-Length @@ -650,7 +833,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): return HttpxBinaryResponseContent(response=raw_response) def _try_transform_vertex_batch_output_to_openai( - self, content: bytes, logging_obj: Optional[LiteLLMLoggingObj] = None + self, + content: bytes, + logging_obj: LiteLLMLoggingObj | None = None, + model: str | None = None, ) -> bytes: """ Try to transform Vertex AI batch output to OpenAI format. @@ -692,7 +878,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # first line is not valid UTF-8/JSON) raises and falls through to the # passthrough below, leaving the content untouched. first_row = json.loads(first_line) - is_vertex_batch_output = ( + is_vertex_batch_output = _is_vertex_embeddings_batch_output_row(first_row) or ( "request" in first_row and "response" in first_row and "processed_time" in first_row @@ -731,11 +917,19 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): output = bytearray() for line in itertools.chain([first_line], lines): try: - openai_output = self._transform_single_vertex_batch_output_to_openai( - vertex_output=json.loads(line), - vertex_gemini_config=vertex_gemini_config, - logging_obj=batch_transform_logging_obj, - mock_httpx_response=mock_httpx_response, + vertex_output_row = json.loads(line) + openai_output = ( + _transform_vertex_embeddings_batch_output_row_to_openai( + vertex_output_row=vertex_output_row, + model=model, + ) + if _is_vertex_embeddings_batch_output_row(vertex_output_row) + else self._transform_single_vertex_batch_output_to_openai( + vertex_output=vertex_output_row, + vertex_gemini_config=vertex_gemini_config, + logging_obj=batch_transform_logging_obj, + mock_httpx_response=mock_httpx_response, + ) ) except Exception: return content @@ -755,30 +949,22 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): vertex_gemini_config: VertexGeminiConfig, logging_obj: Logging, mock_httpx_response: httpx.Response, - ) -> Dict[str, Any]: + ) -> Mapping[str, Any]: """ Transform a single Vertex AI batch output line to OpenAI format. Uses the existing VertexGeminiConfig transformation for the response. """ - # Extract custom_id from request labels (prefer raw for OpenAI round-trip) - request_data = vertex_output.get("request", {}) - labels = request_data.get("labels", {}) or {} - custom_id = _get_litellm_batch_custom_id_from_labels(labels) + custom_id = _get_litellm_batch_custom_id(vertex_output) # Check if there's an error status = vertex_output.get("status", "") has_error = bool(status) if has_error: - return { - "id": f"batch_req_{uuid.uuid4()}", - "custom_id": custom_id, - "response": None, - "error": { - "code": "vertex_ai_error", - "message": status, - }, - } + return _openai_batch_output_row( + custom_id=custom_id, + error={"code": "vertex_ai_error", "message": status}, + ) # Transform successful response using existing transformation vertex_response = vertex_output.get("response", {}) @@ -804,24 +990,13 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): response_dict = transformed_response.model_dump() # Return in OpenAI batch format - return { - "id": f"batch_req_{uuid.uuid4()}", - "custom_id": custom_id, - "response": { - "status_code": 200, - "request_id": response_dict.get("id", ""), - "body": response_dict, - }, - "error": None, - } + return _openai_batch_output_row(custom_id=custom_id, body=response_dict) except Exception as e: - return { - "id": f"batch_req_{uuid.uuid4()}", - "custom_id": custom_id, - "response": None, - "error": { + return _openai_batch_output_row( + custom_id=custom_id, + error={ "code": "transformation_error", "message": f"Failed to transform response: {str(e)}", }, - } + ) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 8c5305ee67b..636a3106617 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -1318,3 +1318,280 @@ class TestConfiguredBucketNameResolution: assert "bucket_name" in OPTIONAL_KWARGS_KEYS params = get_litellm_params(bucket_name="my-legacy-bucket") assert params.get("bucket_name") == "my-legacy-bucket" + + +def _embeddings_entry(**overrides): + entry = { + "custom_id": "request-1", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": "gemini-embedding-2", "input": "hello world"}, + } + entry.update(overrides) + return entry + + +class TestVertexEmbeddingsBatchInputTranslation: + """ + /v1/embeddings batch lines must be translated to Vertex's Gemini Embedding batch + shape, not the generateContent shape. + + Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings + """ + + def test_should_emit_embed_content_request_shape(self): + (row,) = _wrap_entries([_embeddings_entry()]) + + assert row["request"] == {"content": {"parts": [{"text": "hello world"}]}} + assert "contents" not in row["request"] + assert "labels" not in row["request"] + + def test_should_round_trip_custom_id_through_top_level_key(self): + (row,) = _wrap_entries([_embeddings_entry(custom_id="MyRequest-1")]) + + assert row["key"] == "MyRequest-1" + + def test_should_omit_key_when_no_custom_id(self): + entry = _embeddings_entry() + del entry["custom_id"] + + (row,) = _wrap_entries([entry]) + + assert "key" not in row + + def test_should_map_openai_params_to_embed_content_config_sibling(self): + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-001", + "input": "hello world", + "dimensions": 768, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "some_title", + } + ) + ] + ) + + assert row["embed_content_config"] == { + "output_dimensionality": 768, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "some_title", + } + assert "output_dimensionality" not in row["request"] + + def test_should_omit_embed_content_config_when_no_params_given(self): + (row,) = _wrap_entries([_embeddings_entry()]) + + assert "embed_content_config" not in row + + def test_should_translate_multimodal_gcs_input(self): + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-2", + "input": "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + } + ) + ] + ) + + assert row["request"]["content"]["parts"] == [ + { + "file_data": { + "mime_type": "image/jpeg", + "file_uri": "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + } + } + ] + + @pytest.mark.parametrize("url", ["/v1/embeddings", "v1/embeddings", "/v1/embeddings/"]) + def test_should_detect_embeddings_route_variants(self, url): + (row,) = _wrap_entries([_embeddings_entry(url=url)]) + + assert "content" in row["request"] + + def test_should_raise_when_input_missing(self): + with pytest.raises(ValueError, match="`input` is required"): + _wrap_entries([_embeddings_entry(body={"model": "gemini-embedding-2"})]) + + def test_should_keep_chat_completions_lines_on_generate_content_path(self): + (row,) = _wrap_entries( + [ + { + "custom_id": "request-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.0-flash-001", + "messages": [{"role": "user", "content": "Hello"}], + }, + } + ] + ) + + assert row["request"]["contents"] == [ + {"role": "user", "parts": [{"text": "Hello"}]} + ] + assert row["request"]["labels"]["litellm_custom_id"] == "request-1" + assert "key" not in row + + def test_should_translate_each_line_by_its_own_url(self): + chat_row, embeddings_row = _wrap_entries( + [ + { + "custom_id": "chat-1", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.0-flash-001", + "messages": [{"role": "user", "content": "Hello"}], + }, + }, + _embeddings_entry(custom_id="embed-1"), + ] + ) + + assert "contents" in chat_row["request"] + assert "content" in embeddings_row["request"] + + +class TestVertexEmbeddingsBatchOutputTranslation: + """Vertex Gemini Embedding batch output rows must come back as OpenAI batch rows.""" + + def _vertex_embeddings_output_row(self, **overrides): + row = { + "key": "request-1", + "request": {"content": {"parts": [{"text": "hello world"}]}}, + "response": { + "tokenCount": "2", + "embedding": {"values": [-0.015, 0.024]}, + }, + } + row.update(overrides) + return row + + def _transform(self, config, rows, url="https://example.com"): + content = "\n".join(json.dumps(row) for row in rows).encode("utf-8") + result = config.transform_file_content_response( + raw_response=httpx.Response( + status_code=200, + content=content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", url), + ), + logging_obj=MagicMock(), + litellm_params={}, + ) + return [ + json.loads(line) + for line in result.response.content.decode("utf-8").split("\n") + ] + + def test_should_transform_embeddings_output_to_openai_batch_row(self, config): + (result,) = self._transform(config, [self._vertex_embeddings_output_row()]) + + assert result["custom_id"] == "request-1" + assert result["error"] is None + assert result["response"]["status_code"] == 200 + body = result["response"]["body"] + assert body["object"] == "list" + assert body["data"] == [ + {"embedding": [-0.015, 0.024], "index": 0, "object": "embedding"} + ] + assert body["usage"]["prompt_tokens"] == 2 + assert body["usage"]["total_tokens"] == 2 + + def test_should_resolve_model_from_managed_gcs_object_path(self, config): + object_path = urllib.parse.quote( + "litellm-vertex-files/publishers/google/models/gemini-embedding-2/" + "prediction-model-2026-07-29T05:55:52Z/predictions.jsonl", + safe="", + ) + url = f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{object_path}?alt=media" + + (result,) = self._transform( + config, [self._vertex_embeddings_output_row()], url=url + ) + + assert result["response"]["body"]["model"] == "gemini-embedding-2" + + def test_should_surface_failed_embeddings_row_as_error(self, config): + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row( + status="Failed to parse JSON into proto", response={} + ) + ], + ) + + assert result["custom_id"] == "request-1" + assert result["response"] is None + assert result["error"]["code"] == "vertex_ai_error" + assert "Failed to parse JSON into proto" in result["error"]["message"] + + def test_should_transform_every_row_of_a_multi_row_file(self, config): + results = self._transform( + config, + [ + self._vertex_embeddings_output_row(key=f"request-{index}") + for index in range(3) + ], + ) + + assert [result["custom_id"] for result in results] == [ + "request-0", + "request-1", + "request-2", + ] + + def test_should_end_to_end_round_trip_openai_embeddings_batch(self, config): + (vertex_row,) = _wrap_entries( + [ + _embeddings_entry( + custom_id="MyRequest-1", + body={ + "model": "gemini-embedding-2", + "input": "hello world", + "dimensions": 2, + }, + ) + ] + ) + + (result,) = self._transform( + config, + [ + { + **vertex_row, + "status": "", + "processed_time": "2026-07-29T05:55:52.379528Z", + "response": { + "tokenCount": "2", + "embedding": {"values": [-0.015, 0.024]}, + }, + } + ], + ) + + assert result["custom_id"] == "MyRequest-1" + assert result["response"]["body"]["data"][0]["embedding"] == [-0.015, 0.024] + + def test_should_leave_legacy_predict_embeddings_output_untouched(self, config): + legacy_row = { + "instance": {"content": "hello world"}, + "predictions": [ + { + "embeddings": { + "statistics": {"token_count": 2, "truncated": False}, + "values": [0.2], + } + } + ], + "status": "", + } + content = json.dumps(legacy_row).encode("utf-8") + + assert config._try_transform_vertex_batch_output_to_openai(content) == content From e96614a39f45d8b6ffde5a3e2a05e63f6d73541d Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 14:44:25 +0000 Subject: [PATCH 021/504] fix(vertex_ai): put embed config inside the request and read live usage A live Vertex batch run showed the documented "embed_content_config" sibling of "request" is rejected by the API ("unsupported type"), failing the whole job rather than the row; the same fields inside the EmbedContentRequest succeed and honor output_dimensionality. Real output rows also report usage under response.usageMetadata.promptTokenCount, not the documented response.tokenCount, so every row came back with zero tokens. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 29 +++++++------ .../test_vertex_ai_files_transformation.py | 42 ++++++++++++++----- 2 files changed, 48 insertions(+), 23 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 516a3ca7184..1a5d40dabb2 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -82,7 +82,7 @@ _GCP_LABEL_VALUE_MAX_LEN = 63 _CUSTOM_ID_RAW_LABEL_PREFIX = "b32_" _VERTEX_BATCH_KEY_FIELD = "key" _MANAGED_GCS_MODEL_PATH_PATTERN = re.compile(r"publishers/[^/]+/models/([^/?]+)") -_EMBED_CONTENT_CONFIG_FIELD_BY_GEMINI_PARAM = { +_EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = { "outputDimensionality": "output_dimensionality", "taskType": "task_type", "title": "title", @@ -231,10 +231,11 @@ def _transform_vertex_embeddings_batch_output_row_to_openai( output row holding an `/v1/embeddings` response body. Example Vertex jsonl - {"key": "id_1", "request": {...}, "response": {"tokenCount": "2", "embedding": {"values": [-0.015, 0.024]}}} + {"key": "id_1", "request": {...}, "response": {"embedding": {"values": [-0.015, 0.024]}, "usageMetadata": {"promptTokenCount": 2}}} - `tokenCount` is serialized as a string by Vertex (int64 proto field), and the row - carries no `modelVersion`, so the model comes from the batch the row belongs to. + Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as + a fallback. The row carries no `modelVersion`, so the model comes from the batch it + belongs to. """ custom_id = _get_litellm_batch_custom_id(vertex_output_row) status = vertex_output_row.get("status", "") @@ -245,7 +246,8 @@ def _transform_vertex_embeddings_batch_output_row_to_openai( ) vertex_response = vertex_output_row.get("response") or {} - token_count = int(vertex_response.get("tokenCount") or 0) + usage_metadata = vertex_response.get("usageMetadata") or {} + token_count = int(usage_metadata.get("promptTokenCount") or vertex_response.get("tokenCount") or 0) body = EmbeddingResponse( model=model or "", data=[ @@ -296,11 +298,13 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_row( Embedding batch row. Example Vertex jsonl - {"key": "id_1", "request": {"content": {"parts": [{"text": "Hello World"}]}}, "embed_content_config": {"output_dimensionality": 768, "task_type": "RETRIEVAL_DOCUMENT"}} + {"key": "id_1", "request": {"content": {"parts": [{"text": "Hello World"}]}, "output_dimensionality": 768, "task_type": "RETRIEVAL_DOCUMENT"}} Note that `content` is singular (an `EmbedContentRequest`, not a - `GenerateContentRequest`), the per-row config is a sibling of `request` rather than - part of it, and the `custom_id` round-trips through the top-level `key`. + `GenerateContentRequest`) and that the `custom_id` round-trips through the top-level + `key`. The docs put the per-row config in an `embed_content_config` sibling of + `request`, but the API rejects that key outright and fails the whole batch job, so + the config fields go inside the `EmbedContentRequest` itself. API Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings """ @@ -314,17 +318,16 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_row( model=openai_request_body.get("model", ""), optional_params=openai_request_body, ) - embed_content_config = { - config_field: embed_content_request[gemini_param] - for gemini_param, config_field in _EMBED_CONTENT_CONFIG_FIELD_BY_GEMINI_PARAM.items() + embed_request_fields = { + request_field: embed_content_request[gemini_param] + for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM.items() if gemini_param in embed_content_request } custom_id = openai_entry.get("custom_id") return { **({_VERTEX_BATCH_KEY_FIELD: str(custom_id)} if custom_id is not None else {}), - "request": {"content": embed_content_request["content"]}, - **({"embed_content_config": embed_content_config} if embed_content_config else {}), + "request": {"content": embed_content_request["content"], **embed_request_fields}, } diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 636a3106617..3eff3083220 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -1359,7 +1359,11 @@ class TestVertexEmbeddingsBatchInputTranslation: assert "key" not in row - def test_should_map_openai_params_to_embed_content_config_sibling(self): + def test_should_map_openai_params_into_the_embed_content_request(self): + """ + The docs put these in an `embed_content_config` sibling of `request`, but Vertex + rejects that key and fails the whole job, so they belong inside the request. + """ (row,) = _wrap_entries( [ _embeddings_entry( @@ -1374,17 +1378,20 @@ class TestVertexEmbeddingsBatchInputTranslation: ] ) - assert row["embed_content_config"] == { - "output_dimensionality": 768, - "task_type": "RETRIEVAL_DOCUMENT", - "title": "some_title", + assert row == { + "key": "request-1", + "request": { + "content": {"parts": [{"text": "hello world"}]}, + "output_dimensionality": 768, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "some_title", + }, } - assert "output_dimensionality" not in row["request"] - def test_should_omit_embed_content_config_when_no_params_given(self): + def test_should_omit_config_fields_when_no_params_given(self): (row,) = _wrap_entries([_embeddings_entry()]) - assert "embed_content_config" not in row + assert set(row["request"]) == {"content"} def test_should_translate_multimodal_gcs_input(self): (row,) = _wrap_entries( @@ -1465,8 +1472,8 @@ class TestVertexEmbeddingsBatchOutputTranslation: "key": "request-1", "request": {"content": {"parts": [{"text": "hello world"}]}}, "response": { - "tokenCount": "2", "embedding": {"values": [-0.015, 0.024]}, + "usageMetadata": {"promptTokenCount": 2}, }, } row.update(overrides) @@ -1503,6 +1510,21 @@ class TestVertexEmbeddingsBatchOutputTranslation: assert body["usage"]["prompt_tokens"] == 2 assert body["usage"]["total_tokens"] == 2 + def test_should_fall_back_to_documented_token_count_field(self, config): + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row( + response={ + "embedding": {"values": [-0.015, 0.024]}, + "tokenCount": "2", + } + ) + ], + ) + + assert result["response"]["body"]["usage"]["prompt_tokens"] == 2 + def test_should_resolve_model_from_managed_gcs_object_path(self, config): object_path = urllib.parse.quote( "litellm-vertex-files/publishers/google/models/gemini-embedding-2/" @@ -1569,8 +1591,8 @@ class TestVertexEmbeddingsBatchOutputTranslation: "status": "", "processed_time": "2026-07-29T05:55:52.379528Z", "response": { - "tokenCount": "2", "embedding": {"values": [-0.015, 0.024]}, + "usageMetadata": {"promptTokenCount": 2}, }, } ], From 627d2755da56d01a6f5838bec967005027b2f35d Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 15:06:58 +0000 Subject: [PATCH 022/504] fix(vertex_ai): fan array embeddings input out into one vertex row per element Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 221 +++++++++++++----- .../files/test_vertex_ai_files_streaming.py | 5 +- .../test_vertex_ai_files_transformation.py | 180 +++++++++++++- 3 files changed, 339 insertions(+), 67 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 1a5d40dabb2..78a16b002e7 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -66,7 +66,7 @@ from litellm.types.llms.openai import ( PathLike, ) from litellm.types.files import StreamingMediaUploadConfig -from litellm.types.llms.vertex_ai import GcsBucketResponse +from litellm.types.llms.vertex_ai import GcsBucketResponse, GeminiEmbeddingInput from litellm.types.utils import ( Embedding, EmbeddingResponse, @@ -87,6 +87,7 @@ _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = { "taskType": "task_type", "title": "title", } +_VERTEX_BATCH_FANNED_OUT_KEY_PATTERN = re.compile(r"(?P.*)#(?P\d+)/(?P\d+)") def _sanitize_gcp_label_value(value: str) -> str: @@ -222,46 +223,93 @@ def _openai_batch_output_row( } -def _transform_vertex_embeddings_batch_output_row_to_openai( - vertex_output_row: Mapping[str, Any], +def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, int]: + """ + Resolve `(custom_id, index within that custom_id)` for a Vertex batch output row. + + A `/v1/embeddings` entry whose `input` is an array fans out into one Vertex row per + element, tagged `#/` (see `_vertex_batch_embeddings_key`), + so the rows can be reassembled into a single OpenAI response. + """ + key = _get_litellm_batch_custom_id(vertex_output_row) + match = _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN.fullmatch(key) + if match is None or int(match["total"]) < 2: + return key, 0 + return match["custom_id"], int(match["index"]) + + +def _vertex_embeddings_rows_to_openai_batch_output_row( + custom_id: str, + vertex_output_rows: tuple[Mapping[str, Any], ...], model: str | None, ) -> Mapping[str, Any]: """ - Transforms one Vertex Gemini Embedding batch output row into an OpenAI batch - output row holding an `/v1/embeddings` response body. + Transforms the Vertex Gemini Embedding batch output rows belonging to one OpenAI + batch entry into an OpenAI batch output row holding an `/v1/embeddings` response. Example Vertex jsonl {"key": "id_1", "request": {...}, "response": {"embedding": {"values": [-0.015, 0.024]}, "usageMetadata": {"promptTokenCount": 2}}} - Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as - a fallback. The row carries no `modelVersion`, so the model comes from the batch it - belongs to. + An entry that asked for several embeddings at once maps to several rows here, which + become the indexed elements of a single `data` array. One failed element fails the + whole entry, since an OpenAI batch row is either a response or an error. Live rows + report usage under `usageMetadata`; the documented `tokenCount` is kept as a + fallback. Rows carry no `modelVersion`, so the model comes from the batch they + belong to. """ - custom_id = _get_litellm_batch_custom_id(vertex_output_row) - status = vertex_output_row.get("status", "") + status = next((row["status"] for row in vertex_output_rows if row.get("status")), "") if status: return _openai_batch_output_row( custom_id=custom_id, error={"code": "vertex_ai_error", "message": status}, ) - vertex_response = vertex_output_row.get("response") or {} - usage_metadata = vertex_response.get("usageMetadata") or {} - token_count = int(usage_metadata.get("promptTokenCount") or vertex_response.get("tokenCount") or 0) + responses = tuple(row.get("response") or {} for row in vertex_output_rows) + token_count = sum( + int((response.get("usageMetadata") or {}).get("promptTokenCount") or response.get("tokenCount") or 0) + for response in responses + ) body = EmbeddingResponse( model=model or "", data=[ Embedding( - embedding=vertex_response["embedding"]["values"], - index=0, + embedding=response["embedding"]["values"], + index=index, object="embedding", ) + for index, response in enumerate(responses) ], usage=Usage(prompt_tokens=token_count, total_tokens=token_count), ).model_dump() return _openai_batch_output_row(custom_id=custom_id, body=body) +def _transform_vertex_embeddings_batch_output_to_openai( + vertex_output_rows: Iterable[Mapping[str, Any]], + model: str | None, +) -> tuple[Mapping[str, Any], ...]: + """ + Transforms a whole Vertex Gemini Embedding batch output into OpenAI batch output + rows, one per OpenAI batch entry, in the order the entries first appear. + + Rows are grouped rather than mapped one to one because a single entry can fan out + into several Vertex rows, and Vertex returns them in arbitrary order. + """ + keyed_rows = tuple((_split_vertex_batch_key(row), row) for row in vertex_output_rows) + grouped_rows = { + custom_id: tuple(row for _, row in group) + for custom_id, group in itertools.groupby(sorted(keyed_rows, key=lambda kr: kr[0]), key=lambda kr: kr[0][0]) + } + return tuple( + _vertex_embeddings_rows_to_openai_batch_output_row( + custom_id=custom_id, + vertex_output_rows=grouped_rows[custom_id], + model=model, + ) + for custom_id in dict.fromkeys(custom_id for (custom_id, _), _ in keyed_rows) + ) + + def _model_from_managed_gcs_url(url: str) -> str | None: """ Extracts the model from a LiteLLM-managed Vertex batch GCS url. @@ -290,21 +338,49 @@ def _is_embeddings_batch_entry(openai_entry: Mapping[str, Any]) -> bool: return path == "embeddings" or path.endswith("/embeddings") -def _openai_batch_jsonl_entry_to_vertex_embeddings_row( - openai_entry: Mapping[str, Any], -) -> Mapping[str, Any]: +def _openai_embedding_input_elements( + embedding_input: GeminiEmbeddingInput, +) -> tuple[Union[str, List[str]], ...]: """ - Transforms a single OpenAI `/v1/embeddings` batch entry into a Vertex Gemini - Embedding batch row. + Split an OpenAI `input` into the elements that each get their own embedding. + + A string is one embedding, a flat array is one embedding per element, and a nested + array is one combined embedding per inner array, matching the online + `batchEmbedContents` path. + """ + if isinstance(embedding_input, list): + return tuple(embedding_input) + return (embedding_input,) + + +def _vertex_batch_embeddings_key(custom_id: str, index: int, total: int) -> str: + """ + The top-level `key` Vertex echoes back on an embeddings row. + + An entry asking for several embeddings needs several Vertex rows, so its key also + carries the element index and the group size; `_split_vertex_batch_key` reads them + back out. Entries asking for a single embedding keep their bare `custom_id`. + """ + return custom_id if total < 2 else f"{custom_id}#{index}/{total}" + + +def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( + openai_entry: Mapping[str, Any], +) -> tuple[Mapping[str, Any], ...]: + """ + Transforms a single OpenAI `/v1/embeddings` batch entry into Vertex Gemini Embedding + batch rows, one per requested embedding. Example Vertex jsonl {"key": "id_1", "request": {"content": {"parts": [{"text": "Hello World"}]}, "output_dimensionality": 768, "task_type": "RETRIEVAL_DOCUMENT"}} Note that `content` is singular (an `EmbedContentRequest`, not a `GenerateContentRequest`) and that the `custom_id` round-trips through the top-level - `key`. The docs put the per-row config in an `embed_content_config` sibling of - `request`, but the API rejects that key outright and fails the whole batch job, so - the config fields go inside the `EmbedContentRequest` itself. + `key`. An `EmbedContentRequest` returns exactly one vector, so an entry whose `input` + is an array fans out into one row per element and is reassembled on the way back. + The docs put the per-row config in an `embed_content_config` sibling of `request`, + but the API rejects that key outright and fails the whole batch job, so the config + fields go inside the `EmbedContentRequest` itself. API Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings """ @@ -313,37 +389,58 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_row( if embedding_input is None: raise ValueError("`input` is required on /v1/embeddings batch requests, but was not provided") - embed_content_request = transform_openai_input_gemini_embed_content( - input=embedding_input, - model=openai_request_body.get("model", ""), - optional_params=openai_request_body, + elements = _openai_embedding_input_elements(embedding_input) + if not elements: + raise ValueError("`input` on /v1/embeddings batch requests must not be empty") + + embed_content_requests = tuple( + transform_openai_input_gemini_embed_content( + input=element, + model=openai_request_body.get("model", ""), + optional_params=openai_request_body, + ) + for element in elements ) - embed_request_fields = { - request_field: embed_content_request[gemini_param] - for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM.items() - if gemini_param in embed_content_request - } - custom_id = openai_entry.get("custom_id") - return { - **({_VERTEX_BATCH_KEY_FIELD: str(custom_id)} if custom_id is not None else {}), - "request": {"content": embed_content_request["content"], **embed_request_fields}, - } + return tuple( + { + **( + {} + if custom_id is None + else { + _VERTEX_BATCH_KEY_FIELD: _vertex_batch_embeddings_key( + custom_id=str(custom_id), + index=index, + total=len(embed_content_requests), + ) + } + ), + "request": { + "content": embed_content_request["content"], + **{ + request_field: embed_content_request[gemini_param] + for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM.items() + if gemini_param in embed_content_request + }, + }, + } + for index, embed_content_request in enumerate(embed_content_requests) + ) -def _openai_batch_jsonl_entry_to_vertex_wrapped_request( +def _openai_batch_jsonl_entry_to_vertex_rows( openai_entry: Dict[str, Any], map_openai_to_vertex_params: Callable[[Dict[str, Any]], Dict[str, Any]], -) -> Mapping[str, Any]: +) -> tuple[Mapping[str, Any], ...]: """ - Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. + Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to. jsonl body for vertex is {"request": } Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} """ if _is_embeddings_batch_entry(openai_entry): - return _openai_batch_jsonl_entry_to_vertex_embeddings_row(openai_entry) + return _openai_batch_jsonl_entry_to_vertex_embeddings_rows(openai_entry) openai_request_body = openai_entry.get("body") or {} vertex_request_body = _transform_request_body( @@ -361,7 +458,7 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( vertex_request_body["labels"] = {} _set_litellm_batch_custom_id_labels(vertex_request_body["labels"], custom_id) - return {"request": vertex_request_body} + return ({"request": vertex_request_body},) def _iter_stripped_lines(raw_lines: Iterable[Union[str, bytes]]) -> Iterator[str]: @@ -459,10 +556,10 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def _iter_vertex_jsonl_chunks(self) -> Iterator[bytes]: first = True for entry in _iter_openai_jsonl_entries(self._openai_file_content): - wrapped = _openai_batch_jsonl_entry_to_vertex_wrapped_request(entry, self._map_openai_to_vertex_params) - prefix = b"" if first else b"\n" - first = False - yield prefix + json.dumps(wrapped).encode("utf-8") + for wrapped in _openai_batch_jsonl_entry_to_vertex_rows(entry, self._map_openai_to_vertex_params): + prefix = b"" if first else b"\n" + first = False + yield prefix + json.dumps(wrapped).encode("utf-8") def iter_bytes(self) -> Iterator[bytes]: return self._iter_vertex_jsonl_chunks() @@ -914,25 +1011,29 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): request=httpx.Request(method="POST", url="https://example.com"), ) + all_lines = itertools.chain([first_line], lines) + + # Embedding rows are grouped by `custom_id` rather than transformed one at a + # time, since an entry that asked for several embeddings comes back as + # several rows, in arbitrary order. + if _is_vertex_embeddings_batch_output_row(first_row): + openai_outputs = _transform_vertex_embeddings_batch_output_to_openai( + vertex_output_rows=(json.loads(line) for line in all_lines), + model=model, + ) + return b"\n".join(json.dumps(openai_output).encode("utf-8") for openai_output in openai_outputs) + # Transform each row straight into the output buffer, so peak memory # stays at ~one row plus the output. If any row fails, return the # original content unchanged. output = bytearray() - for line in itertools.chain([first_line], lines): + for line in all_lines: try: - vertex_output_row = json.loads(line) - openai_output = ( - _transform_vertex_embeddings_batch_output_row_to_openai( - vertex_output_row=vertex_output_row, - model=model, - ) - if _is_vertex_embeddings_batch_output_row(vertex_output_row) - else self._transform_single_vertex_batch_output_to_openai( - vertex_output=vertex_output_row, - vertex_gemini_config=vertex_gemini_config, - logging_obj=batch_transform_logging_obj, - mock_httpx_response=mock_httpx_response, - ) + openai_output = self._transform_single_vertex_batch_output_to_openai( + vertex_output=json.loads(line), + vertex_gemini_config=vertex_gemini_config, + logging_obj=batch_transform_logging_obj, + mock_httpx_response=mock_httpx_response, ) except Exception: return content diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index 2e3280c0ed1..957fc7dbcf4 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -37,7 +37,7 @@ from litellm.llms.vertex_ai.files.transformation import ( _get_litellm_batch_custom_id_from_labels, _iter_openai_jsonl_entries, _iter_openai_jsonl_lines, - _openai_batch_jsonl_entry_to_vertex_wrapped_request, + _openai_batch_jsonl_entry_to_vertex_rows, ) from litellm.types.llms.openai import CreateFileRequest @@ -84,8 +84,9 @@ def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> st transform, so the streaming path can be checked against it for parity.""" entries = [json.loads(line) for line in content.splitlines() if line.strip()] return "\n".join( - json.dumps(_openai_batch_jsonl_entry_to_vertex_wrapped_request(entry, cfg._map_openai_to_vertex_params)) + json.dumps(row) for entry in entries + for row in _openai_batch_jsonl_entry_to_vertex_rows(entry, cfg._map_openai_to_vertex_params) ) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 3eff3083220..73d1d6eeb5a 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock from litellm.llms.vertex_ai.files.transformation import ( VertexAIFilesConfig, _get_litellm_batch_custom_id_from_labels, - _openai_batch_jsonl_entry_to_vertex_wrapped_request, + _openai_batch_jsonl_entry_to_vertex_rows, _sanitize_gcp_label_value, ) from litellm.types.llms.openai import OpenAIFileObject, HttpxBinaryResponseContent @@ -1054,14 +1054,15 @@ class TestTryTransformDoesNotMutateCallerLoggingObj: def _wrap_entries(openai_jsonl_content): - """Vertex-wrapped requests for a list of OpenAI batch entries, built via the - live single-entry transform that the streaming upload path uses.""" + """Vertex rows for a list of OpenAI batch entries, built via the live + single-entry transform that the streaming upload path uses.""" cfg = VertexAIFilesConfig() return [ - _openai_batch_jsonl_entry_to_vertex_wrapped_request( + row + for entry in openai_jsonl_content + for row in _openai_batch_jsonl_entry_to_vertex_rows( entry, cfg._map_openai_to_vertex_params ) - for entry in openai_jsonl_content ] @@ -1424,6 +1425,86 @@ class TestVertexEmbeddingsBatchInputTranslation: with pytest.raises(ValueError, match="`input` is required"): _wrap_entries([_embeddings_entry(body={"model": "gemini-embedding-2"})]) + def test_should_raise_when_input_empty(self): + with pytest.raises(ValueError, match="must not be empty"): + _wrap_entries( + [_embeddings_entry(body={"model": "gemini-embedding-2", "input": []})] + ) + + def test_should_fan_an_input_array_out_into_one_row_per_element(self): + """ + An `EmbedContentRequest` returns exactly one vector, so an OpenAI entry asking + for several embeddings needs several Vertex rows. + """ + rows = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-001", + "input": ["first", "second"], + "dimensions": 768, + } + ) + ] + ) + + assert rows == [ + { + "key": "request-1#0/2", + "request": { + "content": {"parts": [{"text": "first"}]}, + "output_dimensionality": 768, + }, + }, + { + "key": "request-1#1/2", + "request": { + "content": {"parts": [{"text": "second"}]}, + "output_dimensionality": 768, + }, + }, + ] + + def test_should_keep_the_bare_custom_id_for_single_element_arrays(self): + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={"model": "gemini-embedding-2", "input": ["only one"]} + ) + ] + ) + + assert row["key"] == "request-1" + + def test_should_combine_a_nested_input_into_one_multipart_row(self): + """Nested arrays are the combined-embedding shape, as on the online path.""" + (row,) = _wrap_entries( + [ + _embeddings_entry( + body={ + "model": "gemini-embedding-2", + "input": [ + [ + "a caption", + "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + ] + ], + } + ) + ] + ) + + assert row["key"] == "request-1" + assert row["request"]["content"]["parts"] == [ + {"text": "a caption"}, + { + "file_data": { + "mime_type": "image/jpeg", + "file_uri": "gs://cloud-samples-data/generative-ai/image/benchmark.jpeg", + } + }, + ] + def test_should_keep_chat_completions_lines_on_generate_content_path(self): (row,) = _wrap_entries( [ @@ -1569,6 +1650,95 @@ class TestVertexEmbeddingsBatchOutputTranslation: "request-2", ] + def test_should_reassemble_a_fanned_out_input_array_into_one_row(self, config): + """Vertex returns the rows of one entry in arbitrary order.""" + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row( + key="request-1#1/2", + response={ + "embedding": {"values": [0.3, 0.4]}, + "usageMetadata": {"promptTokenCount": 5}, + }, + ), + self._vertex_embeddings_output_row( + key="request-1#0/2", + response={ + "embedding": {"values": [0.1, 0.2]}, + "usageMetadata": {"promptTokenCount": 3}, + }, + ), + ], + ) + + assert result["custom_id"] == "request-1" + assert result["response"]["body"]["data"] == [ + {"embedding": [0.1, 0.2], "index": 0, "object": "embedding"}, + {"embedding": [0.3, 0.4], "index": 1, "object": "embedding"}, + ] + assert result["response"]["body"]["usage"]["prompt_tokens"] == 8 + + def test_should_keep_fanned_out_entries_apart_and_in_file_order(self, config): + results = self._transform( + config, + [ + self._vertex_embeddings_output_row(key="request-2#0/2"), + self._vertex_embeddings_output_row(key="request-1"), + self._vertex_embeddings_output_row(key="request-2#1/2"), + ], + ) + + assert [result["custom_id"] for result in results] == ["request-2", "request-1"] + assert len(results[0]["response"]["body"]["data"]) == 2 + assert len(results[1]["response"]["body"]["data"]) == 1 + + def test_should_fail_the_whole_entry_when_one_of_its_rows_failed(self, config): + """An OpenAI batch row is either a response or an error, never both.""" + (result,) = self._transform( + config, + [ + self._vertex_embeddings_output_row(key="request-1#0/2"), + self._vertex_embeddings_output_row( + key="request-1#1/2", status="Quota exceeded", response={} + ), + ], + ) + + assert result["custom_id"] == "request-1" + assert result["response"] is None + assert result["error"]["message"] == "Quota exceeded" + + def test_should_end_to_end_round_trip_a_fanned_out_embeddings_batch(self, config): + first_row, second_row = _wrap_entries( + [ + _embeddings_entry( + custom_id="MyRequest-1", + body={ + "model": "gemini-embedding-2", + "input": ["hello world", "goodbye world"], + }, + ) + ] + ) + + (result,) = self._transform( + config, + [ + { + **row, + "status": "", + "response": {"embedding": {"values": values}}, + } + for row, values in ((second_row, [0.3]), (first_row, [0.1])) + ], + ) + + assert result["custom_id"] == "MyRequest-1" + assert [ + embedding["embedding"] for embedding in result["response"]["body"]["data"] + ] == [[0.1], [0.3]] + def test_should_end_to_end_round_trip_openai_embeddings_batch(self, config): (vertex_row,) = _wrap_entries( [ From 3c979f0b471214929b38de0fe61573fc864b7906 Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 15:23:14 +0000 Subject: [PATCH 023/504] test(vertex_ai): cover batch lines without a url staying on the chat path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_vertex_ai_files_transformation.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 73d1d6eeb5a..aef7f8b4684 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -1526,6 +1526,24 @@ class TestVertexEmbeddingsBatchInputTranslation: assert row["request"]["labels"]["litellm_custom_id"] == "request-1" assert "key" not in row + def test_should_keep_lines_without_a_url_on_generate_content_path(self): + """`url` is optional on a batch line, and chat is the shape LiteLLM has always assumed.""" + (row,) = _wrap_entries( + [ + { + "custom_id": "request-1", + "body": { + "model": "gemini-2.0-flash-001", + "messages": [{"role": "user", "content": "Hello"}], + }, + } + ] + ) + + assert row["request"]["contents"] == [ + {"role": "user", "parts": [{"text": "Hello"}]} + ] + def test_should_translate_each_line_by_its_own_url(self): chat_row, embeddings_row = _wrap_entries( [ From bf723fa9c167f48731f68ebe6b3bcab7351f5a83 Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 29 Jul 2026 20:50:27 +0000 Subject: [PATCH 024/504] fix(vertex_ai): percent-encode the custom_id in fanned-out vertex batch keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 29 ++++--- .../test_vertex_ai_files_transformation.py | 75 +++++++++++++++++++ 2 files changed, 92 insertions(+), 12 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 78a16b002e7..90fc5fae082 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -17,7 +17,7 @@ from typing import ( Tuple, Union, ) -from urllib.parse import unquote +from urllib.parse import quote, unquote import httpx from httpx import Headers, Response @@ -87,7 +87,7 @@ _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = { "taskType": "task_type", "title": "title", } -_VERTEX_BATCH_FANNED_OUT_KEY_PATTERN = re.compile(r"(?P.*)#(?P\d+)/(?P\d+)") +_VERTEX_BATCH_FANNED_OUT_KEY_PATTERN = re.compile(r"(?P[^#]*)#(?P\d+)/(?P\d+)") def _sanitize_gcp_label_value(value: str) -> str: @@ -160,7 +160,7 @@ def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, Any]) -> str: """ key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) if key is not None: - return str(key) + return unquote(str(key)) request_data = vertex_output_row.get("request") or {} return _get_litellm_batch_custom_id_from_labels(request_data.get("labels") or {}) @@ -228,14 +228,17 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, Resolve `(custom_id, index within that custom_id)` for a Vertex batch output row. A `/v1/embeddings` entry whose `input` is an array fans out into one Vertex row per - element, tagged `#/` (see `_vertex_batch_embeddings_key`), - so the rows can be reassembled into a single OpenAI response. + element, tagged `#/` (see + `_vertex_batch_embeddings_key`), so the rows can be reassembled into a single OpenAI + response. """ - key = _get_litellm_batch_custom_id(vertex_output_row) - match = _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN.fullmatch(key) - if match is None or int(match["total"]) < 2: - return key, 0 - return match["custom_id"], int(match["index"]) + key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) + if key is None: + return _get_litellm_batch_custom_id(vertex_output_row), 0 + match = _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN.fullmatch(str(key)) + if match is None: + return unquote(str(key)), 0 + return unquote(match["custom_id"]), int(match["index"]) def _vertex_embeddings_rows_to_openai_batch_output_row( @@ -359,9 +362,11 @@ def _vertex_batch_embeddings_key(custom_id: str, index: int, total: int) -> str: An entry asking for several embeddings needs several Vertex rows, so its key also carries the element index and the group size; `_split_vertex_batch_key` reads them - back out. Entries asking for a single embedding keep their bare `custom_id`. + back out. The `custom_id` is percent-encoded so that a customer one ending in + `#/` cannot be mistaken for that tag, which would merge two entries. """ - return custom_id if total < 2 else f"{custom_id}#{index}/{total}" + encoded_custom_id = quote(custom_id, safe="") + return encoded_custom_id if total < 2 else f"{encoded_custom_id}#{index}/{total}" def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index aef7f8b4684..f95a63e4421 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -1476,6 +1476,19 @@ class TestVertexEmbeddingsBatchInputTranslation: assert row["key"] == "request-1" + def test_should_encode_a_custom_id_that_looks_like_a_fan_out_tag(self): + """A customer custom_id ending in `#/` must not read back as fan-out metadata.""" + (row,) = _wrap_entries( + [ + _embeddings_entry( + custom_id="request-1#0/2", + body={"model": "gemini-embedding-2", "input": "hello world"}, + ) + ] + ) + + assert row["key"] == "request-1%230%2F2" + def test_should_combine_a_nested_input_into_one_multipart_row(self): """Nested arrays are the combined-embedding shape, as on the online path.""" (row,) = _wrap_entries( @@ -1711,6 +1724,68 @@ class TestVertexEmbeddingsBatchOutputTranslation: assert len(results[0]["response"]["body"]["data"]) == 2 assert len(results[1]["response"]["body"]["data"]) == 1 + def test_should_not_merge_an_entry_whose_custom_id_looks_like_a_fan_out_tag(self, config): + """`request-1#0/2` is a legal custom_id, and a distinct entry from `request-1`.""" + lookalike_row, plain_row = _wrap_entries( + [ + _embeddings_entry( + custom_id="request-1#0/2", + body={"model": "gemini-embedding-2", "input": "lookalike"}, + ), + _embeddings_entry( + custom_id="request-1", + body={"model": "gemini-embedding-2", "input": "plain"}, + ), + ] + ) + + results = self._transform( + config, + [ + {**row, "status": "", "response": {"embedding": {"values": values}}} + for row, values in ((lookalike_row, [0.1]), (plain_row, [0.2])) + ], + ) + + assert [result["custom_id"] for result in results] == [ + "request-1#0/2", + "request-1", + ] + assert [ + result["response"]["body"]["data"][0]["embedding"] for result in results + ] == [[0.1], [0.2]] + + def test_should_round_trip_a_fan_out_of_a_custom_id_holding_the_separator(self, config): + rows = _wrap_entries( + [ + _embeddings_entry( + custom_id="request#1/1", + body={ + "model": "gemini-embedding-2", + "input": ["first", "second"], + }, + ) + ] + ) + + assert [row["key"] for row in rows] == [ + "request%231%2F1#0/2", + "request%231%2F1#1/2", + ] + + (result,) = self._transform( + config, + [ + {**row, "status": "", "response": {"embedding": {"values": values}}} + for row, values in zip(reversed(rows), ([0.3], [0.1])) + ], + ) + + assert result["custom_id"] == "request#1/1" + assert [ + embedding["embedding"] for embedding in result["response"]["body"]["data"] + ] == [[0.1], [0.3]] + def test_should_fail_the_whole_entry_when_one_of_its_rows_failed(self, config): """An OpenAI batch row is either a response or an error, never both.""" (result,) = self._transform( From 7c56317edf153d61b395f4257476aefdd02f2236 Mon Sep 17 00:00:00 2001 From: Yaroslav Date: Thu, 30 Jul 2026 21:37:28 +0300 Subject: [PATCH 025/504] fix(bedrock): drop toolSpec.strict for Claude Sonnet 5 on Converse (#33196) Bedrock routes Claude Sonnet 5 through the same Anthropic-compatible validator as Opus 4.7/4.8 and Sonnet 4, which rejects toolSpec.strict with 'tools.0.custom.strict: Extra inputs are not permitted'. Set bedrock_converse_supports_strict_tools: false on all six Sonnet 5 entries so the existing gate strips the field, matching the fix shape of #31582 Co-authored-by: Yaroslav Budyanskiy --- ...odel_prices_and_context_window_backup.json | 6 +++++ model_prices_and_context_window.json | 6 +++++ ...edrock_converse_strict_tools_opus_47_48.py | 27 +++++++++++++++---- 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e5afc81b641..5e21a868729 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1759,6 +1759,7 @@ "prompt_cache_min_tokens": 2048 }, "anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -1795,6 +1796,7 @@ "prompt_cache_min_tokens": 1024 }, "global.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -1831,6 +1833,7 @@ "prompt_cache_min_tokens": 1024 }, "us.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -1867,6 +1870,7 @@ "prompt_cache_min_tokens": 1024 }, "eu.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -1903,6 +1907,7 @@ "prompt_cache_min_tokens": 1024 }, "au.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -1939,6 +1944,7 @@ "prompt_cache_min_tokens": 1024 }, "jp.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5cf99ba8bac..7587f71bffc 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1759,6 +1759,7 @@ "prompt_cache_min_tokens": 2048 }, "anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -1795,6 +1796,7 @@ "prompt_cache_min_tokens": 1024 }, "global.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -1831,6 +1833,7 @@ "prompt_cache_min_tokens": 1024 }, "us.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -1867,6 +1870,7 @@ "prompt_cache_min_tokens": 1024 }, "eu.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -1903,6 +1907,7 @@ "prompt_cache_min_tokens": 1024 }, "au.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -1939,6 +1944,7 @@ "prompt_cache_min_tokens": 1024 }, "jp.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py index 791982fc3dc..b02324af0a5 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py @@ -1,9 +1,9 @@ """Regression tests for Bedrock Converse ``toolSpec.strict`` forwarding. -Bedrock Converse routes Claude Opus 4.7/4.8 and Claude Sonnet 4 through an -Anthropic-compatible validator that rejects ``toolSpec.strict`` even though -Anthropic's native API accepts ``strict`` as a top-level tool field. See -BerriAI/litellm#31582. +Bedrock Converse routes Claude Opus 4.7/4.8, Claude Sonnet 4 and Claude +Sonnet 5 through an Anthropic-compatible validator that rejects +``toolSpec.strict`` even though Anthropic's native API accepts ``strict`` +as a top-level tool field. See BerriAI/litellm#31582. """ import pytest @@ -48,12 +48,18 @@ _STRICT_TOOL = [ "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0", "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0", "bedrock/apac.anthropic.claude-sonnet-4-20250514-v1:0", + "anthropic.claude-sonnet-5", + "bedrock/global.anthropic.claude-sonnet-5", + "bedrock/us.anthropic.claude-sonnet-5", + "bedrock/eu.anthropic.claude-sonnet-5", + "bedrock/au.anthropic.claude-sonnet-5", + "bedrock/jp.anthropic.claude-sonnet-5", ], ) def test_bedrock_tools_pt_strict_dropped_for_strict_unsupported_models( model_id: str, ) -> None: - """Opus 4.7/4.8 and Sonnet 4 reject toolSpec.strict and additionalProperties.""" + """Opus 4.7/4.8, Sonnet 4 and Sonnet 5 reject toolSpec.strict and additionalProperties.""" result = _bedrock_tools_pt(_STRICT_TOOL, model=model_id) tool_spec = result[0]["toolSpec"] assert ( @@ -129,6 +135,11 @@ def test_bedrock_converse_supports_strict_tools_helper() -> None: ) is False ) + assert bedrock_converse_supports_strict_tools("anthropic.claude-sonnet-5") is False + assert ( + bedrock_converse_supports_strict_tools("bedrock/us.anthropic.claude-sonnet-5") + is False + ) @pytest.mark.parametrize( @@ -143,6 +154,12 @@ def test_bedrock_converse_supports_strict_tools_helper() -> None: "us.anthropic.claude-sonnet-4-20250514-v1:0", "eu.anthropic.claude-sonnet-4-20250514-v1:0", "apac.anthropic.claude-sonnet-4-20250514-v1:0", + "anthropic.claude-sonnet-5", + "global.anthropic.claude-sonnet-5", + "us.anthropic.claude-sonnet-5", + "eu.anthropic.claude-sonnet-5", + "au.anthropic.claude-sonnet-5", + "jp.anthropic.claude-sonnet-5", ], ) def test_strict_tools_flag_set_in_model_cost_map(cost_map_key: str) -> None: From 3e4669dbc5d31a261e65af8e021d2ad12a2d15c4 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 30 Jul 2026 23:03:19 +0000 Subject: [PATCH 026/504] fix(cost): track OpenAI/Azure web search tool cost per call Adds search_context_cost_per_query pricing for the 82 OpenAI/Azure models that advertise supports_web_search but had none (gpt-5 family, o-series, deep-research at $0.01/call; gpt-4.1 at $0.025/call), so built-in web search is no longer billed as $0. Also counts web_search_call items in Responses output so N searches bill N times instead of once; usage-count providers (gemini, anthropic, xai, vertex) still route through get_cost_for_web_search_request and are unaffected. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_cost_calc/tool_call_cost_tracking.py | 20 +- ...odel_prices_and_context_window_backup.json | 410 ++++++++++++++++++ model_prices_and_context_window.json | 410 ++++++++++++++++++ .../test_tool_call_cost_tracking.py | 85 ++++ 4 files changed, 924 insertions(+), 1 deletion(-) 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 221b1ae6eab..1d8f2a6c965 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 @@ -117,10 +117,28 @@ class StandardBuiltInToolCostTracking: if result is not None: return result - return StandardBuiltInToolCostTracking.get_cost_for_web_search( + per_call_cost = StandardBuiltInToolCostTracking.get_cost_for_web_search( web_search_options=standard_built_in_tools_params.get("web_search_options", None), model_info=model_info, ) + return per_call_cost * StandardBuiltInToolCostTracking._count_web_search_calls(response_object) + + @staticmethod + def _count_web_search_calls(response_object: object) -> int: + """ + Number of web searches to bill for on the per-call pricing path. + + Providers that report a request count in usage (gemini, anthropic, xai, vertex) are handled by + get_cost_for_web_search_request and never reach here. This path prices per call, so it must count + the web_search_call items. Chat-completions responses only expose url_citation annotations with no + count, so they floor to a single billable search. + """ + if isinstance(response_object, ResponsesAPIResponse): + count = sum( + 1 for output_item in response_object.output if getattr(output_item, "type", None) == "web_search_call" + ) + return max(count, 1) + return 1 @staticmethod def _handle_file_search_cost( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5e21a868729..193d2bc2025 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5760,6 +5760,11 @@ "supports_vision": true }, "azure/gpt-5.2-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5791,6 +5796,11 @@ "supports_web_search": true }, "azure/gpt-5.2-pro-2025-12-11": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -6044,6 +6054,11 @@ "supports_vision": true }, "azure/gpt-5.4-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6079,6 +6094,11 @@ "supports_web_search": true }, "azure/gpt-5.4-pro-2026-03-05": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6114,6 +6134,11 @@ "supports_web_search": true }, "azure/gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6159,6 +6184,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6204,6 +6234,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -6249,6 +6284,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_above_272k_tokens": 2e-07, "cache_read_input_token_cost_priority": 2e-07, @@ -6294,6 +6334,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6336,6 +6381,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6378,6 +6428,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "cache_read_input_token_cost_priority": 6.875e-07, @@ -6420,6 +6475,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.1e-07, "cache_read_input_token_cost_above_272k_tokens": 2.2e-07, "cache_read_input_token_cost_priority": 2.75e-07, @@ -6462,6 +6522,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6504,6 +6569,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6546,6 +6616,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "cache_read_input_token_cost_priority": 6.875e-07, @@ -6588,6 +6663,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.1e-07, "cache_read_input_token_cost_above_272k_tokens": 2.2e-07, "cache_read_input_token_cost_priority": 2.75e-07, @@ -6630,6 +6710,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6675,6 +6760,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6717,6 +6807,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6759,6 +6854,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6801,6 +6901,11 @@ "supports_web_search": true }, "azure/us/gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6840,6 +6945,11 @@ "supports_web_search": true }, "azure/eu/gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6879,6 +6989,11 @@ "supports_web_search": true }, "azure/gpt-5.5-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6918,6 +7033,11 @@ "supports_low_reasoning_effort": false }, "azure/gpt-5.5-pro-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6953,6 +7073,11 @@ "supports_web_search": true }, "azure/gpt-5.4-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", @@ -6988,6 +7113,11 @@ "supports_xhigh_reasoning_effort": false }, "azure/gpt-5.4-mini-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", @@ -7023,6 +7153,11 @@ "supports_xhigh_reasoning_effort": false }, "azure/gpt-5.4-nano": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "azure", @@ -7058,6 +7193,11 @@ "supports_xhigh_reasoning_effort": false }, "azure/gpt-5.4-nano-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "azure", @@ -7521,6 +7661,11 @@ "supports_vision": true }, "azure/o3-deep-research": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_token": 1e-05, "litellm_provider": "azure", @@ -21452,6 +21597,11 @@ "supports_tool_choice": true }, "gpt-4.1": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_priority": 8.75e-07, "input_cost_per_token": 2e-06, @@ -21489,6 +21639,11 @@ "supports_web_search": true }, "gpt-4.1-2025-04-14": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -21523,6 +21678,11 @@ "supports_web_search": true }, "gpt-4.1-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_priority": 1.75e-07, "input_cost_per_token": 4e-07, @@ -21560,6 +21720,11 @@ "supports_web_search": true }, "gpt-4.1-mini-2025-04-14": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 4e-07, "input_cost_per_token_batches": 2e-07, @@ -22706,6 +22871,11 @@ "supports_pdf_input": true }, "gpt-5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, @@ -22748,6 +22918,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -22787,6 +22962,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-2025-11-13": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -22826,6 +23006,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-chat-latest": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -22865,6 +23050,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -22905,6 +23095,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-2025-12-11": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -22945,6 +23140,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-chat-latest": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -22983,6 +23183,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.3-chat-latest": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -23021,6 +23226,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -23055,6 +23265,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-pro-2025-12-11": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -23089,6 +23304,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, "cache_creation_input_token_cost_flex": 3.125e-06, @@ -23142,6 +23362,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, "cache_creation_input_token_cost_flex": 3.125e-06, @@ -23195,6 +23420,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 3.125e-06, "cache_creation_input_token_cost_above_272k_tokens": 6.25e-06, "cache_creation_input_token_cost_flex": 1.5625e-06, @@ -23248,6 +23478,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, "cache_creation_input_token_cost_flex": 6.25e-07, @@ -23301,6 +23536,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_flex": 2.5e-07, @@ -23350,6 +23590,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_flex": 2.5e-07, @@ -23399,6 +23644,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.5-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23444,6 +23694,11 @@ "supports_low_reasoning_effort": false }, "gpt-5.5-pro-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23582,6 +23837,11 @@ "supports_vision": true }, "gpt-5.4-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23626,6 +23886,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.4-pro-2026-03-05": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23670,6 +23935,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.4-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "cache_read_input_token_cost_priority": 1.5e-07, @@ -23716,6 +23986,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.4-mini-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "cache_read_input_token_cost_priority": 1.5e-07, @@ -23762,6 +24037,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.4-nano": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_flex": 1e-08, "input_cost_per_token": 2e-07, @@ -23805,6 +24085,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.4-nano-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_flex": 1e-08, "input_cost_per_token": 2e-07, @@ -23848,6 +24133,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", @@ -23884,6 +24174,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-pro-2025-10-06": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", @@ -23920,6 +24215,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-2025-08-07": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, @@ -24032,6 +24332,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -24066,6 +24371,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -24103,6 +24413,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-codex-max": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -24137,6 +24452,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-codex-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_priority": 4.5e-08, "input_cost_per_token": 2.5e-07, @@ -24174,6 +24494,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -24211,6 +24536,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.3-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -24248,6 +24578,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -24290,6 +24625,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-mini-2025-08-07": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -24332,6 +24672,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-nano": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, "input_cost_per_token": 5e-08, @@ -24372,6 +24717,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-nano-2025-08-07": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, "input_cost_per_token": 5e-08, @@ -28548,6 +28898,11 @@ "supports_vision": true }, "o3": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 8.75e-07, @@ -28586,6 +28941,11 @@ "supports_web_search": true }, "o3-2025-04-16": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "openai", @@ -28618,6 +28978,11 @@ "supports_web_search": true }, "o3-deep-research": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_token": 1e-05, "input_cost_per_token_batches": 5e-06, @@ -28652,6 +29017,11 @@ "supports_web_search": true }, "o3-deep-research-2025-06-26": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_token": 1e-05, "input_cost_per_token_batches": 5e-06, @@ -28720,6 +29090,11 @@ "supports_vision": false }, "o3-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -28751,6 +29126,11 @@ "supports_web_search": true }, "o3-pro-2025-06-10": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -28782,6 +29162,11 @@ "supports_web_search": true }, "o4-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_flex": 1.375e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -28807,6 +29192,11 @@ "supports_web_search": true }, "o4-mini-2025-04-16": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "openai", @@ -28826,6 +29216,11 @@ "supports_web_search": true }, "o4-mini-deep-research": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -28860,6 +29255,11 @@ "supports_web_search": true }, "o4-mini-deep-research-2025-06-26": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -43526,6 +43926,11 @@ ] }, "gpt-5-search-api": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -43548,6 +43953,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-search-api-2025-10-14": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7587f71bffc..35bf17675cb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5760,6 +5760,11 @@ "supports_vision": true }, "azure/gpt-5.2-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5791,6 +5796,11 @@ "supports_web_search": true }, "azure/gpt-5.2-pro-2025-12-11": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -6044,6 +6054,11 @@ "supports_vision": true }, "azure/gpt-5.4-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6079,6 +6094,11 @@ "supports_web_search": true }, "azure/gpt-5.4-pro-2026-03-05": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6114,6 +6134,11 @@ "supports_web_search": true }, "azure/gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6159,6 +6184,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6204,6 +6234,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-07, "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -6249,6 +6284,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_above_272k_tokens": 2e-07, "cache_read_input_token_cost_priority": 2e-07, @@ -6294,6 +6334,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6336,6 +6381,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6378,6 +6428,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "cache_read_input_token_cost_priority": 6.875e-07, @@ -6420,6 +6475,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.1e-07, "cache_read_input_token_cost_above_272k_tokens": 2.2e-07, "cache_read_input_token_cost_priority": 2.75e-07, @@ -6462,6 +6522,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6504,6 +6569,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, @@ -6546,6 +6616,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "cache_read_input_token_cost_priority": 6.875e-07, @@ -6588,6 +6663,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.1e-07, "cache_read_input_token_cost_above_272k_tokens": 2.2e-07, "cache_read_input_token_cost_priority": 2.75e-07, @@ -6630,6 +6710,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6675,6 +6760,11 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6717,6 +6807,11 @@ "supports_minimal_reasoning_effort": false }, "azure/eu/gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6759,6 +6854,11 @@ "supports_minimal_reasoning_effort": false }, "azure/gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, @@ -6801,6 +6901,11 @@ "supports_web_search": true }, "azure/us/gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6840,6 +6945,11 @@ "supports_web_search": true }, "azure/eu/gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.38e-06, @@ -6879,6 +6989,11 @@ "supports_web_search": true }, "azure/gpt-5.5-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6918,6 +7033,11 @@ "supports_low_reasoning_effort": false }, "azure/gpt-5.5-pro-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -6953,6 +7073,11 @@ "supports_web_search": true }, "azure/gpt-5.4-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", @@ -6988,6 +7113,11 @@ "supports_xhigh_reasoning_effort": false }, "azure/gpt-5.4-mini-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", @@ -7023,6 +7153,11 @@ "supports_xhigh_reasoning_effort": false }, "azure/gpt-5.4-nano": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "azure", @@ -7058,6 +7193,11 @@ "supports_xhigh_reasoning_effort": false }, "azure/gpt-5.4-nano-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "azure", @@ -7521,6 +7661,11 @@ "supports_vision": true }, "azure/o3-deep-research": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_token": 1e-05, "litellm_provider": "azure", @@ -21527,6 +21672,11 @@ "supports_tool_choice": true }, "gpt-4.1": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_priority": 8.75e-07, "input_cost_per_token": 2e-06, @@ -21564,6 +21714,11 @@ "supports_web_search": true }, "gpt-4.1-2025-04-14": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -21598,6 +21753,11 @@ "supports_web_search": true }, "gpt-4.1-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_priority": 1.75e-07, "input_cost_per_token": 4e-07, @@ -21635,6 +21795,11 @@ "supports_web_search": true }, "gpt-4.1-mini-2025-04-14": { + "search_context_cost_per_query": { + "search_context_size_high": 0.025, + "search_context_size_low": 0.025, + "search_context_size_medium": 0.025 + }, "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 4e-07, "input_cost_per_token_batches": 2e-07, @@ -22781,6 +22946,11 @@ "supports_pdf_input": true }, "gpt-5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, @@ -22823,6 +22993,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -22862,6 +23037,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-2025-11-13": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -22901,6 +23081,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-chat-latest": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -22940,6 +23125,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -22980,6 +23170,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-2025-12-11": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -23020,6 +23215,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-chat-latest": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -23058,6 +23258,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.3-chat-latest": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -23096,6 +23301,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -23130,6 +23340,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-pro-2025-12-11": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2.1e-05, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -23164,6 +23379,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.6": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, "cache_creation_input_token_cost_flex": 3.125e-06, @@ -23217,6 +23437,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-sol": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, "cache_creation_input_token_cost_flex": 3.125e-06, @@ -23270,6 +23495,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-terra": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 3.125e-06, "cache_creation_input_token_cost_above_272k_tokens": 6.25e-06, "cache_creation_input_token_cost_flex": 1.5625e-06, @@ -23323,6 +23553,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-luna": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, "cache_creation_input_token_cost_flex": 6.25e-07, @@ -23376,6 +23611,11 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.5": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_flex": 2.5e-07, @@ -23425,6 +23665,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.5-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_flex": 2.5e-07, @@ -23474,6 +23719,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.5-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23519,6 +23769,11 @@ "supports_low_reasoning_effort": false }, "gpt-5.5-pro-2026-04-23": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23657,6 +23912,11 @@ "supports_vision": true }, "gpt-5.4-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23701,6 +23961,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.4-pro-2026-03-05": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, "input_cost_per_token": 3e-05, @@ -23745,6 +24010,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.4-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "cache_read_input_token_cost_priority": 1.5e-07, @@ -23791,6 +24061,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.4-mini-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, "cache_read_input_token_cost_priority": 1.5e-07, @@ -23837,6 +24112,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.4-nano": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_flex": 1e-08, "input_cost_per_token": 2e-07, @@ -23880,6 +24160,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5.4-nano-2026-03-17": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_flex": 1e-08, "input_cost_per_token": 2e-07, @@ -23923,6 +24208,11 @@ "supports_minimal_reasoning_effort": false }, "gpt-5-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", @@ -23959,6 +24249,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-pro-2025-10-06": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", @@ -23995,6 +24290,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-2025-08-07": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, @@ -24107,6 +24407,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -24141,6 +24446,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -24178,6 +24488,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-codex-max": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -24212,6 +24527,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.1-codex-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_priority": 4.5e-08, "input_cost_per_token": 2.5e-07, @@ -24249,6 +24569,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.2-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -24286,6 +24611,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.3-codex": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, "input_cost_per_token": 1.75e-06, @@ -24323,6 +24653,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -24365,6 +24700,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-mini-2025-08-07": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, @@ -24407,6 +24747,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-nano": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, "input_cost_per_token": 5e-08, @@ -24447,6 +24792,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-nano-2025-08-07": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, "input_cost_per_token": 5e-08, @@ -28623,6 +28973,11 @@ "supports_vision": true }, "o3": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 8.75e-07, @@ -28661,6 +29016,11 @@ "supports_web_search": true }, "o3-2025-04-16": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "openai", @@ -28693,6 +29053,11 @@ "supports_web_search": true }, "o3-deep-research": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_token": 1e-05, "input_cost_per_token_batches": 5e-06, @@ -28727,6 +29092,11 @@ "supports_web_search": true }, "o3-deep-research-2025-06-26": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_token": 1e-05, "input_cost_per_token_batches": 5e-06, @@ -28795,6 +29165,11 @@ "supports_vision": false }, "o3-pro": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -28826,6 +29201,11 @@ "supports_web_search": true }, "o3-pro-2025-06-10": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -28857,6 +29237,11 @@ "supports_web_search": true }, "o4-mini": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_flex": 1.375e-07, "cache_read_input_token_cost_priority": 5e-07, @@ -28882,6 +29267,11 @@ "supports_web_search": true }, "o4-mini-2025-04-16": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "openai", @@ -28901,6 +29291,11 @@ "supports_web_search": true }, "o4-mini-deep-research": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -28935,6 +29330,11 @@ "supports_web_search": true }, "o4-mini-deep-research-2025-06-26": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -43647,6 +44047,11 @@ ] }, "gpt-5-search-api": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -43669,6 +44074,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-search-api-2025-10-14": { + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 24fd3c94ee3..f194db7676a 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -602,5 +602,90 @@ def test_web_search_provider_prefix_fallback_does_not_misprice_non_gemini_model( ) +def _openai_responses_with_web_search_calls(model, num_calls): + from litellm.types.llms.openai import ResponsesAPIResponse + from openai.types.responses.response_function_web_search import ( + ActionSearch, + ResponseFunctionWebSearch, + ) + + output = [ + ResponseFunctionWebSearch( + id=f"ws_{i}", + type="web_search_call", + status="completed", + action=ActionSearch(type="search", query="latest news"), + ) + for i in range(num_calls) + ] + return ResponsesAPIResponse( + id="resp_1", + created_at=0, + model=model, + object="response", + output=output, + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + + +def test_openai_responses_web_search_priced_per_call(local_model_cost_map): + """ + Regression for LIT-5013 bug 1: OpenAI reasoning models (gpt-5 family, o-series, deep-research) + carry supports_web_search but had no search_context_cost_per_query, so get_cost_for_web_search_request + (no openai branch) returned None and the default fallback billed web search as $0. gpt-5-nano now + prices at $0.01 per call, and two web_search_call items in the Responses output must bill 2 x $0.01. + """ + from litellm.types.utils import Usage + + model = "gpt-5-nano" + per_call = litellm.get_model_info(model)["search_context_cost_per_query"][ + "search_context_size_medium" + ] + assert per_call == 0.01 + + response = _openai_responses_with_web_search_calls(model, num_calls=2) + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=response, + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert cost == pytest.approx(2 * per_call), ( + f"gpt-5-nano web search must bill 2 x ${per_call}, got ${cost}" + ) + + +def test_openai_responses_web_search_multiplied_by_call_count(local_model_cost_map): + """ + Regression for LIT-5013 bug 2: web_search_call detection was binary, so a Responses output with + multiple web searches was charged once. gpt-4o-search-preview carries per-call pricing; N calls + must bill N times, and a single call must still bill exactly once. + """ + from litellm.types.utils import Usage + + model = "gpt-4o-search-preview" + per_call = litellm.get_model_info(model)["search_context_cost_per_query"][ + "search_context_size_medium" + ] + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + + for num_calls in (1, 3): + response = _openai_responses_with_web_search_calls(model, num_calls=num_calls) + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=response, + usage=usage, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + assert cost == pytest.approx(num_calls * per_call), ( + f"{num_calls} web searches must bill {num_calls} x ${per_call}, got ${cost}" + ) + + # Note: File search integration test removed due to complex annotation detection logic # The unit tests in test_azure_assistant_cost_tracking.py provide comprehensive coverage From ef614b7b5bcdf94472b876bea64ad17e5dfd4282 Mon Sep 17 00:00:00 2001 From: milan Date: Mon, 3 Aug 2026 14:35:47 +0000 Subject: [PATCH 027/504] refactor(vertex_ai): keep the batch embeddings translation within the LIT002 ceiling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 118 +++++++++++------- 1 file changed, 71 insertions(+), 47 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 9363540fe1b..bbb97a1edc2 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -74,11 +74,11 @@ _GCP_LABEL_VALUE_MAX_LEN = 63 _CUSTOM_ID_RAW_LABEL_PREFIX = "b32_" _VERTEX_BATCH_KEY_FIELD = "key" _MANAGED_GCS_MODEL_PATH_PATTERN = re.compile(r"publishers/[^/]+/models/([^/?]+)") -_EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = { - "outputDimensionality": "output_dimensionality", - "taskType": "task_type", - "title": "title", -} +_EMBED_REQUEST_FIELD_BY_GEMINI_PARAM = ( + ("outputDimensionality", "output_dimensionality"), + ("taskType", "task_type"), + ("title", "title"), +) _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN = re.compile(r"(?P[^#]*)#(?P\d+)/(?P\d+)") @@ -153,12 +153,15 @@ def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, Any]) -> str: key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) if key is not None: return unquote(str(key)) - request_data = vertex_output_row.get("request") or {} - return _get_litellm_batch_custom_id_from_labels(request_data.get("labels") or {}) + request_data = vertex_output_row.get("request") + labels = request_data.get("labels") if isinstance(request_data, Mapping) else None + return _get_litellm_batch_custom_id_from_labels(labels) -def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str: +def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, Any] | None) -> str: """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" + if not labels: + return "unknown" raw = labels.get("litellm_custom_id_raw") if raw: raw_chunks = [str(raw)] @@ -195,7 +198,8 @@ def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, Any]) def _openai_batch_output_row( custom_id: str, body: Mapping[str, Any] | None = None, - error: Mapping[str, str] | None = None, + error_code: str | None = None, + error_message: str = "", ) -> Mapping[str, Any]: """ One row of an OpenAI batch output file. Per the OpenAI Batch spec, failed rows set @@ -211,7 +215,7 @@ def _openai_batch_output_row( "request_id": body.get("id", ""), "body": body, }, - "error": error, + "error": None if error_code is None else {"code": error_code, "message": error_message}, } @@ -233,6 +237,19 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, Any]) -> tuple[str, return unquote(match["custom_id"]), int(match["index"]) +def _embedding_prompt_token_count(vertex_response: Mapping[str, Any]) -> int: + """ + Prompt tokens billed for one Vertex Gemini Embedding batch row. + + Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as + a fallback. + """ + usage_metadata = vertex_response.get("usageMetadata") + if isinstance(usage_metadata, Mapping): + return int(usage_metadata.get("promptTokenCount") or 0) + return int(vertex_response.get("tokenCount") or 0) + + def _vertex_embeddings_rows_to_openai_batch_output_row( custom_id: str, vertex_output_rows: tuple[Mapping[str, Any], ...], @@ -247,23 +264,19 @@ def _vertex_embeddings_rows_to_openai_batch_output_row( An entry that asked for several embeddings at once maps to several rows here, which become the indexed elements of a single `data` array. One failed element fails the - whole entry, since an OpenAI batch row is either a response or an error. Live rows - report usage under `usageMetadata`; the documented `tokenCount` is kept as a - fallback. Rows carry no `modelVersion`, so the model comes from the batch they - belong to. + whole entry, since an OpenAI batch row is either a response or an error. Rows carry + no `modelVersion`, so the model comes from the batch they belong to. """ status = next((row["status"] for row in vertex_output_rows if row.get("status")), "") if status: return _openai_batch_output_row( custom_id=custom_id, - error={"code": "vertex_ai_error", "message": status}, + error_code="vertex_ai_error", + error_message=status, ) - responses = tuple(row.get("response") or {} for row in vertex_output_rows) - token_count = sum( - int((response.get("usageMetadata") or {}).get("promptTokenCount") or response.get("tokenCount") or 0) - for response in responses - ) + responses = tuple(row["response"] for row in vertex_output_rows) + token_count = sum(_embedding_prompt_token_count(response) for response in responses) body = EmbeddingResponse( model=model or "", data=[ @@ -361,6 +374,27 @@ def _vertex_batch_embeddings_key(custom_id: str, index: int, total: int) -> str: return encoded_custom_id if total < 2 else f"{encoded_custom_id}#{index}/{total}" +def _vertex_embeddings_row(key: str | None, embed_content_request: Mapping[str, Any]) -> Mapping[str, Any]: + """ + One Vertex Gemini Embedding batch input row. + + The config fields live inside the `EmbedContentRequest` under their snake_case batch + names, and the OpenAI `custom_id` rides along in the top-level `key` that Vertex + echoes back. + """ + request = { + "content": embed_content_request["content"], + **{ + request_field: embed_content_request[gemini_param] + for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM + if gemini_param in embed_content_request + }, + } + if key is None: + return {"request": request} + return {_VERTEX_BATCH_KEY_FIELD: key, "request": request} + + def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( openai_entry: Mapping[str, Any], ) -> tuple[Mapping[str, Any], ...]: @@ -381,7 +415,9 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( API Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings """ - openai_request_body = openai_entry.get("body") or {} + openai_request_body = openai_entry.get("body") + if not isinstance(openai_request_body, dict): + raise ValueError("`body` is required on /v1/embeddings batch requests, but was not provided") embedding_input = openai_request_body.get("input") if embedding_input is None: raise ValueError("`input` is required on /v1/embeddings batch requests, but was not provided") @@ -400,27 +436,16 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( ) custom_id = openai_entry.get("custom_id") return tuple( - { - **( - {} - if custom_id is None - else { - _VERTEX_BATCH_KEY_FIELD: _vertex_batch_embeddings_key( - custom_id=str(custom_id), - index=index, - total=len(embed_content_requests), - ) - } + _vertex_embeddings_row( + key=None + if custom_id is None + else _vertex_batch_embeddings_key( + custom_id=str(custom_id), + index=index, + total=len(embed_content_requests), ), - "request": { - "content": embed_content_request["content"], - **{ - request_field: embed_content_request[gemini_param] - for gemini_param, request_field in _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM.items() - if gemini_param in embed_content_request - }, - }, - } + embed_content_request=embed_content_request, + ) for index, embed_content_request in enumerate(embed_content_requests) ) @@ -1008,7 +1033,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): request=httpx.Request(method="POST", url="https://example.com"), ) - all_lines = itertools.chain([first_line], lines) + all_lines = itertools.chain((first_line,), lines) # Embedding rows are grouped by `custom_id` rather than transformed one at a # time, since an entry that asked for several embeddings comes back as @@ -1064,7 +1089,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if has_error: return _openai_batch_output_row( custom_id=custom_id, - error={"code": "vertex_ai_error", "message": status}, + error_code="vertex_ai_error", + error_message=status, ) # Transform successful response using existing transformation @@ -1096,8 +1122,6 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): except Exception as e: return _openai_batch_output_row( custom_id=custom_id, - error={ - "code": "transformation_error", - "message": f"Failed to transform response: {e!s}", - }, + error_code="transformation_error", + error_message=f"Failed to transform response: {e!s}", ) From db061d6e3118806dd59340e900af482b15450326 Mon Sep 17 00:00:00 2001 From: ayaangazali Date: Thu, 23 Jul 2026 16:25:13 -0700 Subject: [PATCH 028/504] fix(azure_ai): strip non-OpenAI-spec message fields before request --- litellm/llms/azure_ai/chat/transformation.py | 21 +++++- .../chat/test_azure_ai_transformation.py | 72 +++++++++++++++++++ 2 files changed, 91 insertions(+), 2 deletions(-) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 5540d79f667..683c05cfbaa 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -11,6 +11,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( _audio_or_image_in_message_content, convert_content_list_to_str, + filter_value_from_dict, ) from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj @@ -28,6 +29,13 @@ class AzureFoundryErrorStrings(str, enum.Enum): SET_EXTRA_PARAMETERS_TO_PASS_THROUGH = "Set extra-parameters to 'pass-through'" +NON_OPENAI_SPEC_MESSAGE_FIELDS = ( + "thinking_blocks", + "provider_specific_fields", + "cache_control", +) + + class AzureAIStudioConfig(OpenAIConfig): def get_supported_openai_params(self, model: str) -> list: model_supports_tool_choice = True # azure ai supports this by default @@ -167,10 +175,19 @@ class AzureAIStudioConfig(OpenAIConfig): ) -> list: """ - Azure AI Studio doesn't support content as a list. This handles: - 1. Transforms list content to a string. - 2. If message contains an image or audio, send as is (user-intended) + 1. Strips message fields that are not part of the OpenAI chat-completions + schema (thinking_blocks, provider_specific_fields, cache_control). + Azure AI Foundry backends set additionalProperties=false and reject + these with "Extra inputs are not permitted", which breaks multi-turn + Anthropic-format clients that echo thinking blocks back as history. + 2. Transforms list content to a string. + 3. If message contains an image or audio, send as is (user-intended) """ for message in messages: + message_dict = cast(dict, message) # cast-ok: TypedDict is a runtime dict stripped in place + for field in NON_OPENAI_SPEC_MESSAGE_FIELDS: + filter_value_from_dict(message_dict, field) + # Do nothing if the message contains an image or audio if _audio_or_image_in_message_content(message): continue diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 2e75039139c..80c4355b560 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -262,3 +262,75 @@ def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name(): assert "copilot_mcp_server_name" not in tool assert result["tools"][0]["type"] == "function" assert result["tools"][1]["function"]["name"] == "read_file" + + +def _find_key_anywhere(obj, key: str) -> bool: + if isinstance(obj, dict): + if key in obj: + return True + return any(_find_key_anywhere(v, key) for v in obj.values()) + if isinstance(obj, list): + return any(_find_key_anywhere(item, key) for item in obj) + return False + + +def test_azure_ai_strips_non_openai_spec_message_fields(): + """ + Regression for https://github.com/BerriAI/litellm/issues/33961. + + Azure AI Foundry backends set additionalProperties=false, so any message + field outside the OpenAI chat-completions schema causes a 400 "Extra inputs + are not permitted". Anthropic-format clients (e.g. Claude Code) echo prior + assistant turns back as history carrying thinking_blocks, a nested thought + signature at tool_calls[].function.provider_specific_fields, and Anthropic + cache_control annotations. transform_request must strip all of these before + the request reaches the upstream. + """ + config = AzureAIStudioConfig() + + messages = [ + {"role": "user", "content": "Read a file."}, + { + "role": "assistant", + "content": "I can help.", + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "The user wants me to read a file.", + "signature": "", + "cache_control": {"type": "ephemeral"}, + } + ], + "provider_specific_fields": {"thought_signature": "sig-top"}, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "read_file", + "arguments": "{}", + "provider_specific_fields": {"thought_signature": "sig-nested"}, + }, + } + ], + }, + {"role": "user", "content": "go ahead"}, + ] + + request = config.transform_request( + model="fw-glm-5.2", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + transformed_messages = request["messages"] + + assert not _find_key_anywhere(transformed_messages, "thinking_blocks") + assert not _find_key_anywhere(transformed_messages, "provider_specific_fields") + assert not _find_key_anywhere(transformed_messages, "cache_control") + + assistant_message = transformed_messages[1] + assert assistant_message["content"] == "I can help." + assert assistant_message["tool_calls"][0]["function"]["name"] == "read_file" From 95bc890fcfc157247072752e4005686dd9413a54 Mon Sep 17 00:00:00 2001 From: ayaangazali Date: Thu, 30 Jul 2026 11:42:03 -0700 Subject: [PATCH 029/504] fix(azure_ai): strip non-spec message fields on a copy, not the caller's messages --- litellm/llms/azure_ai/chat/transformation.py | 7 ++- .../chat/test_azure_ai_transformation.py | 50 +++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 683c05cfbaa..067c89214ab 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -1,3 +1,4 @@ +import copy import enum import re from typing import Any, Final, cast @@ -182,9 +183,13 @@ class AzureAIStudioConfig(OpenAIConfig): Anthropic-format clients that echo thinking blocks back as history. 2. Transforms list content to a string. 3. If message contains an image or audio, send as is (user-intended) + + Operates on a deep copy so the caller's messages keep their thinking blocks + and provider metadata, which a fallback to another provider still needs. """ + messages = copy.deepcopy(messages) for message in messages: - message_dict = cast(dict, message) # cast-ok: TypedDict is a runtime dict stripped in place + message_dict = cast(dict, message) # cast-ok: TypedDict is a runtime dict stripped on our copy for field in NON_OPENAI_SPEC_MESSAGE_FIELDS: filter_value_from_dict(message_dict, field) diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 80c4355b560..beb7e9dfab0 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -334,3 +334,53 @@ def test_azure_ai_strips_non_openai_spec_message_fields(): assistant_message = transformed_messages[1] assert assistant_message["content"] == "I can help." assert assistant_message["tool_calls"][0]["function"]["name"] == "read_file" + + +def test_azure_ai_stripping_does_not_mutate_caller_messages(): + """ + The stripping must not touch the caller's messages. LiteLLM reuses the same + message objects when falling back to another provider, so stripping in place + would hand the fallback a conversation history with its thinking blocks and + provider metadata already destroyed. + """ + config = AzureAIStudioConfig() + + messages = [ + {"role": "user", "content": "Read a file."}, + { + "role": "assistant", + "content": "I can help.", + "thinking_blocks": [ + {"type": "thinking", "thinking": "Reading the file.", "signature": "sig"} + ], + "provider_specific_fields": {"thought_signature": "sig-top"}, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "read_file", + "arguments": "{}", + "provider_specific_fields": {"thought_signature": "sig-nested"}, + }, + } + ], + }, + ] + + request = config.transform_request( + model="fw-glm-5.2", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert not _find_key_anywhere(request["messages"], "thinking_blocks") + + original_assistant = messages[1] + assert original_assistant["thinking_blocks"][0]["thinking"] == "Reading the file." + assert original_assistant["provider_specific_fields"] == {"thought_signature": "sig-top"} + assert original_assistant["tool_calls"][0]["function"]["provider_specific_fields"] == { + "thought_signature": "sig-nested" + } From 0c0e1e8374d7e956e65d275ebf5f2f832ec374b9 Mon Sep 17 00:00:00 2001 From: Miles Adkins Date: Wed, 5 Aug 2026 13:21:15 -0500 Subject: [PATCH 030/504] feat(fireworks_ai): translate NIM/vLLM extra params to Fireworks-native args Requests migrated from NIM/vLLM servers carry extras that flow through the extra_body passthrough verbatim, but the Fireworks chat completions API either names them differently or does not accept them at all. Add FireworksAIConfig.map_extra_body_params, invoked from the fireworks chat dispatch, which renames truncate_prompt_tokens to prompt_truncate_len, maps chat_template_kwargs.enable_thinking to reasoning_effort, converts guided_json/guided_grammar/guided_choice to response_format, and drops the remaining extras (min_tokens, stop_token_ids, skip_special_tokens, guided_regex, etc.) with a debug log. Alias and competing-constraint combinations raise BadRequestError. Unrecognized extras keep passing through untouched, as do fireworks-native params like top_k. --- .../llms/fireworks_ai/chat/transformation.py | 162 +++++++++++- litellm/main.py | 6 +- .../test_fireworks_ai_chat_transformation.py | 249 ++++++++++++++++++ 3 files changed, 415 insertions(+), 2 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index a796aa47b70..4740e84d513 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -1,5 +1,5 @@ import json -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Mapping from typing import Any, Final, Literal, cast import httpx @@ -61,6 +61,36 @@ def _extract_fireworks_hidden_params(payload: dict) -> dict: return {**top_level, **per_choice} +def _json_schema_response_format(schema: object) -> Mapping[str, object]: + return {"type": "json_schema", "json_schema": {"schema": schema}} # mutable-ok: JSON request body + + +_NIM_VLLM_STRIP_PARAMS: Final = frozenset( + { + "min_tokens", + "stop_token_ids", + "include_stop_str_in_output", + "skip_special_tokens", + "spaces_between_special_tokens", + "best_of", + "use_beam_search", + "guided_decoding_backend", + "guided_regex", + "add_generation_prompt", + "continue_final_message", + "add_special_tokens", + "detokenize", + "allowed_token_ids", + "bad_words", + } +) + +_EXTRA_BODY_CONSUMED_PARAMS: Final = ( + frozenset({"truncate_prompt_tokens", "chat_template_kwargs", "guided_json", "guided_grammar", "guided_choice"}) + | _NIM_VLLM_STRIP_PARAMS +) + + class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): """ Reference: https://docs.fireworks.ai/api-reference/post-chatcompletions @@ -273,6 +303,136 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): return optional_params + def map_extra_body_params(self, optional_params: Mapping[str, object], model: str) -> dict: # noqa: LIT001 # http handler pops extra_body off the returned dict + extra_body: Final = optional_params.get("extra_body") + if not isinstance(extra_body, dict): + return dict(optional_params) # mutable-ok: JSON request body + + self._validate_extra_body_conflicts(extra_body=extra_body, optional_params=optional_params, model=model) + stripped: Final = tuple(sorted(k for k in extra_body if k in _NIM_VLLM_STRIP_PARAMS)) + if stripped: + verbose_logger.debug( + "fireworks_ai does not support NIM/vLLM params %s for model=%s; dropping them from the request.", + stripped, + model, + ) + promoted: Final = ( + *self._translate_truncate_prompt_tokens(extra_body), + *self._translate_chat_template_kwargs(extra_body, model), + *self._translate_guided_params(extra_body), + ) + remaining: Final = tuple((k, v) for k, v in extra_body.items() if k not in _EXTRA_BODY_CONSUMED_PARAMS) + base: Final = {k: v for k, v in optional_params.items() if k != "extra_body"} # mutable-ok: JSON request body + return { # mutable-ok: JSON request body + **base, + **dict(promoted), # mutable-ok: JSON request body + **({"extra_body": dict(remaining)} if remaining else {}), # mutable-ok: JSON request body + } + + def _validate_extra_body_conflicts( + self, extra_body: Mapping[str, object], optional_params: Mapping[str, object], model: str + ) -> None: + if "truncate_prompt_tokens" in extra_body and ( + "prompt_truncate_len" in extra_body or "prompt_truncate_len" in optional_params + ): + raise litellm.BadRequestError( + message=( + "Fireworks AI chat completions received both `truncate_prompt_tokens` and " + "`prompt_truncate_len`; they are aliases, send only one." + ), + model=model, + llm_provider="fireworks_ai", + ) + chat_template_kwargs: Final = extra_body.get("chat_template_kwargs") + if ( + isinstance(chat_template_kwargs, dict) + and "enable_thinking" in chat_template_kwargs + and ("reasoning_effort" in optional_params or "thinking" in optional_params) + ): + raise litellm.BadRequestError( + message=( + "Fireworks AI chat completions does not support specifying both " + "`chat_template_kwargs.enable_thinking` and `reasoning_effort`/`thinking` in the same request." + ), + model=model, + llm_provider="fireworks_ai", + ) + guided_params: Final = tuple( + k for k in ("guided_json", "guided_grammar", "guided_choice") if extra_body.get(k) is not None + ) + if len(guided_params) > 1: + raise litellm.BadRequestError( + message=( + f"Fireworks AI chat completions received multiple guided decoding params " + f"{guided_params}; send only one." + ), + model=model, + llm_provider="fireworks_ai", + ) + if guided_params and "response_format" in optional_params: + raise litellm.BadRequestError( + message=( + f"Fireworks AI chat completions received both `{guided_params[0]}` and " + "`response_format`; they are competing output constraints, send only one." + ), + model=model, + llm_provider="fireworks_ai", + ) + + @staticmethod + def _translate_truncate_prompt_tokens(extra_body: Mapping[str, object]) -> tuple[tuple[str, object], ...]: + if extra_body.get("truncate_prompt_tokens") is None: + return () + return (("prompt_truncate_len", extra_body["truncate_prompt_tokens"]),) + + def _translate_chat_template_kwargs( + self, extra_body: Mapping[str, object], model: str + ) -> tuple[tuple[str, object], ...]: + chat_template_kwargs: Final = extra_body.get("chat_template_kwargs") + if chat_template_kwargs is None: + return () + if not isinstance(chat_template_kwargs, dict): + raise litellm.BadRequestError( + message="Fireworks AI chat completions expects `chat_template_kwargs` to be an object.", + model=model, + llm_provider="fireworks_ai", + ) + other_keys: Final = tuple(sorted(k for k in chat_template_kwargs if k != "enable_thinking")) + if other_keys: + verbose_logger.debug( + "fireworks_ai does not support chat_template_kwargs keys %s for model=%s; dropping them.", + other_keys, + model, + ) + if "enable_thinking" not in chat_template_kwargs: + return () + if not supports_reasoning(model=model, custom_llm_provider="fireworks_ai"): + verbose_logger.debug( + "fireworks_ai model %r does not support reasoning; dropping chat_template_kwargs.enable_thinking.", + model, + ) + return () + effort: Final = "medium" if chat_template_kwargs["enable_thinking"] else "none" + return (("reasoning_effort", effort),) + + @staticmethod + def _translate_guided_params(extra_body: Mapping[str, object]) -> tuple[tuple[str, object], ...]: + if extra_body.get("guided_json") is not None: + return (("response_format", _json_schema_response_format(extra_body["guided_json"])),) + if extra_body.get("guided_grammar") is not None: + grammar_response_format: Final = { # mutable-ok: JSON request body + "type": "grammar", + "grammar": extra_body["guided_grammar"], + } + return (("response_format", grammar_response_format),) + if extra_body.get("guided_choice") is not None: + choice_schema: Final = { # mutable-ok: JSON request body + "type": "string", + "enum": extra_body["guided_choice"], + } + return (("response_format", _json_schema_response_format(choice_schema)),) + return () + def _transform_tools(self, tools: list[OpenAIChatCompletionToolParam]) -> list[OpenAIChatCompletionToolParam]: for tool in tools: if tool.get("type") != "function": diff --git a/litellm/main.py b/litellm/main.py index f906c78f9ae..660e024b113 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1711,11 +1711,15 @@ def _complete_fireworks_ai( messages: Final = ctx.messages model: Final = ctx.model model_response: Final = ctx.model_response - optional_params: Final = ctx.optional_params provider_config: Final = ctx.provider_config shared_session: Final = ctx.shared_session stream: Final = ctx.stream timeout: Final = ctx.timeout + optional_params: Final = ( + provider_config.map_extra_body_params(optional_params=ctx.optional_params, model=model) + if isinstance(provider_config, litellm.FireworksAIConfig) + else ctx.optional_params + ) try: response: Final = base_llm_http_handler.completion( diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 94945ed4bfb..3fbbc70916a 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1282,3 +1282,252 @@ def test_streaming_surfaces_fireworks_response_fields(): assert surfaced["fireworks_raw_outputs"] == [raw_output] assert surfaced["fireworks_perf_metrics"] == {"prompt-tokens": 5} assert surfaced["fireworks_prompt_token_ids"] == [1, 2, 3] + + +def test_map_extra_body_params_translates_truncate_prompt_tokens(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"truncate_prompt_tokens": 4096}}, _REASONING_MODEL + ) + assert result == {"prompt_truncate_len": 4096} + + +def test_map_extra_body_params_truncate_prompt_tokens_conflicts_with_alias(): + config = FireworksAIConfig() + with pytest.raises(litellm.BadRequestError, match="aliases"): + config.map_extra_body_params( + {"prompt_truncate_len": 2048, "extra_body": {"truncate_prompt_tokens": 4096}}, + _REASONING_MODEL, + ) + with pytest.raises(litellm.BadRequestError, match="aliases"): + config.map_extra_body_params( + {"extra_body": {"truncate_prompt_tokens": 4096, "prompt_truncate_len": 2048}}, + _REASONING_MODEL, + ) + + +def test_map_extra_body_params_chat_template_kwargs_enable_thinking(): + config = FireworksAIConfig() + disabled = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": False}}}, + _REASONING_MODEL, + ) + assert disabled == {"reasoning_effort": "none"} + + enabled = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": True}}}, + _REASONING_MODEL, + ) + assert enabled == {"reasoning_effort": "medium"} + + +def test_map_extra_body_params_chat_template_kwargs_conflicts_with_reasoning_effort(): + config = FireworksAIConfig() + with pytest.raises(litellm.BadRequestError, match="enable_thinking"): + config.map_extra_body_params( + { + "reasoning_effort": "high", + "extra_body": {"chat_template_kwargs": {"enable_thinking": False}}, + }, + _REASONING_MODEL, + ) + + +def test_map_extra_body_params_chat_template_kwargs_conflicts_with_thinking(): + config = FireworksAIConfig() + with pytest.raises(litellm.BadRequestError, match="enable_thinking"): + config.map_extra_body_params( + { + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "extra_body": {"chat_template_kwargs": {"enable_thinking": True}}, + }, + _REASONING_MODEL, + ) + + +def test_map_extra_body_params_chat_template_kwargs_dropped_for_non_reasoning_model(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": False, "custom_flag": 1}}}, + _NON_REASONING_MODEL, + ) + assert result == {} + + +def test_map_extra_body_params_guided_json(): + config = FireworksAIConfig() + schema = {"type": "object", "properties": {"x": {"type": "string"}}} + result = config.map_extra_body_params( + {"extra_body": {"guided_json": schema}}, _REASONING_MODEL + ) + assert result == { + "response_format": {"type": "json_schema", "json_schema": {"schema": schema}} + } + + +def test_map_extra_body_params_guided_grammar_and_choice(): + config = FireworksAIConfig() + grammar = config.map_extra_body_params( + {"extra_body": {"guided_grammar": "root ::= 'hello'"}}, _REASONING_MODEL + ) + assert grammar == { + "response_format": {"type": "grammar", "grammar": "root ::= 'hello'"} + } + + choice = config.map_extra_body_params( + {"extra_body": {"guided_choice": ["yes", "no"]}}, _REASONING_MODEL + ) + assert choice == { + "response_format": { + "type": "json_schema", + "json_schema": {"schema": {"type": "string", "enum": ["yes", "no"]}}, + } + } + + +def test_map_extra_body_params_guided_conflicts_with_response_format(): + config = FireworksAIConfig() + with pytest.raises(litellm.BadRequestError, match="response_format"): + config.map_extra_body_params( + { + "response_format": {"type": "json_object"}, + "extra_body": {"guided_json": {"type": "object"}}, + }, + _REASONING_MODEL, + ) + + +def test_map_extra_body_params_multiple_guided_params_rejected(): + config = FireworksAIConfig() + with pytest.raises(litellm.BadRequestError, match="multiple guided decoding params"): + config.map_extra_body_params( + {"extra_body": {"guided_json": {"type": "object"}, "guided_grammar": "root ::= 'x'"}}, + _REASONING_MODEL, + ) + + +@pytest.mark.parametrize( + "param,value", + [ + ("min_tokens", 10), + ("stop_token_ids", [1, 2]), + ("include_stop_str_in_output", True), + ("skip_special_tokens", False), + ("spaces_between_special_tokens", True), + ("best_of", 2), + ("use_beam_search", True), + ("guided_decoding_backend", "outlines"), + ("guided_regex", "[0-9]+"), + ("add_generation_prompt", True), + ("continue_final_message", True), + ("add_special_tokens", False), + ("detokenize", True), + ("allowed_token_ids", [1]), + ("bad_words", ["foo"]), + ], +) +def test_map_extra_body_params_strips_unsupported_nim_vllm_params(param, value, caplog): + import logging + + config = FireworksAIConfig() + with caplog.at_level(logging.DEBUG): + result = config.map_extra_body_params( + {"extra_body": {param: value}}, _REASONING_MODEL + ) + assert result == {} + assert param in caplog.text + + +def test_map_extra_body_params_preserves_unknown_passthrough(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"top_k": 40, "some_future_param": "x", "truncate_prompt_tokens": 100}}, + _REASONING_MODEL, + ) + assert result == { + "prompt_truncate_len": 100, + "extra_body": {"top_k": 40, "some_future_param": "x"}, + } + + +def test_map_extra_body_params_no_extra_body(): + config = FireworksAIConfig() + assert config.map_extra_body_params({}, _REASONING_MODEL) == {} + unchanged = {"temperature": 0.5, "extra_body": None} + assert config.map_extra_body_params(unchanged, _REASONING_MODEL) == unchanged + + +def test_nim_vllm_extras_translated_end_to_end_in_request_body(): + """ + Passing NIM/vLLM extras to litellm.completion must reach the Fireworks + request body translated, not verbatim: truncate_prompt_tokens becomes + prompt_truncate_len, chat_template_kwargs.enable_thinking becomes + reasoning_effort, min_tokens is dropped, and fireworks-native top_k still + passes through. Asserts on the actual JSON posted to the API, so a revert + of the _complete_fireworks_ai wiring fails this test. + """ + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + model = "accounts/fireworks/models/glm-5p1" + body = { + "id": "chat-1", + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hi"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + raw_response = MagicMock() + raw_response.status_code = 200 + raw_response.headers = {} + raw_response.text = json.dumps(body) + raw_response.json = lambda: body + + client = HTTPHandler() + with patch.object(client, "post", return_value=raw_response) as mock_post: + litellm.completion( + model=f"fireworks_ai/{model}", + messages=[{"role": "user", "content": "hi"}], + api_key="fw-test-key", + client=client, + truncate_prompt_tokens=4096, + chat_template_kwargs={"enable_thinking": False}, + min_tokens=10, + top_k=40, + ) + + request_body = json.loads(mock_post.call_args.kwargs["data"]) + assert request_body["prompt_truncate_len"] == 4096 + assert "truncate_prompt_tokens" not in request_body + assert request_body["reasoning_effort"] == "none" + assert "chat_template_kwargs" not in request_body + assert "min_tokens" not in request_body + assert request_body["top_k"] == 40 + + +def test_in_schema_unsupported_params_still_raise(): + """ + The extras translation channel does not weaken the supported-params gate + for in-schema OpenAI params: store is still rejected with drop_params=False + and dropped with drop_params=True. + """ + with pytest.raises(litellm.UnsupportedParamsError): + litellm.get_optional_params( + model="accounts/fireworks/models/llama-v3-70b-instruct", + custom_llm_provider="fireworks_ai", + drop_params=False, + store=True, + ) + optional_params = litellm.get_optional_params( + model="accounts/fireworks/models/llama-v3-70b-instruct", + custom_llm_provider="fireworks_ai", + drop_params=True, + store=True, + ) + assert "store" not in optional_params From 599283584f2d16448c25a1ee4fbfdda2062eecf3 Mon Sep 17 00:00:00 2001 From: Miles Adkins Date: Wed, 5 Aug 2026 15:09:16 -0500 Subject: [PATCH 031/504] feat(fireworks_ai): drop reasoning_effort=auto to the model default Fireworks rejects reasoning_effort="auto" (accepted set: low, medium, high, xhigh, max, none, adaptive), so OpenAI-compatible clients sending it 400. Omitting the param means model default on Fireworks, which is exactly what auto means on OpenAI's side, so skip it in map_openai_params instead of forwarding. --- litellm/llms/fireworks_ai/chat/transformation.py | 6 ++++-- .../test_fireworks_ai_chat_transformation.py | 16 ++++++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 4740e84d513..a05b160413e 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -295,7 +295,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): optional_params["reasoning_effort"] = "medium" elif value is False: optional_params["reasoning_effort"] = "none" - else: + elif value != "auto": optional_params["reasoning_effort"] = value elif param in supported_openai_params: if value is not None: @@ -303,7 +303,9 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): return optional_params - def map_extra_body_params(self, optional_params: Mapping[str, object], model: str) -> dict: # noqa: LIT001 # http handler pops extra_body off the returned dict + def map_extra_body_params( + self, optional_params: Mapping[str, object], model: str + ) -> dict: # mutable-ok: http handler pops extra_body off the returned dict extra_body: Final = optional_params.get("extra_body") if not isinstance(extra_body, dict): return dict(optional_params) # mutable-ok: JSON request body diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 3fbbc70916a..bbb7fb197d0 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1153,6 +1153,22 @@ def test_reasoning_effort_integer_passthrough(): assert isinstance(result["reasoning_effort"], int) +def test_reasoning_effort_auto_dropped_to_model_default(): + """ + Fireworks rejects reasoning_effort="auto" (accepted set: low/medium/high/ + xhigh/max/none/adaptive). Omitting the param is the model default, which is + exactly what "auto" means on OpenAI's side, so it must not reach the request. + """ + config = FireworksAIConfig() + result = config.map_openai_params( + {"reasoning_effort": "auto"}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert "reasoning_effort" not in result + + def test_transform_response_captures_perf_metrics(): body = { **_BASE_CHAT_COMPLETION_RESPONSE, From 6d80d0509976feb702aad744cb7aff5fa81d3f54 Mon Sep 17 00:00:00 2001 From: Miles Adkins Date: Wed, 5 Aug 2026 15:56:33 -0500 Subject: [PATCH 032/504] fix(fireworks_ai): align extras translation with the API gateway matrix min_tokens is accepted natively by the Fireworks API (verified live), so stop stripping it and let it pass through extra_body. Add the NIM-specific include_reasoning and nvext keys to the strip set. enable_thinking=true now omits reasoning_effort (model default) instead of forcing medium, matching the gateway translation and preserving default-off models' behavior; enable_thinking=false still maps to none. --- litellm/llms/fireworks_ai/chat/transformation.py | 8 +++++--- .../test_fireworks_ai_chat_transformation.py | 16 ++++++++++------ 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index a05b160413e..b5c82129f45 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -67,7 +67,6 @@ def _json_schema_response_format(schema: object) -> Mapping[str, object]: _NIM_VLLM_STRIP_PARAMS: Final = frozenset( { - "min_tokens", "stop_token_ids", "include_stop_str_in_output", "skip_special_tokens", @@ -82,6 +81,8 @@ _NIM_VLLM_STRIP_PARAMS: Final = frozenset( "detokenize", "allowed_token_ids", "bad_words", + "include_reasoning", + "nvext", } ) @@ -414,8 +415,9 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): model, ) return () - effort: Final = "medium" if chat_template_kwargs["enable_thinking"] else "none" - return (("reasoning_effort", effort),) + if chat_template_kwargs["enable_thinking"]: + return () + return (("reasoning_effort", "none"),) @staticmethod def _translate_guided_params(extra_body: Mapping[str, object]) -> tuple[tuple[str, object], ...]: diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index bbb7fb197d0..d25bbd69b91 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1334,7 +1334,7 @@ def test_map_extra_body_params_chat_template_kwargs_enable_thinking(): {"extra_body": {"chat_template_kwargs": {"enable_thinking": True}}}, _REASONING_MODEL, ) - assert enabled == {"reasoning_effort": "medium"} + assert enabled == {} def test_map_extra_body_params_chat_template_kwargs_conflicts_with_reasoning_effort(): @@ -1425,7 +1425,6 @@ def test_map_extra_body_params_multiple_guided_params_rejected(): @pytest.mark.parametrize( "param,value", [ - ("min_tokens", 10), ("stop_token_ids", [1, 2]), ("include_stop_str_in_output", True), ("skip_special_tokens", False), @@ -1440,6 +1439,8 @@ def test_map_extra_body_params_multiple_guided_params_rejected(): ("detokenize", True), ("allowed_token_ids", [1]), ("bad_words", ["foo"]), + ("include_reasoning", False), + ("nvext", {"verbosity": 1}), ], ) def test_map_extra_body_params_strips_unsupported_nim_vllm_params(param, value, caplog): @@ -1478,9 +1479,10 @@ def test_nim_vllm_extras_translated_end_to_end_in_request_body(): Passing NIM/vLLM extras to litellm.completion must reach the Fireworks request body translated, not verbatim: truncate_prompt_tokens becomes prompt_truncate_len, chat_template_kwargs.enable_thinking becomes - reasoning_effort, min_tokens is dropped, and fireworks-native top_k still - passes through. Asserts on the actual JSON posted to the API, so a revert - of the _complete_fireworks_ai wiring fails this test. + reasoning_effort, include_reasoning is dropped, and min_tokens and + fireworks-native top_k still pass through. Asserts on the actual JSON + posted to the API, so a revert of the _complete_fireworks_ai wiring + fails this test. """ from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -1515,6 +1517,7 @@ def test_nim_vllm_extras_translated_end_to_end_in_request_body(): truncate_prompt_tokens=4096, chat_template_kwargs={"enable_thinking": False}, min_tokens=10, + include_reasoning=False, top_k=40, ) @@ -1523,7 +1526,8 @@ def test_nim_vllm_extras_translated_end_to_end_in_request_body(): assert "truncate_prompt_tokens" not in request_body assert request_body["reasoning_effort"] == "none" assert "chat_template_kwargs" not in request_body - assert "min_tokens" not in request_body + assert "include_reasoning" not in request_body + assert request_body["min_tokens"] == 10 assert request_body["top_k"] == 40 From 431f61b4f7c20b8f722f30c42c279edd19fe6a2d Mon Sep 17 00:00:00 2001 From: Miles Adkins Date: Wed, 5 Aug 2026 16:02:47 -0500 Subject: [PATCH 033/504] fix(fireworks_ai): prefer native values silently on extras conflicts Align with the API gateway translation: instead of raising BadRequestError on alias or competing-constraint conflicts, the explicit Fireworks-native param wins and the NIM/vLLM extra is dropped with a debug log. Covers truncate_prompt_tokens vs prompt_truncate_len, chat_template_kwargs enable_thinking vs reasoning_effort/thinking, guided_* vs response_format (including response_format nested in an explicit extra_body, which the previous conflict check missed), and multiple guided_* params (priority order json, grammar, choice). Malformed non-object chat_template_kwargs is also dropped with a log instead of raising. --- .../llms/fireworks_ai/chat/transformation.py | 108 +++++++---------- .../test_fireworks_ai_chat_transformation.py | 111 +++++++++++------- 2 files changed, 107 insertions(+), 112 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index b5c82129f45..3bacb3cd28e 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -311,7 +311,6 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): if not isinstance(extra_body, dict): return dict(optional_params) # mutable-ok: JSON request body - self._validate_extra_body_conflicts(extra_body=extra_body, optional_params=optional_params, model=model) stripped: Final = tuple(sorted(k for k in extra_body if k in _NIM_VLLM_STRIP_PARAMS)) if stripped: verbose_logger.debug( @@ -320,9 +319,9 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): model, ) promoted: Final = ( - *self._translate_truncate_prompt_tokens(extra_body), - *self._translate_chat_template_kwargs(extra_body, model), - *self._translate_guided_params(extra_body), + *self._translate_truncate_prompt_tokens(extra_body, optional_params), + *self._translate_chat_template_kwargs(extra_body, optional_params, model), + *self._translate_guided_params(extra_body, optional_params), ) remaining: Final = tuple((k, v) for k, v in extra_body.items() if k not in _EXTRA_BODY_CONSUMED_PARAMS) base: Final = {k: v for k, v in optional_params.items() if k != "extra_body"} # mutable-ok: JSON request body @@ -332,74 +331,32 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): **({"extra_body": dict(remaining)} if remaining else {}), # mutable-ok: JSON request body } - def _validate_extra_body_conflicts( - self, extra_body: Mapping[str, object], optional_params: Mapping[str, object], model: str - ) -> None: - if "truncate_prompt_tokens" in extra_body and ( - "prompt_truncate_len" in extra_body or "prompt_truncate_len" in optional_params - ): - raise litellm.BadRequestError( - message=( - "Fireworks AI chat completions received both `truncate_prompt_tokens` and " - "`prompt_truncate_len`; they are aliases, send only one." - ), - model=model, - llm_provider="fireworks_ai", - ) - chat_template_kwargs: Final = extra_body.get("chat_template_kwargs") - if ( - isinstance(chat_template_kwargs, dict) - and "enable_thinking" in chat_template_kwargs - and ("reasoning_effort" in optional_params or "thinking" in optional_params) - ): - raise litellm.BadRequestError( - message=( - "Fireworks AI chat completions does not support specifying both " - "`chat_template_kwargs.enable_thinking` and `reasoning_effort`/`thinking` in the same request." - ), - model=model, - llm_provider="fireworks_ai", - ) - guided_params: Final = tuple( - k for k in ("guided_json", "guided_grammar", "guided_choice") if extra_body.get(k) is not None - ) - if len(guided_params) > 1: - raise litellm.BadRequestError( - message=( - f"Fireworks AI chat completions received multiple guided decoding params " - f"{guided_params}; send only one." - ), - model=model, - llm_provider="fireworks_ai", - ) - if guided_params and "response_format" in optional_params: - raise litellm.BadRequestError( - message=( - f"Fireworks AI chat completions received both `{guided_params[0]}` and " - "`response_format`; they are competing output constraints, send only one." - ), - model=model, - llm_provider="fireworks_ai", - ) - @staticmethod - def _translate_truncate_prompt_tokens(extra_body: Mapping[str, object]) -> tuple[tuple[str, object], ...]: + def _translate_truncate_prompt_tokens( + extra_body: Mapping[str, object], optional_params: Mapping[str, object] + ) -> tuple[tuple[str, object], ...]: if extra_body.get("truncate_prompt_tokens") is None: return () + if "prompt_truncate_len" in extra_body or "prompt_truncate_len" in optional_params: + verbose_logger.debug( + "fireworks_ai ignoring truncate_prompt_tokens; explicit prompt_truncate_len takes precedence." + ) + return () return (("prompt_truncate_len", extra_body["truncate_prompt_tokens"]),) def _translate_chat_template_kwargs( - self, extra_body: Mapping[str, object], model: str + self, extra_body: Mapping[str, object], optional_params: Mapping[str, object], model: str ) -> tuple[tuple[str, object], ...]: chat_template_kwargs: Final = extra_body.get("chat_template_kwargs") if chat_template_kwargs is None: return () if not isinstance(chat_template_kwargs, dict): - raise litellm.BadRequestError( - message="Fireworks AI chat completions expects `chat_template_kwargs` to be an object.", - model=model, - llm_provider="fireworks_ai", + verbose_logger.debug( + "fireworks_ai dropping chat_template_kwargs for model=%s; expected an object, got %s.", + model, + type(chat_template_kwargs).__name__, ) + return () other_keys: Final = tuple(sorted(k for k in chat_template_kwargs if k != "enable_thinking")) if other_keys: verbose_logger.debug( @@ -409,6 +366,11 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): ) if "enable_thinking" not in chat_template_kwargs: return () + if "reasoning_effort" in optional_params or "thinking" in optional_params: + verbose_logger.debug( + "fireworks_ai ignoring chat_template_kwargs.enable_thinking; explicit reasoning_effort/thinking takes precedence." + ) + return () if not supports_reasoning(model=model, custom_llm_provider="fireworks_ai"): verbose_logger.debug( "fireworks_ai model %r does not support reasoning; dropping chat_template_kwargs.enable_thinking.", @@ -420,7 +382,19 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): return (("reasoning_effort", "none"),) @staticmethod - def _translate_guided_params(extra_body: Mapping[str, object]) -> tuple[tuple[str, object], ...]: + def _translate_guided_params( + extra_body: Mapping[str, object], optional_params: Mapping[str, object] + ) -> tuple[tuple[str, object], ...]: + has_guided: Final = any( + extra_body.get(key) is not None for key in ("guided_json", "guided_grammar", "guided_choice") + ) + if not has_guided: + return () + if "response_format" in optional_params or "response_format" in extra_body: + verbose_logger.debug( + "fireworks_ai ignoring guided decoding params; explicit response_format takes precedence." + ) + return () if extra_body.get("guided_json") is not None: return (("response_format", _json_schema_response_format(extra_body["guided_json"])),) if extra_body.get("guided_grammar") is not None: @@ -429,13 +403,11 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): "grammar": extra_body["guided_grammar"], } return (("response_format", grammar_response_format),) - if extra_body.get("guided_choice") is not None: - choice_schema: Final = { # mutable-ok: JSON request body - "type": "string", - "enum": extra_body["guided_choice"], - } - return (("response_format", _json_schema_response_format(choice_schema)),) - return () + choice_schema: Final = { # mutable-ok: JSON request body + "type": "string", + "enum": extra_body["guided_choice"], + } + return (("response_format", _json_schema_response_format(choice_schema)),) def _transform_tools(self, tools: list[OpenAIChatCompletionToolParam]) -> list[OpenAIChatCompletionToolParam]: for tool in tools: diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index d25bbd69b91..48d868b5846 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1308,18 +1308,19 @@ def test_map_extra_body_params_translates_truncate_prompt_tokens(): assert result == {"prompt_truncate_len": 4096} -def test_map_extra_body_params_truncate_prompt_tokens_conflicts_with_alias(): +def test_map_extra_body_params_truncate_prompt_tokens_native_wins(): config = FireworksAIConfig() - with pytest.raises(litellm.BadRequestError, match="aliases"): - config.map_extra_body_params( - {"prompt_truncate_len": 2048, "extra_body": {"truncate_prompt_tokens": 4096}}, - _REASONING_MODEL, - ) - with pytest.raises(litellm.BadRequestError, match="aliases"): - config.map_extra_body_params( - {"extra_body": {"truncate_prompt_tokens": 4096, "prompt_truncate_len": 2048}}, - _REASONING_MODEL, - ) + top_level = config.map_extra_body_params( + {"prompt_truncate_len": 2048, "extra_body": {"truncate_prompt_tokens": 4096}}, + _REASONING_MODEL, + ) + assert top_level == {"prompt_truncate_len": 2048} + + nested = config.map_extra_body_params( + {"extra_body": {"truncate_prompt_tokens": 4096, "prompt_truncate_len": 2048}}, + _REASONING_MODEL, + ) + assert nested == {"extra_body": {"prompt_truncate_len": 2048}} def test_map_extra_body_params_chat_template_kwargs_enable_thinking(): @@ -1337,28 +1338,29 @@ def test_map_extra_body_params_chat_template_kwargs_enable_thinking(): assert enabled == {} -def test_map_extra_body_params_chat_template_kwargs_conflicts_with_reasoning_effort(): +def test_map_extra_body_params_chat_template_kwargs_native_reasoning_effort_wins(): config = FireworksAIConfig() - with pytest.raises(litellm.BadRequestError, match="enable_thinking"): - config.map_extra_body_params( - { - "reasoning_effort": "high", - "extra_body": {"chat_template_kwargs": {"enable_thinking": False}}, - }, - _REASONING_MODEL, - ) + result = config.map_extra_body_params( + { + "reasoning_effort": "high", + "extra_body": {"chat_template_kwargs": {"enable_thinking": False}}, + }, + _REASONING_MODEL, + ) + assert result == {"reasoning_effort": "high"} -def test_map_extra_body_params_chat_template_kwargs_conflicts_with_thinking(): +def test_map_extra_body_params_chat_template_kwargs_native_thinking_wins(): config = FireworksAIConfig() - with pytest.raises(litellm.BadRequestError, match="enable_thinking"): - config.map_extra_body_params( - { - "thinking": {"type": "enabled", "budget_tokens": 4096}, - "extra_body": {"chat_template_kwargs": {"enable_thinking": True}}, - }, - _REASONING_MODEL, - ) + thinking = {"type": "enabled", "budget_tokens": 4096} + result = config.map_extra_body_params( + { + "thinking": thinking, + "extra_body": {"chat_template_kwargs": {"enable_thinking": True}}, + }, + _REASONING_MODEL, + ) + assert result == {"thinking": thinking} def test_map_extra_body_params_chat_template_kwargs_dropped_for_non_reasoning_model(): @@ -1370,6 +1372,15 @@ def test_map_extra_body_params_chat_template_kwargs_dropped_for_non_reasoning_mo assert result == {} +def test_map_extra_body_params_non_dict_chat_template_kwargs_dropped(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": "enable_thinking"}}, + _REASONING_MODEL, + ) + assert result == {} + + def test_map_extra_body_params_guided_json(): config = FireworksAIConfig() schema = {"type": "object", "properties": {"x": {"type": "string"}}} @@ -1401,25 +1412,37 @@ def test_map_extra_body_params_guided_grammar_and_choice(): } -def test_map_extra_body_params_guided_conflicts_with_response_format(): +def test_map_extra_body_params_guided_native_response_format_wins(): config = FireworksAIConfig() - with pytest.raises(litellm.BadRequestError, match="response_format"): - config.map_extra_body_params( - { - "response_format": {"type": "json_object"}, - "extra_body": {"guided_json": {"type": "object"}}, - }, - _REASONING_MODEL, - ) + top_level = config.map_extra_body_params( + { + "response_format": {"type": "json_object"}, + "extra_body": {"guided_json": {"type": "object"}}, + }, + _REASONING_MODEL, + ) + assert top_level == {"response_format": {"type": "json_object"}} + + nested_format = {"type": "json_object"} + nested = config.map_extra_body_params( + {"extra_body": {"guided_json": {"type": "object"}, "response_format": nested_format}}, + _REASONING_MODEL, + ) + assert nested == {"extra_body": {"response_format": nested_format}} -def test_map_extra_body_params_multiple_guided_params_rejected(): +def test_map_extra_body_params_multiple_guided_params_priority_order(): config = FireworksAIConfig() - with pytest.raises(litellm.BadRequestError, match="multiple guided decoding params"): - config.map_extra_body_params( - {"extra_body": {"guided_json": {"type": "object"}, "guided_grammar": "root ::= 'x'"}}, - _REASONING_MODEL, - ) + result = config.map_extra_body_params( + {"extra_body": {"guided_grammar": "root ::= 'x'", "guided_json": {"type": "object"}}}, + _REASONING_MODEL, + ) + assert result == { + "response_format": { + "type": "json_schema", + "json_schema": {"schema": {"type": "object"}}, + } + } @pytest.mark.parametrize( From 80d8e952280580e21346b9c200b2d0035d55dfc8 Mon Sep 17 00:00:00 2001 From: Scott Wilson Date: Mon, 3 Aug 2026 16:34:14 -0400 Subject: [PATCH 034/504] fix(responses): unwrap object-form tool_choice before calling the Responses API Clients send tool_choice as {"type": "auto"} (Cursor on chat completions, Claude Code's Anthropic tool_choice shape). validate_chat_completion_tool_choice recognized that shape but returned it verbatim, and the chat -> Responses API bridge only normalized {"type": "function"}, so the wrapper reached OpenAI and the whole call failed with: Invalid value: 'auto'. Supported values are: 'code_interpreter', ..., 'web_search_preview', ... (param: tool_choice.type) That broke every tool call, web search included, on responses-mode models. Unwrap {"type": "auto"|"none"|"required"} to the bare string at both layers: the chat completions validation boundary where the shape is first accepted, and the Responses API bridge that owns the Responses tool_choice contract. No OpenAI surface accepts the object form for these values, so the previous passthrough only deferred the 400 to the provider. --- .../transformation.py | 2 + litellm/utils.py | 12 +-- .../test_validate_tool_choice.py | 15 +-- ...responses_transformation_transformation.py | 91 ++++++++++++++++++- 4 files changed, 105 insertions(+), 15 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index f31e228e456..d3c5290bd9a 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -169,6 +169,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if not isinstance(tool_choice, dict): return tool_choice choice_type: Final = tool_choice.get("type") + if isinstance(choice_type, str) and choice_type in ("auto", "none", "required"): + return choice_type if choice_type not in ("function", "custom"): return tool_choice if isinstance(tool_choice.get("name"), str) and tool_choice.get("name"): diff --git a/litellm/utils.py b/litellm/utils.py index d93c88e05a0..f4a9ee19e4f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7577,17 +7577,13 @@ def validate_chat_completion_tool_choice( Prevents user errors like: https://github.com/BerriAI/litellm/issues/7483 """ - from litellm.types.llms.openai import ( - ChatCompletionToolChoiceObjectParam, - ChatCompletionToolChoiceStringValues, - ) - if tool_choice is None or isinstance(tool_choice, str): return tool_choice elif isinstance(tool_choice, dict): - # Handle Cursor IDE format: {"type": "auto"} -> return as-is - if tool_choice.get("type") in ["auto", "none", "required"] and "function" not in tool_choice: - return tool_choice + # Handle Cursor IDE format: {"type": "auto"} -> unwrap to the bare string + tool_choice_type = tool_choice.get("type") + if tool_choice_type in ("auto", "none", "required") and "function" not in tool_choice: + return tool_choice_type # Standard OpenAI format: {"type": "function", "function": {...}} if tool_choice.get("type") is None or tool_choice.get("function") is None: diff --git a/tests/litellm_utils_tests/test_validate_tool_choice.py b/tests/litellm_utils_tests/test_validate_tool_choice.py index 8150403c145..c3f80f31864 100644 --- a/tests/litellm_utils_tests/test_validate_tool_choice.py +++ b/tests/litellm_utils_tests/test_validate_tool_choice.py @@ -28,12 +28,15 @@ def test_validate_tool_choice_standard_dict(): def test_validate_tool_choice_cursor_format(): - """Test Cursor IDE format: {"type": "auto"} -> {"type": "auto"}.""" - assert validate_chat_completion_tool_choice({"type": "auto"}) == {"type": "auto"} - assert validate_chat_completion_tool_choice({"type": "none"}) == {"type": "none"} - assert validate_chat_completion_tool_choice({"type": "required"}) == { - "type": "required" - } + """Cursor IDE format {"type": "auto"} must be unwrapped to the bare string. + + No OpenAI surface accepts the object form of these values. Forwarding it + verbatim makes the provider reject the call with + "Invalid value: 'auto' ... param: tool_choice.type". + """ + assert validate_chat_completion_tool_choice({"type": "auto"}) == "auto" + assert validate_chat_completion_tool_choice({"type": "none"}) == "none" + assert validate_chat_completion_tool_choice({"type": "required"}) == "required" def test_validate_tool_choice_invalid_dict(): diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index b8bd5c951ee..70563aa880a 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -2474,7 +2474,9 @@ def test_map_optional_params_tool_choice_chat_nested_to_responses_api(): {"type": "function", "name": "foo", "function": {"name": "bar"}}, {"type": "function", "name": "foo"}, ), - ({"type": "required"}, {"type": "required"}), + ({"type": "auto"}, "auto"), + ({"type": "none"}, "none"), + ({"type": "required"}, "required"), ( {"type": "custom", "custom": {"name": "ApplyPatch"}}, {"type": "custom", "name": "ApplyPatch"}, @@ -3400,3 +3402,90 @@ def test_output_item_done_with_stream_map_keeps_empty_delta(): ) assert chunk.choices[0].delta.tool_calls is None assert chunk.choices[0].finish_reason is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_choice,expected_wire_tool_choice", + [ + ({"type": "auto"}, "auto"), + ({"type": "none"}, "none"), + ({"type": "required"}, "required"), + ("auto", "auto"), + ({"type": "function", "function": {"name": "get_weather"}}, {"type": "function", "name": "get_weather"}), + ], +) +async def test_acompletion_bridge_normalizes_tool_choice_on_the_wire(tool_choice, expected_wire_tool_choice): + """Object-wrapped tool_choice must never reach /v1/responses. + + Clients (Cursor, Claude Code via /v1/messages) send ``{"type": "auto"}``. + The Responses API only accepts a hosted-tool name in ``tool_choice.type``, + so forwarding the wrapper verbatim fails the whole call with + ``Invalid value: 'auto' ... param: tool_choice.type`` -- which broke every + tool call, including web search, on responses-mode models. + """ + from unittest.mock import AsyncMock + + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + responses_payload = { + "id": "resp_bridge_tool_choice", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "gpt-5.5", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + "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, + } + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.text = json.dumps(responses_payload) + mock_response.headers = httpx.Headers({}) + mock_response.json.return_value = responses_payload + + with patch.object(AsyncHTTPHandler, "post", new_callable=AsyncMock) as mock_post: + mock_post.return_value = mock_response + + await litellm.acompletion( + model="openai/responses/gpt-5.5", + messages=[{"role": "user", "content": "what is the DJIA today"}], + api_key="fake-api-key", + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + tool_choice=tool_choice, + ) + + mock_post.assert_called_once() + post_kwargs = mock_post.call_args.kwargs + request_body = post_kwargs["json"] if "json" in post_kwargs else json.loads(post_kwargs["data"]) + assert request_body["tool_choice"] == expected_wire_tool_choice From ebf6167d8acf48499e294ecf3a7642b4112913eb Mon Sep 17 00:00:00 2001 From: Scott Wilson Date: Mon, 3 Aug 2026 16:50:58 -0400 Subject: [PATCH 035/504] fix(anthropic): stop emitting empty thinking blocks on the Responses adapter OpenAI emits a reasoning output item on every reasoning turn, but only emits reasoning_summary_text deltas when a summary was requested and actually produced. The Anthropic /v1/messages Responses stream adapter opened the thinking content block eagerly on response.output_item.added, so a summary-less reasoning item surfaced as {"type": "thinking", "thinking": ""}. Clients persist that in their session transcript and replay it on the next turn; an Anthropic model then rejects the request with "each thinking block must contain thinking", which is what users hit when a resumed session falls back to the default Anthropic model. Open the thinking block on the first non-empty summary delta instead, and only emit content_block_stop for items that actually have an open block. --- .../responses_adapters/streaming_iterator.py | 83 +++++++------------ ...t_responses_adapters_streaming_iterator.py | 79 +++++++++++++++++- 2 files changed, 106 insertions(+), 56 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index f12dd979338..c588e791cd9 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -68,6 +68,19 @@ class AnthropicResponsesStreamWrapper: self._current_block_index += 1 return self._current_block_index + def _open_block(self, item_id: str | None, content_block: dict[str, Any]) -> int: + block_idx = self._next_block_index() + if item_id: + self._item_id_to_block_index[item_id] = block_idx + self._chunk_queue.append( + { + "type": "content_block_start", + "index": block_idx, + "content_block": content_block, + } + ) + return block_idx + def _process_event(self, event: Any) -> None: """Convert one Responses API event into zero or more Anthropic chunks queued for emission.""" event_type = getattr(event, "type", None) @@ -93,47 +106,22 @@ class AnthropicResponsesStreamWrapper: item_id = getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item_type == "message": - block_idx = self._next_block_index() - if item_id: - self._item_id_to_block_index[item_id] = block_idx - self._chunk_queue.append( - { - "type": "content_block_start", - "index": block_idx, - "content_block": {"type": "text", "text": ""}, - } - ) + self._open_block(item_id, {"type": "text", "text": ""}) elif item_type == "function_call": call_id: Final = ( getattr(item, "call_id", None) or (item.get("call_id") if isinstance(item, dict) else None) or "" ) name = getattr(item, "name", None) or (item.get("name") if isinstance(item, dict) else None) or "" - block_idx = self._next_block_index() if item_id: - self._item_id_to_block_index[item_id] = block_idx self._pending_tool_ids[item_id] = call_id - self._chunk_queue.append( + self._open_block( + item_id, { - "type": "content_block_start", - "index": block_idx, - "content_block": { - "type": "tool_use", - "id": call_id, - "name": name, - "input": {}, - }, - } - ) - elif item_type == "reasoning": - block_idx = self._next_block_index() - if item_id: - self._item_id_to_block_index[item_id] = block_idx - self._chunk_queue.append( - { - "type": "content_block_start", - "index": block_idx, - "content_block": {"type": "thinking", "thinking": ""}, - } + "type": "tool_use", + "id": call_id, + "name": name, + "input": {}, + }, ) return @@ -146,16 +134,7 @@ class AnthropicResponsesStreamWrapper: # Some providers (e.g. LMStudio) skip response.output_item.added, # so no text block is open yet; synthesize content_block_start # instead of emitting a delta with index -1 - block_idx = self._next_block_index() - if item_id: - self._item_id_to_block_index[item_id] = block_idx - self._chunk_queue.append( - { - "type": "content_block_start", - "index": block_idx, - "content_block": {"type": "text", "text": ""}, - } - ) + block_idx = self._open_block(item_id, {"type": "text", "text": ""}) self._chunk_queue.append( { "type": "content_block_delta", @@ -169,11 +148,11 @@ class AnthropicResponsesStreamWrapper: if event_type == "response.reasoning_summary_text.delta": item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") - block_idx = ( - self._item_id_to_block_index.get(item_id, self._current_block_index) - if item_id - else self._current_block_index - ) + block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index + if block_idx < 0: + if not delta: + return + block_idx = self._open_block(item_id, {"type": "thinking", "thinking": ""}) self._chunk_queue.append( { "type": "content_block_delta", @@ -207,11 +186,9 @@ class AnthropicResponsesStreamWrapper: item_id = ( getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item else None ) - block_idx = ( - self._item_id_to_block_index.get(item_id, self._current_block_index) - if item_id - else self._current_block_index - ) + block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index + if block_idx < 0: + return self._chunk_queue.append( { "type": "content_block_stop", diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index 73b58e71009..b1ae865fde1 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -76,6 +76,78 @@ class TestProcessEventResponseCreatedGuard: assert len(message_starts) == 1 +class TestReasoningItemWithoutSummaryText: + """Regression: a reasoning item whose summary never produces text must not + surface as a thinking content block. + + OpenAI emits ``response.output_item.added`` with ``type: "reasoning"`` on + every reasoning turn, but only emits + ``response.reasoning_summary_text.delta`` when a summary was requested and + the model actually produced one. Eagerly opening the block on + ``output_item.added`` left ``{"type": "thinking", "thinking": ""}`` in the + assistant turn, which clients persist in their session transcript. Replaying + that transcript against an Anthropic model (what ``claude --resume`` does + once the resumed session falls back to the default Anthropic model) fails + with:: + + 400 invalid_request_error - messages.2.content.0.thinking: + each thinking block must contain thinking + + So the thinking block is opened on the first non-empty summary delta. + """ + + @staticmethod + def _gpt_turn(reasoning_summary_deltas: list) -> list: + return [ + {"type": "response.created"}, + {"type": "response.output_item.added", "item": {"type": "reasoning", "id": "rs_1"}}, + *( + {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": delta} + for delta in reasoning_summary_deltas + ), + {"type": "response.output_item.done", "item": {"type": "reasoning", "id": "rs_1"}}, + {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}}, + {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hello"}, + {"type": "response.output_item.done", "item": {"type": "message", "id": "msg_1"}}, + ] + + def test_reasoning_without_summary_emits_no_thinking_block(self): + chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=[])) + + assert not [ + c for c in chunks if c["type"] == "content_block_start" and c["content_block"]["type"] == "thinking" + ] + assert [(c["type"], c.get("index")) for c in chunks[1:]] == [ + ("content_block_start", 0), + ("content_block_delta", 0), + ("content_block_stop", 0), + ] + assert chunks[1]["content_block"] == {"type": "text", "text": ""} + + def test_reasoning_with_only_empty_summary_deltas_emits_no_thinking_block(self): + chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=["", ""])) + + assert not [c for c in chunks if c["type"] == "content_block_delta" and c["delta"]["type"] == "thinking_delta"] + assert not [ + c for c in chunks if c["type"] == "content_block_start" and c["content_block"]["type"] == "thinking" + ] + + def test_reasoning_with_summary_text_still_emits_a_thinking_block(self): + chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=["Weigh", "ing options"])) + + assert [(c["type"], c.get("index")) for c in chunks[1:]] == [ + ("content_block_start", 0), + ("content_block_delta", 0), + ("content_block_delta", 0), + ("content_block_stop", 0), + ("content_block_start", 1), + ("content_block_delta", 1), + ("content_block_stop", 1), + ] + assert chunks[1]["content_block"] == {"type": "thinking", "thinking": ""} + assert "".join(c["delta"]["thinking"] for c in chunks[2:4]) == "Weighing options" + + class TestProcessEventTextDeltaWithoutOutputItemAdded: """Streams that skip response.output_item.added (e.g. LMStudio) must still open a text block before any delta and never emit index -1.""" @@ -110,12 +182,13 @@ class TestProcessEventTextDeltaWithoutOutputItemAdded: "type": "response.output_item.added", "item": {"type": "reasoning", "id": "rs_1"}, }, + {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": "hm"}, {"type": "response.output_text.delta", "item_id": "m1", "delta": "Hi"}, ] ) - assert chunks[1]["type"] == "content_block_start" - assert chunks[1]["content_block"] == {"type": "text", "text": ""} - assert [c["index"] for c in chunks[1:]] == [1, 1] + assert chunks[2]["type"] == "content_block_start" + assert chunks[2]["content_block"] == {"type": "text", "text": ""} + assert [c["index"] for c in chunks[2:]] == [1, 1] def test_process_event_registered_item_id_does_not_synthesize_start(self): chunks = _process_all( From 889c1f584a6bda2d1812c1142117b5d1eed01932 Mon Sep 17 00:00:00 2001 From: Scott Wilson Date: Wed, 5 Aug 2026 23:23:12 -0400 Subject: [PATCH 036/504] test(responses): annotate the tool_choice bridge test signature --- ...extras_litellm_responses_transformation_transformation.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 70563aa880a..c2484970ed9 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3415,7 +3415,10 @@ def test_output_item_done_with_stream_map_keeps_empty_delta(): ({"type": "function", "function": {"name": "get_weather"}}, {"type": "function", "name": "get_weather"}), ], ) -async def test_acompletion_bridge_normalizes_tool_choice_on_the_wire(tool_choice, expected_wire_tool_choice): +async def test_acompletion_bridge_normalizes_tool_choice_on_the_wire( + tool_choice: str | dict[str, object], + expected_wire_tool_choice: str | dict[str, str], +) -> None: """Object-wrapped tool_choice must never reach /v1/responses. Clients (Cursor, Claude Code via /v1/messages) send ``{"type": "auto"}``. From 4a601c49a60d34d12810bd0372b062dca57d34b7 Mon Sep 17 00:00:00 2001 From: Miles Adkins Date: Thu, 6 Aug 2026 10:46:48 -0500 Subject: [PATCH 037/504] feat(fireworks_ai): full chat_template_kwargs parity with the gateway Map the remaining gateway-documented effort keys: thinking as an alias for enable_thinking (enable_thinking wins when both are present), reasoning_budget to an integer reasoning_effort (skipped when thinking is explicitly off), and low_effort=true to reasoning_effort=low (budget wins when both are set). guided_json and guided_choice response_format wrappers now include the name field (response and choice) to match the gateway wire shape. --- .../llms/fireworks_ai/chat/transformation.py | 47 ++++++++---- .../test_fireworks_ai_chat_transformation.py | 72 ++++++++++++++++++- 2 files changed, 104 insertions(+), 15 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 3bacb3cd28e..6b763be0bfe 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -61,8 +61,32 @@ def _extract_fireworks_hidden_params(payload: dict) -> dict: return {**top_level, **per_choice} -def _json_schema_response_format(schema: object) -> Mapping[str, object]: - return {"type": "json_schema", "json_schema": {"schema": schema}} # mutable-ok: JSON request body +def _json_schema_response_format(schema: object, name: str) -> Mapping[str, object]: + return {"type": "json_schema", "json_schema": {"name": name, "schema": schema}} # mutable-ok: JSON request body + + +_EFFORT_KWARG_KEYS: Final = frozenset({"enable_thinking", "thinking", "reasoning_budget", "low_effort"}) + + +def _bool_from_kwargs(kwargs: Mapping[str, object], keys: tuple[str, ...]) -> bool | None: + for key in keys: + value = kwargs.get(key) + if isinstance(value, bool): + return value + return None + + +def _effort_from_chat_template_kwargs(kwargs: Mapping[str, object]) -> object: + enable_thinking: Final = _bool_from_kwargs(kwargs, ("enable_thinking", "thinking")) + if enable_thinking is False: + return "none" + budget: Final = kwargs.get("reasoning_budget") + if isinstance(budget, (int, float)) and not isinstance(budget, bool) and budget > 0: + return int(budget) + low_effort: Final = _bool_from_kwargs(kwargs, ("low_effort",)) + if low_effort is True: + return "low" + return None _NIM_VLLM_STRIP_PARAMS: Final = frozenset( @@ -357,29 +381,28 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): type(chat_template_kwargs).__name__, ) return () - other_keys: Final = tuple(sorted(k for k in chat_template_kwargs if k != "enable_thinking")) + other_keys: Final = tuple(sorted(k for k in chat_template_kwargs if k not in _EFFORT_KWARG_KEYS)) if other_keys: verbose_logger.debug( "fireworks_ai does not support chat_template_kwargs keys %s for model=%s; dropping them.", other_keys, model, ) - if "enable_thinking" not in chat_template_kwargs: - return () if "reasoning_effort" in optional_params or "thinking" in optional_params: verbose_logger.debug( - "fireworks_ai ignoring chat_template_kwargs.enable_thinking; explicit reasoning_effort/thinking takes precedence." + "fireworks_ai ignoring chat_template_kwargs; explicit reasoning_effort/thinking takes precedence." ) return () + effort: Final = _effort_from_chat_template_kwargs(chat_template_kwargs) + if effort is None: + return () if not supports_reasoning(model=model, custom_llm_provider="fireworks_ai"): verbose_logger.debug( - "fireworks_ai model %r does not support reasoning; dropping chat_template_kwargs.enable_thinking.", + "fireworks_ai model %r does not support reasoning; dropping chat_template_kwargs effort keys.", model, ) return () - if chat_template_kwargs["enable_thinking"]: - return () - return (("reasoning_effort", "none"),) + return (("reasoning_effort", effort),) @staticmethod def _translate_guided_params( @@ -396,7 +419,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): ) return () if extra_body.get("guided_json") is not None: - return (("response_format", _json_schema_response_format(extra_body["guided_json"])),) + return (("response_format", _json_schema_response_format(extra_body["guided_json"], "response")),) if extra_body.get("guided_grammar") is not None: grammar_response_format: Final = { # mutable-ok: JSON request body "type": "grammar", @@ -407,7 +430,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): "type": "string", "enum": extra_body["guided_choice"], } - return (("response_format", _json_schema_response_format(choice_schema)),) + return (("response_format", _json_schema_response_format(choice_schema, "choice")),) def _transform_tools(self, tools: list[OpenAIChatCompletionToolParam]) -> list[OpenAIChatCompletionToolParam]: for tool in tools: diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 48d868b5846..e1b5d457205 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1338,6 +1338,66 @@ def test_map_extra_body_params_chat_template_kwargs_enable_thinking(): assert enabled == {} +def test_map_extra_body_params_chat_template_kwargs_thinking_alias(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"thinking": False}}}, + _REASONING_MODEL, + ) + assert result == {"reasoning_effort": "none"} + + +def test_map_extra_body_params_chat_template_kwargs_enable_thinking_wins_over_thinking(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": True, "thinking": False}}}, + _REASONING_MODEL, + ) + assert result == {} + + +def test_map_extra_body_params_chat_template_kwargs_reasoning_budget(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"reasoning_budget": 512}}}, + _REASONING_MODEL, + ) + assert result == {"reasoning_effort": 512} + + +def test_map_extra_body_params_chat_template_kwargs_budget_ignored_when_thinking_off(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"enable_thinking": False, "reasoning_budget": 512}}}, + _REASONING_MODEL, + ) + assert result == {"reasoning_effort": "none"} + + +def test_map_extra_body_params_chat_template_kwargs_low_effort(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"low_effort": True}}}, + _REASONING_MODEL, + ) + assert result == {"reasoning_effort": "low"} + + budget_wins = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"low_effort": True, "reasoning_budget": 256}}}, + _REASONING_MODEL, + ) + assert budget_wins == {"reasoning_effort": 256} + + +def test_map_extra_body_params_chat_template_kwargs_effort_keys_dropped_for_non_reasoning_model(): + config = FireworksAIConfig() + result = config.map_extra_body_params( + {"extra_body": {"chat_template_kwargs": {"reasoning_budget": 512, "low_effort": True}}}, + _NON_REASONING_MODEL, + ) + assert result == {} + + def test_map_extra_body_params_chat_template_kwargs_native_reasoning_effort_wins(): config = FireworksAIConfig() result = config.map_extra_body_params( @@ -1388,7 +1448,10 @@ def test_map_extra_body_params_guided_json(): {"extra_body": {"guided_json": schema}}, _REASONING_MODEL ) assert result == { - "response_format": {"type": "json_schema", "json_schema": {"schema": schema}} + "response_format": { + "type": "json_schema", + "json_schema": {"name": "response", "schema": schema}, + } } @@ -1407,7 +1470,10 @@ def test_map_extra_body_params_guided_grammar_and_choice(): assert choice == { "response_format": { "type": "json_schema", - "json_schema": {"schema": {"type": "string", "enum": ["yes", "no"]}}, + "json_schema": { + "name": "choice", + "schema": {"type": "string", "enum": ["yes", "no"]}, + }, } } @@ -1440,7 +1506,7 @@ def test_map_extra_body_params_multiple_guided_params_priority_order(): assert result == { "response_format": { "type": "json_schema", - "json_schema": {"schema": {"type": "object"}}, + "json_schema": {"name": "response", "schema": {"type": "object"}}, } } From 1d8a642e0683e13be122c532436e0919d6c540f5 Mon Sep 17 00:00:00 2001 From: heathriel Date: Wed, 22 Jul 2026 08:41:01 -0700 Subject: [PATCH 038/504] fix(fireworks_ai): support router slugs via routers/ prefix Bare fireworks_ai/ only resolved to accounts/fireworks/models/, so Fireworks routers (served at accounts/fireworks/routers/, e.g. glm-latest and firerouter) could not be reached without passing the full resource id. Add a shared resolve_fireworks_resource_name helper that maps an explicit routers/ or models/ segment to the right resource path, keeps the existing -fast router heuristic, and defaults bare slugs to models/ for backward compatibility. Wire it into both the chat and text-completion transforms, which had drifted (completion lacked router handling entirely) --- .../llms/fireworks_ai/chat/transformation.py | 18 ++++---- litellm/llms/fireworks_ai/common_utils.py | 11 +++++ .../fireworks_ai/completion/transformation.py | 7 +-- .../test_fireworks_ai_chat_transformation.py | 43 ++++++++++++++++++ ..._fireworks_ai_completion_transformation.py | 34 ++++++++++++++ .../test_fireworks_ai_common_utils.py | 45 +++++++++++++++++++ type-discipline-budget.json | 2 +- 7 files changed, 146 insertions(+), 14 deletions(-) create mode 100644 tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py create mode 100644 tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index a796aa47b70..26f0caefacd 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -39,7 +39,11 @@ from ...openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig, ) -from ..common_utils import FireworksAIException, FireworksAIMixin +from ..common_utils import ( + FireworksAIException, + FireworksAIMixin, + resolve_fireworks_resource_name, +) def _extract_fireworks_hidden_params(payload: dict) -> dict: @@ -459,12 +463,10 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): litellm_params: dict, headers: dict, ) -> dict: - if not model.startswith("accounts/") and "#" not in model: - if model.endswith("-fast"): - model = f"accounts/fireworks/routers/{model}" - else: - model = f"accounts/fireworks/models/{model}" - messages = self._transform_messages_helper(messages=messages, model=model, litellm_params=litellm_params) + resolved_model: Final = resolve_fireworks_resource_name(model) + messages = self._transform_messages_helper( + messages=messages, model=resolved_model, litellm_params=litellm_params + ) if "tools" in optional_params and optional_params["tools"] is not None: tools: Final = self._transform_tools(tools=optional_params["tools"]) optional_params["tools"] = tools @@ -478,7 +480,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): "include_usage": True, } return super().transform_request( - model=model, + model=resolved_model, messages=messages, optional_params=optional_params, litellm_params=litellm_params, diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index 143dd151027..e07e7a26f9e 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -29,6 +29,17 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None: return None +def resolve_fireworks_resource_name(model: str) -> str: + stripped: Final = model.removeprefix("fireworks_ai/") + if stripped.startswith("accounts/") or "#" in stripped: + return stripped + if stripped.startswith(("routers/", "models/")): + return f"accounts/fireworks/{stripped}" + if stripped.endswith("-fast"): + return f"accounts/fireworks/routers/{stripped}" + return f"accounts/fireworks/models/{stripped}" + + class FireworksAIMixin: """ Common Base Config functions across Fireworks AI Endpoints diff --git a/litellm/llms/fireworks_ai/completion/transformation.py b/litellm/llms/fireworks_ai/completion/transformation.py index c141e097d3a..c460510f39c 100644 --- a/litellm/llms/fireworks_ai/completion/transformation.py +++ b/litellm/llms/fireworks_ai/completion/transformation.py @@ -4,7 +4,7 @@ from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUser from ...base_llm.completion.transformation import BaseTextCompletionConfig from ...openai.completion.utils import _transform_prompt -from ..common_utils import FireworksAIMixin +from ..common_utils import FireworksAIMixin, resolve_fireworks_resource_name class FireworksAITextCompletionConfig(FireworksAIMixin, BaseTextCompletionConfig): @@ -50,11 +50,8 @@ class FireworksAITextCompletionConfig(FireworksAIMixin, BaseTextCompletionConfig ) -> dict: prompt: Final = _transform_prompt(messages=messages) - if not model.startswith("accounts/") and "#" not in model: - model = f"accounts/fireworks/models/{model}" - data: Final = { - "model": model, + "model": resolve_fireworks_resource_name(model), "prompt": prompt, **optional_params, } diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 94945ed4bfb..87908ef60c3 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1282,3 +1282,46 @@ def test_streaming_surfaces_fireworks_response_fields(): assert surfaced["fireworks_raw_outputs"] == [raw_output] assert surfaced["fireworks_perf_metrics"] == {"prompt-tokens": 5} assert surfaced["fireworks_prompt_token_ids"] == [1, 2, 3] + + +def test_transform_request_routes_router_slug(): + config = FireworksAIConfig() + + data = config.transform_request( + model="routers/glm-latest", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert data["model"] == "accounts/fireworks/routers/glm-latest" + + +def test_transform_request_bare_slug_stays_model(): + config = FireworksAIConfig() + + data = config.transform_request( + model="glm-4p6", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert data["model"] == "accounts/fireworks/models/glm-4p6" + + +def test_transform_request_direct_route_passthrough(): + config = FireworksAIConfig() + model = "accounts/fireworks/models/qwen2p5-coder-7b#accounts/gitlab/deployments/2fb7764c" + + data = config.transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert data["model"] == model diff --git a/tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py b/tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py new file mode 100644 index 00000000000..996f1fd975b --- /dev/null +++ b/tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py @@ -0,0 +1,34 @@ +import os +import sys + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.fireworks_ai.completion.transformation import ( + FireworksAITextCompletionConfig, +) + + +def test_transform_text_completion_request_routes_router_slug(): + config = FireworksAITextCompletionConfig() + + data = config.transform_text_completion_request( + model="routers/glm-latest", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + headers={}, + ) + + assert data["model"] == "accounts/fireworks/routers/glm-latest" + + +def test_transform_text_completion_request_bare_slug_stays_model(): + config = FireworksAITextCompletionConfig() + + data = config.transform_text_completion_request( + model="glm-4p6", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + headers={}, + ) + + assert data["model"] == "accounts/fireworks/models/glm-4p6" diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py new file mode 100644 index 00000000000..4af395baf41 --- /dev/null +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py @@ -0,0 +1,45 @@ +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.llms.fireworks_ai.common_utils import resolve_fireworks_resource_name + + +@pytest.mark.parametrize( + "model, expected", + [ + ("routers/glm-latest", "accounts/fireworks/routers/glm-latest"), + ("routers/firerouter", "accounts/fireworks/routers/firerouter"), + ("fireworks_ai/routers/glm-latest", "accounts/fireworks/routers/glm-latest"), + ("models/glm-4p6", "accounts/fireworks/models/glm-4p6"), + ("fireworks_ai/models/glm-4p6", "accounts/fireworks/models/glm-4p6"), + ("glm-4p6", "accounts/fireworks/models/glm-4p6"), + ("fireworks_ai/glm-4p6", "accounts/fireworks/models/glm-4p6"), + ("kimi-k2p6-fast", "accounts/fireworks/routers/kimi-k2p6-fast"), + ( + "accounts/fireworks/routers/glm-latest", + "accounts/fireworks/routers/glm-latest", + ), + ( + "accounts/fireworks/models/glm-4p6", + "accounts/fireworks/models/glm-4p6", + ), + ( + "fireworks_ai/accounts/fireworks/routers/glm-latest", + "accounts/fireworks/routers/glm-latest", + ), + ( + "accounts/fireworks/models/qwen2p5-coder-7b#accounts/gitlab/deployments/2fb7764c", + "accounts/fireworks/models/qwen2p5-coder-7b#accounts/gitlab/deployments/2fb7764c", + ), + ( + "glm-4p6#accounts/gitlab/deployments/2fb7764c", + "glm-4p6#accounts/gitlab/deployments/2fb7764c", + ), + ], +) +def test_resolve_fireworks_resource_name(model, expected): + assert resolve_fireworks_resource_name(model) == expected diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ab8198304bb..d9038e20df9 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -30,6 +30,6 @@ "limit": 16783 }, "LIT011": { - "limit": 5602 + "limit": 5599 } } From f5d98c0b8ce15164f25b880258fb88c24f03baeb Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 6 Aug 2026 14:25:11 -0700 Subject: [PATCH 039/504] feat(ui): migrate playground chat controls toward shadcn Continue the Playground Chat Ant Design/Tremor migration: shared MultiSelect, upload validation with semantic file inputs, collapsible message widgets, and AdditionalModelSettings on Base UI controls --- .../components/chat_ui/A2AMetrics.tsx | 237 +++++++------ .../chat_ui/AdditionalModelSettings.tsx | 224 +++++++----- .../components/chat_ui/ChatImageUpload.tsx | 87 +++-- .../components/chat_ui/ChatMessageBubble.tsx | 8 +- .../playground/components/chat_ui/ChatUI.tsx | 251 +++++++------- .../chat_ui/CodeInterpreterOutput.tsx | 215 ++++++------ .../chat_ui/CodeInterpreterTool.tsx | 27 +- .../components/chat_ui/EndpointSelector.tsx | 14 +- .../components/chat_ui/FilePreviewCard.tsx | 17 +- .../chat_ui/ResponsesImageUpload.tsx | 83 +++-- .../chat_ui/SearchResultsDisplay.tsx | 168 ++++----- .../components/chat_ui/SessionManagement.tsx | 73 ++-- .../chat_ui/uploadValidation.test.ts | 78 +++++ .../components/chat_ui/uploadValidation.ts | 97 ++++++ .../src/app/(dashboard)/playground/page.tsx | 10 +- .../components/chat_ui/MCPEventsDisplay.tsx | 323 ++++++++---------- .../components/chat_ui/ReasoningContent.tsx | 117 +++---- .../components/chat_ui/ResponseMetrics.tsx | 131 +++---- .../guardrails/GuardrailSelector.tsx | 13 +- .../src/components/llm_calls/fetch_models.tsx | 20 +- .../components/policies/PolicySelector.tsx | 13 +- .../src/components/shared/MultiSelect.tsx | 119 +++++++ .../src/components/shared/SearchSelect.tsx | 4 +- .../components/tag_management/TagSelector.tsx | 17 +- .../src/components/ui/combobox.tsx | 7 +- .../VectorStoreSelector.tsx | 17 +- 26 files changed, 1414 insertions(+), 956 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/uploadValidation.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/uploadValidation.ts create mode 100644 ui/litellm-dashboard/src/components/shared/MultiSelect.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/A2AMetrics.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/A2AMetrics.tsx index 004a513f061..6ddfe1442f5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/A2AMetrics.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/A2AMetrics.tsx @@ -1,17 +1,19 @@ import React, { useState } from "react"; -import { Tooltip, Button } from "antd"; import { - CheckCircleOutlined, - ClockCircleOutlined, - LoadingOutlined, - ExclamationCircleOutlined, - CopyOutlined, - DownOutlined, - RightOutlined, - LinkOutlined, - FileTextOutlined, - RobotOutlined, -} from "@ant-design/icons"; + Bot, + CheckCircle, + ChevronDown, + ChevronRight, + CircleAlert, + Clock, + Copy, + FileText, + Link, + LoaderCircle, +} from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; export interface A2ATaskMetadata { taskId?: string; @@ -21,7 +23,7 @@ export interface A2ATaskMetadata { timestamp?: string; message?: string; }; - metadata?: Record; + metadata?: Record; } interface A2AMetricsProps { @@ -33,15 +35,15 @@ interface A2AMetricsProps { const getStatusIcon = (state?: string) => { switch (state) { case "completed": - return ; + return ; case "working": case "submitted": - return ; + return ; case "failed": case "canceled": - return ; + return ; default: - return ; + return ; } }; @@ -91,7 +93,7 @@ const A2AMetrics: React.FC = ({ a2aMetadata, timeToFirstToken,
{/* A2A Metadata Header */}
- + A2A Metadata
@@ -109,28 +111,33 @@ const A2AMetrics: React.FC = ({ a2aMetadata, timeToFirstToken, {/* Timestamp */} {formattedTime && ( - - - + + }> + {formattedTime} - + + {status?.timestamp} )} {/* Latency */} {totalLatency !== undefined && ( - - - + + }> + {(totalLatency / 1000).toFixed(2)}s - + + Total latency )} {/* Time to first token */} {timeToFirstToken !== undefined && ( - - TTFT: {(timeToFirstToken / 1000).toFixed(2)}s + + }> + TTFT: {(timeToFirstToken / 1000).toFixed(2)}s + + Time to first token )}
@@ -139,95 +146,133 @@ const A2AMetrics: React.FC = ({ a2aMetadata, timeToFirstToken,
{/* Task ID */} {taskId && ( - - copyToClipboard(taskId)} + + copyToClipboard(taskId)} + aria-label={`Copy task ID ${taskId}`} + /> + } > - + Task: {truncateId(taskId)} - - + + + Click to copy: {taskId} )} {/* Context/Session ID */} {contextId && ( - - copyToClipboard(contextId)} + + copyToClipboard(contextId)} + aria-label={`Copy session ID ${contextId}`} + /> + } > - + Session: {truncateId(contextId)} - - + + + Click to copy: {contextId} )} {/* Details toggle */} {(metadata || status?.message) && ( - + + + } + > + {showDetails ? : } + Details + + )}
{/* Expandable details panel */} - {showDetails && ( -
- {/* Status message */} - {status?.message && ( -
- Status Message: - {status.message} -
- )} + + +
+ {/* Status message */} + {status?.message && ( +
+ Status Message: + {status.message} +
+ )} - {/* Full IDs */} - {taskId && ( -
- Task ID: - - {taskId} - - copyToClipboard(taskId)} - /> -
- )} + {/* Full IDs */} + {taskId && ( +
+ Task ID: + + {taskId} + + +
+ )} - {contextId && ( -
- Session ID: - - {contextId} - - copyToClipboard(contextId)} - /> -
- )} + {contextId && ( +
+ Session ID: + + {contextId} + + +
+ )} - {/* Metadata fields */} - {metadata && Object.keys(metadata).length > 0 && ( -
- Custom Metadata: -
-                {JSON.stringify(metadata, null, 2)}
-              
-
- )} -
- )} + {/* Metadata fields */} + {metadata && Object.keys(metadata).length > 0 && ( +
+ Custom Metadata: +
+                  {JSON.stringify(metadata, null, 2)}
+                
+
+ )} +
+ + ); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx index d4320110c4c..4deb7051954 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx @@ -1,7 +1,10 @@ -import { InfoCircleOutlined } from "@ant-design/icons"; -import { Text } from "@tremor/react"; -import { Checkbox, InputNumber, Popover, Slider, Tooltip, Typography } from "antd"; -import React, { useEffect, useState } from "react"; +import { Info } from "lucide-react"; +import React, { useEffect, useId, useState } from "react"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Input } from "@/components/ui/input"; +import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { cn } from "@/lib/cva.config"; interface AdditionalModelSettingsProps { temperature?: number; @@ -17,6 +20,10 @@ interface AdditionalModelSettingsProps { showAdvancedParams?: boolean; } +function clamp(value: number, min: number, max: number): number { + return Math.min(max, Math.max(min, value)); +} + const AdditionalModelSettings: React.FC = ({ temperature = 1.0, maxTokens = 2048, @@ -36,7 +43,12 @@ const AdditionalModelSettings: React.FC = ({ const [localTemperature, setLocalTemperature] = useState(temperature); const [localMaxTokens, setLocalMaxTokens] = useState(maxTokens); - // Sync local state with props when they change + const streamingId = useId(); + const advancedId = useId(); + const fallbacksId = useId(); + const temperatureId = useId(); + const maxTokensId = useId(); + useEffect(() => { setLocalTemperature(temperature); }, [temperature]); @@ -45,21 +57,18 @@ const AdditionalModelSettings: React.FC = ({ setLocalMaxTokens(maxTokens); }, [maxTokens]); - const handleTemperatureChange = (value: number | null) => { - const newValue = value ?? 1.0; + const handleTemperatureChange = (value: number) => { + const newValue = clamp(Number.isFinite(value) ? value : 1.0, 0, 2); setLocalTemperature(newValue); onTemperatureChange?.(newValue); }; - const handleMaxTokensChange = (value: number | null) => { - const newValue = value ?? 1000; + const handleMaxTokensChange = (value: number) => { + const newValue = clamp(Number.isFinite(value) ? Math.round(value) : 1000, 1, 32768); setLocalMaxTokens(newValue); onMaxTokensChange?.(newValue); }; - const disabledOpacity = useAdvancedParams ? 1 : 0.4; - const disabledTextColor = useAdvancedParams ? "text-gray-700" : "text-gray-400"; - const handleUseAdvancedParamsChange = (checked: boolean) => { if (onUseAdvancedParamsChange) { onUseAdvancedParamsChange(checked); @@ -68,129 +77,176 @@ const AdditionalModelSettings: React.FC = ({ } }; + const disabledTextColor = useAdvancedParams ? "text-gray-700" : "text-gray-400"; + return ( -
+
{onStreamingChange && ( -
- onStreamingChange(e.target.checked)}> - Stream responses - - - +
+ onStreamingChange(checked === true)} + aria-label="Stream responses" + /> + + + + + + + Streams the answer token by token. Uncheck to send a non-streaming request and render the full response at + once. +
)} {showAdvancedParams && ( - handleUseAdvancedParamsChange(e.target.checked)}> - Use Advanced Parameters - +
+ handleUseAdvancedParamsChange(checked === true)} + aria-label="Use Advanced Parameters" + /> + +
)} {onMockTestFallbacksChange && ( -
- onMockTestFallbacksChange(e.target.checked)}> - Simulate failure to test fallbacks - - - - Causes the first request to fail so the router tries fallbacks (if configured). Use this to verify - your fallback setup. - - - Behavior can differ when keys, teams, or router settings are configured.{" "} - - Learn more - - -
- } - > - +
+ onMockTestFallbacksChange(checked === true)} + aria-label="Simulate failure to test fallbacks" + /> + + + + + + +

+ Causes the first request to fail so the router tries fallbacks (if configured). Use this to verify your + fallback setup. +

+

+ Behavior can differ when keys, teams, or router settings are configured.{" "} + + Learn more + +

+
)} {showAdvancedParams && ( -
+
-
+
- Temperature - - + + + + + + + Controls randomness. Lower values make output more deterministic, higher values more creative. +
- handleTemperatureChange(Number(event.target.value))} />
- handleTemperatureChange(Number(event.target.value))} /> +
+ 0 + 1.0 + 2.0 +
-
+
- Max Tokens - - + + + + + + + Maximum number of tokens to generate in the response. +
- handleMaxTokensChange(Number(event.target.value))} />
- handleMaxTokensChange(Number(event.target.value))} /> +
+ 1 + 32768 +
)} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUpload.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUpload.tsx index 55527d997ac..6f210118281 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUpload.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUpload.tsx @@ -1,43 +1,70 @@ -import React from "react"; -import { Upload, Tooltip } from "antd"; -import { PaperClipOutlined } from "@ant-design/icons"; - -const { Dragger } = Upload; +import React, { useId, useRef } from "react"; +import { Paperclip } from "lucide-react"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { Button } from "@/components/ui/button"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { CHAT_ATTACHMENT_ACCEPT, validateChatAttachment } from "./uploadValidation"; interface ChatImageUploadProps { chatUploadedImage: File | null; chatImagePreviewUrl: string | null; - onImageUpload: (file: File) => false; + onImageUpload: (file: File) => void; onRemoveImage: () => void; + disabled?: boolean; } -const ChatImageUpload: React.FC = ({ - chatUploadedImage, - chatImagePreviewUrl, - onImageUpload, - onRemoveImage, -}) => { +const ChatImageUpload: React.FC = ({ chatUploadedImage, onImageUpload, disabled = false }) => { + const inputRef = useRef(null); + const inputId = useId(); + + if (chatUploadedImage) { + return null; + } + + const handleFileChange = (event: React.ChangeEvent) => { + const file = event.target.files?.[0]; + event.target.value = ""; + if (!file) { + return; + } + const result = validateChatAttachment(file); + if (!result.ok) { + NotificationsManager.error(result.error); + return; + } + onImageUpload(file); + }; + return ( <> - {/* Subtle upload button - only show when no image */} - {!chatUploadedImage && ( - - - - - - )} + variant="ghost" + size="icon-sm" + disabled={disabled} + aria-label="Attach image or PDF" + className="text-gray-400 hover:text-gray-600" + onClick={() => inputRef.current?.click()} + /> + } + > + + + Attach image or PDF + ); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx index 8e71017a7b5..c438b4982bf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx @@ -41,9 +41,9 @@ function ChatMessageBubble({ const isUser = message.role === "user"; return ( -
+
{/* Header: role icon + name + model badge */} -
+
{message.role} {message.role === "assistant" && message.model && ( - + {message.model} )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index 57ff7906eda..5edcbe84aa8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -66,7 +66,16 @@ import { MessageType } from "@/components/chat_ui/types"; import { useCodeInterpreter } from "../../hooks/useCodeInterpreter"; import { useChatHistory } from "../../hooks/useChatHistory"; import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; +import { SearchSelect } from "@/components/shared/SearchSelect"; +import { Select as ShadcnSelect, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; +import { + AUDIO_ACCEPT, + IMAGE_EDIT_ACCEPT, + validateAudioFile, + validateChatAttachment, + validateImageEditFile, +} from "./uploadValidation"; const { TextArea } = Input; const { Dragger } = Upload; @@ -177,6 +186,8 @@ const ChatUI: React.FC = ({ const [selectedModel, setSelectedModel] = useState(simplified ? fixedModel : undefined); const [showCustomModelInput, setShowCustomModelInput] = useState(false); const [modelInfo, setModelInfo] = useState([]); + const [isLoadingModels, setIsLoadingModels] = useState(false); + const [modelLoadError, setModelLoadError] = useState(false); const [agentInfo, setAgentInfo] = useState([]); const [selectedAgent, setSelectedAgent] = useState(undefined); const debouncedSetSelectedModel = useDebouncedCallback((value: string) => setSelectedModel(value), { @@ -388,17 +399,17 @@ const ChatUI: React.FC = ({ ]); useEffect(() => { - let userApiKey = apiKeySource === "session" ? accessToken : apiKey; - if (!userApiKey || !token || !userRole || !userID) { + const userApiKey = apiKeySource === "session" ? accessToken : apiKey.trim(); + if (!userApiKey) { + setModelInfo([]); + setModelLoadError(false); return; } - // Fetch model info and set the default selected model (skip in simplified mode; we use fixedModel) const loadModels = async () => { + setIsLoadingModels(true); + setModelLoadError(false); try { - if (!userApiKey) { - return; - } const uniqueModels = await fetchAvailableModels(userApiKey); setModelInfo(uniqueModels); @@ -412,6 +423,10 @@ const ChatUI: React.FC = ({ } } catch (error) { console.error("Error fetching model info:", error); + setModelInfo([]); + setModelLoadError(true); + } finally { + setIsLoadingModels(false); } }; @@ -419,7 +434,7 @@ const ChatUI: React.FC = ({ loadModels(); } loadMCPServers(); - }, [accessToken, userID, userRole, apiKeySource, apiKey, token, simplified]); + }, [accessToken, apiKeySource, apiKey, simplified]); // Load tools when MCP direct mode has a server (or toolset) selected useEffect(() => { @@ -494,13 +509,35 @@ const ChatUI: React.FC = ({ } }; - const handleImageUpload = (file: File) => { - setUploadedImages((prev) => [...prev, file]); + const createBlobPreviewUrl = (file: File): string => { const rawPreviewUrl = URL.createObjectURL(file); - // Sanitize: only allow blob: URLs to prevent XSS via img src injection. - const previewUrl = rawPreviewUrl.startsWith("blob:") ? rawPreviewUrl : ""; - setImagePreviewUrls((prev) => [...prev, previewUrl]); - return false; // Prevent default upload behavior + return rawPreviewUrl.startsWith("blob:") ? rawPreviewUrl : ""; + }; + + const handleImageFiles = (files: File[]) => { + let nextCount = uploadedImages.length; + const accepted: File[] = []; + const previews: string[] = []; + for (const file of files) { + const result = validateImageEditFile(file, nextCount); + if (!result.ok) { + NotificationsManager.error(result.error); + continue; + } + accepted.push(file); + previews.push(createBlobPreviewUrl(file)); + nextCount += 1; + } + if (accepted.length === 0) { + return; + } + setUploadedImages((prev) => [...prev, ...accepted]); + setImagePreviewUrls((prev) => [...prev, ...previews]); + }; + + const handleImageUpload = (file: File): false => { + handleImageFiles([file]); + return false; }; const handleRemoveImage = (index: number) => { @@ -519,11 +556,14 @@ const ChatUI: React.FC = ({ setImagePreviewUrls([]); }; - const handleResponsesImageUpload = (file: File): false => { + const handleResponsesImageUpload = (file: File): void => { + const result = validateChatAttachment(file); + if (!result.ok) { + NotificationsManager.error(result.error); + return; + } setResponsesUploadedImage(file); - const previewUrl = URL.createObjectURL(file); - setResponsesImagePreviewUrl(previewUrl); - return false; // Prevent default upload behavior + setResponsesImagePreviewUrl(createBlobPreviewUrl(file)); }; const handleRemoveResponsesImage = () => { @@ -534,11 +574,14 @@ const ChatUI: React.FC = ({ setResponsesImagePreviewUrl(null); }; - const handleChatImageUpload = (file: File): false => { + const handleChatImageUpload = (file: File): void => { + const result = validateChatAttachment(file); + if (!result.ok) { + NotificationsManager.error(result.error); + return; + } setChatUploadedImage(file); - const previewUrl = URL.createObjectURL(file); - setChatImagePreviewUrl(previewUrl); - return false; // Prevent default upload behavior + setChatImagePreviewUrl(createBlobPreviewUrl(file)); }; const handleRemoveChatImage = () => { @@ -550,8 +593,13 @@ const ChatUI: React.FC = ({ }; const handleAudioUpload = (file: File): false => { + const result = validateAudioFile(file); + if (!result.ok) { + NotificationsManager.error(result.error); + return false; + } setUploadedAudio(file); - return false; // Prevent default upload behavior + return false; }; const handleRemoveAudio = () => { @@ -1002,8 +1050,12 @@ const ChatUI: React.FC = ({ const onModelChange = (value: string) => { setSelectedModel(value); - setShowCustomModelInput(value === "custom"); + + const model = modelInfo.find((option) => option.model_group === value); + if (model?.mode) { + setEndpointType(getEndpointType(model.mode)); + } }; // Check if the selected model is a chat model @@ -1020,35 +1072,43 @@ const ChatUI: React.FC = ({ }; const supportsStreamingToggle = endpointType === EndpointType.CHAT || endpointType === EndpointType.RESPONSES; + let modelEmptyText = "No models available for this key"; + if (modelLoadError) { + modelEmptyText = "Unable to load models for this key"; + } else if (apiKeySource === "custom" && !apiKey.trim()) { + modelEmptyText = "Enter a Virtual Key to load models"; + } const antIcon = ; return ( -
- -
+
+ +
{/* Left Sidebar with Controls - hidden in simplified mode */} {!simplified && ( -
+
Configurations
Virtual Key Source - { + onValueChange={(value) => { setSelectedVoice(value); sessionStorage.setItem("selectedVoice", value); }} - style={{ width: "100%" }} - className="rounded-md" - options={OPEN_AI_VOICE_SELECT_OPTIONS} - /> + > + + + + + {OPEN_AI_VOICE_SELECT_OPTIONS.map((voice) => ( + + {voice.label} + + ))} + +
)} @@ -1212,46 +1280,20 @@ const ChatUI: React.FC = ({ )} - setSelectedAgent(value)} + onValueChange={(value) => setSelectedAgent(value)} options={agentInfo.map((agent) => ({ value: agent.agent_name, label: agent.agent_name || agent.agent_id, - key: agent.agent_id, + sublabel: agent.agent_card_params?.description, }))} - style={{ width: "100%" }} - showSearch={true} - className="rounded-md" - optionLabelProp="label" - > - {agentInfo.map((agent) => ( - -
- {agent.agent_name || agent.agent_id} - {agent.agent_card_params?.description && ( - {agent.agent_card_params.description} - )} -
-
- ))} - + /> {agentInfo.length === 0 && ( No agents found. Create agents via /v1/agents endpoint. @@ -1697,7 +1720,7 @@ const ChatUI: React.FC = ({ )} {/* Main Chat Area */} -
+
{endpointType === EndpointType.REALTIME ? ( = ({ /> ) : ( <> -
+
{simplified ? "Chat" : "Test Key"} -
+
= ({ )}
-
+
{chatHistory.length === 0 && (
@@ -1788,18 +1811,18 @@ const ChatUI: React.FC = ({
-
+
{/* Image Upload Section for Image Edits */} {endpointType === EndpointType.IMAGE_EDITS && (
{uploadedImages.length === 0 ? ( - +

Click or drag images to upload

- Support for PNG, JPG, JPEG formats. Multiple images supported. + Support for PNG, JPG, JPEG, GIF, WebP. Multiple images supported.

) : ( @@ -1840,12 +1863,12 @@ const ChatUI: React.FC = ({ { - const files = Array.from(e.target.files || []); - files.forEach((file) => handleImageUpload(file)); + handleImageFiles(Array.from(e.target.files || [])); + e.target.value = ""; }} />
@@ -1858,11 +1881,7 @@ const ChatUI: React.FC = ({ {endpointType === EndpointType.TRANSCRIPTION && (
{!uploadedAudio ? ( - +

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx index c27273a116d..8dc503f369c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx @@ -1,15 +1,10 @@ -import React, { useState, useEffect } from "react"; -import { Collapse, Spin } from "antd"; -import { - CodeOutlined, - DownloadOutlined, - FileImageOutlined, - FileTextOutlined, - LoadingOutlined, -} from "@ant-design/icons"; +import React, { useEffect, useState } from "react"; +import { Code, Download, FileImage, FileText, Loader2 } from "lucide-react"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; +import { Button } from "@/components/ui/button"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; interface ContainerFileCitation { type: "container_file_citation"; @@ -27,48 +22,60 @@ interface CodeInterpreterOutputProps { accessToken: string; } -const CodeInterpreterOutput: React.FC = ({ - code, - containerId, - annotations = [], - accessToken, -}) => { +const IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".gif"] as const; + +function isImageFilename(filename: string | undefined): boolean { + if (!filename) { + return false; + } + const lower = filename.toLowerCase(); + return IMAGE_EXTENSIONS.some((ext) => lower.endsWith(ext)); +} + +const CodeInterpreterOutput: React.FC = ({ code, annotations = [], accessToken }) => { const [imageUrls, setImageUrls] = useState>({}); const [loadingImages, setLoadingImages] = useState>({}); + const [codeOpen, setCodeOpen] = useState(false); const proxyBaseUrl = getProxyBaseUrl(); - // Fetch images from container files API useEffect(() => { + const createdUrls: string[] = []; + let cancelled = false; + const fetchImages = async () => { for (const annotation of annotations) { - const isImage = - annotation.filename?.toLowerCase().endsWith(".png") || - annotation.filename?.toLowerCase().endsWith(".jpg") || - annotation.filename?.toLowerCase().endsWith(".jpeg") || - annotation.filename?.toLowerCase().endsWith(".gif"); + if (!isImageFilename(annotation.filename) || !annotation.container_id || !annotation.file_id) { + continue; + } - if (isImage && annotation.container_id && annotation.file_id) { + if (!cancelled) { setLoadingImages((prev) => ({ ...prev, [annotation.file_id]: true })); + } - try { - // Fetch image content from container files API - const response = await fetch( - `${proxyBaseUrl}/v1/containers/${annotation.container_id}/files/${annotation.file_id}/content`, - { - headers: { - [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, - }, + try { + const response = await fetch( + `${proxyBaseUrl}/v1/containers/${annotation.container_id}/files/${annotation.file_id}/content`, + { + headers: { + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, }, - ); + }, + ); - if (response.ok) { - const blob = await response.blob(); - const url = URL.createObjectURL(blob); + if (response.ok) { + const blob = await response.blob(); + const url = URL.createObjectURL(blob); + createdUrls.push(url); + if (!cancelled) { setImageUrls((prev) => ({ ...prev, [annotation.file_id]: url })); + } else { + URL.revokeObjectURL(url); } - } catch (error) { - console.error("Error fetching image:", error); - } finally { + } + } catch (error) { + console.error("Error fetching image:", error); + } finally { + if (!cancelled) { setLoadingImages((prev) => ({ ...prev, [annotation.file_id]: false })); } } @@ -76,12 +83,12 @@ const CodeInterpreterOutput: React.FC = ({ }; if (annotations.length > 0 && accessToken) { - fetchImages(); + void fetchImages(); } - // Cleanup URLs on unmount return () => { - Object.values(imageUrls).forEach((url) => URL.revokeObjectURL(url)); + cancelled = true; + createdUrls.forEach((url) => URL.revokeObjectURL(url)); }; }, [annotations, accessToken, proxyBaseUrl]); @@ -112,22 +119,8 @@ const CodeInterpreterOutput: React.FC = ({ } }; - // Separate images and other files - const imageAnnotations = annotations.filter( - (a) => - a.filename?.toLowerCase().endsWith(".png") || - a.filename?.toLowerCase().endsWith(".jpg") || - a.filename?.toLowerCase().endsWith(".jpeg") || - a.filename?.toLowerCase().endsWith(".gif"), - ); - - const fileAnnotations = annotations.filter( - (a) => - !a.filename?.toLowerCase().endsWith(".png") && - !a.filename?.toLowerCase().endsWith(".jpg") && - !a.filename?.toLowerCase().endsWith(".jpeg") && - !a.filename?.toLowerCase().endsWith(".gif"), - ); + const imageAnnotations = annotations.filter((a) => isImageFilename(a.filename)); + const fileAnnotations = annotations.filter((a) => !isImageFilename(a.filename)); if (!code && annotations.length === 0) { return null; @@ -135,44 +128,46 @@ const CodeInterpreterOutput: React.FC = ({ return (
- {/* Executed Code - Collapsible */} {code && ( - - Python Code Executed - - ), - children: ( - - {code} - - ), - }, - ]} - /> + + + } + > + + Python Code Executed + + +
+ + {code} + +
+
+
)} - {/* Generated Images */} {imageAnnotations.map((annotation) => ( -
+
{loadingImages[annotation.file_id] ? ( -
- } /> +
+
) : imageUrls[annotation.file_id] ? ( @@ -180,42 +175,48 @@ const CodeInterpreterOutput: React.FC = ({ {annotation.filename -
- - {annotation.filename} +
+ + - + + Download +
) : ( -
+
Image not available
)}
))} - {/* Download Links for Other Files */} {fileAnnotations.length > 0 && (
{fileAnnotations.map((annotation) => ( - +
)} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterTool.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterTool.tsx index d2682e3a7a8..e5744ac8e38 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterTool.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterTool.tsx @@ -1,8 +1,8 @@ import React from "react"; -import { Switch, Tooltip } from "antd"; import MessageManager from "@/components/molecules/message_manager"; -import { CodeOutlined, InfoCircleOutlined, ExclamationCircleOutlined } from "@ant-design/icons"; -import { Text } from "@tremor/react"; +import { Code, Info, TriangleAlert } from "lucide-react"; +import { Switch } from "@/components/ui/switch"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; interface CodeInterpreterToolProps { accessToken: string; @@ -49,25 +49,30 @@ const CodeInterpreterTool: React.FC = ({
- - Code Interpreter - - + + Code Interpreter + + + + + + Run Python code to generate files, charts, and analyze data. Container is created automatically. +
{!isOpenAI && (
- +
Code Interpreter is currently only supported for OpenAI models. = ({ endpointType, onEndpointChange, className }) => { return (
- + { return { label: `${guardrail.guardrail_name}`, value: guardrail.guardrail_name, }; })} - optionFilterProp="label" - showSearch - style={{ width: "100%" }} />
); diff --git a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx index 0de98330c2e..24f1e038f85 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx @@ -7,6 +7,13 @@ export interface ModelGroup { mode?: string; } +interface AvailableModel { + model_group?: string | null; + model_name?: string | null; + id?: string | null; + mode?: string | null; +} + /** * Fetches available models using modelHubCall and formats them for the selection dropdown. */ @@ -15,14 +22,15 @@ export const fetchAvailableModels = async (accessToken: string): Promise 0) { - const models: ModelGroup[] = fetchedModels.data.map((item: any) => ({ - model_group: item.model_group, // Display the model_group to the user - mode: item?.mode, // Save the mode for auto-selection of endpoint type - })); + const models: ModelGroup[] = fetchedModels.data + .map((item: AvailableModel) => ({ + model_group: item.model_group || item.id || item.model_name || "", + mode: item.mode || undefined, + })) + .filter((model: ModelGroup) => model.model_group !== ""); - // Sort models alphabetically by label models.sort((a, b) => a.model_group.localeCompare(b.model_group)); - return models; + return Array.from(new Map(models.map((model) => [model.model_group, model])).values()); } return []; } catch (error) { diff --git a/ui/litellm-dashboard/src/components/policies/PolicySelector.tsx b/ui/litellm-dashboard/src/components/policies/PolicySelector.tsx index 132538d439f..1e565d938d7 100644 --- a/ui/litellm-dashboard/src/components/policies/PolicySelector.tsx +++ b/ui/litellm-dashboard/src/components/policies/PolicySelector.tsx @@ -1,7 +1,7 @@ import React, { useEffect, useState } from "react"; -import { Select } from "antd"; import { Policy } from "./types"; import { getPoliciesList } from "../networking"; +import { MultiSelect } from "@/components/shared/MultiSelect"; /** Prefix for policy version IDs in request body; must match backend POLICY_VERSION_ID_PREFIX. */ export const POLICY_VERSION_ID_PREFIX = "policy_"; @@ -80,22 +80,17 @@ const PolicySelector: React.FC = ({ }; return ( -
- ({ label: tag.name, value: tag.name, - title: tag.description || tag.name, + description: tag.description || undefined, }))} - optionFilterProp="label" - tokenSeparators={[","]} - maxTagCount="responsive" - allowClear - style={{ width: "100%" }} /> ); }; diff --git a/ui/litellm-dashboard/src/components/ui/combobox.tsx b/ui/litellm-dashboard/src/components/ui/combobox.tsx index 2854928140e..541ad8bb25c 100644 --- a/ui/litellm-dashboard/src/components/ui/combobox.tsx +++ b/ui/litellm-dashboard/src/components/ui/combobox.tsx @@ -83,10 +83,14 @@ function ComboboxContent({ sideOffset = 6, align = "start", alignOffset = 0, + collisionAvoidance, anchor, ...props }: ComboboxPrimitive.Popup.Props & - Pick) { + Pick< + ComboboxPrimitive.Positioner.Props, + "side" | "align" | "sideOffset" | "alignOffset" | "collisionAvoidance" | "anchor" + >) { return ( diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.tsx index 2642b74e492..d80996f6a28 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.tsx @@ -1,7 +1,7 @@ import React, { useEffect, useState } from "react"; -import { Select } from "antd"; import { VectorStore } from "./types"; import { vectorStoreListCall } from "../networking"; +import { MultiSelect } from "@/components/shared/MultiSelect"; interface VectorStoreSelectorProps { onChange: (selectedVectorStores: string[]) => void; value?: string[]; @@ -43,24 +43,19 @@ const VectorStoreSelector: React.FC = ({ }, [accessToken]); return ( -
- setApiKey(event.target.value)} + value={apiKey} + /> +
)}
-
- - Custom Proxy Base URL - +
+ {proxySettings?.LITELLM_UI_API_DOC_BASE_URL && !customProxyBaseUrl && ( )} {customProxyBaseUrl && ( )}
- { - setCustomProxyBaseUrl(value); - sessionStorage.setItem("customProxyBaseUrl", value); - }} - value={customProxyBaseUrl} - icon={ApiOutlined} - /> +
+ + { + setCustomProxyBaseUrl(event.target.value); + sessionStorage.setItem("customProxyBaseUrl", event.target.value); + }} + /> +
{customProxyBaseUrl && ( - API calls will be sent to: {customProxyBaseUrl} +

API calls will be sent to: {customProxyBaseUrl}

)}
- - Endpoint Type - + { setEndpointType(value); - // Clear model/agent selection when switching endpoint type setSelectedModel(undefined); setSelectedAgent(undefined); setShowCustomModelInput(false); setSelectedMCPDirectTool(undefined); - // For MCP direct mode, require single server (clear __all__ or multiple) if (value === EndpointType.MCP) { setSelectedMCPServers((prev) => (prev.length === 1 && prev[0] !== "__all__" ? prev : [])); } @@ -1194,13 +1279,12 @@ const ChatUI: React.FC = ({ className="mb-4" /> - {/* Voice Selector for Speech Endpoint */} {endpointType === EndpointType.SPEECH && (
- - + + { @@ -1222,7 +1306,6 @@ const ChatUI: React.FC = ({
)} - {/* Session Management Component */} = ({ />
- {/* Model Selector - shown when NOT using A2A Agents or MCP direct mode */} {endpointType !== EndpointType.A2A_AGENTS && endpointType !== EndpointType.MCP && (
- +
- Select Model + {isChatModel() || supportsStreamingToggle ? ( - + + } + > + + + +
Model Settings
= ({ streamingEnabled={streamingEnabled} onStreamingChange={supportsStreamingToggle ? setStreamingEnabled : undefined} /> - } - title="Model Settings" - trigger="click" - placement="right" - > -
= ({ ]} /> {showCustomModelInput && ( - debouncedSetSelectedModel(event.target.value)} /> )}
)} - {/* Agent Selector - shown ONLY for A2A Agents endpoint */} {endpointType === EndpointType.A2A_AGENTS && (
- - Select Agent - + = ({ }))} /> {agentInfo.length === 0 && ( - +

No agents found. Create agents via /v1/agents endpoint. - +

)}
)}
- - Tags - + = ({ />
- {/* MCP Server Selection */}
- - +
+
)} - {/* BYOK credential status for selected servers */} {selectedMCPServers.length > 0 && !selectedMCPServers.includes("__all__") && selectedMCPServers.some((serverId) => { @@ -1593,28 +1571,31 @@ const ChatUI: React.FC = ({ return (
- {serverName} requires your API key +

{serverName} requires your API key

{server.has_user_credential ? (
- - Connected + + Connected
) : ( - + )}
); @@ -1624,23 +1605,21 @@ const ChatUI: React.FC = ({
- - Vector Store - - Select vector store(s) to use for this LLM API call. You can set up your vector store{" "} - - here - - . - - } - > - +
+
= ({
- - Guardrails - - Select guardrail(s) to use for this LLM API call. You can set up your guardrails{" "} - - here - - . - - } - > - +
+
= ({
- - Policies - - Select policy/policies to apply to this LLM API call. Policies define which guardrails are - applied based on conditions. You can set up your policies{" "} - - here - - . - - } - > - +
+
= ({ />
- {/* Code Interpreter Toggle - Only for Responses endpoint */} {endpointType === EndpointType.RESPONSES && (
= ({
)} - {/* Main Chat Area */}
{endpointType === EndpointType.REALTIME ? ( = ({ ) : ( <>
- {simplified ? "Chat" : "Test Key"} +

{simplified ? "Chat" : "Test Key"}

- + {!simplified && ( - setIsGetCodeModalVisible(true)} - className="bg-gray-100 hover:bg-gray-200 text-gray-700 border-gray-300" - icon={CodeOutlined} - > + )}
{chatHistory.length === 0 && ( -
- - Start a conversation, generate an image, or handle audio +
+
)} @@ -1772,29 +1739,26 @@ const ChatUI: React.FC = ({
))} - {/* Show MCP events during loading if no assistant message exists yet */} {isLoading && mcpEvents.length > 0 && (endpointType === EndpointType.RESPONSES || endpointType === EndpointType.CHAT) && chatHistory.length > 0 && chatHistory[chatHistory.length - 1].role === "user" && ( -
+
-
+
- +
Assistant
@@ -1804,27 +1768,34 @@ const ChatUI: React.FC = ({ )} {isLoading && ( -
- +
+
)}
- {/* Image Upload Section for Image Edits */} {endpointType === EndpointType.IMAGE_EDITS && (
{uploadedImages.length === 0 ? ( - -

- -

-

Click or drag images to upload

-

+

Click or drag images to upload

+

Support for PNG, JPG, JPEG, GIF, WebP. Multiple images supported.

-
+ { + handleImageFiles(Array.from(event.target.files || [])); + event.target.value = ""; + }} + /> + ) : (
{uploadedImages.map((file, index) => ( @@ -1841,76 +1812,83 @@ const ChatUI: React.FC = ({ } })()} alt={`Upload preview ${index + 1}`} - className="max-w-32 max-h-32 rounded-md border border-gray-200 object-cover" + className="max-h-32 max-w-32 rounded-md border border-gray-200 object-cover" /> - + +
))} - {/* Add more images button */} -
document.getElementById("additional-image-upload")?.click()} - > -
- -

Add more

-
+
+
)}
)} - {/* Audio Upload Section for Transcriptions */} {endpointType === EndpointType.TRANSCRIPTION && (
{!uploadedAudio ? ( - -

- -

-

Click or drag audio file to upload

-

+

Click or drag audio file to upload

+

Support for MP3, MP4, MPEG, MPGA, M4A, WAV, WEBM formats. Max file size: 25 MB.

-
+ { + const file = event.target.files?.[0]; + if (file) handleAudioUpload(file); + event.target.value = ""; + }} + /> + ) : ( -
-
- +
+
+
- + + Remove +
)}
)} - {/* Show file previews above input when files are uploaded */} {endpointType === EndpointType.RESPONSES && responsesUploadedImage && ( = ({ /> )} - {/* Code Interpreter indicator and sample prompts when enabled */} {endpointType === EndpointType.RESPONSES && codeInterpreter.enabled && (
-
+
{isLoading ? ( <> - - Running Python code... +
- {/* Sample prompts - only show when not loading */} {!isLoading && (
{[ @@ -1961,7 +1938,8 @@ const ChatUI: React.FC = ({ ].map((prompt, idx) => (
)} - {/* Suggested prompts - show when chat is empty and not loading (skip for MCP - uses structured form) */} {chatHistory.length === 0 && !isLoading && endpointType !== EndpointType.MCP && ( -
+
{(endpointType === EndpointType.A2A_AGENTS ? ["What can you help me with?", "Tell me about yourself", "What tasks can you perform?"] : ["Write me a poem", "Explain quantum computing", "Draft a polite email requesting a meeting"] @@ -1982,7 +1959,7 @@ const ChatUI: React.FC = ({ + + + + {codeInterpreter.enabled + ? "Code Interpreter enabled (click to disable)" + : "Enable Code Interpreter"} + )}
- {/* Middle: input field or MCP structured form */} {endpointType === EndpointType.MCP && selectedMCPServers.length === 1 && selectedMCPServers[0] !== "__all__" && selectedMCPDirectTool ? ( -
+
{(() => { const rawSel = selectedMCPServers[0]; - let toolPool: any[] = []; + let toolPool: { name: string }[] = []; if (rawSel.startsWith("toolset:")) { const toolsetId = rawSel.slice("toolset:".length); const toolset = mcpToolsets.find((t) => t.toolset_id === toolsetId); @@ -2060,82 +2043,51 @@ const ChatUI: React.FC = ({ } else { toolPool = serverToolsMap[rawSel] || []; } - const mcpTool = toolPool.find((t: any) => t.name === selectedMCPDirectTool); + const mcpTool = toolPool.find((t) => t.name === selectedMCPDirectTool); return mcpTool ? ( ) : ( -
+
Loading tool schema...
); })()}
) : ( -