From 71ee1a852a02787097c025d71f24202fdfc6b8c0 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 25 Jun 2026 17:35:15 -0700 Subject: [PATCH 01/16] fix(proxy/client): redact api key from key/info client error messages (#31342) * fix(proxy/client): redact api key from key/info client error messages The keys management client builds GET /key/info?key= and lets the requests HTTPError propagate. str(HTTPError) renders the failing request URL verbatim ("... for url: .../key/info?key=sk-..."), so any caller that logs the exception leaks the full key; the 401 branch leaked the same way through UnauthorizedError(str(orig_exception)) Redact both branches with the existing redact_secrets helper so the secret-bearing query param is scrubbed to ?REDACTED while the status code, reason, and response object are preserved. Server-side responses already mask the key, so this closes the remaining client-side surface * fix: preserve key info unauthorized response --------- Co-authored-by: Cursor Agent --- litellm/proxy/client/exceptions.py | 20 +++- litellm/proxy/client/keys.py | 9 +- tests/test_litellm/proxy/client/test_keys.py | 96 +++++++++++++++++++- 3 files changed, 118 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/client/exceptions.py b/litellm/proxy/client/exceptions.py index fffd1b78b7e..c4089381e30 100644 --- a/litellm/proxy/client/exceptions.py +++ b/litellm/proxy/client/exceptions.py @@ -2,18 +2,30 @@ from typing import Union import requests +from litellm.litellm_core_utils.secret_redaction import redact_string + + +def _redact_orig_exception( + orig_exception: Union[requests.exceptions.HTTPError, str], +) -> Union[requests.exceptions.HTTPError, str]: + if isinstance(orig_exception, requests.exceptions.HTTPError): + return requests.exceptions.HTTPError( + redact_string(str(orig_exception)), response=orig_exception.response + ) + return redact_string(str(orig_exception)) + class UnauthorizedError(Exception): """Exception raised when the API returns a 401 Unauthorized response.""" def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]): - self.orig_exception = orig_exception - super().__init__(str(orig_exception)) + self.orig_exception = _redact_orig_exception(orig_exception) + super().__init__(str(self.orig_exception)) class NotFoundError(Exception): """Exception raised when the API returns a 404 Not Found response or indicates a resource was not found.""" def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]): - self.orig_exception = orig_exception - super().__init__(str(orig_exception)) + self.orig_exception = _redact_orig_exception(orig_exception) + super().__init__(str(self.orig_exception)) diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index d8687cbad16..845b49d1581 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -2,6 +2,8 @@ from typing import Any, Dict, List, Optional, Union import requests +from litellm.litellm_core_utils.secret_redaction import redact_string + from .exceptions import UnauthorizedError @@ -314,6 +316,9 @@ class KeysManagementClient: response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: + redacted_message = redact_string(str(e)) if e.response.status_code == 401: - raise UnauthorizedError(e) - raise + raise UnauthorizedError(e) from None + raise requests.exceptions.HTTPError( + redacted_message, response=e.response + ) from None diff --git a/tests/test_litellm/proxy/client/test_keys.py b/tests/test_litellm/proxy/client/test_keys.py index 136408e01ac..620daefb39e 100644 --- a/tests/test_litellm/proxy/client/test_keys.py +++ b/tests/test_litellm/proxy/client/test_keys.py @@ -1,5 +1,6 @@ import os import sys +import traceback import pytest import requests @@ -11,7 +12,7 @@ sys.path.insert( import responses -from litellm.proxy.client.exceptions import UnauthorizedError +from litellm.proxy.client.exceptions import NotFoundError, UnauthorizedError from litellm.proxy.client.keys import KeysManagementClient @@ -420,3 +421,96 @@ def test_info_server_error(client): ) with pytest.raises(requests.exceptions.HTTPError): client.info(key="test-key") + + +LEAKY_KEY = "sk-1234567890abcdefghijklmnop" + + +def _render_full_traceback(exc: BaseException) -> str: + return "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)) + + +@responses.activate +def test_info_not_found_redacts_key_everywhere(client): + """A 404 must not echo the raw key embedded in the request URL. + + Covers str(exc) and the rendered traceback, since the chain through + __cause__ / __context__ is what logging.exception() and the default + excepthook print. + """ + responses.add( + responses.GET, + f"{client._base_url}/key/info?key={LEAKY_KEY}", + status=404, + json={"error": {"message": "Key not found", "code": "404"}}, + ) + with pytest.raises(requests.exceptions.HTTPError) as excinfo: + client.info(key=LEAKY_KEY) + + exc = excinfo.value + assert LEAKY_KEY not in str(exc) + assert "REDACTED" in str(exc) + assert LEAKY_KEY not in _render_full_traceback(exc) + assert exc.__cause__ is None and exc.__suppress_context__ + assert exc.response is not None + assert exc.response.status_code == 404 + assert exc.request is not None + # Known residual: the live request URL still carries the key, since the + # response is preserved so callers keep status_code / text. str(exc) and the + # traceback are scrubbed; the URL-borne key is the root issue tracked in + # LIT-4013 (move the lookup key out of the query string server-side). + assert LEAKY_KEY in exc.response.request.url + + +@responses.activate +def test_info_unauthorized_redacts_key_everywhere(client): + """A 401 surfaced as UnauthorizedError must not echo the raw key in the + message, the retained original, or the rendered traceback chain.""" + responses.add( + responses.GET, + f"{client._base_url}/key/info?key={LEAKY_KEY}", + status=401, + json={"error": "Unauthorized"}, + ) + with pytest.raises(UnauthorizedError) as excinfo: + client.info(key=LEAKY_KEY) + + exc = excinfo.value + assert LEAKY_KEY not in str(exc) + assert "REDACTED" in str(exc) + assert LEAKY_KEY not in str(exc.orig_exception) + assert LEAKY_KEY not in _render_full_traceback(exc) + assert exc.__cause__ is None and exc.__suppress_context__ + assert isinstance(exc.orig_exception, requests.exceptions.HTTPError) + assert exc.orig_exception.response is not None + assert exc.orig_exception.response.status_code == 401 + + +def _http_error_with_key(prefix: str, status: int) -> requests.exceptions.HTTPError: + resp = requests.Response() + resp.status_code = status + return requests.exceptions.HTTPError( + f"{prefix} for url: http://x/key/info?key={LEAKY_KEY}", response=resp + ) + + +def test_unauthorized_error_redacts_wrapped_key(): + """UnauthorizedError scrubs the key in str(exc) and in the retained + orig_exception, while preserving the response for structured access.""" + wrapped = UnauthorizedError( + _http_error_with_key("401 Client Error: Unauthorized", 401) + ) + assert LEAKY_KEY not in str(wrapped) + assert "REDACTED" in str(wrapped) + assert LEAKY_KEY not in str(wrapped.orig_exception) + assert wrapped.orig_exception.response.status_code == 401 + + +def test_not_found_error_redacts_wrapped_key(): + """NotFoundError scrubs the key in str(exc) and in the retained + orig_exception, while preserving the response for structured access.""" + wrapped = NotFoundError(_http_error_with_key("404 Client Error: Not Found", 404)) + assert LEAKY_KEY not in str(wrapped) + assert "REDACTED" in str(wrapped) + assert LEAKY_KEY not in str(wrapped.orig_exception) + assert wrapped.orig_exception.response.status_code == 404 From 9203488578471390c2a3aebb305288ebb0fe0526 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 25 Jun 2026 17:45:37 -0700 Subject: [PATCH 02/16] feat(spend): store litellm_call_id on spend logs for DB-to-trace correlation (#31344) * feat(spend): store litellm_call_id on spend logs for DB-to-trace correlation Successful spend logs keyed request_id to the provider response id while tracing uses x-litellm-call-id, so a DB row could not be correlated with its trace; this only worked for failures, where request_id already fell back to the call id. Add a nullable litellm_call_id column to LiteLLM_SpendLogs, populate it in get_logging_payload, and surface it in the spend logs read endpoints so correlation works both directions for successful calls Fixes LIT-3868 * chore: sync schema.prisma copies from root * test(spend): cover cache-hit and missing-response-id paths for litellm_call_id Lock the intended behavior surfaced in review: on a cache hit request_id gets the uniqueness suffix while litellm_call_id stays the raw call id, and when the provider returns no id request_id falls back to the call id so both columns match. Both assertions fail when the populate line is reverted * test(spend): ignore litellm_call_id in spend logs payload comparisons get_logging_payload now always writes litellm_call_id, so the full-payload comparisons in test_spend_management_endpoints.py saw an unexpected key and failed. litellm_call_id is a per-request runtime uuid like request_id, which is already ignored, so add it to ignored_keys * test(logging): ignore litellm_call_id in gcs pubsub spend logs comparison The gcs pubsub spend logs payload comparison flags any key present in the actual payload but absent from the golden snapshot. get_logging_payload now always emits litellm_call_id, a per-request runtime uuid like request_id which is already ignored, so add it to ignored_keys * refactor(spend): store litellm_call_id in spend log metadata, drop column Switch DB-to-trace correlation off a dedicated column and onto the existing metadata JSON, avoiding a schema migration entirely. litellm_call_id is now written into spend log metadata (already selected and re-hydrated on the read paths) instead of a new LiteLLM_SpendLogs column, so the three schema.prisma copies and the migration are reverted and the read SELECTs go back to their original form. Correlation is queryable via metadata->>'litellm_call_id' Trade-off: an unindexed JSON lookup rather than an indexed column; acceptable for this use case and removes all migration risk * refactor(spend): thread litellm_call_id into _get_spend_logs_metadata Set litellm_call_id beside the other computed metadata values inside _get_spend_logs_metadata rather than mutating clean_metadata back in the caller, matching how applied_guardrails, cost_breakdown and the rest are threaded. No behavior change; the value still comes from kwargs with a litellm_params fallback --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 1 + .../spend_tracking/spend_tracking_utils.py | 7 ++ .../test_gcs_pub_sub.py | 1 + .../test_spend_management_endpoints.py | 1 + .../test_spend_tracking_utils.py | 109 ++++++++++++++++++ 5 files changed, 119 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7db0156bbe0..be5d7a2db78 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3181,6 +3181,7 @@ class SpendLogsMetadata(TypedDict): dict ] # special param to log k,v pairs to spendlogs for a call requester_ip_address: Optional[str] + litellm_call_id: Optional[str] applied_guardrails: Optional[List[str]] mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 8d89ff4a1ff..f6e0303e349 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -79,6 +79,7 @@ def _get_spend_logs_metadata( cold_storage_object_key: Optional[str] = None, litellm_overhead_time_ms: Optional[float] = None, cost_breakdown: Optional[CostBreakdown] = None, + litellm_call_id: Optional[str] = None, ) -> SpendLogsMetadata: if metadata is None: return SpendLogsMetadata( @@ -109,6 +110,7 @@ def _get_spend_logs_metadata( attempted_retries=None, max_retries=None, cost_breakdown=None, + litellm_call_id=litellm_call_id, ) verbose_proxy_logger.debug( "getting payload for SpendLogs, available keys in metadata: " @@ -133,6 +135,7 @@ def _get_spend_logs_metadata( clean_metadata["cold_storage_object_key"] = cold_storage_object_key clean_metadata["litellm_overhead_time_ms"] = litellm_overhead_time_ms clean_metadata["cost_breakdown"] = cost_breakdown + clean_metadata["litellm_call_id"] = litellm_call_id return clean_metadata @@ -383,6 +386,10 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs if standard_logging_payload is not None else None ), + litellm_call_id=cast( + Optional[str], + kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"), + ), ) special_usage_fields = ["completion_tokens", "prompt_tokens", "total_tokens"] diff --git a/tests/logging_callback_tests/test_gcs_pub_sub.py b/tests/logging_callback_tests/test_gcs_pub_sub.py index f4cc9735177..c37a2e3f65d 100644 --- a/tests/logging_callback_tests/test_gcs_pub_sub.py +++ b/tests/logging_callback_tests/test_gcs_pub_sub.py @@ -31,6 +31,7 @@ verbose_logger.setLevel(logging.DEBUG) ignored_keys = [ "request_id", + "metadata.litellm_call_id", "session_id", "startTime", "endTime", diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 0b583129591..6d0b4509b45 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -342,6 +342,7 @@ def test_ui_view_request_response_forbids_non_admin_without_db(client, monkeypat ignored_keys = [ "request_id", + "metadata.litellm_call_id", "session_id", "startTime", "endTime", diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index e305054d075..874e0654a1f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2120,3 +2120,112 @@ def test_get_logging_payload_failure_without_recovered_usage_is_zero(): ) assert payload["total_tokens"] == 0 + + +def test_get_logging_payload_sets_litellm_call_id_for_correlation(): + """LIT-3868: a successful spend log must carry the x-litellm-call-id (the + trace id) in its metadata, distinct from request_id, which stays the + provider response id. Without this there is no way to correlate a DB row + with its trace for a successful call. + """ + provider_response_id = "chatcmpl-e6e6f3e9-c392-404e-9a71-5361c79d8470" + trace_call_id = "c6a77556-19ce-4406-b287-53f5fb4b2b55" + + kwargs = { + "model": "openai/gpt-4o-mini", + "call_type": "acompletion", + "litellm_call_id": trace_call_id, + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + } + response_obj = { + "id": provider_response_id, + "usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12}, + } + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + metadata = json.loads(payload["metadata"]) + + assert payload["request_id"] == provider_response_id + assert metadata["litellm_call_id"] == trace_call_id + assert metadata["litellm_call_id"] != payload["request_id"] + + +def test_get_logging_payload_litellm_call_id_falls_back_to_litellm_params(): + """litellm_call_id may only be present in litellm_params; it must still land + in the spend log metadata so correlation works on that path too. + """ + trace_call_id = "fallback-7a1c-42d9-9f0e-2b6c5d4e3f21" + kwargs = { + "model": "openai/gpt-4o-mini", + "call_type": "acompletion", + "litellm_params": { + "litellm_call_id": trace_call_id, + "metadata": {"user_api_key": "sk-test"}, + }, + } + response_obj = { + "id": "chatcmpl-abc123", + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + + assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id + + +def test_get_logging_payload_litellm_call_id_when_response_has_no_id(): + """When the provider returns no id, request_id falls back to the call id, so + request_id and the metadata call id hold the same value and correlation + still resolves. + """ + trace_call_id = "noid-5b2e-4c7a-9d10-3f8a1c2b4e6d" + kwargs = { + "model": "openai/gpt-4o-mini", + "call_type": "acompletion", + "litellm_call_id": trace_call_id, + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + } + response_obj = { + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + } + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + + assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id + assert payload["request_id"] == trace_call_id + + +def test_get_logging_payload_cache_hit_keeps_raw_litellm_call_id(): + """On a cache hit request_id is suffixed to stay unique, but the metadata + litellm_call_id stays the raw trace id so the row still points at its trace. + """ + trace_call_id = "cache-9a1c-42d9-9f0e-2b6c5d4e3f21" + kwargs = { + "model": "openai/gpt-4o-mini", + "call_type": "acompletion", + "litellm_call_id": trace_call_id, + "cache_hit": True, + "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, + } + response_obj = { + "id": "chatcmpl-cache-src", + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + now = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now + ) + + assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id + assert "_cache_hit" in payload["request_id"] + assert json.loads(payload["metadata"])["litellm_call_id"] != payload["request_id"] From 97008bad29be6ed0539860807377200082048ef7 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 25 Jun 2026 18:17:54 -0700 Subject: [PATCH 03/16] chore(deps): bump deps (#31377) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * bump: version 0.1.43 → 0.1.44 * uv lock --- enterprise/pyproject.toml | 4 ++-- pyproject.toml | 2 +- uv.lock | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index b032942427c..66f6aeb7abc 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.43" +version = "0.1.44" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.43" +version = "0.1.44" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/pyproject.toml b/pyproject.toml index 6a6a47f540e..6e99d81f8f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,7 +63,7 @@ proxy = [ "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.26.0,<2.0", "litellm-proxy-extras==0.4.74", - "litellm-enterprise==0.1.43", + "litellm-enterprise==0.1.44", "RestrictedPython>=8.1,<9.0", "rich>=13.9.4,<14.0", "polars>=1.38.1,<2.0", diff --git a/uv.lock b/uv.lock index 917dff39e38..da44ad25715 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-06-22T19:05:55.080417Z" +exclude-newer = "2026-06-23T00:31:52.495979Z" exclude-newer-span = "P3D" [manifest] @@ -3597,7 +3597,7 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.43" +version = "0.1.44" source = { editable = "enterprise" } [[package]] From 6cc9ea2538be4ad6681bb6b597b8529847bc10da Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 25 Jun 2026 18:27:18 -0700 Subject: [PATCH 04/16] fix(cost-map): retarget mistral-medium-latest to Medium 3.5 and add date-pinned aliases (#31373) * fix(cost-map): retarget mistral-medium-latest to Medium 3.5 and add date-pinned aliases Mistral repointed the rolling mistral-medium-latest alias from Medium 3.1 to Medium 3.5, but the static cost map still carried Medium 3.1 specs, showing wrong pricing/context in the model hub and undercharging spend by about 3.75x (LIT-3883). Update mistral/mistral-medium-latest to Medium 3.5 ($1.50/$7.50 per 1M, 256K context, reasoning + vision), add the bare date-pinned aliases mistral/mistral-medium-2604 (Medium 3.5) and mistral/mistral-medium-2508 (Medium 3.1) that match Mistral's real API model ids, and add supports_reasoning to mistral/mistral-medium-3-5. Apply every change to both model_prices_and_context_window.json and the bundled litellm/model_prices_and_context_window_backup.json so the two stay in sync, and extend the regression tests to lock the resolved get_model_info values and the main/backup parity for all touched models. * test(cost-map): force local cost map in mistral-medium-latest resolution test get_model_info reads litellm.model_cost, which is fetched from the remote main branch at import time when LITELLM_LOCAL_MODEL_COST_MAP is unset. Until this PR lands on main, that remote map still carries the pre-merge Medium 3.1 pricing, so the assertion was only passing when the remote fetch happened to fail and fell back to the bundled backup. Force the local cost map (the same fixture pattern the other get_model_info tests use) so the alias resolution is verified deterministically against the in-repo file. --- ...odel_prices_and_context_window_backup.json | 36 +++++++- model_prices_and_context_window.json | 36 +++++++- .../test_mistral_medium_3_5_model_metadata.py | 87 ++++++++++++++----- 3 files changed, 134 insertions(+), 25 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1393c683860..d44fc654a56 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25793,7 +25793,7 @@ "supports_response_schema": true, "supports_tool_choice": true }, - "mistral/mistral-medium-latest": { + "mistral/mistral-medium-2508": { "input_cost_per_token": 4e-07, "litellm_provider": "mistral", "max_input_tokens": 131072, @@ -25801,12 +25801,45 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2e-06, + "source": "https://mistral.ai/news/mistral-medium-3", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true }, + "mistral/mistral-medium-2604": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-medium-latest": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/mistral-medium-3-1-2508": { "input_cost_per_token": 4e-07, "litellm_provider": "mistral", @@ -25833,6 +25866,7 @@ "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9867874cb65..a174c1b5efd 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25965,7 +25965,7 @@ "supports_response_schema": true, "supports_tool_choice": true }, - "mistral/mistral-medium-latest": { + "mistral/mistral-medium-2508": { "input_cost_per_token": 4e-07, "litellm_provider": "mistral", "max_input_tokens": 131072, @@ -25973,12 +25973,45 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2e-06, + "source": "https://mistral.ai/news/mistral-medium-3", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true }, + "mistral/mistral-medium-2604": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "mistral/mistral-medium-latest": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/mistral-medium-3-1-2508": { "input_cost_per_token": 4e-07, "litellm_provider": "mistral", @@ -26005,6 +26038,7 @@ "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true diff --git a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py index 496b0276b87..7cc05d6e30a 100644 --- a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py +++ b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py @@ -3,19 +3,45 @@ from pathlib import Path import pytest +import litellm from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" -@pytest.mark.parametrize("model", ["mistral/mistral-medium-3-5"]) -def test_mistral_medium_3_5_model_info(model): - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) +MEDIUM_3_5_MODELS = ( + "mistral/mistral-medium-3-5", + "mistral/mistral-medium-2604", + "mistral/mistral-medium-latest", +) - info = model_cost.get(model) - assert ( - info is not None - ), f"{model} not found in model_prices_and_context_window.json" +SYNCED_MODELS = MEDIUM_3_5_MODELS + ( + "mistral/mistral-medium-2508", + "mistral/mistral-medium-3-1-2508", +) + + +def _load(path): + with open(path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force get_model_info to resolve against the in-repo cost map instead of the + remote one fetched at import time, which still carries the pre-merge pricing.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model", MEDIUM_3_5_MODELS) +def test_medium_3_5_specs(model): + info = _load(MAIN_PATH).get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" assert info["litellm_provider"] == "mistral" assert info["mode"] == "chat" @@ -27,10 +53,11 @@ def test_mistral_medium_3_5_model_info(model): assert info["max_output_tokens"] == 262144 assert info["max_tokens"] == 262144 + assert info["supports_reasoning"] is True + assert info["supports_vision"] is True assert info["supports_function_calling"] is True assert info["supports_response_schema"] is True assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True assert info["supports_assistant_prefill"] is True routed_model, provider, _, _ = get_llm_provider(model=model) @@ -38,18 +65,32 @@ def test_mistral_medium_3_5_model_info(model): assert provider == "mistral" -def test_mistral_medium_3_5_backup_matches_main(): - """Ensure the bundled model cost map stays in sync with the canonical file.""" - repo_root = Path(__file__).parents[2] - main_path = repo_root / "model_prices_and_context_window.json" - backup_path = repo_root / "litellm" / "model_prices_and_context_window_backup.json" +def test_mistral_medium_latest_resolves_to_medium_3_5(local_model_cost_map): + """LIT-3883: the -latest alias was retargeted to Medium 3.5; get_model_info must + return the 3.5 pricing/context/reasoning, not the stale Medium 3.1 values.""" + info = litellm.get_model_info(model="mistral/mistral-medium-latest") - with open(main_path) as f: - main_cost = json.load(f) - with open(backup_path) as f: - backup_cost = json.load(f) + assert info["input_cost_per_token"] == 1.5e-06 + assert info["output_cost_per_token"] == 7.5e-06 + assert info["max_input_tokens"] == 262144 + assert info["supports_reasoning"] is True - for model in ("mistral/mistral-medium-3-5",): - assert backup_cost.get(model) == main_cost.get( - model - ), f"{model} differs between main and backup model cost maps" + +def test_mistral_medium_2508_keeps_medium_3_1_specs(): + """The date-pinned 2508 alias is Medium 3.1 and must not inherit 3.5 pricing.""" + info = _load(MAIN_PATH).get("mistral/mistral-medium-2508") + assert info is not None, "mistral/mistral-medium-2508 missing from cost map" + + assert info["input_cost_per_token"] == 4e-07 + assert info["output_cost_per_token"] == 2e-06 + assert info["max_input_tokens"] == 131072 + assert info.get("supports_reasoning") is not True + + +@pytest.mark.parametrize("model", SYNCED_MODELS) +def test_backup_matches_main(model): + """Ensure the bundled (backup) cost map stays in sync with the canonical file.""" + main_cost = _load(MAIN_PATH) + backup_cost = _load(BACKUP_PATH) + + assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" From bdafc9a00842856b2e2e1785dd75a09a27d4b65d Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Thu, 25 Jun 2026 18:42:59 -0700 Subject: [PATCH 05/16] feat(ocr): thin Rust OCR Python bridge (#31368) * feat(ocr): thin Rust OCR Python bridge * refactor(rust): group provider routing helpers --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/ai-gateway/Cargo.toml | 4 +- .../crates/ai-gateway/src/constants.rs | 1 + .../ai-gateway/src/integrations/README.md | 127 ++++ .../src/integrations/custom_guardrail/mod.rs | 468 ++++++++++++++ .../integrations/custom_guardrail/types.rs | 110 ++++ .../src/integrations/custom_logger.rs | 24 - .../src/integrations/custom_logger/mod.rs | 317 +++++++++ .../src/integrations/custom_logger/types.rs | 194 ++++++ .../mod.rs} | 110 ++-- .../litellm_python_proxy_api/types.rs | 72 +++ .../crates/ai-gateway/src/integrations/mod.rs | 4 +- .../ai-gateway/src/integrations/types.rs | 81 --- litellm-rust/crates/ai-gateway/src/io/ocr.rs | 407 +----------- .../crates/ai-gateway/src/io/realtime.rs | 6 +- litellm-rust/crates/ai-gateway/src/lib.rs | 9 +- .../crates/ai-gateway/src/ocr/client.rs | 14 + .../src/{io => }/ocr/common_utils.rs | 3 +- .../crates/ai-gateway/src/ocr/handler.rs | 71 ++ .../crates/ai-gateway/src/ocr/hooks.rs | 329 ++++++++++ litellm-rust/crates/ai-gateway/src/ocr/mod.rs | 25 + .../crates/ai-gateway/src/ocr/prepare.rs | 57 ++ .../crates/ai-gateway/src/ocr/tests.rs | 610 ++++++++++++++++++ .../crates/ai-gateway/src/ocr/types.rs | 57 ++ .../ai-gateway/src/realtime/streaming.rs | 108 ++-- .../ai-gateway/src/routes/realtime/mod.rs | 2 +- litellm-rust/crates/core/Cargo.toml | 3 + .../crates/core/src/call_lifecycle/README.md | 167 +++++ .../crates/core/src/call_lifecycle/mod.rs | 414 ++++++++++++ .../crates/core/src/call_lifecycle/types.rs | 75 +++ litellm-rust/crates/core/src/lib.rs | 2 + .../crates/core/src/routing_utils/README.md | 7 + .../crates/core/src/routing_utils/mod.rs | 1 + .../crates/core/src/routing_utils/provider.rs | 77 +++ litellm-rust/crates/python-bridge/src/lib.rs | 14 +- litellm/__init__.py | 2 +- litellm/ocr/main.py | 137 ++-- litellm/rust_bridge/__init__.py | 3 +- .../rust_bridge.py => rust_bridge/ocr.py} | 99 ++- tests/test_litellm/ocr/test_rust_bridge.py | 168 +++-- 40 files changed, 3583 insertions(+), 797 deletions(-) create mode 100644 litellm-rust/crates/ai-gateway/src/integrations/README.md create mode 100644 litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs create mode 100644 litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/types.rs delete mode 100644 litellm-rust/crates/ai-gateway/src/integrations/custom_logger.rs create mode 100644 litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs create mode 100644 litellm-rust/crates/ai-gateway/src/integrations/custom_logger/types.rs rename litellm-rust/crates/ai-gateway/src/integrations/{litellm_python_proxy_api.rs => litellm_python_proxy_api/mod.rs} (68%) create mode 100644 litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/types.rs create mode 100644 litellm-rust/crates/ai-gateway/src/ocr/client.rs rename litellm-rust/crates/ai-gateway/src/{io => }/ocr/common_utils.rs (99%) create mode 100644 litellm-rust/crates/ai-gateway/src/ocr/handler.rs create mode 100644 litellm-rust/crates/ai-gateway/src/ocr/hooks.rs create mode 100644 litellm-rust/crates/ai-gateway/src/ocr/mod.rs create mode 100644 litellm-rust/crates/ai-gateway/src/ocr/prepare.rs create mode 100644 litellm-rust/crates/ai-gateway/src/ocr/tests.rs create mode 100644 litellm-rust/crates/ai-gateway/src/ocr/types.rs create mode 100644 litellm-rust/crates/core/src/call_lifecycle/README.md create mode 100644 litellm-rust/crates/core/src/call_lifecycle/mod.rs create mode 100644 litellm-rust/crates/core/src/call_lifecycle/types.rs create mode 100644 litellm-rust/crates/core/src/routing_utils/README.md create mode 100644 litellm-rust/crates/core/src/routing_utils/mod.rs create mode 100644 litellm-rust/crates/core/src/routing_utils/provider.rs rename litellm/{ocr/rust_bridge.py => rust_bridge/ocr.py} (52%) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ce86a0ee6ac..9bffe9f9ec6 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -673,6 +673,7 @@ dependencies = [ "serde", "serde_json", "thiserror 2.0.18", + "tokio", ] [[package]] diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 2f414159158..4055be36785 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -25,7 +25,7 @@ futures-util.workspace = true serde_json.workspace = true base64.workspace = true axum = { workspace = true, features = ["ws"], optional = true } -serde = { workspace = true, optional = true } +serde.workspace = true subtle = { workspace = true, optional = true } # sha2 hashes the master key into user_api_key_hash (matches the proxy's # SHA-256 hash_token) so the plaintext credential never enters a log payload. @@ -34,7 +34,7 @@ pyo3 = { workspace = true, features = ["auto-initialize"], optional = true } [features] default = [] -server = ["dep:axum", "dep:subtle", "dep:serde", "dep:sha2"] +server = ["dep:axum", "dep:subtle", "dep:sha2"] # Build the gateway's config from the proxy YAML via an embedded Python # interpreter (links libpython; requires `litellm` importable at runtime). python-config = ["dep:pyo3"] diff --git a/litellm-rust/crates/ai-gateway/src/constants.rs b/litellm-rust/crates/ai-gateway/src/constants.rs index 3116a4c9932..109b648f5db 100644 --- a/litellm-rust/crates/ai-gateway/src/constants.rs +++ b/litellm-rust/crates/ai-gateway/src/constants.rs @@ -26,4 +26,5 @@ pub(crate) const DEFAULT_MAX_BATCH_SIZE: usize = 256; pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500; /// Provider attributed to realtime sessions in the logging payload. +#[cfg(feature = "server")] pub(crate) const DEFAULT_PROVIDER: &str = "openai"; diff --git a/litellm-rust/crates/ai-gateway/src/integrations/README.md b/litellm-rust/crates/ai-gateway/src/integrations/README.md new file mode 100644 index 00000000000..16a162dac57 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/README.md @@ -0,0 +1,127 @@ +# LiteLLM Rust integrations + +This directory contains Rust-native equivalents of LiteLLM integration hooks. +The first supported surfaces are terminal custom loggers and pre/during-call +custom guardrails. + +## File layout + +Every integration is a folder: + +- `mod.rs` contains the implementation, trait, runner, or adapter +- `types.rs` contains the integration-local request, response, error, and future + types + +Do not add new flat integration files such as `custom_logger.rs`. Shared wire +contracts that are used by multiple integrations can stay in +`integrations/types.rs`. + +Call ordering and lifecycle timing live in `litellm-core/src/call_lifecycle`. +Call-type modules, such as OCR, adapt their request and response shapes into +that generic lifecycle runner. + +## CustomLogger + +Implement `CustomLogger` when Rust code needs to observe terminal success or +failure events. Method names intentionally match Python `CustomLogger` names. + +```rust +use litellm_ai_gateway::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails, +}; + +struct RecordingLogger; + +impl CustomLogger for RecordingLogger { + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: &'a CallbackValue, + timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + let model = &model_call_details.model; + let provider = &model_call_details.custom_llm_provider; + let call_type = model_call_details.call_type.to_string(); + let request_id = model_call_details.request_id.as_deref(); + let response_object = &response_obj.object; + let duration = timing.end_time - timing.start_time; + let standard_payload = model_call_details.standard_logging_payload.as_ref(); + + Ok(()) + }) + } + + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: Option<&'a CallbackValue>, + timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + let error = model_call_details.failure_error.as_ref(); + let response_object = response_obj.map(|value| value.object.as_str()); + let duration = timing.end_time - timing.start_time; + + Ok(()) + }) + } +} +``` + +Use `CustomLoggerRunner` to fan out terminal events to configured loggers. The +runner is a no-op when no loggers are configured, which is the expected fast +path for requests without callbacks. + +## CustomGuardrail + +Implement `CustomGuardrail` when Rust code needs to run pre-call or native +during-call checks. Method names intentionally match Python `CustomGuardrail` +entrypoints inherited from Python `CustomLogger`. + +```rust +use litellm_ai_gateway::integrations::custom_guardrail::{ + CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailEventHook, + GuardrailFuture, GuardrailRequest, +}; + +struct BlocklistedPromptGuardrail; + +impl CustomGuardrail for BlocklistedPromptGuardrail { + fn guardrail_name(&self) -> &str { + "blocklisted-prompt" + } + + fn supported_event_hooks(&self) -> &[GuardrailEventHook] { + &[GuardrailEventHook::PreCall] + } + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + if request.data.to_string().contains("blocked phrase") { + return Ok(GuardrailDecision::Block( + litellm_ai_gateway::integrations::custom_guardrail::GuardrailError::blocked( + "blocked phrase detected", + ), + )); + } + Ok(GuardrailDecision::Allow(request)) + }) + } +} +``` + +Use `CustomGuardrailRunner::run_pre_call` for `pre_call` guardrails and +`CustomGuardrailRunner::run_during_call` for `during_call` guardrails. A +`GuardrailDecision::Mask` continues with modified request data. +`GuardrailDecision::Block` short-circuits the provider call. + +## Current boundary + +These are Rust-only primitives. Python callback and guardrail adapters are a +separate layer that should implement these Rust traits instead of changing the +runner interfaces. diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs new file mode 100644 index 00000000000..e5d4ce3a708 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/mod.rs @@ -0,0 +1,468 @@ +//! Rust mirror of Python `CustomGuardrail` entrypoints used by the proxy. +//! +//! This module is intentionally Rust-only: Python/PyO3 adapters are a later +//! layer that should implement this trait rather than changing the runner. + +use std::future::Future; +use std::sync::Arc; + +use crate::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails, +}; + +pub mod types; + +pub use types::{ + GuardrailContext, GuardrailDecision, GuardrailDispatchReport, GuardrailError, + GuardrailEventHook, GuardrailFuture, GuardrailRequest, +}; + +pub trait CustomGuardrail: Send + Sync { + fn guardrail_name(&self) -> &str; + + fn supported_event_hooks(&self) -> &[GuardrailEventHook]; + + /// Python 1:1 name: `async_pre_call_hook(user_api_key_dict, cache, data, call_type)`. + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { Ok(GuardrailDecision::Allow(request)) }) + } + + /// Python 1:1 name: `async_moderation_hook(data, user_api_key_dict, call_type)`. + fn async_moderation_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { Ok(GuardrailDecision::Allow(request)) }) + } +} + +pub struct CustomGuardrailRunner { + guardrails: Vec>, +} + +impl CustomGuardrailRunner { + pub fn new(guardrails: Vec>) -> Self { + Self { guardrails } + } + + pub fn is_empty(&self) -> bool { + self.guardrails.is_empty() + } + + pub async fn run_pre_call( + &self, + context: &GuardrailContext, + request: GuardrailRequest, + ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> { + self.run_hook(GuardrailEventHook::PreCall, context, request) + .await + } + + pub async fn run_during_call( + &self, + context: &GuardrailContext, + request: GuardrailRequest, + ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> { + self.run_hook(GuardrailEventHook::DuringCall, context, request) + .await + } + + pub async fn run_before_provider( + &self, + event_hook: GuardrailEventHook, + context: &GuardrailContext, + request: GuardrailRequest, + provider: F, + ) -> Result + where + F: FnOnce(GuardrailRequest) -> Fut, + Fut: Future>, + { + let (request, _) = self.run_hook(event_hook, context, request).await?; + provider(request).await + } + + pub async fn run_pre_call_with_failure_logging( + &self, + context: &GuardrailContext, + request: GuardrailRequest, + logger_runner: &CustomLoggerRunner, + model_call_details: &ModelCallDetails, + timing: CallbackTiming, + ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> { + match self.run_pre_call(context, request).await { + Ok(result) => Ok(result), + Err(error) => { + let failure_details = model_call_details.clone().with_failure_error(LoggingError { + message: error.message.clone(), + kind: error.kind.clone(), + }); + let response_obj = CallbackValue::new( + "guardrail_error", + serde_json::json!({ + "message": error.message, + "kind": error.kind, + }), + ); + logger_runner + .async_log_failure_event(&failure_details, Some(&response_obj), timing) + .await; + Err(error) + } + } + } + + async fn run_hook( + &self, + event_hook: GuardrailEventHook, + context: &GuardrailContext, + mut request: GuardrailRequest, + ) -> Result<(GuardrailRequest, GuardrailDispatchReport), GuardrailError> { + if self.guardrails.is_empty() { + return Ok((request, GuardrailDispatchReport::default())); + } + + let mut report = GuardrailDispatchReport::default(); + for guardrail in &self.guardrails { + if !self.should_run(guardrail.as_ref(), event_hook, context) { + continue; + } + + report.invoked += 1; + let decision = match event_hook { + GuardrailEventHook::PreCall => { + guardrail + .async_pre_call_hook(context, request.clone()) + .await? + } + GuardrailEventHook::DuringCall => { + guardrail + .async_moderation_hook(context, request.clone()) + .await? + } + }; + match decision.into_request() { + Ok(next_request) => request = next_request, + Err(error) => return Err(error), + } + } + + Ok((request, report)) + } + + fn should_run( + &self, + guardrail: &dyn CustomGuardrail, + event_hook: GuardrailEventHook, + context: &GuardrailContext, + ) -> bool { + let supports_hook = guardrail.supported_event_hooks().contains(&event_hook); + let selected = context.selected_guardrails.is_empty() + || context + .selected_guardrails + .iter() + .any(|name| name == guardrail.guardrail_name()); + supports_hook && selected + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::integrations::custom_logger::{CallType, CallbackValue, CustomLogger, LogFuture}; + use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload}; + use serde_json::json; + use std::sync::Mutex; + + #[derive(Clone)] + enum TestDecision { + Allow, + Mask, + Block, + } + + struct RecordingCustomGuardrail { + name: String, + hooks: Vec, + decision: TestDecision, + calls: Mutex>, + } + + impl RecordingCustomGuardrail { + fn new(name: &str, hooks: Vec, decision: TestDecision) -> Self { + Self { + name: name.to_string(), + hooks, + decision, + calls: Mutex::new(Vec::new()), + } + } + + fn calls(&self) -> Vec<&'static str> { + self.calls.lock().unwrap().clone() + } + + fn decision(&self, mut request: GuardrailRequest) -> GuardrailDecision { + match self.decision { + TestDecision::Allow => GuardrailDecision::Allow(request), + TestDecision::Mask => { + request.data["masked"] = json!(true); + GuardrailDecision::Mask(request) + } + TestDecision::Block => { + GuardrailDecision::Block(GuardrailError::blocked("blocked by guardrail")) + } + } + } + } + + impl CustomGuardrail for RecordingCustomGuardrail { + fn guardrail_name(&self) -> &str { + &self.name + } + + fn supported_event_hooks(&self) -> &[GuardrailEventHook] { + &self.hooks + } + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + self.calls.lock().unwrap().push("async_pre_call_hook"); + Ok(self.decision(request)) + }) + } + + fn async_moderation_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + self.calls.lock().unwrap().push("async_moderation_hook"); + Ok(self.decision(request)) + }) + } + } + + #[tokio::test] + async fn pre_call_dispatches_to_async_pre_call_hook() { + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "pre", + vec![GuardrailEventHook::PreCall], + TestDecision::Allow, + )); + let runner = CustomGuardrailRunner::new(vec![guardrail.clone()]); + let context = + GuardrailContext::new(CallType::Ocr).with_selected_guardrails(vec!["pre".to_string()]); + let request = GuardrailRequest::new(json!({"messages": ["hello"]})); + + let (result, report) = runner + .run_pre_call(&context, request) + .await + .expect("guardrail allows request"); + + assert_eq!(report.invoked, 1); + assert_eq!(result.data["messages"], json!(["hello"])); + assert_eq!(guardrail.calls(), vec!["async_pre_call_hook"]); + } + + #[tokio::test] + async fn during_call_dispatches_to_async_moderation_hook() { + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "during", + vec![GuardrailEventHook::DuringCall], + TestDecision::Allow, + )); + let runner = CustomGuardrailRunner::new(vec![guardrail.clone()]); + let context = GuardrailContext::new(CallType::Completion) + .with_selected_guardrails(vec!["during".to_string()]); + let request = GuardrailRequest::new(json!({"prompt": "hello"})); + + let (_result, report) = runner + .run_during_call(&context, request) + .await + .expect("guardrail allows request"); + + assert_eq!(report.invoked, 1); + assert_eq!(guardrail.calls(), vec!["async_moderation_hook"]); + } + + #[tokio::test] + async fn mask_decision_continues_with_updated_request() { + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "masker", + vec![GuardrailEventHook::PreCall], + TestDecision::Mask, + )); + let runner = CustomGuardrailRunner::new(vec![guardrail]); + let context = GuardrailContext::new(CallType::Ocr); + let request = GuardrailRequest::new(json!({"document": "secret"})); + + let (result, report) = runner + .run_pre_call(&context, request) + .await + .expect("mask continues"); + + assert_eq!(report.invoked, 1); + assert_eq!(result.data["masked"], json!(true)); + } + + #[tokio::test] + async fn block_decision_short_circuits_and_logs_failure() { + struct RecordingFailureLogger { + errors: Mutex>, + } + + impl CustomLogger for RecordingFailureLogger { + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + _response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.errors.lock().unwrap().push( + model_call_details + .failure_error + .as_ref() + .map(|error| error.kind.clone()) + .unwrap_or_default(), + ); + Ok(()) + }) + } + } + + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "blocker", + vec![GuardrailEventHook::PreCall], + TestDecision::Block, + )); + let guardrail_runner = CustomGuardrailRunner::new(vec![guardrail]); + let logger = Arc::new(RecordingFailureLogger { + errors: Mutex::new(Vec::new()), + }); + let logger_runner = CustomLoggerRunner::new(vec![logger.clone()]); + let context = GuardrailContext::new(CallType::Ocr); + let details = ModelCallDetails::from_standard_logging_payload(StandardLoggingPayload { + id: "req_ocr".to_string(), + litellm_call_id: "req_ocr".to_string(), + call_type: "ocr".to_string(), + model: "mistral-ocr-latest".to_string(), + custom_llm_provider: "mistral".to_string(), + response_cost: 0.0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + start_time: 1.0, + end_time: 1.0, + stream: false, + metadata: StandardLoggingMetadata::default(), + messages: None, + }); + + let err = guardrail_runner + .run_pre_call_with_failure_logging( + &context, + GuardrailRequest::new(json!({"document": "bad"})), + &logger_runner, + &details, + CallbackTiming::new(1.0, 2.0), + ) + .await + .expect_err("guardrail blocks request"); + + assert_eq!(err.kind, "GuardrailBlocked"); + assert_eq!( + logger.errors.lock().unwrap().as_slice(), + ["GuardrailBlocked"] + ); + } + + #[tokio::test] + async fn block_decision_short_circuits_later_guardrails_and_provider_work() { + let blocking_guardrail = Arc::new(RecordingCustomGuardrail::new( + "blocker", + vec![GuardrailEventHook::PreCall], + TestDecision::Block, + )); + let later_guardrail = Arc::new(RecordingCustomGuardrail::new( + "later", + vec![GuardrailEventHook::PreCall], + TestDecision::Allow, + )); + let runner = + CustomGuardrailRunner::new(vec![blocking_guardrail.clone(), later_guardrail.clone()]); + let provider_called = Arc::new(Mutex::new(false)); + let provider_called_for_closure = provider_called.clone(); + + let result = runner + .run_before_provider( + GuardrailEventHook::PreCall, + &GuardrailContext::new(CallType::Completion), + GuardrailRequest::new(json!({"prompt": "blocked"})), + move |_request| async move { + *provider_called_for_closure.lock().unwrap() = true; + Ok("provider response") + }, + ) + .await; + + assert!(result.is_err()); + assert_eq!(blocking_guardrail.calls(), vec!["async_pre_call_hook"]); + assert_eq!(later_guardrail.calls(), Vec::<&'static str>::new()); + assert!(!*provider_called.lock().unwrap()); + } + + #[tokio::test] + async fn run_before_provider_returns_provider_guardrail_error_directly() { + let guardrail = Arc::new(RecordingCustomGuardrail::new( + "allow", + vec![GuardrailEventHook::PreCall], + TestDecision::Allow, + )); + let runner = CustomGuardrailRunner::new(vec![guardrail]); + + let result = runner + .run_before_provider( + GuardrailEventHook::PreCall, + &GuardrailContext::new(CallType::Completion), + GuardrailRequest::new(json!({"prompt": "allowed"})), + |_request| async move { + Err::<&'static str, GuardrailError>(GuardrailError::blocked( + "provider-side guardrail error", + )) + }, + ) + .await; + + let err = result.expect_err("provider error is returned directly"); + assert_eq!(err.kind, "GuardrailBlocked"); + assert_eq!(err.message, "provider-side guardrail error"); + } + + #[tokio::test] + async fn no_guardrails_fast_path_dispatches_nothing() { + let runner = CustomGuardrailRunner::new(Vec::new()); + let context = GuardrailContext::new(CallType::Ocr); + let request = GuardrailRequest::new(json!({"document": "ok"})); + + let (result, report) = runner + .run_pre_call(&context, request) + .await + .expect("no guardrails allow request"); + + assert!(runner.is_empty()); + assert_eq!(report, GuardrailDispatchReport::default()); + assert_eq!(result.data["document"], json!("ok")); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/types.rs new file mode 100644 index 00000000000..825e56cc0d7 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_guardrail/types.rs @@ -0,0 +1,110 @@ +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; + +use serde_json::Value; + +use crate::integrations::custom_logger::CallType; + +pub type GuardrailFuture<'a> = + Pin> + Send + 'a>>; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum GuardrailEventHook { + PreCall, + DuringCall, +} + +impl GuardrailEventHook { + pub fn as_str(&self) -> &'static str { + match self { + Self::PreCall => "pre_call", + Self::DuringCall => "during_call", + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct GuardrailError { + pub message: String, + pub kind: String, +} + +impl GuardrailError { + pub fn blocked(message: impl Into) -> Self { + Self { + message: message.into(), + kind: "GuardrailBlocked".to_string(), + } + } +} + +impl std::fmt::Display for GuardrailError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}: {}", self.kind, self.message) + } +} + +impl std::error::Error for GuardrailError {} + +#[derive(Clone, Debug)] +pub struct GuardrailContext { + pub call_type: CallType, + pub selected_guardrails: Vec, + pub metadata: HashMap, + pub user_api_key_hash: Option, + pub user_api_key_user_id: Option, + pub user_api_key_team_id: Option, + pub trace_parent: Option, +} + +impl GuardrailContext { + pub fn new(call_type: CallType) -> Self { + Self { + call_type, + selected_guardrails: Vec::new(), + metadata: HashMap::new(), + user_api_key_hash: None, + user_api_key_user_id: None, + user_api_key_team_id: None, + trace_parent: None, + } + } + + pub fn with_selected_guardrails(mut self, selected_guardrails: Vec) -> Self { + self.selected_guardrails = selected_guardrails; + self + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct GuardrailRequest { + pub data: Value, +} + +impl GuardrailRequest { + pub fn new(data: Value) -> Self { + Self { data } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub enum GuardrailDecision { + Allow(GuardrailRequest), + Mask(GuardrailRequest), + Block(GuardrailError), +} + +impl GuardrailDecision { + pub(super) fn into_request(self) -> Result { + match self { + Self::Allow(request) | Self::Mask(request) => Ok(request), + Self::Block(error) => Err(error), + } + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct GuardrailDispatchReport { + pub invoked: usize, +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_logger.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger.rs deleted file mode 100644 index 53b599d8c98..00000000000 --- a/litellm-rust/crates/ai-gateway/src/integrations/custom_logger.rs +++ /dev/null @@ -1,24 +0,0 @@ -//! The `CustomLogger` trait — the Rust mirror of Python -//! `litellm/integrations/custom_logger.py::CustomLogger`. -//! -//! Synchronous (no `async_trait`): callbacks are O(1) enqueue-and-return so the -//! realtime splice never blocks on a logger. Default bodies are no-ops so a -//! logger can implement only the events it cares about. - -use crate::integrations::types::{LogError, LoggingError, StandardLoggingPayload}; - -pub trait CustomLogger: Send + Sync { - /// Record a successful call. Default: no-op. - fn log_success_event(&self, _payload: &StandardLoggingPayload) -> Result<(), LogError> { - Ok(()) - } - - /// Record a failed call. Default: no-op. - fn log_failure_event( - &self, - _payload: &StandardLoggingPayload, - _error: &LoggingError, - ) -> Result<(), LogError> { - Ok(()) - } -} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs new file mode 100644 index 00000000000..792717dacfc --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs @@ -0,0 +1,317 @@ +//! The `CustomLogger` trait — the Rust mirror of Python +//! `litellm/integrations/custom_logger.py::CustomLogger`. +//! +//! The Python-named async terminal methods are the public Rust callback shape. + +use std::sync::Arc; + +pub mod types; + +pub use types::{ + CallType, CallbackDispatchReport, CallbackTiming, CallbackValue, LogError, LogFuture, + LoggingError, ModelCallDetails, +}; + +pub trait CustomLogger: Send + Sync { + /// Python 1:1 name: `async_log_success_event(model_call_details, response_obj, start_time, end_time)`. + fn async_log_success_event<'a>( + &'a self, + _model_call_details: &'a ModelCallDetails, + _response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async { Ok(()) }) + } + + /// Python 1:1 name: `async_log_failure_event(model_call_details, response_obj, start_time, end_time)`. + fn async_log_failure_event<'a>( + &'a self, + _model_call_details: &'a ModelCallDetails, + _response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async { Ok(()) }) + } +} + +pub struct CustomLoggerRunner { + loggers: Vec>, +} + +impl CustomLoggerRunner { + pub fn new(loggers: Vec>) -> Self { + Self { loggers } + } + + pub fn is_empty(&self) -> bool { + self.loggers.is_empty() + } + + pub async fn async_log_success_event( + &self, + model_call_details: &ModelCallDetails, + response_obj: &CallbackValue, + timing: CallbackTiming, + ) -> CallbackDispatchReport { + if self.loggers.is_empty() { + return CallbackDispatchReport::default(); + } + + let mut report = CallbackDispatchReport::default(); + for logger in &self.loggers { + report.invoked += 1; + if let Err(err) = logger + .async_log_success_event(model_call_details, response_obj, timing) + .await + { + report.dropped += 1; + eprintln!("litellm-ai-gateway: async_log_success_event dropped: {err}"); + } + } + report + } + + pub async fn async_log_failure_event( + &self, + model_call_details: &ModelCallDetails, + response_obj: Option<&CallbackValue>, + timing: CallbackTiming, + ) -> CallbackDispatchReport { + if self.loggers.is_empty() { + return CallbackDispatchReport::default(); + } + + let mut report = CallbackDispatchReport::default(); + for logger in &self.loggers { + report.invoked += 1; + if let Err(err) = logger + .async_log_failure_event(model_call_details, response_obj, timing) + .await + { + report.dropped += 1; + eprintln!("litellm-ai-gateway: async_log_failure_event dropped: {err}"); + } + } + report + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload}; + use serde_json::json; + use std::sync::Mutex; + + #[derive(Clone, Debug, PartialEq)] + struct RecordedEvent { + hook: &'static str, + model: String, + provider: String, + call_type: String, + request_id: Option, + litellm_call_id: Option, + user_id: Option, + response_object: Option, + error_kind: Option, + start_time: f64, + end_time: f64, + standard_logging_model: Option, + } + + #[derive(Default)] + struct RecordingCustomLogger { + events: Mutex>, + } + + impl RecordingCustomLogger { + fn events(&self) -> Vec { + self.events.lock().unwrap().clone() + } + } + + impl CustomLogger for RecordingCustomLogger { + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: &'a CallbackValue, + timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push(RecordedEvent { + hook: "async_log_success_event", + model: model_call_details.model.clone(), + provider: model_call_details.custom_llm_provider.clone(), + call_type: model_call_details.call_type.to_string(), + request_id: model_call_details.request_id.clone(), + litellm_call_id: model_call_details.litellm_call_id.clone(), + user_id: model_call_details.metadata.user_api_key_user_id.clone(), + response_object: Some(response_obj.object.clone()), + error_kind: None, + start_time: timing.start_time, + end_time: timing.end_time, + standard_logging_model: model_call_details + .standard_logging_payload + .as_ref() + .map(|payload| payload.model.clone()), + }); + Ok(()) + }) + } + + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: Option<&'a CallbackValue>, + timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push(RecordedEvent { + hook: "async_log_failure_event", + model: model_call_details.model.clone(), + provider: model_call_details.custom_llm_provider.clone(), + call_type: model_call_details.call_type.to_string(), + request_id: model_call_details.request_id.clone(), + litellm_call_id: model_call_details.litellm_call_id.clone(), + user_id: model_call_details.metadata.user_api_key_user_id.clone(), + response_object: response_obj.map(|value| value.object.clone()), + error_kind: model_call_details + .failure_error + .as_ref() + .map(|error| error.kind.clone()), + start_time: timing.start_time, + end_time: timing.end_time, + standard_logging_model: model_call_details + .standard_logging_payload + .as_ref() + .map(|payload| payload.model.clone()), + }); + Ok(()) + }) + } + } + + fn payload(call_type: &str, model: &str, provider: &str) -> StandardLoggingPayload { + StandardLoggingPayload { + id: format!("req_{call_type}"), + litellm_call_id: format!("call_{call_type}"), + call_type: call_type.to_string(), + model: model.to_string(), + custom_llm_provider: provider.to_string(), + response_cost: 0.25, + prompt_tokens: 3, + completion_tokens: 4, + total_tokens: 7, + start_time: 10.0, + end_time: 11.5, + stream: false, + metadata: StandardLoggingMetadata { + user_api_key_hash: Some("hash".to_string()), + user_api_key_user_id: Some("user".to_string()), + user_api_key_team_id: Some("team".to_string()), + ..Default::default() + }, + messages: Some(json!([{"role": "user", "content": "read this"}])), + } + } + + #[tokio::test] + async fn rust_custom_logger_reads_success_payload_for_ocr() { + let logger = Arc::new(RecordingCustomLogger::default()); + let runner = CustomLoggerRunner::new(vec![logger.clone()]); + let details = ModelCallDetails::from_standard_logging_payload(payload( + "ocr", + "mistral-ocr-latest", + "mistral", + )); + let response = CallbackValue::new("ocr", json!({"pages": [{"markdown": "ok"}]})); + let report = runner + .async_log_success_event(&details, &response, CallbackTiming::new(10.0, 11.5)) + .await; + + assert_eq!(report.invoked, 1); + assert_eq!(report.dropped, 0); + assert_eq!( + logger.events(), + vec![RecordedEvent { + hook: "async_log_success_event", + model: "mistral-ocr-latest".to_string(), + provider: "mistral".to_string(), + call_type: "ocr".to_string(), + request_id: Some("req_ocr".to_string()), + litellm_call_id: Some("call_ocr".to_string()), + user_id: Some("user".to_string()), + response_object: Some("ocr".to_string()), + error_kind: None, + start_time: 10.0, + end_time: 11.5, + standard_logging_model: Some("mistral-ocr-latest".to_string()), + }] + ); + } + + #[tokio::test] + async fn rust_custom_logger_reads_failure_payload_for_non_ocr_call_type() { + let logger = Arc::new(RecordingCustomLogger::default()); + let runner = CustomLoggerRunner::new(vec![logger.clone()]); + let details = ModelCallDetails::from_standard_logging_payload(payload( + "acompletion", + "gpt-4.1-mini", + "openai", + )) + .with_failure_error(LoggingError { + message: "provider failed".to_string(), + kind: "ProviderError".to_string(), + }); + let response = CallbackValue::new("error", json!({"message": "provider failed"})); + let report = runner + .async_log_failure_event(&details, Some(&response), CallbackTiming::new(2.0, 3.0)) + .await; + + assert_eq!(report.invoked, 1); + assert_eq!(report.dropped, 0); + assert_eq!( + logger.events(), + vec![RecordedEvent { + hook: "async_log_failure_event", + model: "gpt-4.1-mini".to_string(), + provider: "openai".to_string(), + call_type: "acompletion".to_string(), + request_id: Some("req_acompletion".to_string()), + litellm_call_id: Some("call_acompletion".to_string()), + user_id: Some("user".to_string()), + response_object: Some("error".to_string()), + error_kind: Some("ProviderError".to_string()), + start_time: 2.0, + end_time: 3.0, + standard_logging_model: Some("gpt-4.1-mini".to_string()), + }] + ); + } + + #[tokio::test] + async fn no_callback_fast_path_dispatches_nothing() { + let runner = CustomLoggerRunner::new(Vec::new()); + let details = ModelCallDetails::new("mistral-ocr-latest", "mistral", CallType::Ocr); + let response = CallbackValue::new("ocr", json!({})); + + let report = runner + .async_log_success_event(&details, &response, CallbackTiming::new(1.0, 1.5)) + .await; + + assert!(runner.is_empty()); + assert_eq!(report, CallbackDispatchReport::default()); + } + + #[test] + fn with_standard_logging_payload_keeps_top_level_fields_in_sync() { + let details = ModelCallDetails::new("old-model", "old-provider", CallType::Completion) + .with_standard_logging_payload(payload("ocr", "mistral-ocr-latest", "mistral")); + + assert_eq!(details.model, "mistral-ocr-latest"); + assert_eq!(details.custom_llm_provider, "mistral"); + assert_eq!(details.call_type, CallType::Ocr); + assert_eq!(details.request_id, Some("req_ocr".to_string())); + assert_eq!(details.litellm_call_id, Some("call_ocr".to_string())); + } +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/types.rs new file mode 100644 index 00000000000..ba7d67bd46e --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/custom_logger/types.rs @@ -0,0 +1,194 @@ +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; + +use serde_json::Value; + +use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload}; + +pub type LogFuture<'a> = Pin> + Send + 'a>>; + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct CallbackDispatchReport { + pub invoked: usize, + pub dropped: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CallType { + Ocr, + Realtime, + Completion, + Acompletion, + ChatCompletion, + Other(String), +} + +impl CallType { + pub fn as_str(&self) -> &str { + match self { + Self::Ocr => "ocr", + Self::Realtime => "realtime", + Self::Completion => "completion", + Self::Acompletion => "acompletion", + Self::ChatCompletion => "chat_completion", + Self::Other(value) => value.as_str(), + } + } +} + +impl From<&str> for CallType { + fn from(value: &str) -> Self { + match value { + "ocr" => Self::Ocr, + "realtime" => Self::Realtime, + "completion" => Self::Completion, + "acompletion" => Self::Acompletion, + "chat_completion" => Self::ChatCompletion, + other => Self::Other(other.to_string()), + } + } +} + +impl std::fmt::Display for CallType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct CallbackTiming { + pub start_time: f64, + pub end_time: f64, +} + +impl CallbackTiming { + pub fn new(start_time: f64, end_time: f64) -> Self { + Self { + start_time, + end_time, + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct CallbackValue { + pub object: String, + pub value: Value, +} + +impl CallbackValue { + pub fn new(object: impl Into, value: Value) -> Self { + Self { + object: object.into(), + value, + } + } +} + +#[derive(Clone, Debug)] +pub struct ModelCallDetails { + pub model: String, + pub custom_llm_provider: String, + pub call_type: CallType, + pub metadata: StandardLoggingMetadata, + pub extra_metadata: HashMap, + pub request_id: Option, + pub litellm_call_id: Option, + pub response_cost: Option, + pub standard_logging_payload: Option, + pub failure_error: Option, +} + +impl ModelCallDetails { + pub fn new( + model: impl Into, + custom_llm_provider: impl Into, + call_type: CallType, + ) -> Self { + Self { + model: model.into(), + custom_llm_provider: custom_llm_provider.into(), + call_type, + metadata: StandardLoggingMetadata::default(), + extra_metadata: HashMap::new(), + request_id: None, + litellm_call_id: None, + response_cost: None, + standard_logging_payload: None, + failure_error: None, + } + } + + pub fn from_standard_logging_payload(payload: StandardLoggingPayload) -> Self { + let request_id = Some(payload.id.clone()); + let litellm_call_id = Some(payload.litellm_call_id.clone()); + let response_cost = Some(payload.response_cost); + let metadata = payload.metadata.clone(); + Self { + model: payload.model.clone(), + custom_llm_provider: payload.custom_llm_provider.clone(), + call_type: CallType::from(payload.call_type.as_str()), + metadata, + extra_metadata: HashMap::new(), + request_id, + litellm_call_id, + response_cost, + standard_logging_payload: Some(payload), + failure_error: None, + } + } + + pub fn with_standard_logging_payload(mut self, payload: StandardLoggingPayload) -> Self { + self.model = payload.model.clone(); + self.custom_llm_provider = payload.custom_llm_provider.clone(); + self.call_type = CallType::from(payload.call_type.as_str()); + self.request_id = Some(payload.id.clone()); + self.litellm_call_id = Some(payload.litellm_call_id.clone()); + self.response_cost = Some(payload.response_cost); + self.metadata = payload.metadata.clone(); + self.standard_logging_payload = Some(payload); + self + } + + pub fn with_failure_error(mut self, error: LoggingError) -> Self { + self.failure_error = Some(error); + self + } +} + +#[derive(Clone, Debug)] +pub struct LoggingError { + pub message: String, + pub kind: String, +} + +#[derive(Clone, Debug)] +pub struct LogError { + pub message: String, + pub kind: String, +} + +impl LogError { + pub fn channel_full() -> Self { + Self { + message: "logging channel is full; dropping record".to_string(), + kind: "ChannelFull".to_string(), + } + } + + pub fn channel_closed() -> Self { + Self { + message: "logging channel is closed; worker has shut down".to_string(), + kind: "ChannelClosed".to_string(), + } + } +} + +impl std::fmt::Display for LogError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}: {}", self.kind, self.message) + } +} + +impl std::error::Error for LogError {} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api.rs b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/mod.rs similarity index 68% rename from litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api.rs rename to litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/mod.rs index 165a90d8dbe..3dad18cb7a3 100644 --- a/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api.rs +++ b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/mod.rs @@ -1,7 +1,8 @@ //! A `CustomLogger` that ships finished events to the LiteLLM Python proxy's //! `/v1/rust_control_plane/logs` endpoint. //! -//! The callback path is non-blocking: `log_success_event` / `log_failure_event` +//! The callback path is non-blocking: `async_log_success_event` / +//! `async_log_failure_event` //! build a `LogRecord` and `try_send` it onto a bounded channel, returning a //! `LogError` (never panicking, never awaiting) if the channel is full or the //! worker has gone away. A spawned background worker drains the channel, batches @@ -15,54 +16,14 @@ use reqwest::Client; use tokio::sync::mpsc::{self, Receiver, Sender}; use tokio::time::interval; -use crate::constants::{ - DEFAULT_CHANNEL_CAPACITY, DEFAULT_FLUSH_INTERVAL_MS, DEFAULT_MAX_BATCH_SIZE, - DEFAULT_PROXY_BASE_URL, RUST_CONTROL_PLANE_LOGS_PATH, -}; -use crate::integrations::custom_logger::CustomLogger; -use crate::integrations::types::{ - CallbackLogsRequest, LogError, LogRecord, LoggingError, StandardLoggingPayload, +use crate::constants::{DEFAULT_PROXY_BASE_URL, RUST_CONTROL_PLANE_LOGS_PATH}; +use crate::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLogger, LogError, LogFuture, LoggingError, + ModelCallDetails, }; +use types::{CallbackLogsRequest, EgressTunables, LogRecord}; -/// Egress worker tunables. Each field defaults to the matching `DEFAULT_*` const -/// in `crate::constants` and is overridable via an env var (read once at logger -/// construction). -struct EgressTunables { - channel_capacity: usize, - max_batch_size: usize, - flush_interval: Duration, -} - -impl EgressTunables { - fn from_env() -> Self { - Self { - channel_capacity: env_positive( - "LITELLM_LOG_CHANNEL_CAPACITY", - DEFAULT_CHANNEL_CAPACITY, - ), - max_batch_size: env_positive("LITELLM_LOG_BATCH_SIZE", DEFAULT_MAX_BATCH_SIZE), - flush_interval: Duration::from_millis(env_positive( - "LITELLM_LOG_FLUSH_INTERVAL_MS", - DEFAULT_FLUSH_INTERVAL_MS, - )), - } - } -} - -/// Parse a positive integer env var, falling back to `default` on missing, -/// unparseable, or non-positive values. Generic over the integer type so one -/// helper serves both the `usize` capacities and the `u64` interval. -fn env_positive(name: &str, default: T) -> T -where - T: std::str::FromStr + PartialOrd + From, -{ - let zero = T::from(0u8); - std::env::var(name) - .ok() - .and_then(|value| value.trim().parse::().ok()) - .filter(|n| *n > zero) - .unwrap_or(default) -} +pub mod types; /// Ships realtime logging events to the LiteLLM Python proxy. pub struct LiteLLMPythonProxyAPILogger { @@ -118,23 +79,50 @@ impl LiteLLMPythonProxyAPILogger { } impl CustomLogger for LiteLLMPythonProxyAPILogger { - fn log_success_event(&self, payload: &StandardLoggingPayload) -> Result<(), LogError> { - self.enqueue(LogRecord { - status: "success".to_string(), - payload: payload.clone(), - error: None, + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + _response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + if let Some(payload) = &model_call_details.standard_logging_payload { + self.enqueue(LogRecord { + status: "success".to_string(), + payload: payload.clone(), + error: None, + })?; + } + Ok(()) }) } - fn log_failure_event( - &self, - payload: &StandardLoggingPayload, - error: &LoggingError, - ) -> Result<(), LogError> { - self.enqueue(LogRecord { - status: "failure".to_string(), - payload: payload.clone(), - error: Some(format!("{}: {}", error.kind, error.message)), + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + _response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + if let Some(payload) = &model_call_details.standard_logging_payload { + let fallback_error; + let error = match &model_call_details.failure_error { + Some(error) => error, + None => { + fallback_error = LoggingError { + message: "callback failure event".to_string(), + kind: "CallbackFailure".to_string(), + }; + &fallback_error + } + }; + self.enqueue(LogRecord { + status: "failure".to_string(), + payload: payload.clone(), + error: Some(format!("{}: {}", error.kind, error.message)), + })?; + } + Ok(()) }) } } diff --git a/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/types.rs new file mode 100644 index 00000000000..481a437747f --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/integrations/litellm_python_proxy_api/types.rs @@ -0,0 +1,72 @@ +use std::time::Duration; + +use serde::Serialize; + +use crate::constants::{ + DEFAULT_CHANNEL_CAPACITY, DEFAULT_FLUSH_INTERVAL_MS, DEFAULT_MAX_BATCH_SIZE, +}; +use crate::integrations::types::StandardLoggingPayload; + +#[derive(Serialize)] +pub struct CallbackLogsRequest { + pub records: Vec, +} + +#[derive(Serialize)] +pub struct CallbackLogRecord { + pub status: String, + pub standard_logging_payload: StandardLoggingPayload, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +#[derive(Clone, Debug)] +pub struct LogRecord { + pub status: String, + pub payload: StandardLoggingPayload, + pub error: Option, +} + +impl LogRecord { + pub fn into_callback_record(self) -> CallbackLogRecord { + CallbackLogRecord { + status: self.status, + standard_logging_payload: self.payload, + error: self.error, + } + } +} + +pub(super) struct EgressTunables { + pub channel_capacity: usize, + pub max_batch_size: usize, + pub flush_interval: Duration, +} + +impl EgressTunables { + pub fn from_env() -> Self { + Self { + channel_capacity: env_positive( + "LITELLM_LOG_CHANNEL_CAPACITY", + DEFAULT_CHANNEL_CAPACITY, + ), + max_batch_size: env_positive("LITELLM_LOG_BATCH_SIZE", DEFAULT_MAX_BATCH_SIZE), + flush_interval: Duration::from_millis(env_positive( + "LITELLM_LOG_FLUSH_INTERVAL_MS", + DEFAULT_FLUSH_INTERVAL_MS, + )), + } + } +} + +fn env_positive(name: &str, default: T) -> T +where + T: std::str::FromStr + PartialOrd + From, +{ + let zero = T::from(0u8); + std::env::var(name) + .ok() + .and_then(|value| value.trim().parse::().ok()) + .filter(|n| *n > zero) + .unwrap_or(default) +} diff --git a/litellm-rust/crates/ai-gateway/src/integrations/mod.rs b/litellm-rust/crates/ai-gateway/src/integrations/mod.rs index 8799be0c040..c62f1821ef8 100644 --- a/litellm-rust/crates/ai-gateway/src/integrations/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/integrations/mod.rs @@ -1,10 +1,12 @@ //! Pure-Rust logging integrations. Names map 1:1 to Python //! `litellm/integrations/`: +//! - [`custom_guardrail::CustomGuardrail`] — the guardrail callback trait //! - [`custom_logger::CustomLogger`] — the callback trait //! - [`litellm_python_proxy_api::LiteLLMPythonProxyAPILogger`] — ships events -//! to the Python proxy's `/v1/callbacks/logs` endpoint +//! to the Python proxy's `/v1/rust_control_plane/logs` endpoint //! - [`types`] — the typed `StandardLoggingPayload` wire contract +pub mod custom_guardrail; pub mod custom_logger; pub mod litellm_python_proxy_api; pub mod types; diff --git a/litellm-rust/crates/ai-gateway/src/integrations/types.rs b/litellm-rust/crates/ai-gateway/src/integrations/types.rs index d61a1f816a7..34dce93d8e0 100644 --- a/litellm-rust/crates/ai-gateway/src/integrations/types.rs +++ b/litellm-rust/crates/ai-gateway/src/integrations/types.rs @@ -28,68 +28,6 @@ pub struct RequestMetadata { pub user_api_key_team_id: Option, } -/// A logging-callback failure (e.g. a custom logger raised). Mirrors the Python -/// failure-event shape: a message plus an exception kind/class name. -#[derive(Clone, Debug)] -pub struct LoggingError { - pub message: String, - pub kind: String, -} - -/// A non-fatal error returned by a `CustomLogger` when it cannot enqueue an -/// event (channel full or the background worker has shut down). -#[derive(Clone, Debug)] -pub struct LogError { - pub message: String, - pub kind: String, -} - -impl LogError { - pub fn channel_full() -> Self { - Self { - message: "logging channel is full; dropping record".to_string(), - kind: "ChannelFull".to_string(), - } - } - - pub fn channel_closed() -> Self { - Self { - message: "logging channel is closed; worker has shut down".to_string(), - kind: "ChannelClosed".to_string(), - } - } -} - -impl std::fmt::Display for LogError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}: {}", self.kind, self.message) - } -} - -impl std::error::Error for LogError {} - -/// Batch wrapper — the top-level request body. -/// Matches Python `CallbackLogsRequest { records: list[CallbackLogRecord] }`. -#[derive(Serialize)] -pub struct CallbackLogsRequest { - pub records: Vec, -} - -/// One finished logging event. -/// Matches `CallbackLogRecord { status, standard_logging_payload, error? }`. -#[derive(Serialize)] -pub struct CallbackLogRecord { - /// "success" | "failure". On "failure", `error` (or payload.error_str) - /// becomes the replayed exception string. - pub status: String, - - pub standard_logging_payload: StandardLoggingPayload, - - /// Only meaningful when status == "failure". Omitted on success. - #[serde(skip_serializing_if = "Option::is_none")] - pub error: Option, -} - /// The self-describing payload. Field names are the EXACT JSON keys the Python /// replay path + spend-logs builder read. #[derive(Clone, Debug, Serialize)] @@ -143,22 +81,3 @@ pub struct StandardLoggingMetadata { #[serde(skip_serializing_if = "Option::is_none")] pub spend_logs_metadata: Option>, } - -/// The unit handed to a `CustomLogger` sink: a finished payload plus its status -/// and (on failure) the replayed error string. -#[derive(Clone, Debug)] -pub struct LogRecord { - pub status: String, - pub payload: StandardLoggingPayload, - pub error: Option, -} - -impl LogRecord { - pub fn into_callback_record(self) -> CallbackLogRecord { - CallbackLogRecord { - status: self.status, - standard_logging_payload: self.payload, - error: self.error, - } - } -} diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr.rs b/litellm-rust/crates/ai-gateway/src/io/ocr.rs index 35e511fa982..55e02839c4e 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr.rs +++ b/litellm-rust/crates/ai-gateway/src/io/ocr.rs @@ -1,406 +1 @@ -//! End-to-end OCR orchestration. -//! -//! Owns supported OCR provider calls so the Python side stays a thin bridge: -//! resolve the API key, build the URL + body via the pure transforms, POST it, -//! and normalize the response. The HTTP client is built once and reused. - -use std::sync::OnceLock; -use std::time::Duration; - -use litellm_core::error::CoreError; -use litellm_core::ocr::transformation::{OcrAuthStrategy, OcrResponseHandling}; -use litellm_core::CoreResult; -use serde_json::{Map, Value}; - -mod common_utils; - -use common_utils::{ - convert_document_url_to_data_uri, has_header, ocr_provider_config, poll_document_intelligence, - string_headers, truncate_error_body, -}; - -/// OCR over large documents can take a while; bound it generously rather than -/// hanging forever on an unresponsive upstream. The client-level limit is the -/// outer ceiling; callers can tighten it per request via ``run_ocr``'s ``timeout``. -const OCR_TIMEOUT_SECS: u64 = 600; - -/// Process-wide async HTTP client (connection pool + TLS reused across calls). -/// -/// The Python fallback path uses LiteLLM's standard `BaseLLMHTTPHandler`. This -/// Rust path is opt-in and owns end-to-end OCR I/O, so it cannot call the -/// Python handler directly; keep this route-scoped until litellm-rust has a -/// shared HTTP abstraction. -fn http_client() -> &'static reqwest::Client { - static CLIENT: OnceLock = OnceLock::new(); - CLIENT.get_or_init(|| { - reqwest::Client::builder() - .timeout(Duration::from_secs(OCR_TIMEOUT_SECS)) - .build() - .expect("failed to build reqwest client") - }) -} - -fn upstream_headers( - headers: &[(String, String)], - auth_strategy: OcrAuthStrategy, - api_key: Option<&str>, -) -> Vec<(String, String)> { - let auth_header = api_key.map(|api_key| match auth_strategy { - OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), - OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()), - }); - auth_header - .into_iter() - .chain(headers.iter().cloned()) - .collect() -} - -pub struct OcrRequest<'a> { - pub model: &'a str, - pub document: Value, - pub api_key: Option<&'a str>, - pub api_base: Option<&'a str>, - pub custom_llm_provider: &'a str, - pub extra_headers: Option>, - pub optional_params: Map, - pub timeout: Option, -} - -/// Perform an OCR call end to end and return the normalized response as -/// JSON (the shape the Python `OCRResponse` model expects). -/// -/// Async: intended to be awaited directly by the Python bridge's async entrypoint. -pub async fn ocr(request: OcrRequest<'_>) -> CoreResult { - let model = request.model; - let config = ocr_provider_config(request.custom_llm_provider, model) - .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.to_string()))?; - let env_lookup = |key: &str| std::env::var(key).ok(); - - let headers = string_headers(request.extra_headers)?; - let auth_strategy = config.auth_strategy(); - let api_key = (!has_header(&headers, auth_strategy.header_name())) - .then(|| config.resolve_api_key(request.api_key, &env_lookup)) - .transpose()?; - let url = config.complete_url( - request.api_base, - model, - &request.optional_params, - &env_lookup, - )?; - let filtered_params = config.map_ocr_params(&request.optional_params); - let document = if config.requires_data_uri_document() { - convert_document_url_to_data_uri(request.document).await? - } else { - request.document - }; - let body = config - .transform_ocr_request(model, document, filtered_params)? - .data; - let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref()); - - let mut request_builder = http_client().post(&url).json(&body); - for (key, value) in &upstream_headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - - let response = request_builder - .send() - .await - .map_err(|err| CoreError::Network(err.to_string()))?; - - let status = response.status(); - if config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll - && status.as_u16() == 202 - { - let operation_url = response - .headers() - .get("operation-location") - .and_then(|value| value.to_str().ok()) - .map(str::to_string) - .ok_or_else(|| { - CoreError::InvalidResponse( - "Azure Document Intelligence returned 202 but no Operation-Location header found" - .to_string(), - ) - })?; - let response_json = - poll_document_intelligence(&operation_url, &url, &upstream_headers, request.timeout) - .await?; - return Ok(config - .transform_ocr_response(model, response_json)? - .into_json()); - } - - let text = response - .text() - .await - .map_err(|err| CoreError::Network(err.to_string()))?; - - if !status.is_success() { - return Err(CoreError::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }); - } - - let response_json: Value = serde_json::from_str(&text) - .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?; - - Ok(config - .transform_ocr_response(model, response_json)? - .into_json()) -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - use tokio::net::{TcpListener, TcpStream}; - - async fn read_http_headers(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); - if request.windows(4).any(|window| window == b"\r\n\r\n") { - break; - } - } - String::from_utf8(request).expect("request is utf8") - } - - #[test] - fn truncate_error_body_passes_short_strings_through() { - let body = "Unauthorized"; - assert_eq!(truncate_error_body(body), "Unauthorized"); - } - - #[test] - fn truncate_error_body_caps_long_payloads() { - let body = "x".repeat(306); - let truncated = truncate_error_body(&body); - - assert!(truncated.ends_with("... (truncated)")); - let prefix_chars = truncated - .strip_suffix("... (truncated)") - .expect("truncated marker present") - .chars() - .count(); - assert_eq!(prefix_chars, 256); - } - - #[test] - fn truncate_error_body_does_not_split_multibyte_chars() { - let body = "é".repeat(266); - let truncated = truncate_error_body(&body); - assert!(truncated.is_char_boundary(truncated.len())); - } - - #[test] - fn ocr_dispatch_supports_migrated_providers() { - assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some()); - assert!(ocr_provider_config("azure_ai", "pixtral-12b-2409") - .expect("azure ai config resolves") - .requires_data_uri_document()); - assert_eq!( - ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read") - .expect("document intelligence config resolves") - .response_handling(), - OcrResponseHandling::AzureDocumentIntelligencePoll - ); - assert!(ocr_provider_config("vertex_ai", "deepseek-ocr-maas") - .expect("vertex deepseek config resolves") - .supported_ocr_params() - .contains(&"temperature")); - assert!(ocr_provider_config("openai", "gpt-4o").is_none()); - } - - #[test] - fn string_headers_accepts_string_values() { - let headers = json!({ - "x-trace-id": "trace-1" - }) - .as_object() - .unwrap() - .clone(); - - assert_eq!( - string_headers(Some(headers)).expect("string headers accepted"), - vec![("x-trace-id".to_string(), "trace-1".to_string())] - ); - } - - #[test] - fn auth_header_detection_is_case_insensitive() { - let headers = vec![ - ("x-trace-id".to_string(), "trace-1".to_string()), - ("authorization".to_string(), "Bearer sk-test".to_string()), - ]; - - assert!(has_header(&headers, "authorization")); - - let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())]; - assert!(has_header(&headers, "authorization")); - - let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())]; - assert!(!has_header(&headers, "authorization")); - } - - #[tokio::test] - async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("test listener binds"); - let addr = listener.local_addr().expect("listener has local addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts one request"); - let request = read_http_headers(&mut socket).await; - - let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#; - let response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - response_body.len(), - response_body - ); - socket - .write_all(response.as_bytes()) - .await - .expect("writes response"); - request - }); - - let mut headers = Map::new(); - headers.insert( - "Authorization".to_string(), - Value::String("Bearer sk-from-python".to_string()), - ); - headers.insert( - "x-trace-id".to_string(), - Value::String("trace-1".to_string()), - ); - - let response = ocr(OcrRequest { - model: "mistral-ocr-latest", - document: json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }), - api_key: Some("sk-for-rust-fallback"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: "mistral", - extra_headers: Some(headers), - optional_params: Map::new(), - timeout: Some(Duration::from_secs(5)), - }) - .await - .expect("ocr request succeeds"); - - assert_eq!(response["pages"][0]["markdown"], "ok"); - - let request = server.await.expect("server task completes"); - let authorization_count = request - .lines() - .filter(|line| line.to_ascii_lowercase().starts_with("authorization:")) - .count(); - assert_eq!(authorization_count, 1, "{request}"); - assert!( - request.contains("authorization: Bearer sk-from-python") - || request.contains("Authorization: Bearer sk-from-python"), - "{request}" - ); - } - - #[tokio::test] - async fn document_intelligence_poll_uses_resolved_subscription_key() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("test listener binds"); - let addr = listener.local_addr().expect("listener has local addr"); - let operation_url = format!("http://{addr}/operations/1"); - - let server = tokio::spawn(async move { - let (mut post_socket, _) = listener.accept().await.expect("accepts post request"); - let post_request = read_http_headers(&mut post_socket).await; - let post_response = format!( - "HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n" - ); - post_socket - .write_all(post_response.as_bytes()) - .await - .expect("writes post response"); - - let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request"); - let poll_request = read_http_headers(&mut poll_socket).await; - let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#; - let poll_response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - response_body.len(), - response_body - ); - poll_socket - .write_all(poll_response.as_bytes()) - .await - .expect("writes poll response"); - (post_request, poll_request) - }); - - let response = ocr(OcrRequest { - model: "prebuilt-read", - document: json!({ - "type": "document_url", - "document_url": "https://example.com/doc.pdf" - }), - api_key: Some("di-key"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: "azure_ai/doc-intelligence", - extra_headers: None, - optional_params: Map::new(), - timeout: Some(Duration::from_secs(5)), - }) - .await - .expect("document intelligence request succeeds"); - - assert_eq!(response["pages"][0]["markdown"], "ok"); - - let (post_request, poll_request) = server.await.expect("server task completes"); - assert!( - post_request - .to_ascii_lowercase() - .contains("ocp-apim-subscription-key: di-key"), - "{post_request}" - ); - assert!( - poll_request - .to_ascii_lowercase() - .contains("ocp-apim-subscription-key: di-key"), - "{poll_request}" - ); - } - - #[test] - fn string_headers_rejects_non_string_values() { - let headers = json!({ - "x-retry-count": 3 - }) - .as_object() - .unwrap() - .clone(); - - let err = string_headers(Some(headers)).expect_err("non-string header rejected"); - assert_eq!( - err, - CoreError::InvalidRequest( - "OCR extra_headers.x-retry-count must be a string, got number".to_string() - ) - ); - } -} +pub use crate::ocr::{ocr, OcrRequest}; diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs index 4047de5cb26..40a38c1579a 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -1,8 +1,8 @@ //! End-to-end OpenAI realtime invocation. //! -//! The host-facing entry point, mirroring `crate::io::ocr::run_ocr`: open the -//! WebSocket to OpenAI, then splice a client realtime stream to the upstream, -//! driving typed events through the pure `OPENAI_REALTIME_CONFIG` transforms. +//! The host-facing entry point opens the WebSocket to OpenAI, then splices a +//! client realtime stream to the upstream, driving typed events through the pure +//! `OPENAI_REALTIME_CONFIG` transforms. //! Network, auth header, key resolution, and wire (de)serialization live here so //! the `transformation` module stays pure and typed. //! diff --git a/litellm-rust/crates/ai-gateway/src/lib.rs b/litellm-rust/crates/ai-gateway/src/lib.rs index 6c04fbb7626..d8ef7bb5ba1 100644 --- a/litellm-rust/crates/ai-gateway/src/lib.rs +++ b/litellm-rust/crates/ai-gateway/src/lib.rs @@ -3,15 +3,16 @@ //! Two layers, split by feature so the Python `cdylib` can depend on the I/O //! without pulling in the HTTP server: //! -//! - [`io`]: all network I/O (OCR HTTP call, realtime WebSocket splice, the -//! pre-warmed realtime pool). Always available — no feature required. The -//! Python bridge links this for `run_ocr`. +//! - Call-type modules such as [`ocr`]: provider transforms, lifecycle hooks, +//! and provider I/O. Always available — no feature required. +//! - [`io`]: compatibility exports and realtime WebSocket splice helpers. //! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling //! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway` //! binary turns on. The `python-config` feature additionally pulls in [`python`] //! for the load-time config reader. pub mod io; +pub mod ocr; /// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and /// the `python-config` reader, so it is available without either feature. @@ -27,9 +28,7 @@ pub mod state; // Realtime request logging. Only the server serves realtime, so these are // `server`-gated; `io::realtime` exposes the generic `observe` hook while the // collector and callback fan-out live here. -#[cfg(feature = "server")] mod constants; -#[cfg(feature = "server")] pub mod integrations; #[cfg(feature = "server")] mod realtime; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/client.rs b/litellm-rust/crates/ai-gateway/src/ocr/client.rs new file mode 100644 index 00000000000..79cc7816227 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/client.rs @@ -0,0 +1,14 @@ +use std::sync::OnceLock; +use std::time::Duration; + +const OCR_TIMEOUT_SECS: u64 = 600; + +pub(super) fn http_client() -> &'static reqwest::Client { + static CLIENT: OnceLock = OnceLock::new(); + CLIENT.get_or_init(|| { + reqwest::Client::builder() + .timeout(Duration::from_secs(OCR_TIMEOUT_SECS)) + .build() + .expect("failed to build reqwest client") + }) +} diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs similarity index 99% rename from litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs rename to litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index ee0d86c3000..d4b4d9338e7 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -18,7 +18,7 @@ use litellm_core::providers::vertex_ai::ocr::transformation::{ VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG, }; -use super::http_client; +use super::client::http_client; const ERROR_BODY_MAX_CHARS: usize = 256; const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120; @@ -42,7 +42,6 @@ pub(super) fn ocr_provider_config( "azure_ai" if is_azure_document_intelligence_model(model) => { Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG) } - "azure_ai/doc-intelligence" => Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG), "azure_ai" => Some(&AZURE_AI_OCR_CONFIG), "vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG), "vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG), diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs new file mode 100644 index 00000000000..4d93c2a25db --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -0,0 +1,71 @@ +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrResponseHandling; +use litellm_core::CoreResult; +use serde_json::Value; + +use super::client::http_client; +use super::common_utils::{poll_document_intelligence, truncate_error_body}; +use super::types::ProviderOcrRequest; + +pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult { + let mut request_builder = http_client().post(&request.url).json(&request.body); + for (key, value) in &request.upstream_headers { + request_builder = request_builder.header(key, value); + } + if let Some(duration) = request.timeout { + request_builder = request_builder.timeout(duration); + } + + let response = request_builder + .send() + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + + let status = response.status(); + if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll + && status.as_u16() == 202 + { + let operation_url = response + .headers() + .get("operation-location") + .and_then(|value| value.to_str().ok()) + .map(str::to_string) + .ok_or_else(|| { + CoreError::InvalidResponse( + "Azure Document Intelligence returned 202 but no Operation-Location header found" + .to_string(), + ) + })?; + let response_json = poll_document_intelligence( + &operation_url, + &request.url, + &request.upstream_headers, + request.timeout, + ) + .await?; + return Ok(request + .config + .transform_ocr_response(&request.model, response_json)? + .into_json()); + } + + let text = response + .text() + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + + if !status.is_success() { + return Err(CoreError::Http { + status: status.as_u16(), + body: truncate_error_body(&text), + }); + } + + let response_json: Value = serde_json::from_str(&text) + .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?; + + Ok(request + .config + .transform_ocr_response(&request.model, response_json)? + .into_json()) +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs new file mode 100644 index 00000000000..6be74ed2714 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -0,0 +1,329 @@ +use std::future::Future; +use std::pin::Pin; + +use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrAuthStrategy; +use litellm_core::CoreResult; +use serde_json::{json, Map, Value}; + +use super::common_utils::{ + convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers, +}; +use super::types::{PreparedOcrRequest, ProviderOcrRequest}; +use crate::integrations::custom_guardrail::{ + CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest, +}; +use crate::integrations::custom_logger::{ + CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails, +}; +use crate::integrations::types::{ + RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload, +}; + +pub(crate) struct OcrLifecycleHooks { + logger_runner: CustomLoggerRunner, + guardrail_runner: CustomGuardrailRunner, + request_metadata: RequestMetadata, +} + +type OcrFuture<'a, T> = Pin> + Send + 'a>>; +type OcrLogFuture<'a> = Pin + Send + 'a>>; + +impl OcrLifecycleHooks { + pub(crate) fn new( + logger_runner: CustomLoggerRunner, + guardrail_runner: CustomGuardrailRunner, + request_metadata: RequestMetadata, + ) -> Self { + Self { + logger_runner, + guardrail_runner, + request_metadata, + } + } + + async fn run_pre_call_guardrails( + &self, + request: PreparedOcrRequest, + ) -> CoreResult { + if self.guardrail_runner.is_empty() { + return Ok(request); + } + + let context = guardrail_context(&self.request_metadata); + let guardrail_request = GuardrailRequest::new(json!({ + "model": request.model, + "custom_llm_provider": request.custom_llm_provider, + "document": request.document, + "optional_params": request.optional_params, + })); + let (guardrail_request, _) = self + .guardrail_runner + .run_pre_call(&context, guardrail_request) + .await + .map_err(guardrail_error_to_core_error)?; + let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?; + Ok(PreparedOcrRequest { + document, + optional_params, + ..request + }) + } + + async fn prepare_provider_request( + &self, + request: PreparedOcrRequest, + ) -> CoreResult { + let config = ocr_provider_config(&request.custom_llm_provider, &request.model) + .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?; + let env_lookup = |key: &str| std::env::var(key).ok(); + let headers = string_headers(request.extra_headers)?; + let auth_strategy = config.auth_strategy(); + let api_key = (!has_header(&headers, auth_strategy.header_name())) + .then(|| config.resolve_api_key(request.api_key.as_deref(), &env_lookup)) + .transpose()?; + let url = config.complete_url( + request.api_base.as_deref(), + &request.model, + &request.optional_params, + &env_lookup, + )?; + let filtered_params = config.map_ocr_params(&request.optional_params); + let model = request.model.clone(); + let custom_llm_provider = request.custom_llm_provider.clone(); + let document = if config.requires_data_uri_document() { + convert_document_url_to_data_uri(request.document).await? + } else { + request.document + }; + let body = config + .transform_ocr_request(&request.model, document, filtered_params)? + .data; + let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref()); + let body = self + .run_during_call_guardrails(&model, &custom_llm_provider, &url, body) + .await?; + Ok(ProviderOcrRequest { + model, + config, + url, + body, + upstream_headers, + timeout: request.timeout, + }) + } + + async fn run_during_call_guardrails( + &self, + model: &str, + custom_llm_provider: &str, + url: &str, + body: Value, + ) -> CoreResult { + if self.guardrail_runner.is_empty() { + return Ok(body); + } + + let context = guardrail_context(&self.request_metadata); + let guardrail_request = GuardrailRequest::new(json!({ + "model": model, + "custom_llm_provider": custom_llm_provider, + "url": url, + "body": body, + })); + let (guardrail_request, _) = self + .guardrail_runner + .run_during_call(&context, guardrail_request) + .await + .map_err(guardrail_error_to_core_error)?; + parse_ocr_during_call_guardrail_request(guardrail_request) + } + + fn standard_logging_payload( + &self, + context: &CallLifecycleContext, + timing: &CallLifecycleTiming, + ) -> StandardLoggingPayload { + StandardLoggingPayload { + id: context.litellm_call_id.clone(), + litellm_call_id: context.litellm_call_id.clone(), + call_type: context.call_type.clone(), + model: context.model.clone(), + custom_llm_provider: context.custom_llm_provider.clone(), + response_cost: 0.0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + start_time: timing.start_time, + end_time: timing.end_time, + stream: false, + metadata: StandardLoggingMetadata { + user_api_key_hash: self.request_metadata.user_api_key_hash.clone(), + user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(), + user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(), + ..Default::default() + }, + messages: None, + } + } +} + +impl CallLifecycleHooks for OcrLifecycleHooks { + type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>; + type DuringCallFuture<'a> = OcrFuture<'a, ProviderOcrRequest>; + type SuccessFuture<'a> = OcrLogFuture<'a>; + type FailureFuture<'a> = OcrLogFuture<'a>; + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: PreparedOcrRequest, + ) -> Self::PreCallFuture<'a> { + Box::pin(async move { self.run_pre_call_guardrails(request).await }) + } + + fn async_during_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: PreparedOcrRequest, + ) -> Self::DuringCallFuture<'a> { + Box::pin(async move { self.prepare_provider_request(request).await }) + } + + fn async_log_success_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + response: &'a Value, + timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a> { + Box::pin(async move { + if self.logger_runner.is_empty() { + return; + } + let response_obj = CallbackValue::new("ocr", response.clone()); + self.logger_runner + .async_log_success_event( + &ModelCallDetails::from_standard_logging_payload( + self.standard_logging_payload(context, timing), + ), + &response_obj, + CallbackTiming::new(timing.start_time, timing.end_time), + ) + .await; + }) + } + + fn async_log_failure_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + error: &'a CoreError, + timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a> { + Box::pin(async move { + if self.logger_runner.is_empty() { + return; + } + let logging_error = LoggingError { + message: error.to_string(), + kind: core_error_kind(error).to_string(), + }; + let response_obj = CallbackValue::new( + "error", + json!({ + "message": logging_error.message, + "kind": logging_error.kind, + }), + ); + self.logger_runner + .async_log_failure_event( + &ModelCallDetails::from_standard_logging_payload( + self.standard_logging_payload(context, timing), + ) + .with_failure_error(logging_error), + Some(&response_obj), + CallbackTiming::new(timing.start_time, timing.end_time), + ) + .await; + }) + } +} + +fn upstream_headers( + headers: &[(String, String)], + auth_strategy: OcrAuthStrategy, + api_key: Option<&str>, +) -> Vec<(String, String)> { + api_key + .map(|api_key| match auth_strategy { + OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), + OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()), + }) + .into_iter() + .chain(headers.iter().cloned()) + .collect() +} + +fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { + GuardrailContext { + call_type: CallType::Ocr, + selected_guardrails: Vec::new(), + metadata: std::collections::HashMap::new(), + user_api_key_hash: metadata.user_api_key_hash.clone(), + user_api_key_user_id: metadata.user_api_key_user_id.clone(), + user_api_key_team_id: metadata.user_api_key_team_id.clone(), + trace_parent: None, + } +} + +fn parse_ocr_pre_call_guardrail_request( + request: GuardrailRequest, +) -> CoreResult<(Value, Map)> { + let Value::Object(mut data) = request.data else { + return Err(CoreError::InvalidRequest( + "OCR pre_call guardrail must return an object".to_string(), + )); + }; + let document = data.remove("document").ok_or_else(|| { + CoreError::InvalidRequest("OCR pre_call guardrail removed document".to_string()) + })?; + let optional_params = match data.remove("optional_params") { + Some(Value::Object(params)) => params, + Some(_) => { + return Err(CoreError::InvalidRequest( + "OCR pre_call guardrail optional_params must be an object".to_string(), + )) + } + None => Map::new(), + }; + Ok((document, optional_params)) +} + +fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> CoreResult { + let Value::Object(mut data) = request.data else { + return Err(CoreError::InvalidRequest( + "OCR during_call guardrail must return an object".to_string(), + )); + }; + data.remove("body").ok_or_else(|| { + CoreError::InvalidRequest("OCR during_call guardrail removed body".to_string()) + }) +} + +fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError { + CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message)) +} + +fn core_error_kind(error: &CoreError) -> &'static str { + match error { + CoreError::Auth(_) => "AuthError", + CoreError::InvalidProvider(_) => "InvalidProvider", + CoreError::InvalidRequest(_) => "InvalidRequest", + CoreError::InvalidType { .. } => "InvalidType", + CoreError::MissingField(_) => "MissingField", + CoreError::Http { .. } => "HttpError", + CoreError::InvalidResponse(_) => "InvalidResponse", + CoreError::Network(_) => "NetworkError", + CoreError::Routing(_) => "RoutingError", + } +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs new file mode 100644 index 00000000000..b54ee39b21d --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -0,0 +1,25 @@ +use litellm_core::call_lifecycle::CallLifecycle; +use litellm_core::CoreResult; +use serde_json::Value; + +mod client; +mod common_utils; +mod handler; +mod hooks; +mod prepare; +mod types; + +pub use types::OcrRequest; + +use handler::execute_ocr_provider_call; +use prepare::{prepare_ocr_call, PreparedOcrCall}; + +pub async fn ocr(request: OcrRequest<'_>) -> CoreResult { + let PreparedOcrCall { request, hooks } = prepare_ocr_call(request); + CallLifecycle::default() + .run_request(request, &hooks, execute_ocr_provider_call) + .await +} + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs new file mode 100644 index 00000000000..5a4b350a4c4 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs @@ -0,0 +1,57 @@ +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider}; + +use super::hooks::OcrLifecycleHooks; +use super::types::{OcrRequest, PreparedOcrRequest}; +use crate::integrations::custom_guardrail::CustomGuardrailRunner; +use crate::integrations::custom_logger::CustomLoggerRunner; + +pub(crate) struct PreparedOcrCall { + pub(crate) request: PreparedOcrRequest, + pub(crate) hooks: OcrLifecycleHooks, +} + +pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { + let call_id = request + .litellm_call_id + .map(str::to_string) + .unwrap_or_else(new_ocr_call_id); + let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) + .unwrap_or(CustomLlmProvider { + model: request.model, + custom_llm_provider: "mistral", + }); + let model = provider_info.model.to_string(); + let custom_llm_provider = provider_info.custom_llm_provider.to_string(); + + PreparedOcrCall { + request: PreparedOcrRequest { + model, + custom_llm_provider, + litellm_call_id: call_id, + document: request.document, + api_key: request.api_key.map(str::to_string), + api_base: request.api_base.map(str::to_string), + extra_headers: request.extra_headers, + optional_params: request.optional_params, + timeout: request.timeout, + }, + hooks: OcrLifecycleHooks::new( + CustomLoggerRunner::new(request.callbacks), + CustomGuardrailRunner::new(request.guardrails), + request.request_metadata, + ), + } +} + +fn new_ocr_call_id() -> String { + static COUNTER: AtomicU64 = AtomicU64::new(1); + let sequence = COUNTER.fetch_add(1, Ordering::Relaxed); + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_nanos()) + .unwrap_or(0); + format!("ocr-{timestamp}-{sequence}") +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs new file mode 100644 index 00000000000..35747dc6985 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -0,0 +1,610 @@ +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrResponseHandling; +use serde_json::{json, Map, Value}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; + +use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body}; +use super::{ocr, OcrRequest}; +use crate::integrations::custom_guardrail::{ + CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook, + GuardrailFuture, GuardrailRequest, +}; +use crate::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails, +}; +use crate::integrations::types::RequestMetadata; + +async fn read_http_headers(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + String::from_utf8(request).expect("request is utf8") +} + +async fn read_http_request(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break request.len(); + } + request.extend_from_slice(&buffer[..n]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while request.len().saturating_sub(header_end) < content_length { + let n = socket.read(&mut buffer).await.expect("reads body"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + } + String::from_utf8(request).expect("request is utf8") +} + +#[derive(Clone, Debug, PartialEq)] +struct RecordedLogEvent { + hook: &'static str, + model: String, + call_type: String, + user_id: Option, + response_object: Option, + error_kind: Option, +} + +#[derive(Default)] +struct RecordingOcrLogger { + events: Mutex>, +} + +impl RecordingOcrLogger { + fn events(&self) -> Vec { + self.events.lock().unwrap().clone() + } +} + +impl CustomLogger for RecordingOcrLogger { + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push(RecordedLogEvent { + hook: "async_log_success_event", + model: model_call_details.model.clone(), + call_type: model_call_details.call_type.to_string(), + user_id: model_call_details.metadata.user_api_key_user_id.clone(), + response_object: Some(response_obj.object.clone()), + error_kind: None, + }); + Ok(()) + }) + } + + fn async_log_failure_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push(RecordedLogEvent { + hook: "async_log_failure_event", + model: model_call_details.model.clone(), + call_type: model_call_details.call_type.to_string(), + user_id: model_call_details.metadata.user_api_key_user_id.clone(), + response_object: response_obj.map(|value| value.object.clone()), + error_kind: model_call_details + .failure_error + .as_ref() + .map(|error| error.kind.clone()), + }); + Ok(()) + }) + } +} + +struct RecordingOcrGuardrail { + hooks: Vec, + events: Mutex>, + block_pre_call: bool, +} + +impl RecordingOcrGuardrail { + fn new(hooks: Vec) -> Self { + Self { + hooks, + events: Mutex::new(Vec::new()), + block_pre_call: false, + } + } + + fn blocking_pre_call() -> Self { + Self { + hooks: vec![GuardrailEventHook::PreCall], + events: Mutex::new(Vec::new()), + block_pre_call: true, + } + } + + fn events(&self) -> Vec<&'static str> { + self.events.lock().unwrap().clone() + } +} + +impl CustomGuardrail for RecordingOcrGuardrail { + fn guardrail_name(&self) -> &str { + "recording-ocr-guardrail" + } + + fn supported_event_hooks(&self) -> &[GuardrailEventHook] { + &self.hooks + } + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + mut request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("async_pre_call_hook"); + if self.block_pre_call { + return Ok(GuardrailDecision::Block(GuardrailError::blocked( + "blocked before provider", + ))); + } + request.data["document"]["guarded_pre"] = json!(true); + Ok(GuardrailDecision::Mask(request)) + }) + } + + fn async_moderation_hook<'a>( + &'a self, + _context: &'a GuardrailContext, + mut request: GuardrailRequest, + ) -> GuardrailFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("async_moderation_hook"); + request.data["body"]["guarded_during"] = json!(true); + Ok(GuardrailDecision::Mask(request)) + }) + } +} + +#[test] +fn truncate_error_body_passes_short_strings_through() { + let body = "Unauthorized"; + assert_eq!(truncate_error_body(body), "Unauthorized"); +} + +#[test] +fn truncate_error_body_caps_long_payloads() { + let body = "x".repeat(306); + let truncated = truncate_error_body(&body); + + assert!(truncated.ends_with("... (truncated)")); + let prefix_chars = truncated + .strip_suffix("... (truncated)") + .expect("truncated marker present") + .chars() + .count(); + assert_eq!(prefix_chars, 256); +} + +#[test] +fn truncate_error_body_does_not_split_multibyte_chars() { + let body = "é".repeat(266); + let truncated = truncate_error_body(&body); + assert!(truncated.is_char_boundary(truncated.len())); +} + +#[test] +fn ocr_dispatch_supports_migrated_providers() { + assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some()); + assert!(ocr_provider_config("azure_ai", "pixtral-12b-2409") + .expect("azure ai config resolves") + .requires_data_uri_document()); + assert_eq!( + ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read") + .expect("document intelligence config resolves") + .response_handling(), + OcrResponseHandling::AzureDocumentIntelligencePoll + ); + assert!(ocr_provider_config("vertex_ai", "deepseek-ocr-maas") + .expect("vertex deepseek config resolves") + .supported_ocr_params() + .contains(&"temperature")); + assert!(ocr_provider_config("openai", "gpt-4o").is_none()); +} + +#[test] +fn string_headers_accepts_string_values() { + let headers = json!({ + "x-trace-id": "trace-1" + }) + .as_object() + .unwrap() + .clone(); + + assert_eq!( + string_headers(Some(headers)).expect("string headers accepted"), + vec![("x-trace-id".to_string(), "trace-1".to_string())] + ); +} + +#[test] +fn auth_header_detection_is_case_insensitive() { + let headers = vec![ + ("x-trace-id".to_string(), "trace-1".to_string()), + ("authorization".to_string(), "Bearer sk-test".to_string()), + ]; + + assert!(has_header(&headers, "authorization")); + + let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())]; + assert!(has_header(&headers, "authorization")); + + let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())]; + assert!(!has_header(&headers, "authorization")); +} + +#[tokio::test] +async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts one request"); + let request = read_http_request(&mut socket).await; + let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + request + }); + + let logger = Arc::new(RecordingOcrLogger::default()); + let guardrail = Arc::new(RecordingOcrGuardrail::new(vec![ + GuardrailEventHook::PreCall, + GuardrailEventHook::DuringCall, + ])); + let response = ocr(OcrRequest { + model: "mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-test"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("mistral"), + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + callbacks: vec![logger.clone()], + guardrails: vec![guardrail.clone()], + request_metadata: RequestMetadata { + user_api_key_user_id: Some("user-1".to_string()), + ..Default::default() + }, + litellm_call_id: Some("ocr-call-1"), + }) + .await + .expect("ocr request succeeds"); + + assert_eq!(response["pages"][0]["markdown"], "ok"); + assert_eq!( + guardrail.events(), + vec!["async_pre_call_hook", "async_moderation_hook"] + ); + assert_eq!( + logger.events(), + vec![RecordedLogEvent { + hook: "async_log_success_event", + model: "mistral-ocr-latest".to_string(), + call_type: "ocr".to_string(), + user_id: Some("user-1".to_string()), + response_object: Some("ocr".to_string()), + error_kind: None, + }] + ); + + let request = server.await.expect("server task completes"); + assert!(request.contains(r#""guarded_pre":true"#), "{request}"); + assert!(request.contains(r#""guarded_during":true"#), "{request}"); +} + +#[tokio::test] +async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts one request"); + let _request = read_http_request(&mut socket).await; + let response_body = "provider failed"; + let response = format!( + "HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + }); + + let logger = Arc::new(RecordingOcrLogger::default()); + let err = ocr(OcrRequest { + model: "mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-test"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("mistral"), + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + callbacks: vec![logger.clone()], + guardrails: Vec::new(), + request_metadata: RequestMetadata::default(), + litellm_call_id: Some("ocr-call-2"), + }) + .await + .expect_err("provider error propagates"); + + assert!(matches!(err, CoreError::Http { status: 500, .. })); + server.await.expect("server task completes"); + assert_eq!( + logger.events(), + vec![RecordedLogEvent { + hook: "async_log_failure_event", + model: "mistral-ocr-latest".to_string(), + call_type: "ocr".to_string(), + user_id: None, + response_object: Some("error".to_string()), + error_kind: Some("HttpError".to_string()), + }] + ); +} + +#[tokio::test] +async fn ocr_lifecycle_pre_call_block_skips_provider_socket() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + let logger = Arc::new(RecordingOcrLogger::default()); + let guardrail = Arc::new(RecordingOcrGuardrail::blocking_pre_call()); + + let err = ocr(OcrRequest { + model: "mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-test"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("mistral"), + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_millis(100)), + callbacks: vec![logger.clone()], + guardrails: vec![guardrail.clone()], + request_metadata: RequestMetadata::default(), + litellm_call_id: Some("ocr-call-3"), + }) + .await + .expect_err("guardrail blocks request"); + + assert!(matches!(err, CoreError::InvalidRequest(_))); + assert_eq!(guardrail.events(), vec!["async_pre_call_hook"]); + assert_eq!( + logger.events(), + vec![RecordedLogEvent { + hook: "async_log_failure_event", + model: "mistral-ocr-latest".to_string(), + call_type: "ocr".to_string(), + user_id: None, + response_object: Some("error".to_string()), + error_kind: Some("InvalidRequest".to_string()), + }] + ); + let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await; + assert!(accepted.is_err(), "provider socket should not be touched"); +} + +#[tokio::test] +async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts one request"); + let request = read_http_headers(&mut socket).await; + let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + request + }); + + let mut headers = Map::new(); + headers.insert( + "Authorization".to_string(), + Value::String("Bearer sk-from-python".to_string()), + ); + headers.insert( + "x-trace-id".to_string(), + Value::String("trace-1".to_string()), + ); + + let response = ocr(OcrRequest { + model: "mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-for-rust-fallback"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("mistral"), + extra_headers: Some(headers), + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: RequestMetadata::default(), + litellm_call_id: None, + }) + .await + .expect("ocr request succeeds"); + + assert_eq!(response["pages"][0]["markdown"], "ok"); + + let request = server.await.expect("server task completes"); + let authorization_count = request + .lines() + .filter(|line| line.to_ascii_lowercase().starts_with("authorization:")) + .count(); + assert_eq!(authorization_count, 1, "{request}"); + assert!( + request.contains("authorization: Bearer sk-from-python") + || request.contains("Authorization: Bearer sk-from-python"), + "{request}" + ); +} + +#[tokio::test] +async fn document_intelligence_poll_uses_resolved_subscription_key() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + let operation_url = format!("http://{addr}/operations/1"); + + let server = tokio::spawn(async move { + let (mut post_socket, _) = listener.accept().await.expect("accepts post request"); + let post_request = read_http_headers(&mut post_socket).await; + let post_response = format!( + "HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n" + ); + post_socket + .write_all(post_response.as_bytes()) + .await + .expect("writes post response"); + + let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request"); + let poll_request = read_http_headers(&mut poll_socket).await; + let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#; + let poll_response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + poll_socket + .write_all(poll_response.as_bytes()) + .await + .expect("writes poll response"); + (post_request, poll_request) + }); + + let response = ocr(OcrRequest { + model: "doc-intelligence/prebuilt-read", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("di-key"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("azure_ai"), + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: RequestMetadata::default(), + litellm_call_id: None, + }) + .await + .expect("document intelligence request succeeds"); + + assert_eq!(response["pages"][0]["markdown"], "ok"); + + let (post_request, poll_request) = server.await.expect("server task completes"); + assert!( + post_request + .to_ascii_lowercase() + .contains("ocp-apim-subscription-key: di-key"), + "{post_request}" + ); + assert!( + poll_request + .to_ascii_lowercase() + .contains("ocp-apim-subscription-key: di-key"), + "{poll_request}" + ); +} + +#[test] +fn string_headers_rejects_non_string_values() { + let headers = json!({ + "x-retry-count": 3 + }) + .as_object() + .unwrap() + .clone(); + + let err = string_headers(Some(headers)).expect_err("non-string header rejected"); + assert_eq!( + err, + CoreError::InvalidRequest( + "OCR extra_headers.x-retry-count must be a string, got number".to_string() + ) + ); +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/types.rs b/litellm-rust/crates/ai-gateway/src/ocr/types.rs new file mode 100644 index 00000000000..bde734a4dd1 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/ocr/types.rs @@ -0,0 +1,57 @@ +use std::sync::Arc; +use std::time::Duration; + +use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest}; +use litellm_core::ocr::transformation::OcrProviderConfig; +use serde_json::{Map, Value}; + +use crate::integrations::custom_guardrail::CustomGuardrail; +use crate::integrations::custom_logger::CustomLogger; +use crate::integrations::types::RequestMetadata; + +pub struct OcrRequest<'a> { + pub model: &'a str, + pub document: Value, + pub api_key: Option<&'a str>, + pub api_base: Option<&'a str>, + pub custom_llm_provider: Option<&'a str>, + pub extra_headers: Option>, + pub optional_params: Map, + pub timeout: Option, + pub callbacks: Vec>, + pub guardrails: Vec>, + pub request_metadata: RequestMetadata, + pub litellm_call_id: Option<&'a str>, +} + +pub(crate) struct PreparedOcrRequest { + pub(crate) model: String, + pub(crate) custom_llm_provider: String, + pub(crate) litellm_call_id: String, + pub(crate) document: Value, + pub(crate) api_key: Option, + pub(crate) api_base: Option, + pub(crate) extra_headers: Option>, + pub(crate) optional_params: Map, + pub(crate) timeout: Option, +} + +impl CallLifecycleRequest for PreparedOcrRequest { + fn lifecycle_context(&self) -> CallLifecycleContext { + CallLifecycleContext::new( + "ocr", + self.model.clone(), + self.custom_llm_provider.clone(), + self.litellm_call_id.clone(), + ) + } +} + +pub(crate) struct ProviderOcrRequest { + pub(crate) model: String, + pub(crate) config: &'static dyn OcrProviderConfig, + pub(crate) url: String, + pub(crate) body: Value, + pub(crate) upstream_headers: Vec<(String, String)>, + pub(crate) timeout: Option, +} diff --git a/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs index 34c82897808..c32e727de54 100644 --- a/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs +++ b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs @@ -13,7 +13,9 @@ use litellm_core::realtime::types::RealtimeEvent; use serde_json::Value; use crate::constants::DEFAULT_PROVIDER; -use crate::integrations::custom_logger::CustomLogger; +use crate::integrations::custom_logger::{ + CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails, +}; use crate::integrations::types::{ RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload, Usage, }; @@ -183,30 +185,45 @@ impl RealTimeStreaming { /// Finish the session: stamp the end time and fan the payload out to every /// callback. On a logger enqueue error we bump a non-fatal counter (the /// realtime session has already ended; a dropped log must never propagate). - pub fn log_messages(&mut self, status: SessionStatus) { + pub async fn log_messages(&mut self, status: SessionStatus) { self.end_time = epoch_seconds(); let payload = self.build_payload(); + let timing = CallbackTiming::new(payload.start_time, payload.end_time); + let runner = CustomLoggerRunner::new(self.callbacks.clone()); match status { SessionStatus::Success => { - for callback in &self.callbacks { - if let Err(err) = callback.log_success_event(&payload) { - self.dropped += 1; - eprintln!("litellm-ai-gateway: log_success_event dropped: {err}"); - } - } + let response = CallbackValue::new("realtime", serde_json::Value::Null); + let report = runner + .async_log_success_event( + &ModelCallDetails::from_standard_logging_payload(payload), + &response, + timing, + ) + .await; + self.dropped += report.dropped as u64; } SessionStatus::Failure => { - let error = crate::integrations::types::LoggingError { + let error = LoggingError { message: "realtime session ended in failure".to_string(), kind: "RealtimeSessionError".to_string(), }; - for callback in &self.callbacks { - if let Err(err) = callback.log_failure_event(&payload, &error) { - self.dropped += 1; - eprintln!("litellm-ai-gateway: log_failure_event dropped: {err}"); - } - } + let response = CallbackValue::new( + "error", + serde_json::json!({ + "message": error.message, + "kind": error.kind, + }), + ); + let report = runner + .async_log_failure_event( + &ModelCallDetails::from_standard_logging_payload(payload) + .with_failure_error(error), + Some(&response), + timing, + ) + .await; + self.dropped += report.dropped as u64; } } } @@ -215,7 +232,8 @@ impl RealTimeStreaming { #[cfg(test)] mod tests { use super::*; - use crate::integrations::types::{LogError, LoggingError}; + use crate::integrations::custom_logger::LogError; + use crate::integrations::custom_logger::LogFuture; use std::sync::atomic::{AtomicU64, Ordering}; fn event(raw: &str) -> RealtimeEvent { @@ -231,17 +249,28 @@ mod tests { } impl CustomLogger for CapturingLogger { - fn log_success_event(&self, payload: &StandardLoggingPayload) -> Result<(), LogError> { - self.calls.fetch_add(1, Ordering::SeqCst); - *self.last_model.lock().unwrap() = Some(payload.model.clone()); - self.last_total_tokens - .store(payload.total_tokens, Ordering::SeqCst); - Ok(()) + fn async_log_success_event<'a>( + &'a self, + model_call_details: &'a ModelCallDetails, + _response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async move { + let payload = model_call_details + .standard_logging_payload + .as_ref() + .expect("standard logging payload"); + self.calls.fetch_add(1, Ordering::SeqCst); + *self.last_model.lock().unwrap() = Some(payload.model.clone()); + self.last_total_tokens + .store(payload.total_tokens, Ordering::SeqCst); + Ok(()) + }) } } - #[test] - fn observe_accumulates_model_and_tokens_then_logs() { + #[tokio::test] + async fn observe_accumulates_model_and_tokens_then_logs() { let logger = Arc::new(CapturingLogger::default()); let callbacks: Vec> = vec![logger.clone()]; let mut streaming = RealTimeStreaming::new( @@ -284,7 +313,7 @@ mod tests { Some("hash123") ); - streaming.log_messages(SessionStatus::Success); + streaming.log_messages(SessionStatus::Success).await; assert_eq!(logger.calls.load(Ordering::SeqCst), 1); assert_eq!( logger.last_model.lock().unwrap().as_deref(), @@ -324,19 +353,26 @@ mod tests { /// A logger whose enqueue always fails should bump the dropped counter, not /// panic or propagate. - #[test] - fn failing_logger_bumps_dropped_counter() { + #[tokio::test] + async fn failing_logger_bumps_dropped_counter() { struct FailingLogger; impl CustomLogger for FailingLogger { - fn log_success_event(&self, _p: &StandardLoggingPayload) -> Result<(), LogError> { - Err(LogError::channel_full()) + fn async_log_success_event<'a>( + &'a self, + _model_call_details: &'a ModelCallDetails, + _response_obj: &'a CallbackValue, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async { Err(LogError::channel_full()) }) } - fn log_failure_event( - &self, - _p: &StandardLoggingPayload, - _e: &LoggingError, - ) -> Result<(), LogError> { - Err(LogError::channel_closed()) + + fn async_log_failure_event<'a>( + &'a self, + _model_call_details: &'a ModelCallDetails, + _response_obj: Option<&'a CallbackValue>, + _timing: CallbackTiming, + ) -> LogFuture<'a> { + Box::pin(async { Err(LogError::channel_closed()) }) } } let callbacks: Vec> = vec![Arc::new(FailingLogger)]; @@ -346,7 +382,7 @@ mod tests { "gpt-realtime".to_string(), RequestMetadata::default(), ); - streaming.log_messages(SessionStatus::Success); + streaming.log_messages(SessionStatus::Success).await; assert_eq!(streaming.dropped(), 1); } } diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs index 899ad73829f..c3f929f5f0b 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs @@ -162,5 +162,5 @@ async fn bridge( } else { SessionStatus::Failure }; - collector.log_messages(status); + collector.log_messages(status).await; } diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 1881bcfa602..9bd4634cc2a 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -10,3 +10,6 @@ rand.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true + +[dev-dependencies] +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/litellm-rust/crates/core/src/call_lifecycle/README.md b/litellm-rust/crates/core/src/call_lifecycle/README.md new file mode 100644 index 00000000000..692e249ef27 --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/README.md @@ -0,0 +1,167 @@ +# Call lifecycle + +`litellm_core::call_lifecycle` is the shared execution wrapper for LiteLLM call +types migrated to Rust. It owns lifecycle ordering, phase timing, and trace +observer calls. It must not know about OCR, chat, messages, responses, +completions, provider auth, request transforms, or response normalization. + +Call-type modules own their domain behavior. For example, OCR owns document +payloads, OCR provider transforms, safe document fetch, guardrail payload shape, +callback payload shape, and provider HTTP execution. + +## Runtime order + +Every wrapped call runs in this order: + +1. `async_pre_call_hook` +2. `async_during_call_hook` +3. provider call +4. `async_log_success_event` or `async_log_failure_event` + +`async_pre_call_hook` receives the initial LiteLLM request shape. It is where +pre-call custom guardrails run. + +`async_during_call_hook` converts the initial request into the provider-ready +request. It is where provider config selection, parameter mapping, auth/header +resolution, request transforms, and during-call guardrails belong. + +The provider call receives only the provider-ready request. It should execute +I/O and call the provider response transform. + +Success and failure callbacks receive `CallLifecycleTiming`. Callback failures +must not replace the original provider or guardrail result. + +## Trace contract + +The lifecycle runner records: + +- full call start and end time +- `pre_call` phase timing +- `during_call` phase timing +- `provider_call` phase timing +- `success_callback` phase timing +- `failure_callback` phase timing + +`CallLifecycleObserver` receives phase start and end events. The default +observer is a no-op. Future OTEL support should implement this observer instead +of editing OCR, chat, messages, responses, completions, or provider modules. + +## Required shape + +Each migrated call type should use this folder shape: + +```text +litellm-rust/crates/ai-gateway/src// + mod.rs # thin public entrypoint + types.rs # public request, prepared request, provider request, response types + prepare.rs # model/provider/callback/guardrail setup + hooks.rs # CallLifecycleHooks implementation + handler.rs # provider I/O and response normalization + tests.rs # call-type lifecycle and handler tests +``` + +Provider transforms can live in `litellm-rust/crates/core/src/providers/...`. +Shared call-type helpers can live beside the call type, but generic lifecycle +code stays in this folder. + +## Core API + +The prepared request implements `CallLifecycleRequest`: + +```rust +impl CallLifecycleRequest for PreparedMessagesRequest { + fn lifecycle_context(&self) -> CallLifecycleContext { + CallLifecycleContext::new( + "messages", + self.model.clone(), + self.custom_llm_provider.clone(), + self.litellm_call_id.clone(), + ) + } +} +``` + +The call-type hooks implement `CallLifecycleHooks`: + +```rust +impl CallLifecycleHooks< + PreparedMessagesRequest, + ProviderMessagesRequest, + MessagesResponse, +> for MessagesLifecycleHooks { + fn async_pre_call_hook(...) { + // run pre-call custom guardrails against the LiteLLM request shape + } + + fn async_during_call_hook(...) { + // map params, validate env, transform request, run during-call guardrails + } + + fn async_log_success_event(...) { + // call async_log_success_event on configured custom loggers + } + + fn async_log_failure_event(...) { + // call async_log_failure_event without swallowing the original error + } +} +``` + +The public entrypoint stays thin: + +```rust +pub async fn messages(request: MessagesRequest<'_>) -> CoreResult { + let PreparedMessagesCall { request, hooks } = prepare_messages_call(request)?; + + CallLifecycle::default() + .run_request(request, &hooks, execute_messages_provider_call) + .await +} +``` + +Use `run_request` for new call types. Keep `run` available only for specialized +tests or existing code that already has a `CallLifecycleContext`. + +## Adding a new call type + +1. Add `/types.rs` + +Define the public request accepted by the bridge, the prepared request used by +the lifecycle runner, and the provider request consumed by the handler. + +2. Implement `CallLifecycleRequest` + +Return `call_type`, `model`, `custom_llm_provider`, and `litellm_call_id`. +Do not put provider-specific logic here. + +3. Add `/prepare.rs` + +Resolve model/provider once, generate or preserve `litellm_call_id`, construct +callback and guardrail runners, and return `PreparedCall`. + +4. Add `/hooks.rs` + +Implement `CallLifecycleHooks`. Put pre-call guardrail payload construction, +provider config selection, param mapping, request transform, during-call +guardrail payload construction, and callback payload construction here. + +5. Add `/handler.rs` + +Execute the provider request and normalize the provider response. Do not repeat +provider-specific transforms here; call the provider config. + +6. Add tests + +Cover hook order, success callback payload, failure callback payload, pre-call +guardrail blocking before provider I/O, during-call body mutation, and provider +error mapping. + +## Review checklist + +- Core lifecycle has no call-type or provider-specific branches +- Public call-type entrypoint only prepares and calls `run_request` +- Provider behavior lives behind provider config/transformation code +- Hook method names map to the Python custom logger and guardrail concepts +- Phase timing is recorded once in lifecycle, not separately per call type +- Callback failures never hide the original provider or guardrail error +- Tests prove the provider socket is not touched when pre-call guardrails block diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs new file mode 100644 index 00000000000..d9b68a1b726 --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -0,0 +1,414 @@ +use std::future::Future; +use std::time::{Instant, SystemTime, UNIX_EPOCH}; + +use crate::{CoreError, CoreResult}; + +pub mod types; + +pub use types::{ + CallLifecycleContext, CallLifecyclePhase, CallLifecyclePhaseTiming, CallLifecycleRequest, + CallLifecycleTiming, +}; + +pub trait CallLifecycleHooks: Send + Sync { + type PreCallFuture<'a>: Future> + Send + 'a + where + Self: 'a, + InitialReq: 'a, + ProviderReq: 'a, + Resp: 'a; + + type DuringCallFuture<'a>: Future> + Send + 'a + where + Self: 'a, + InitialReq: 'a, + ProviderReq: 'a, + Resp: 'a; + + type SuccessFuture<'a>: Future + Send + 'a + where + Self: 'a, + Resp: 'a; + + type FailureFuture<'a>: Future + Send + 'a + where + Self: 'a; + + fn async_pre_call_hook<'a>( + &'a self, + context: &'a CallLifecycleContext, + request: InitialReq, + ) -> Self::PreCallFuture<'a>; + + fn async_during_call_hook<'a>( + &'a self, + context: &'a CallLifecycleContext, + request: InitialReq, + ) -> Self::DuringCallFuture<'a>; + + fn async_log_success_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + response: &'a Resp, + timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a>; + + fn async_log_failure_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + error: &'a CoreError, + timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a>; +} + +pub trait CallLifecycleObserver: Send + Sync { + fn on_phase_start(&self, _context: &CallLifecycleContext, _phase: CallLifecyclePhase) {} + + fn on_phase_end(&self, _context: &CallLifecycleContext, _timing: &CallLifecyclePhaseTiming) {} +} + +#[derive(Default)] +pub struct NoopCallLifecycleObserver; + +impl CallLifecycleObserver for NoopCallLifecycleObserver {} + +pub struct CallLifecycle<'a> { + observer: &'a dyn CallLifecycleObserver, +} + +impl<'a> CallLifecycle<'a> { + pub fn new(observer: &'a dyn CallLifecycleObserver) -> Self { + Self { observer } + } + + pub async fn run_request( + &self, + request: InitialReq, + hooks: &Hooks, + provider_call: ProviderCall, + ) -> CoreResult + where + InitialReq: CallLifecycleRequest, + Hooks: CallLifecycleHooks, + ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, + ProviderFuture: Future>, + { + let context = request.lifecycle_context(); + self.run(context, request, hooks, provider_call).await + } + + pub async fn run( + &self, + context: CallLifecycleContext, + request: InitialReq, + hooks: &Hooks, + provider_call: ProviderCall, + ) -> CoreResult + where + Hooks: CallLifecycleHooks, + ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, + ProviderFuture: Future>, + { + let call_start = epoch_seconds(); + let mut phases = Vec::new(); + + let pre_call = self.start_phase(&context, CallLifecyclePhase::PreCall); + let request = match hooks.async_pre_call_hook(&context, request).await { + Ok(request) => { + phases.push(self.finish_phase(&context, pre_call)); + request + } + Err(error) => { + phases.push(self.finish_phase(&context, pre_call)); + self.log_failure(&context, hooks, &error, call_start, &mut phases) + .await; + return Err(error); + } + }; + + let during_call = self.start_phase(&context, CallLifecyclePhase::DuringCall); + let provider_request = match hooks.async_during_call_hook(&context, request).await { + Ok(request) => { + phases.push(self.finish_phase(&context, during_call)); + request + } + Err(error) => { + phases.push(self.finish_phase(&context, during_call)); + self.log_failure(&context, hooks, &error, call_start, &mut phases) + .await; + return Err(error); + } + }; + + let provider_phase = self.start_phase(&context, CallLifecyclePhase::ProviderCall); + let result = provider_call(provider_request).await; + phases.push(self.finish_phase(&context, provider_phase)); + + match &result { + Ok(response) => { + let success_phase = self.start_phase(&context, CallLifecyclePhase::SuccessCallback); + let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone()); + hooks + .async_log_success_event(&context, response, &timing) + .await; + phases.push(self.finish_phase(&context, success_phase)); + } + Err(error) => { + self.log_failure(&context, hooks, error, call_start, &mut phases) + .await; + } + } + + result + } + + async fn log_failure( + &self, + context: &CallLifecycleContext, + hooks: &Hooks, + error: &CoreError, + call_start: f64, + phases: &mut Vec, + ) where + Hooks: CallLifecycleHooks, + { + let failure_phase = self.start_phase(context, CallLifecyclePhase::FailureCallback); + let timing = CallLifecycleTiming::new(call_start, epoch_seconds(), phases.clone()); + hooks.async_log_failure_event(context, error, &timing).await; + phases.push(self.finish_phase(context, failure_phase)); + } + + fn start_phase(&self, context: &CallLifecycleContext, phase: CallLifecyclePhase) -> PhaseStart { + self.observer.on_phase_start(context, phase); + PhaseStart { + phase, + start_time: epoch_seconds(), + started_at: Instant::now(), + } + } + + fn finish_phase( + &self, + context: &CallLifecycleContext, + phase_start: PhaseStart, + ) -> CallLifecyclePhaseTiming { + let timing = CallLifecyclePhaseTiming { + phase: phase_start.phase, + start_time: phase_start.start_time, + end_time: epoch_seconds(), + duration: phase_start.started_at.elapsed(), + }; + self.observer.on_phase_end(context, &timing); + timing + } +} + +impl Default for CallLifecycle<'static> { + fn default() -> Self { + static OBSERVER: NoopCallLifecycleObserver = NoopCallLifecycleObserver; + Self::new(&OBSERVER) + } +} + +struct PhaseStart { + phase: CallLifecyclePhase, + start_time: f64, + started_at: Instant, +} + +fn epoch_seconds() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or(0.0) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::pin::Pin; + use std::sync::Mutex; + + type BoxFuture<'a, T> = Pin + Send + 'a>>; + + #[derive(Default)] + struct RecordingHooks { + events: Mutex>, + } + + struct RecordingRequest(String); + + impl CallLifecycleRequest for RecordingRequest { + fn lifecycle_context(&self) -> CallLifecycleContext { + CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1") + } + } + + impl RecordingHooks { + fn events(&self) -> Vec<&'static str> { + self.events.lock().unwrap().clone() + } + } + + impl CallLifecycleHooks for RecordingHooks { + type PreCallFuture<'a> = BoxFuture<'a, CoreResult>; + type DuringCallFuture<'a> = BoxFuture<'a, CoreResult>; + type SuccessFuture<'a> = BoxFuture<'a, ()>; + type FailureFuture<'a> = BoxFuture<'a, ()>; + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: String, + ) -> Self::PreCallFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("pre_call"); + Ok(format!("{request}:pre")) + }) + } + + fn async_during_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: String, + ) -> Self::DuringCallFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("during_call"); + Ok(format!("{request}:during")) + }) + } + + fn async_log_success_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _response: &'a String, + timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a> { + Box::pin(async move { + assert!(timing.end_time >= timing.start_time); + assert_eq!(timing.phases.len(), 3); + self.events.lock().unwrap().push("success"); + }) + } + + fn async_log_failure_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _error: &'a CoreError, + _timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("failure"); + }) + } + } + + impl CallLifecycleHooks for RecordingHooks { + type PreCallFuture<'a> = BoxFuture<'a, CoreResult>; + type DuringCallFuture<'a> = BoxFuture<'a, CoreResult>; + type SuccessFuture<'a> = BoxFuture<'a, ()>; + type FailureFuture<'a> = BoxFuture<'a, ()>; + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: RecordingRequest, + ) -> Self::PreCallFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("pre_call"); + Ok(RecordingRequest(format!("{}:pre", request.0))) + }) + } + + fn async_during_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: RecordingRequest, + ) -> Self::DuringCallFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("during_call"); + Ok(format!("{}:during", request.0)) + }) + } + + fn async_log_success_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _response: &'a String, + _timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("success"); + }) + } + + fn async_log_failure_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _error: &'a CoreError, + _timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("failure"); + }) + } + } + + #[tokio::test] + async fn lifecycle_runs_hooks_around_provider_call() { + let hooks = RecordingHooks::default(); + let response = CallLifecycle::default() + .run( + CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"), + "request".to_string(), + &hooks, + |request| async move { + assert_eq!(request, "request:pre:during"); + Ok("response".to_string()) + }, + ) + .await + .expect("call succeeds"); + + assert_eq!(response, "response"); + assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]); + } + + #[tokio::test] + async fn lifecycle_logs_failure_when_provider_fails() { + let hooks = RecordingHooks::default(); + let error = CallLifecycle::default() + .run( + CallLifecycleContext::new("ocr", "mistral-ocr-latest", "mistral", "call_1"), + "request".to_string(), + &hooks, + |_request| async move { + Err::(CoreError::Network("provider down".to_string())) + }, + ) + .await + .expect_err("call fails"); + + assert_eq!(error, CoreError::Network("provider down".to_string())); + assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]); + } + + #[tokio::test] + async fn lifecycle_can_run_any_request_with_embedded_context() { + let hooks = RecordingHooks::default(); + let response = CallLifecycle::default() + .run_request( + RecordingRequest("request".to_string()), + &hooks, + |request| async move { + assert_eq!(request, "request:pre:during"); + Ok("response".to_string()) + }, + ) + .await + .expect("call succeeds"); + + assert_eq!(response, "response"); + assert_eq!(hooks.events(), vec!["pre_call", "during_call", "success"]); + } +} diff --git a/litellm-rust/crates/core/src/call_lifecycle/types.rs b/litellm-rust/crates/core/src/call_lifecycle/types.rs new file mode 100644 index 00000000000..8819c8830d2 --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/types.rs @@ -0,0 +1,75 @@ +use std::time::Duration; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CallLifecycleContext { + pub call_type: String, + pub model: String, + pub custom_llm_provider: String, + pub litellm_call_id: String, +} + +impl CallLifecycleContext { + pub fn new( + call_type: impl Into, + model: impl Into, + custom_llm_provider: impl Into, + litellm_call_id: impl Into, + ) -> Self { + Self { + call_type: call_type.into(), + model: model.into(), + custom_llm_provider: custom_llm_provider.into(), + litellm_call_id: litellm_call_id.into(), + } + } +} + +pub trait CallLifecycleRequest { + fn lifecycle_context(&self) -> CallLifecycleContext; +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CallLifecyclePhase { + PreCall, + DuringCall, + ProviderCall, + SuccessCallback, + FailureCallback, +} + +impl CallLifecyclePhase { + pub fn as_str(self) -> &'static str { + match self { + Self::PreCall => "pre_call", + Self::DuringCall => "during_call", + Self::ProviderCall => "provider_call", + Self::SuccessCallback => "success_callback", + Self::FailureCallback => "failure_callback", + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct CallLifecyclePhaseTiming { + pub phase: CallLifecyclePhase, + pub start_time: f64, + pub end_time: f64, + pub duration: Duration, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct CallLifecycleTiming { + pub start_time: f64, + pub end_time: f64, + pub phases: Vec, +} + +impl CallLifecycleTiming { + pub fn new(start_time: f64, end_time: f64, phases: Vec) -> Self { + Self { + start_time, + end_time, + phases, + } + } +} diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 2ac479cc725..555a04ce853 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,7 +1,9 @@ +pub mod call_lifecycle; pub mod error; pub mod ocr; pub mod providers; pub mod realtime; pub mod router; +pub mod routing_utils; pub use error::{CoreError, CoreResult}; diff --git a/litellm-rust/crates/core/src/routing_utils/README.md b/litellm-rust/crates/core/src/routing_utils/README.md new file mode 100644 index 00000000000..8585c18e421 --- /dev/null +++ b/litellm-rust/crates/core/src/routing_utils/README.md @@ -0,0 +1,7 @@ +# Routing Utils + +Shared helpers for deciding how a LiteLLM model routes to an LLM provider. +Keep provider-name parsing, explicit `custom_llm_provider` handling, and model-prefix normalization here. +Do not put deployment selection or load-balancing logic here; that belongs in `router`. +Do not put provider HTTP transformation logic here; that belongs in `providers`. +Helpers in this folder should be deterministic and easy to unit test without network calls. diff --git a/litellm-rust/crates/core/src/routing_utils/mod.rs b/litellm-rust/crates/core/src/routing_utils/mod.rs new file mode 100644 index 00000000000..8336397f870 --- /dev/null +++ b/litellm-rust/crates/core/src/routing_utils/mod.rs @@ -0,0 +1 @@ +pub mod provider; diff --git a/litellm-rust/crates/core/src/routing_utils/provider.rs b/litellm-rust/crates/core/src/routing_utils/provider.rs new file mode 100644 index 00000000000..6333eedebfc --- /dev/null +++ b/litellm-rust/crates/core/src/routing_utils/provider.rs @@ -0,0 +1,77 @@ +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CustomLlmProvider<'a> { + pub model: &'a str, + pub custom_llm_provider: &'a str, +} + +pub fn get_custom_llm_provider<'a>( + model: &'a str, + custom_llm_provider: Option<&'a str>, +) -> Option> { + if let Some(custom_llm_provider) = custom_llm_provider.filter(|provider| !provider.is_empty()) { + return Some(CustomLlmProvider { + model: strip_custom_llm_provider_prefix(model, custom_llm_provider), + custom_llm_provider, + }); + } + + let (custom_llm_provider, model) = model.split_once('/')?; + if custom_llm_provider.is_empty() || model.is_empty() { + return None; + } + Some(CustomLlmProvider { + model, + custom_llm_provider, + }) +} + +fn strip_custom_llm_provider_prefix<'a>(model: &'a str, custom_llm_provider: &str) -> &'a str { + model + .strip_prefix(custom_llm_provider) + .and_then(|model| model.strip_prefix('/')) + .unwrap_or(model) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn gets_custom_llm_provider_from_model_prefix() { + assert_eq!( + get_custom_llm_provider("mistral/mistral-ocr-latest", None), + Some(CustomLlmProvider { + model: "mistral-ocr-latest", + custom_llm_provider: "mistral", + }) + ); + assert_eq!( + get_custom_llm_provider("azure_ai/doc-intelligence/prebuilt-layout", None), + Some(CustomLlmProvider { + model: "doc-intelligence/prebuilt-layout", + custom_llm_provider: "azure_ai", + }) + ); + assert_eq!(get_custom_llm_provider("mistral-ocr-latest", None), None); + assert_eq!(get_custom_llm_provider("/model", None), None); + assert_eq!(get_custom_llm_provider("provider/", None), None); + } + + #[test] + fn explicit_custom_llm_provider_strips_matching_model_prefix() { + assert_eq!( + get_custom_llm_provider("mistral/mistral-ocr-latest", Some("mistral")), + Some(CustomLlmProvider { + model: "mistral-ocr-latest", + custom_llm_provider: "mistral", + }) + ); + assert_eq!( + get_custom_llm_provider("mistral/mistral-ocr-latest", Some("vertex_ai")), + Some(CustomLlmProvider { + model: "mistral/mistral-ocr-latest", + custom_llm_provider: "vertex_ai", + }) + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 82024e1bf47..946a99f990c 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -96,7 +96,6 @@ fn ocr( optional_params: Option>, timeout_seconds: Option, ) -> PyResult> { - let custom_llm_provider = custom_llm_provider.unwrap_or_else(|| "mistral".to_string()); let (document, extra_headers, optional_params, timeout) = marshal_inputs( py, document, @@ -111,10 +110,14 @@ fn ocr( document, api_key: api_key.as_deref(), api_base: api_base.as_deref(), - custom_llm_provider: &custom_llm_provider, + custom_llm_provider: custom_llm_provider.as_deref(), extra_headers, optional_params, timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, })) }); @@ -138,7 +141,6 @@ fn aocr( optional_params: Option>, timeout_seconds: Option, ) -> PyResult> { - let custom_llm_provider = custom_llm_provider.unwrap_or_else(|| "mistral".to_string()); let (document, extra_headers, optional_params, timeout) = marshal_inputs( py, document, @@ -153,10 +155,14 @@ fn aocr( document, api_key: api_key.as_deref(), api_base: api_base.as_deref(), - custom_llm_provider: &custom_llm_provider, + custom_llm_provider: custom_llm_provider.as_deref(), extra_headers, optional_params, timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, }) .await .map_err(core_error_to_pyerr)?; diff --git a/litellm/__init__.py b/litellm/__init__.py index 9650dc12c97..5ae5942f32f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1394,7 +1394,7 @@ from .skills.main import ( ) from .containers.main import * from .ocr.main import * -from .ocr.rust_bridge import use_litellm_rust +from .rust_bridge.ocr import use_litellm_rust from .rag.main import * from .sandbox.main import * from .search.main import * diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 6a196d41768..93b3a892659 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -19,13 +19,7 @@ from litellm.constants import request_timeout from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.ocr.rust_bridge import ( - RustAocr, - RustOcr, - load_rust_aocr, - load_rust_ocr, - rust_ocr_enabled, -) +from litellm.rust_bridge import ocr as rust_ocr_bridge from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -60,27 +54,10 @@ class _PreparedRustOCRCall: _RUST_OCR_PROVIDERS = { "mistral", "azure_ai", - "azure_ai/doc-intelligence", "vertex_ai", } -def _timeout_to_seconds( - timeout: Union[float, httpx.Timeout] | None, -) -> float | None: - """Convert the Python OCR timeout to a single seconds value for the Rust bridge. - - The Rust HTTP client takes one duration; ``httpx.Timeout`` carries separate - connect/read/write/pool values, so pick the read deadline as the closest - analog to a total-request timeout. - """ - if timeout is None: - return None - if isinstance(timeout, httpx.Timeout): - return timeout.read - return float(timeout) - - def _prepare_ocr_request( model: str, document: dict[str, Any], @@ -218,13 +195,9 @@ def _rust_bridge_api_base( ) -> str | None: if prepared_request.api_base is not None: return prepared_request.api_base - if prepared_request.custom_llm_provider == "azure_ai/doc-intelligence": - return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") if prepared_request.custom_llm_provider == "azure_ai": - if ( - "doc-intelligence" in prepared_request.model - or "documentintelligence" in prepared_request.model - ): + model = prepared_request.model.lower() + if "doc-intelligence" in model or "documentintelligence" in model: return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") return resolve_secret("AZURE_AI_API_BASE") return None @@ -278,57 +251,53 @@ def _prepare_rust_ocr_call( def _run_rust_ocr( - rust_ocr: RustOcr, prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], -) -> OCRResponse: - """Run the Mistral OCR call through the Rust bridge and wrap the result. - - Resolves the key the same way the Python path does so secret-manager backends - (AWS/Azure/GCP/Vault) work; the Rust bridge's own fallback only reads the - process environment. The request that Rust actually sends (resolved URL and - headers) is mirrored into pre_call so logs match the wire. Dependencies are - injected so this stays unit-testable without patching module globals. - """ +) -> OCRResponse | None: + if rust_ocr_bridge.load_rust_ocr() is None: + return None prepared = _prepare_rust_ocr_call( prepared_request=prepared_request, resolve_api_key=resolve_api_key, ) - return OCRResponse.model_validate( - rust_ocr( - model=prepared_request.model, - document=cast(dict[str, object], prepared_request.document), - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout), - ) + rust_response = rust_ocr_bridge.ocr( + model=prepared_request.model, + document=prepared_request.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout=prepared_request.effective_timeout, ) + if rust_response is None: + return None + return OCRResponse.model_validate(rust_response) async def _run_rust_aocr( - rust_aocr: RustAocr, prepared_request: _PreparedOCRRequest, resolve_api_key: Callable[[str], str | None], -) -> OCRResponse: +) -> OCRResponse | None: + if rust_ocr_bridge.load_rust_aocr() is None: + return None prepared = _prepare_rust_ocr_call( prepared_request=prepared_request, resolve_api_key=resolve_api_key, ) - return OCRResponse.model_validate( - await rust_aocr( - model=prepared_request.model, - document=cast(dict[str, object], prepared_request.document), - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout), - ) + rust_response = await rust_ocr_bridge.aocr( + model=prepared_request.model, + document=prepared_request.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout=prepared_request.effective_timeout, ) + if rust_response is None: + return None + return OCRResponse.model_validate(rust_response) @client @@ -427,21 +396,19 @@ async def aocr( {"model": model, "custom_llm_provider": custom_llm_provider} ) - if _rust_ocr_supported(prepared) and rust_ocr_enabled(): - rust_aocr = load_rust_aocr() - if rust_aocr is None: + if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled(): + from litellm.secret_managers.main import get_secret_str + + rust_response = await _run_rust_aocr( + prepared_request=prepared, + resolve_api_key=get_secret_str, + ) + if rust_response is None: verbose_logger.debug( "Async Rust OCR bridge unavailable; falling back to Python path" ) else: - from litellm.secret_managers.main import get_secret_str - - response = await _run_rust_aocr( - rust_aocr=rust_aocr, - prepared_request=prepared, - resolve_api_key=get_secret_str, - ) - return response + return rust_response response = base_llm_http_handler.ocr( model=prepared.model, @@ -704,21 +671,19 @@ def ocr( {"model": model, "custom_llm_provider": custom_llm_provider} ) - # Optional Rust path: hand supported OCR calls to the Rust bridge. - if _rust_ocr_supported(prepared) and rust_ocr_enabled(): - rust_ocr = load_rust_ocr() - if rust_ocr is None: + if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled(): + from litellm.secret_managers.main import get_secret_str + + rust_response = _run_rust_ocr( + prepared_request=prepared, + resolve_api_key=get_secret_str, + ) + if rust_response is None: verbose_logger.debug( "Rust OCR bridge unavailable; falling back to Python path" ) else: - from litellm.secret_managers.main import get_secret_str - - return _run_rust_ocr( - rust_ocr=rust_ocr, - prepared_request=prepared, - resolve_api_key=get_secret_str, - ) + return rust_response response = base_llm_http_handler.ocr( model=prepared.model, diff --git a/litellm/rust_bridge/__init__.py b/litellm/rust_bridge/__init__.py index ec89e3b65b4..3da5b98449b 100644 --- a/litellm/rust_bridge/__init__.py +++ b/litellm/rust_bridge/__init__.py @@ -4,5 +4,6 @@ from litellm.rust_bridge.loader import ( get_native_bridge, native_bridge_available, ) +from litellm.rust_bridge.ocr import use_litellm_rust -__all__ = ["get_native_bridge", "native_bridge_available"] +__all__ = ["get_native_bridge", "native_bridge_available", "use_litellm_rust"] diff --git a/litellm/ocr/rust_bridge.py b/litellm/rust_bridge/ocr.py similarity index 52% rename from litellm/ocr/rust_bridge.py rename to litellm/rust_bridge/ocr.py index 631fd4c63c5..36a088b6b1a 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/rust_bridge/ocr.py @@ -1,30 +1,21 @@ -""" -Optional Rust-backed OCR path. - -Enable with ``litellm.use_litellm_rust()``; the sync ``litellm.ocr()`` entrypoint -then routes supported Mistral calls through the compiled ``litellm.rust_bridge._native`` -extension, which performs the whole OCR call (URL, headers, HTTP, parse) in Rust. - -No module-level ``litellm`` imports keep this a leaf so ``litellm/ocr/main.py`` -can import it statically without forming an import cycle. -""" +"""Thin Python wrapper for the native Rust OCR bridge.""" from __future__ import annotations import os -from typing import Awaitable, Final, Protocol, cast +from typing import Any, Awaitable, Final, Protocol, Union, cast + +import httpx class RustOcr(Protocol): - """Signature of the compiled Rust OCR entrypoint.""" - def __call__( self, model: str, document: dict[str, object], api_key: str | None, api_base: str | None, - custom_llm_provider: str, + custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, @@ -33,15 +24,13 @@ class RustOcr(Protocol): class RustAocr(Protocol): - """Signature of the compiled ``litellm_python_bridge.aocr`` entrypoint.""" - def __call__( self, model: str, document: dict[str, object], api_key: str | None, api_base: str | None, - custom_llm_provider: str, + custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, @@ -50,7 +39,7 @@ class RustAocr(Protocol): class _Unset: - """Sentinel type so ``ocr=None`` can clear a prior injection while omission preserves it.""" + pass _UNSET: Final[_Unset] = _Unset() @@ -76,12 +65,6 @@ def use_litellm_rust( ocr: RustOcr | None | _Unset = _UNSET, aocr: RustAocr | None | _Unset = _UNSET, ) -> None: - """Route supported OCR calls through the packaged Rust extension. - - ``ocr`` and ``aocr`` inject bridge callables; when omitted the compiled - extension is loaded on demand and any previously injected bridge is - preserved. Pass ``None`` explicitly to clear a prior injection. - """ global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl _rust_ocr_enabled = enabled if not isinstance(ocr, _Unset): @@ -91,17 +74,10 @@ def use_litellm_rust( def rust_ocr_enabled() -> bool: - """Whether the Rust OCR path has been turned on via ``use_litellm_rust()``.""" return _rust_ocr_enabled def load_rust_ocr() -> RustOcr | None: - """Return the Rust OCR callable, or ``None`` when no bridge is available. - - Prefers an injected implementation, otherwise loads the compiled - ``litellm.rust_bridge._native`` extension; a missing extension yields ``None`` so - the caller can fall back to the Python path instead of hard-failing. - """ if _rust_ocr_impl is not None: return _rust_ocr_impl from litellm.rust_bridge import get_native_bridge @@ -113,7 +89,6 @@ def load_rust_ocr() -> RustOcr | None: def load_rust_aocr() -> RustAocr | None: - """Return the async Rust OCR callable, or ``None`` when unavailable.""" if _rust_aocr_impl is not None: return _rust_aocr_impl from litellm.rust_bridge import get_native_bridge @@ -122,3 +97,63 @@ def load_rust_aocr() -> RustAocr | None: if native_bridge is None: return None return cast(RustAocr, getattr(native_bridge, "aocr", None)) + + +def _timeout_to_seconds(timeout: Union[float, httpx.Timeout] | None) -> float | None: + if timeout is None: + return None + if isinstance(timeout, httpx.Timeout): + return timeout.read + return float(timeout) + + +def ocr( + *, + model: str, + document: dict[str, Any], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, Any] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, +) -> dict[str, object] | None: + rust_ocr = load_rust_ocr() + if rust_ocr is None: + return None + return rust_ocr( + model=model, + document=cast(dict[str, object], document), + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=cast(dict[str, object] | None, extra_headers), + optional_params=optional_params, + timeout_seconds=_timeout_to_seconds(timeout), + ) + + +async def aocr( + *, + model: str, + document: dict[str, Any], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, Any] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, +) -> dict[str, object] | None: + rust_aocr = load_rust_aocr() + if rust_aocr is None: + return None + return await rust_aocr( + model=model, + document=cast(dict[str, object], document), + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=cast(dict[str, object] | None, extra_headers), + optional_params=optional_params, + timeout_seconds=_timeout_to_seconds(timeout), + ) diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 7e23e441f50..acad249a2bb 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -1,4 +1,4 @@ -"""Tests for the optional Rust-backed OCR path (``litellm/ocr/rust_bridge.py``).""" +"""Tests for the optional Rust-backed OCR path.""" import importlib import builtins @@ -15,7 +15,7 @@ from litellm.llms.base_llm.ocr.transformation import OCRResponse # function onto `litellm.ocr` and shadows the submodule, so import the modules # explicitly via importlib rather than attribute traversal. ocr_main = importlib.import_module("litellm.ocr.main") -rust_bridge = importlib.import_module("litellm.ocr.rust_bridge") +rust_bridge = importlib.import_module("litellm.rust_bridge.ocr") rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") MODEL = "mistral/mistral-ocr-latest" @@ -49,7 +49,7 @@ class RecordingBridge: document: dict[str, object], api_key: str | None, api_base: str | None, - custom_llm_provider: str, + custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, @@ -81,7 +81,7 @@ class RecordingAsyncBridge: document: dict[str, object], api_key: str | None, api_base: str | None, - custom_llm_provider: str, + custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, @@ -108,7 +108,7 @@ class RaisingBridge: document: dict[str, object], api_key: str | None, api_base: str | None, - custom_llm_provider: str, + custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, @@ -123,7 +123,7 @@ class RaisingAsyncBridge: document: dict[str, object], api_key: str | None, api_base: str | None, - custom_llm_provider: str, + custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, @@ -366,17 +366,78 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): def test_timeout_to_seconds_handles_float_timeout_and_none(): - assert ocr_main._timeout_to_seconds(12.5) == 12.5 - assert ocr_main._timeout_to_seconds(None) is None - assert ocr_main._timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 + assert rust_bridge._timeout_to_seconds(12.5) == 12.5 + assert rust_bridge._timeout_to_seconds(None) is None + assert rust_bridge._timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 -def test_run_rust_ocr_forwards_args_and_wraps_response(): +def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): + bridge = RecordingBridge() + + litellm.use_litellm_rust(True, ocr=bridge) + response = rust_bridge.ocr( + model="mistral-ocr-latest", + document=DOCUMENT, + api_key="sk-test", + api_base="https://proxy.internal", + custom_llm_provider="mistral", + extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"}, + optional_params={"include_image_base64": True, "pages": [0]}, + timeout=12.5, + ) + + assert response == FAKE_OCR_RESPONSE + call = bridge.calls[0] + assert call == { + "model": "mistral-ocr-latest", + "document": DOCUMENT, + "api_key": "sk-test", + "api_base": "https://proxy.internal", + "custom_llm_provider": "mistral", + "extra_headers": { + "Authorization": "Bearer sk-test", + "x-trace-id": "trace-1", + }, + "optional_params": {"include_image_base64": True, "pages": [0]}, + "timeout_seconds": 12.5, + } + + +@pytest.mark.asyncio +async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): + bridge = RecordingAsyncBridge() + + litellm.use_litellm_rust(True, aocr=bridge) + response = await rust_bridge.aocr( + model="mistral-ocr-maas", + document=DOCUMENT, + api_key=None, + api_base=None, + custom_llm_provider="vertex_ai", + extra_headers=None, + optional_params={"vertex_project": "project-1"}, + timeout=httpx.Timeout(30.0, read=42.0), + ) + + assert response == FAKE_OCR_RESPONSE + assert bridge.calls[0] == { + "model": "mistral-ocr-maas", + "document": DOCUMENT, + "api_key": None, + "api_base": None, + "custom_llm_provider": "vertex_ai", + "extra_headers": None, + "optional_params": {"vertex_project": "project-1"}, + "timeout_seconds": 42.0, + } + + +def test_run_rust_ocr_prepares_request_and_wraps_response(): bridge = RecordingBridge() logging_obj = RecordingLogging() + litellm.use_litellm_rust(True, ocr=bridge) response = ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request( logging_obj=logging_obj, api_base="https://proxy.internal", @@ -389,8 +450,7 @@ def test_run_rust_ocr_forwards_args_and_wraps_response(): assert isinstance(response, OCRResponse) assert response.pages[0].markdown == "hello world" - call = bridge.calls[0] - assert call == { + assert bridge.calls[0] == { "model": "mistral-ocr-latest", "document": DOCUMENT, "api_key": "sk-test", @@ -406,12 +466,10 @@ def test_run_rust_ocr_forwards_args_and_wraps_response(): def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): - """No explicit api_key: the resolver (get_secret_str in production) supplies it, - so secret-manager backends (AWS/Azure/GCP/Vault) work like the Python path.""" bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request(api_key=None, timeout=None), resolve_api_key=lambda name: ( "sk-from-vault" if name == "MISTRAL_API_KEY" else None @@ -421,16 +479,34 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): assert bridge.calls[0]["api_key"] == "sk-from-vault" +def test_run_rust_ocr_prefers_explicit_key_over_resolver(): + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + + def _resolver(name: str) -> str | None: + raise AssertionError(f"resolver should not be called for {name}") + + ocr_main._run_rust_ocr( + prepared_request=build_prepared_request( + api_key="sk-explicit", + timeout=None, + ), + resolve_api_key=_resolver, + ) + + assert bridge.calls[0]["api_key"] == "sk-explicit" + + def test_run_rust_ocr_uses_provider_api_key_env_var(): bridge = RecordingBridge() resolver_calls = [] + litellm.use_litellm_rust(True, ocr=bridge) def _resolver(name): resolver_calls.append(name) return "sk-provider-env" ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request( provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"), model="provider-ocr-model", @@ -446,9 +522,9 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", @@ -472,6 +548,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) def _resolver(name: str) -> str | None: return { @@ -480,7 +557,6 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana }.get(name) ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", @@ -495,9 +571,9 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request( custom_llm_provider="azure_ai", model="pixtral-12b-2409", @@ -514,12 +590,12 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) ocr_main._run_rust_ocr( - rust_ocr=bridge, prepared_request=build_prepared_request( - custom_llm_provider="azure_ai/doc-intelligence", - model="prebuilt-layout", + custom_llm_provider="azure_ai", + model="doc-intelligence/prebuilt-layout", api_base=None, timeout=None, ), @@ -533,30 +609,12 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): assert bridge.calls[0]["api_base"] == "https://document-intelligence.example.com" -def test_run_rust_ocr_prefers_explicit_key_over_resolver(): - bridge = RecordingBridge() - resolver_calls = [] - - def _resolver(name): - resolver_calls.append(name) - return "sk-from-vault" - - ocr_main._run_rust_ocr( - rust_ocr=bridge, - prepared_request=build_prepared_request(api_key="sk-explicit", timeout=None), - resolve_api_key=_resolver, - ) - - assert bridge.calls[0]["api_key"] == "sk-explicit" - assert resolver_calls == [] # resolver never consulted when a key is supplied - - def test_run_rust_ocr_runs_pre_call_logging(): - """The Rust shortcut must run pre_call so callbacks and spend tracking fire.""" logging_obj = RecordingLogging() + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) ocr_main._run_rust_ocr( - rust_ocr=RecordingBridge(), prepared_request=build_prepared_request( logging_obj=logging_obj, api_base="https://api.mistral.ai/v1", @@ -573,7 +631,6 @@ def test_run_rust_ocr_runs_pre_call_logging(): complete_input = additional_args["complete_input_dict"] assert complete_input["document"] == DOCUMENT assert complete_input["include_image_base64"] is True - # The logged request mirrors what Rust sends: resolved URL + headers. assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr" assert additional_args["headers"] == { "Authorization": "Bearer sk-test", @@ -594,7 +651,6 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): assert response.pages[0].markdown == "hello world" assert len(fake_bridge.calls) == 1 call = fake_bridge.calls[0] - # Provider prefix is stripped before reaching the bridge. assert call["model"] == "mistral-ocr-latest" assert call["document"] == DOCUMENT assert call["api_key"] == "sk-test" @@ -603,7 +659,6 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): "Authorization": "Bearer sk-test", "x-trace-id": "trace-1", } - # Raw OCR params ride along in optional_params; Rust filters to supported keys. assert call["optional_params"].get("include_image_base64") is True @@ -621,6 +676,19 @@ def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge): assert fake_bridge.calls[0]["custom_llm_provider"] == "azure_ai" +def test_ocr_rust_path_converts_file_document_before_bridge(fake_bridge): + response = litellm.ocr( + model=MODEL, + document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"}, + api_key="sk-test", + ) + + assert isinstance(response, OCRResponse) + document = fake_bridge.calls[0]["document"] + assert document["type"] == "document_url" + assert document["document_url"].startswith("data:application/pdf;base64,") + + def test_ocr_exception_type_uses_resolved_provider_context( monkeypatch: pytest.MonkeyPatch, ): @@ -694,12 +762,10 @@ def test_ocr_forwards_timeout_to_rust(fake_bridge): def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): - """When no explicit timeout is given, the library default (request_timeout) - must still be forwarded so the Rust path matches the Python path's deadline.""" - from litellm.constants import request_timeout - litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + from litellm.constants import request_timeout + assert fake_bridge.calls[0]["timeout_seconds"] == float(request_timeout) @@ -717,7 +783,7 @@ def test_ocr_does_not_route_to_rust_when_disabled(): def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): """Rust enabled but no bridge available (no injected impl, no compiled wheel): ocr() must degrade to the Python HTTP handler instead of raising.""" - monkeypatch.setattr(ocr_main, "load_rust_ocr", lambda: None) + monkeypatch.setattr(rust_bridge, "load_rust_ocr", lambda: None) litellm.use_litellm_rust(True) # enabled, but load_rust_ocr() returns None in CI captured = {} From f16af8853b4c3f4052c29ce5680abc40f726569a Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 25 Jun 2026 18:49:15 -0700 Subject: [PATCH 06/16] feat(mcp): opt-in least-privilege default for team key MCP access (#31380) * feat(mcp): add require_key_mcp_access_defined to stop keys inheriting team MCP servers By default a virtual key that grants no MCP servers of its own inherits its team's full MCP server list. The new general_settings flag require_key_mcp_access_defined (default false) flips this so the team list acts purely as a ceiling: a key reaches only the servers it grants explicitly (or via an access group), and inherits none. This mirrors the existing require_end_user_mcp_access_defined setting. The default is unchanged, so existing deployments keep today's behavior until they opt in. The no-mcp-servers sentinel and key access-group grants are unaffected. * docs(mcp): note require_key_mcp_access_defined effect in resolver docstring --- .../mcp_server/auth/user_api_key_auth_mcp.py | 15 ++- .../auth/test_user_api_key_auth_mcp.py | 98 +++++++++++++++++++ 2 files changed, 111 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 19f508e6af9..b7ac6a8c325 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -624,7 +624,9 @@ class MCPRequestHandler: Permission hierarchy (all rules are intersections): 1. Get allowed servers from key permissions - 2. Get allowed servers from team permissions (key inherits from team, or intersection) + 2. Get allowed servers from team permissions (key inherits from team, or + intersection; or inherits nothing when require_key_mcp_access_defined + is enabled, making the team a ceiling rather than a default) 3. Get allowed servers from end_user permissions (intersected if set) 4. Get allowed servers from agent permissions (intersected if set) 5. Get allowed servers from org permissions — org acts as a ceiling: if the org @@ -677,7 +679,16 @@ class MCPRequestHandler: if not team_set: base = key_set # no team restriction elif not key_set: - base = team_set # key has no own perms → inherits team + # A key that grants no MCP servers of its own inherits the + # team's by default. With require_key_mcp_access_defined the + # team is a ceiling rather than a default, so the key must + # grant servers explicitly (or via an access group) to reach + # any — it inherits none. + base = ( + set() + if general_settings.get("require_key_mcp_access_defined", False) + else team_set + ) else: base = key_set & team_set # both restrict → intersect diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 20fd5c1d86a..b3b0e8adcf6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -170,6 +170,104 @@ class TestMCPRequestHandler: mock_key_servers.assert_called_once_with(user_api_key_auth) mock_team_servers.assert_called_once_with(user_api_key_auth) + @pytest.mark.parametrize( + "require_key_mcp_access_defined,expected", + [ + # Default (flag off): a key with no MCP scope of its own inherits + # the team's servers. + (False, ["team_server1", "team_server2"]), + # Flag on: the team is a ceiling, not a default — the key inherits + # nothing and must grant servers explicitly. + (True, []), + ], + ) + async def test_require_key_mcp_access_defined_gates_team_inheritance( + self, require_key_mcp_access_defined, expected + ): + """The require_key_mcp_access_defined general setting flips an empty key + from inheriting its team's MCP servers (default) to inheriting none.""" + auth = UserAPIKeyAuth( + api_key="test-key", user_id="test-user", team_id="test-team" + ) + with ( + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + new_callable=AsyncMock, + return_value=[], + ), + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + new_callable=AsyncMock, + return_value=["team_server1", "team_server2"], + ), + patch.object( + MCPRequestHandler, + "_get_key_access_group_mcp_server_extras", + new_callable=AsyncMock, + return_value=[], + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"require_key_mcp_access_defined": require_key_mcp_access_defined}, + ), + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + + assert sorted(result) == sorted(expected) + + @pytest.mark.parametrize( + "key_servers,grants,expected,scenario", + [ + # Explicit key subset is still honored under the flag (intersected + # with the team ceiling) — the flag only removes empty-key inheritance. + (["team_server1"], [], ["team_server1"], "explicit_subset_survives"), + # An access-group grant is the escape hatch: it surfaces even though + # the key inherits nothing from the team. + ([], ["granted_server"], ["granted_server"], "access_group_grant_survives"), + ], + ) + async def test_require_key_mcp_access_defined_preserves_explicit_grants( + self, key_servers, grants, expected, scenario + ): + """With require_key_mcp_access_defined on, a key still reaches servers it + grants explicitly or via an access group — only blanket team inheritance + is removed.""" + auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + access_group_ids=["grp"] if grants else [], + ) + with ( + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + new_callable=AsyncMock, + return_value=key_servers, + ), + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + new_callable=AsyncMock, + return_value=["team_server1", "team_server2"], + ), + patch.object( + MCPRequestHandler, + "_get_key_access_group_mcp_server_extras", + new_callable=AsyncMock, + return_value=grants, + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"require_key_mcp_access_defined": True}, + ), + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(auth) + + assert sorted(result) == sorted(expected) + @pytest.mark.parametrize("team_servers", [[], ["team_server1", "team_server2"]]) async def test_no_mcp_servers_sentinel_returns_empty(self, team_servers): """A key scoped to the no-mcp-servers sentinel resolves to zero servers, From 997c7a267622d73f3a7a3e29656e78e483b0b187 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 25 Jun 2026 19:22:01 -0700 Subject: [PATCH 07/16] chore(ci): main into internal_staging (reconcile OCR hotfix history; unblocks #31384) (#31390) * docs(readme): add Deploy on AWS/GCP with Terraform section Adds a quickstart for the two published Terraform modules on the public registry (BerriAI/litellm/aws and BerriAI/litellm/google). Copy-paste main.tf for each cloud, the one-time GCP Artifact Registry remote-repo command, and pointers to the registry pages for the full input surface. Sits inside the Get Started section, between the gateway/SDK table and Run in Developer Mode -- where someone scanning the README for "how do I deploy this" will land. Co-Authored-By: Claude Opus 4.7 * docs(readme): add 1-click deploy buttons for AWS + GCP GCP gets the real 1-click: Open in Cloud Shell badge that clones the repo and walks through `terraform apply` via the existing DeployStack tutorial (already shipped at terraform/litellm/gcp/examples/default/ TUTORIAL.md). User just picks a project. AWS gets a soft 1-click: a Launch in AWS CloudShell badge that opens an in-browser, already-authenticated shell. User runs four commands (clone + cd + cp tfvars + terraform apply) once inside. There's no native AWS deeplink that pre-clones a repo + runs a tutorial -- CFN "Launch Stack" + CodeBuild would be needed for that, and that's a separate piece of work. Co-Authored-By: Claude Opus 4.7 * docs(readme): move AWS + GCP deploy buttons next to Render button * docs(readme): unify deploy button sizes and badge styles * docs(readme): bump deploy button height to 48 to match Render/Railway * docs(readme): bump AWS/GCP badge height to compensate for SVG padding * docs(readme): bump AWS/GCP badge height to 72 * docs(readme): bump AWS/GCP badge height to 84 * fix(readme): make deploy buttons same height (48px) https://claude.ai/code/session_01MxQRMHSDXbqJh74rF86UBc * docs(readme): flag GCP project ID substitution in image_registry * docs(readme): equalize deploy button heights and fix Cloud Shell button font GitHub rewrites an image's height attribute to "height: auto; max-height: Npx", which only caps and never stretches, so each image renders at its intrinsic height. The AWS/GCP shields badges are intrinsically 28px while the Render/Railway buttons are 40px, leaving the row uneven regardless of the height="48" we set. Replace the two shields badges with committed 40px PNGs so all four header buttons render at the same 40px. Also swap the Cloud Shell button from open-btn.svg to open-btn.png. The SVG renders its label as live text with font-family "Roboto, Sans" and no generic fallback; since neither font exists in GitHub's render environment, the text fell back to a serif (Times New Roman). The PNG bakes in the correct typeface. * docs(readme): collapse Railway deploy anchor to a single line The Railway button wrapped its img across indented lines, so the anchor contained leading and trailing whitespace. GitHub underlines link content, rendering that whitespace as a small blue underline beside the button. Put the anchor on one line like the other three buttons so there is no inner whitespace to underline. * Add Claude Fable 5 cost map entries as a data-only hotfix Backports only the model map changes from #30064 so deployments on released litellm versions pick up Fable 5 pricing, context window, and the adaptive thinking flag through the hosted cost map fetch without upgrading. Includes the supports_sampling_params flag on the 28 Fable 5 / Opus 4.7 / Opus 4.8 entries (ignored by released code, read by the gating that ships with the next release) and the matching one-line schema declaration so the map validation test passes. https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * feat: make rust OCR async-first * docs: clarify rust provider call flow * docs: clarify OCR provider transform contract * docs: note Tokio route contract * fix: address OCR bridge review comments * docs: bound rust OCR HTTP exception * feat: generate rust providers from registry * chore: move rust provider registry into core * chore: source rust providers from endpoint registry * fix: satisfy OCR lint budget * fix: reduce OCR basedpyright argument errors * fix: address OCR greptile feedback * fix: align rust OCR request preparation * fix: resolve OCR CodeQL alerts * fix: avoid duplicate Rust OCR authorization header * ci: rerun CircleCI --------- Co-authored-by: shin-berri Co-authored-by: Yassin Kortam Co-authored-by: Claude Opus 4.7 Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Krrish Dholakia Co-authored-by: Ishaan Jaff Co-authored-by: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> From 6e3540856cd1b3527f58e6bd1649ef2a6cb321e4 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 25 Jun 2026 19:53:16 -0700 Subject: [PATCH 08/16] fix(vertex): preserve Gemini Embedding 2 usageMetadata for cost tracking (#31354) * fix(vertex): preserve Gemini Embedding 2 usageMetadata for cost tracking * style(vertex): apply ruff format to batch_embed_content_transformation * fix(vertex): bill files/ image refs in Gemini embedContent at per-image rate Resolved files/... references whose mime type is an image were not detected by _is_image_element, so image_count stayed 0 and generic_cost_per_token fell back to the text token rate instead of input_cost_per_image. Thread the resolved_files mapping into the usage builder so resolved image references are counted and billed per image. Also modernize the _flatten_input return annotation to satisfy the ruff UP006 strict gate. * fix(vertex): bill Gemini embedding audio per-second and stop video+audio double-billing Audio-only embedContent responses set audio_tokens, but generic_cost_per_token only charges audio via input_cost_per_audio_token. gemini-embedding-2 prices audio via input_cost_per_audio_per_second, so spend stayed at $0. Plumb a new audio_length_seconds field through PromptTokensDetailsWrapper, parse it in _parse_prompt_tokens_details, and bill it from _calculate_input_cost. The vertex embedding transformation derives audio_length_seconds from audio_tokens using the documented 32 tokens/sec Gemini rate. The 1-token text floor that protects video billing only fired when no other modality was billable, but audio presence flipped that flag, leaving text_tokens at zero for video+audio responses. generic_cost_per_token then rewrote text_tokens to prompt_tokens minus audio_tokens (the video token count), charging video tokens as text on top of the per-second video cost. The rewrite trigger is text_tokens == 0 and image_count == 0; align the floor with that trigger and ignore audio_tokens. --------- Co-authored-by: Cursor Agent --- .../litellm_core_utils/llm_cost_calc/utils.py | 18 ++ .../batch_embed_content_handler.py | 2 + .../batch_embed_content_transformation.py | 151 +++++++++++- litellm/types/utils.py | 5 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 1 + ...test_batch_embed_content_transformation.py | 219 +++++++++++++++++- 6 files changed, 383 insertions(+), 13 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 19b711f6748..d2c97a5d6e7 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -485,6 +485,7 @@ class PromptTokensDetailsResult(TypedDict): character_count: int image_count: int video_length_seconds: float + audio_length_seconds: float def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: @@ -535,6 +536,13 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: ) or 0.0 ) + audio_length_seconds = ( + cast( + Optional[float], + getattr(usage.prompt_tokens_details, "audio_length_seconds", 0), + ) + or 0.0 + ) return PromptTokensDetailsResult( cache_hit_tokens=cache_hit_tokens, @@ -546,6 +554,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: character_count=character_count, image_count=image_count, video_length_seconds=float(video_length_seconds), + audio_length_seconds=float(audio_length_seconds), ) @@ -667,6 +676,14 @@ def _calculate_input_cost( prompt_tokens_details["video_length_seconds"], ) + ### AUDIO LENGTH COST + if prompt_tokens_details["audio_length_seconds"]: + prompt_cost += calculate_cost_component( + model_info, + "input_cost_per_audio_per_second", + prompt_tokens_details["audio_length_seconds"], + ) + return prompt_cost @@ -743,6 +760,7 @@ def generic_cost_per_token( character_count=0, image_count=0, video_length_seconds=0.0, + audio_length_seconds=0.0, ) if usage.prompt_tokens_details: prompt_tokens_details = _parse_prompt_tokens_details(usage) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 165dac24903..53f5fb464df 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -274,6 +274,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, + resolved_files=resolved_files, ) else: _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore @@ -377,6 +378,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, + resolved_files=resolved_files, ) else: _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index ba6e6f0c056..27ca1bd92a4 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -4,7 +4,10 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc Why separate file? Make it easy to see how transformation works """ -from typing import Dict, List, Optional, Tuple +from collections.abc import Mapping +from typing import Dict, List, Optional, Sequence, Tuple + +from pydantic import TypeAdapter, ValidationError from litellm.types.llms.vertex_ai import ( BlobType, @@ -13,10 +16,17 @@ from litellm.types.llms.vertex_ai import ( FileDataType, GeminiEmbeddingInput, PartType, + PromptTokensDetails, + UsageMetadata, VertexAIBatchEmbeddingsRequestBody, VertexAIBatchEmbeddingsResponseObject, ) -from litellm.types.utils import Embedding, EmbeddingResponse, Usage +from litellm.types.utils import ( + Embedding, + EmbeddingResponse, + PromptTokensDetailsWrapper, + Usage, +) from litellm.utils import get_formatted_prompt, token_counter SUPPORTED_EMBEDDING_MIME_TYPES = { @@ -294,11 +304,133 @@ def transform_openai_input_gemini_embed_content( return request_body +_IMAGE_MIME_TYPES = frozenset({"image/png", "image/jpeg"}) +_VIDEO_TOKENS_PER_SECOND = 258.0 +_AUDIO_TOKENS_PER_SECOND = 32.0 +_usage_metadata_adapter = TypeAdapter(UsageMetadata) + + +def _parse_usage_metadata(raw_usage_metadata: object) -> Optional[UsageMetadata]: + if not isinstance(raw_usage_metadata, dict): + return None + try: + return _usage_metadata_adapter.validate_python(raw_usage_metadata) + except ValidationError: + return None + + +def _flatten_input(input: GeminiEmbeddingInput) -> tuple[str, ...]: + if isinstance(input, str): + return (input,) + return tuple( + sub + for element in input + for sub in (element if isinstance(element, list) else [element]) + ) + + +def _is_image_element( + element: str, + resolved_files: Mapping[str, Mapping[str, str]], +) -> bool: + if element.startswith("data:") and ";base64," in element: + try: + mime_type, _ = _parse_data_url(element) + except ValueError: + return False + return mime_type in _IMAGE_MIME_TYPES + if _is_gcs_url(element): + try: + return _infer_mime_type_from_gcs_url(element) in _IMAGE_MIME_TYPES + except ValueError: + return False + if _is_file_reference(element): + file_info = resolved_files.get(element) + return file_info is not None and file_info.get("mime_type") in _IMAGE_MIME_TYPES + return False + + +def _count_input_images( + input: GeminiEmbeddingInput, + resolved_files: Mapping[str, Mapping[str, str]], +) -> int: + return sum( + 1 + for element in _flatten_input(input) + if _is_image_element(element, resolved_files) + ) + + +def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int: + return sum( + detail["tokenCount"] for detail in details if detail["modality"] == modality + ) + + +def _fallback_usage(input: GeminiEmbeddingInput, model: str) -> Usage: + if _is_multimodal_input(input): + return Usage(prompt_tokens=0, total_tokens=0) + input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") + prompt_tokens = token_counter(model=model, text=input_text) + return Usage(prompt_tokens=prompt_tokens, total_tokens=prompt_tokens) + + +def _usage_from_embed_content_response( + input: GeminiEmbeddingInput, + model: str, + raw_usage_metadata: object, + resolved_files: Mapping[str, Mapping[str, str]], +) -> Usage: + usage_metadata = _parse_usage_metadata(raw_usage_metadata) + if usage_metadata is None: + return _fallback_usage(input, model) + + prompt_tokens = usage_metadata.get("promptTokenCount", 0) + total_tokens = usage_metadata.get("totalTokenCount") or prompt_tokens + + details: Sequence[PromptTokensDetails] = ( + usage_metadata.get("promptTokensDetails") or () + ) + text_tokens = _tokens_for_modality(details, "TEXT") + audio_tokens = _tokens_for_modality(details, "AUDIO") + video_tokens = _tokens_for_modality(details, "VIDEO") + image_count = _count_input_images(input, resolved_files) + + video_length_seconds = ( + video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0 + ) + audio_length_seconds = ( + audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0 + ) + + # generic_cost_per_token rewrites text_tokens to the full prompt minus + # other modalities when both text_tokens and image_count are zero. For + # video, that misallocates video tokens to text; a 1-token floor sidesteps + # the rewrite and keeps billing on input_cost_per_video_per_second. + needs_video_text_floor = ( + video_length_seconds > 0 and text_tokens == 0 and image_count == 0 + ) + resolved_text_tokens = 1 if needs_video_text_floor else text_tokens + + return Usage( + prompt_tokens=prompt_tokens, + total_tokens=total_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=resolved_text_tokens, + audio_tokens=audio_tokens, + image_count=image_count, + video_length_seconds=video_length_seconds, + audio_length_seconds=audio_length_seconds, + ), + ) + + def process_embed_content_response( input: GeminiEmbeddingInput, model_response: EmbeddingResponse, model: str, response_json: dict, + resolved_files: Mapping[str, Mapping[str, str]] | None = None, ) -> EmbeddingResponse: """ Process Gemini embedContent response (single embedding for multimodal input). @@ -308,6 +440,8 @@ def process_embed_content_response( model_response: EmbeddingResponse to populate model: Model name response_json: Raw JSON response from embedContent endpoint + resolved_files: Mapping of file references (files/abc) to {mime_type, uri}, + used to bill resolved image references at the per-image rate Returns: EmbeddingResponse with single embedding @@ -327,14 +461,11 @@ def process_embed_content_response( model_response.data = [openai_embedding] model_response.model = model - - if _is_multimodal_input(input): - prompt_tokens = 0 - else: - input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") - prompt_tokens = token_counter(model=model, text=input_text) - model_response.usage = Usage( - prompt_tokens=prompt_tokens, total_tokens=prompt_tokens + model_response.usage = _usage_from_embed_content_response( + input=input, + model=model, + raw_usage_metadata=response_json.get("usageMetadata"), + resolved_files=resolved_files or {}, ) return model_response diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 24d6e84fba7..c3f99b1d18b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1539,6 +1539,9 @@ class PromptTokensDetailsWrapper( video_length_seconds: Optional[float] = None """Length of videos sent to the model. Used for Vertex AI multimodal embeddings.""" + audio_length_seconds: Optional[float] = None + """Length of audio sent to the model. Used for multimodal embeddings priced per audio-second.""" + cache_creation_tokens: Optional[int] = None """Number of cache creation tokens sent to the model. Used for Anthropic prompt caching.""" @@ -1553,6 +1556,8 @@ class PromptTokensDetailsWrapper( del self.image_count if self.video_length_seconds is None: del self.video_length_seconds + if self.audio_length_seconds is None: + del self.audio_length_seconds if self.web_search_requests is None: del self.web_search_requests if self.cache_creation_tokens is None: diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 7f3d5a959a1..a5d1934237f 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -951,6 +951,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): "character_count": 0, "image_count": 0, "video_length_seconds": 0.0, + "audio_length_seconds": 0.0, } model_info: ModelInfo = {} diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index bb4e6c67e9e..86b3f0976ab 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -10,9 +10,11 @@ Covers: import pytest +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( _build_part_for_input, _is_multimodal_input, + process_embed_content_response, process_response, transform_openai_input_gemini_content, transform_openai_input_gemini_embed_content, @@ -72,7 +74,9 @@ class TestBuildPartForInput: assert part["file_data"]["file_uri"] == GCS_URL def test_file_reference_resolved(self): - resolved = {"files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"}} + resolved = { + "files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"} + } part = _build_part_for_input("files/abc", resolved_files=resolved) assert part["file_data"] is not None assert part["file_data"]["mime_type"] == "image/jpeg" @@ -94,7 +98,9 @@ class TestTransformOpenaiInputGeminiContent: def test_multiple_texts(self): result = transform_openai_input_gemini_content( - input=["hello", "world"], model="gemini-embedding-2-preview", optional_params={} + input=["hello", "world"], + model="gemini-embedding-2-preview", + optional_params={}, ) assert len(result["requests"]) == 2 assert result["requests"][0]["content"]["parts"][0]["text"] == "hello" @@ -109,7 +115,10 @@ class TestTransformOpenaiInputGeminiContent: ) assert len(result["requests"]) == 2 # First request is text - assert result["requests"][0]["content"]["parts"][0]["text"] == "The food was delicious" + assert ( + result["requests"][0]["content"]["parts"][0]["text"] + == "The food was delicious" + ) # Second request is image assert result["requests"][1]["content"]["parts"][0]["inline_data"] is not None @@ -288,3 +297,207 @@ class TestProcessResponse: model="gemini-embedding-2-preview", optional_params={}, ) + + +class TestProcessEmbedContentResponseUsage: + """Gemini Embedding 2 embedContent usageMetadata must drive spend. + + Regression for multimodal calls recording prompt_tokens=0 / spend=$0. + """ + + MODEL = "gemini-embedding-2" + + def test_multimodal_image_preserves_usage_metadata(self): + response_json = { + "embedding": {"values": [0.1, 0.2, 0.3]}, + "usageMetadata": { + "promptTokenCount": 258, + "totalTokenCount": 258, + "promptTokensDetails": [{"modality": "IMAGE", "tokenCount": 258}], + }, + } + result = process_embed_content_response( + input=[IMAGE_DATA_URI], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 258 + assert result.usage.total_tokens == 258 + assert result.usage.prompt_tokens_details.image_count == 1 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost > 0 + + def test_text_modality_detail_populated(self): + response_json = { + "embedding": {"values": [0.1, 0.2]}, + "usageMetadata": { + "promptTokenCount": 12, + "totalTokenCount": 12, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}], + }, + } + result = process_embed_content_response( + input="a short caption", + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 12 + assert result.usage.prompt_tokens_details.text_tokens == 12 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost > 0 + + def test_video_modality_derives_seconds_and_text_floor(self): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 516, + "totalTokenCount": 516, + "promptTokensDetails": [{"modality": "VIDEO", "tokenCount": 516}], + }, + } + result = process_embed_content_response( + input=["gs://bucket/clip.mp4"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 516 + assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx( + 2.0 + ) + assert result.usage.prompt_tokens_details.text_tokens == 1 + + def test_missing_usage_metadata_does_not_estimate_from_base64(self): + response_json = {"embedding": {"values": [0.1, 0.2]}} + result = process_embed_content_response( + input=[IMAGE_DATA_URI], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 0 + assert result.usage.total_tokens == 0 + + def test_missing_usage_metadata_text_falls_back_to_token_counter(self): + response_json = {"embedding": {"values": [0.1, 0.2]}} + result = process_embed_content_response( + input="hello world this is plain text", + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens > 0 + + def test_file_reference_image_billed_per_image_not_text(self): + """files/... image refs must bill per-image, not at the text token rate.""" + response_json = { + "embedding": {"values": [0.1, 0.2, 0.3]}, + "usageMetadata": { + "promptTokenCount": 258, + "totalTokenCount": 258, + "promptTokensDetails": [{"modality": "IMAGE", "tokenCount": 258}], + }, + } + result = process_embed_content_response( + input=["files/img123"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + resolved_files={ + "files/img123": { + "mime_type": "image/png", + "uri": "https://example.com/img123", + } + }, + ) + assert result.usage.prompt_tokens_details.image_count == 1 + assert result.usage.prompt_tokens_details.text_tokens == 0 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(0.00012) + + def test_file_reference_non_image_not_counted_as_image(self): + """A files/... ref resolving to a non-image mime must not be image-counted.""" + response_json = { + "embedding": {"values": [0.1, 0.2]}, + "usageMetadata": { + "promptTokenCount": 64, + "totalTokenCount": 64, + "promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 64}], + }, + } + result = process_embed_content_response( + input=["files/clip1"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + resolved_files={ + "files/clip1": { + "mime_type": "audio/mpeg", + "uri": "https://example.com/clip1", + } + }, + ) + assert result.usage.prompt_tokens_details.image_count == 0 + assert result.usage.prompt_tokens_details.audio_tokens == 64 + assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx( + 2.0 + ) + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(2.0 * 0.00016) + + def test_video_plus_audio_does_not_double_bill_text(self): + """Video+audio responses must not get video tokens reassigned to text.""" + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 580, + "totalTokenCount": 580, + "promptTokensDetails": [ + {"modality": "VIDEO", "tokenCount": 516}, + {"modality": "AUDIO", "tokenCount": 64}, + ], + }, + } + result = process_embed_content_response( + input=["gs://bucket/clip.mp4"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens_details.text_tokens == 1 + assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx( + 2.0 + ) + assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx( + 2.0 + ) + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + # 1 floor text token at 2e-7 + 2s of video at 7.9e-4 + 2s of audio at 1.6e-4 + assert prompt_cost == pytest.approx(1 * 2e-7 + 2 * 0.00079 + 2 * 0.00016) From 1a4009caf4b49a9884c7529a017e1cf7b52cf58b Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 25 Jun 2026 20:05:42 -0700 Subject: [PATCH 09/16] chore: remove CI section (#31376) We now require all checks --- .github/pull_request_template.md | 32 ++++++++------------------------ 1 file changed, 8 insertions(+), 24 deletions(-) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 9658baeb89a..12ad124fa20 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,17 +1,17 @@ ## Relevant issues - + ## Linear ticket - + ## Pre-Submission checklist **Please complete all items before asking a LiteLLM maintainer to review your PR** - [ ] I have added meaningful tests -- [ ] My PR passes all unit tests on [`make test-unit`](https://docs.litellm.ai/docs/extras/contributing_code) +- [ ] My PR passes all CI/CD checks (e.g., lint, format, unit tests) - [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem - [ ] I have requested a Greptile review by commenting `@greptileai` and received a **Confidence Score of at least 4/5** before requesting a maintainer review @@ -19,29 +19,13 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slack (#pr-review)](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA). -## CI (LiteLLM team) - -> **CI status guideline:** -> -> - 50-55 passing tests: main is stable with minor issues. -> - 45-49 passing tests: acceptable but needs attention -> - <= 40 passing tests: unstable; be careful with your merges and assess the risk. - -- [ ] **Branch creation CI run** - Link: - -- [ ] **CI run for the last commit** - Link: - -- [ ] **Merge / cherry-pick CI run** - Links: - ## Screenshots / Proof of Fix - + ## Type From 7680cedf42fcd75f96cbfc7cd2df727ca0f39384 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 25 Jun 2026 20:30:04 -0700 Subject: [PATCH 10/16] test(logging): regression coverage for streaming /v1/messages OpenAI Responses spend logs (#31388) * test(logging): cover streaming /v1/messages OpenAI Responses spend logs The #28595 fix added unit tests that call _handle_anthropic_messages_response_logging directly, but nothing exercises the streaming wiring that actually regressed: a streaming /v1/messages call cross-routed to the OpenAI Responses backend whose success handler took the no-op async_log_stream_event path and dropped the SpendLogs row. Add an end-to-end test that drives litellm.anthropic_messages(stream=True) with a mocked upstream Responses SSE and asserts async_log_success_event fires with non-zero cost and call_type anthropic_messages, plus a key-gated live counterpart. * test(logging): exercise stream deltas and assert single success log Address review on the streaming bridge regression test: emit output_item.added plus text deltas before response.completed so it covers mid-stream delta handling rather than only end-of-stream success logging, assert at least one content_block_delta surfaces, restore litellm.callbacks via monkeypatch instead of leaking global state, and assert async_log_success_event fires exactly once. * test(logging): drop live network test from mock-only suite Greptile flagged that tests/test_litellm only permits mock tests; network calls belong in tests/e2e. Remove the key-gated live counterpart and keep the deterministic mocked test as the regression guard. The live verification stays in the PR description as the proof of fix. --- .../test_litellm_logging.py | 136 ++++++++++++++++++ 1 file changed, 136 insertions(+) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index f0db0409bd7..5a8ac0313a7 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -11,7 +11,9 @@ sys.path.insert( import time +import litellm from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import set_callbacks from litellm.types.utils import ModelResponse, TextCompletionResponse @@ -3408,6 +3410,140 @@ def test_handle_anthropic_messages_response_logging_degrades_on_unparseable_resp assert result.usage.prompt_tokens == 4 # type: ignore[attr-defined] +class _SuccessCapturingLogger(CustomLogger): + """Records the success payload. success_payload is populated only in + async_log_success_event, so it stays None when the buggy no-op + async_log_stream_event path runs for streaming.""" + + def __init__(self): + super().__init__() + self.success_payload = None + self.success_calls = 0 + self.stream_event_calls = 0 + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.success_calls += 1 + self.success_payload = kwargs.get("standard_logging_object") + + async def async_log_stream_event(self, kwargs, response_obj, start_time, end_time): + self.stream_event_calls += 1 + + +def _responses_stream_sse_bytes(): + """A full Responses stream: an opened message item, two text deltas, then the + terminal response.completed carrying usage. Exercises mid-stream delta handling + in addition to end-of-stream success logging.""" + import json + + events = [ + { + "type": "response.output_item.added", + "output_index": 0, + "item": { + "id": "msg-1", + "type": "message", + "status": "in_progress", + "role": "assistant", + "content": [], + }, + }, + { + "type": "response.output_text.delta", + "item_id": "msg-1", + "output_index": 0, + "content_index": 0, + "delta": "hello ", + "sequence_number": 2, + }, + { + "type": "response.output_text.delta", + "item_id": "msg-1", + "output_index": 0, + "content_index": 0, + "delta": "world", + "sequence_number": 3, + }, + { + "type": "response.completed", + "sequence_number": 4, + "response": _responses_api_response_with_text("hello world").model_dump(), + }, + ] + return [f"data: {json.dumps(e)}\n\n".encode("utf-8") for e in events] + + +def _fake_streaming_responses_http_response(): + sse_chunks = _responses_stream_sse_bytes() + + async def aiter_bytes(*args, **kwargs): + for chunk in sse_chunks: + yield chunk + + resp = MagicMock() + resp.status_code = 200 + resp.headers = {} + resp.aiter_bytes = aiter_bytes + return resp + + +def _chunk_text(chunk): + if isinstance(chunk, (bytes, bytearray)): + return chunk.decode("utf-8", "ignore") + return str(chunk) + + +async def _drain_until_logged(logger, max_iter=30): + for _ in range(max_iter): + if logger.success_payload is not None: + break + await asyncio.sleep(0.1) + + +@pytest.mark.asyncio +async def test_streaming_anthropic_messages_openai_bridge_fires_success_logging( + monkeypatch, +): + """Regression for #28595 / #28943. The existing tests above call + _handle_anthropic_messages_response_logging directly; they do not cover the + streaming wiring that originally broke. Drive a real streaming + anthropic_messages call routed to the OpenAI Responses backend (upstream SSE + mocked) and assert the bridge surfaces delta chunks and fires success logging + exactly once with real cost. On the broken version the stream ran but only the + no-op async_log_stream_event was called, so success_payload stayed None and the + SpendLogs row never landed.""" + logger = _SuccessCapturingLogger() + monkeypatch.setattr(litellm, "callbacks", [logger]) + + chunks = [] + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=AsyncMock(return_value=_fake_streaming_responses_http_response()), + ): + stream = await litellm.anthropic_messages( + model="openai/gpt-4o", + api_key="sk-test-28595", + messages=[{"role": "user", "content": "ping"}], + max_tokens=16, + stream=True, + ) + async for chunk in stream: # logging fires on stream end; must drain fully + chunks.append(chunk) + + await _drain_until_logged(logger) + + assert chunks, "stream yielded no chunks" + assert any("content_block_delta" in _chunk_text(c) for c in chunks), ( + "no delta chunks surfaced; the streaming text deltas were not forwarded" + ) + assert logger.success_payload is not None, ( + "async_log_success_event never fired for streaming /v1/messages -> openai " + "Responses bridge; the no-op stream path dropped the spend row" + ) + assert logger.success_calls == 1, "bridge call must log success exactly once" + assert logger.success_payload["response_cost"] > 0 + assert logger.success_payload["call_type"] == "anthropic_messages" + + def test_failure_handler_records_recovered_partial_spend(logging_obj): """A stream interrupted mid-flight still billed the provider for the chunks already delivered. When the router stashes that recovered usage as From 7eacdd5258e9dda6b75493c22e3063854c4dc9e4 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 26 Jun 2026 09:30:28 +0530 Subject: [PATCH 11/16] chore: litellm oss staging 250626 (#31305) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(anthropic): support Bearer auth for custom api_base endpoints (Fixes #30926) * style: format common_utils.py with black * fix(anthropic): extract api_base from litellm_params in batches/files validate_environment * fix(anthropic): scope Bearer key check to custom api_base endpoints * fix(streaming): reset Anthropic message_start cursor (output_tokens=1) when no message_delta arrives The Anthropic streaming protocol emits `message_start.usage.output_tokens=1` as a placeholder cursor; the real cumulative output count only arrives in the final `message_delta` event. When a stream is cancelled before `message_delta` lands (common for thinking models on long-tail prompts), ChunkProcessor._calculate_usage_per_chunk's last-wins accumulator left completion_tokens stuck at 1. Because 1 is truthy, the `completion_tokens or token_counter(text=...)` fallback in calculate_usage() never fired, and requests were billed for 1 output token even when several thousand tokens of text had actually streamed. Fix: track whether any chunk's completion_tokens exceeded 1 (saw_non_cursor_completion). If the only update we saw was the cursor, reset completion_tokens to 0 so the text-based fallback estimates from the real completion content. Legitimate 1-token completions (model returns "Yes." etc.) are unaffected in practice — token_counter on a 1-token completion_output also yields ~1, so billing stays approximately correct. Tests: - TestAnthropicCursorBug (6 cases) — pins the post-fix behavior - TestNonAnthropicStreamingIntact (2 cases) — guards against regression on providers without the cursor pattern All 8 new tests pass; 9 existing streaming_chunk_builder_utils tests still pass. * fix(streaming): scope cursor reset to anthropic provider + recognize message_delta arrival Addresses both Greptile P2 threads on PR #30420: CLASS A — Anthropic-specific heuristic was applied globally ============================================================ The `completion_tokens == 1 and not saw_non_cursor_completion` reset lived in provider-neutral `streaming_chunk_builder_utils.py`. Any non-Anthropic provider that legitimately reports completion_tokens=1 in a single usage chunk (perfectly normal for short OpenAI / Bedrock / Vertex single-token replies with stream_options.include_usage=true) would have its value silently rewritten to 0 and re-billed via token_counter — producing a different number than what the provider actually charged. Fix: gate the reset on `custom_llm_provider == "anthropic"`, resolved from the first chunk's `_hidden_params` (the same field set by streaming_handler.py:722 on the live path). Unknown / missing provider is treated as non-Anthropic and skips the reset, so newer providers and custom plugins are also safe by default. CLASS B — `saw_non_cursor_completion` missed legitimate single-token replies ============================================================ Previous condition was `usage_chunk_dict["completion_tokens"] > 1`, which never fires for an Anthropic stream where the model legitimately emits exactly one output token (e.g., "Yes."). Anthropic still sends message_start (output_tokens=1, the cursor) AND message_delta (output_tokens=1, the real value) — same value, but two distinct usage events. The old check couldn't tell that apart from a cancelled stream where only message_start landed. Fix: track `completion_usage_updates` and flip `saw_non_cursor_completion` when EITHER (1) the value exceeds 1 (definitely not a placeholder), OR (2) we've seen >=2 completion-bearing usage events (positive evidence that message_delta arrived). Cancelled cursor-only streams still have exactly one event and still hit the reset; cache chunks with completion_tokens=0 don't count toward the threshold. Tests ============================================================ - _make_chunk now sets `_hidden_params["custom_llm_provider"]` (default "anthropic") so the gate is exercised by every existing test — none of them needed assertion changes besides the legitimate-single- token case, which now expects exactly 1 (was a fuzzy 0..3 range). - New: test_anthropic_cache_only_chunks_after_message_start_still_resets - New: test_non_anthropic_provider_completion_tokens_one_not_reset - New: test_unknown_provider_completion_tokens_one_not_reset 11/11 tests pass. * chore: add Co-authored-by trailer for attribution Co-authored-by: songkuan-zheng * fix(anthropic): preserve messages cache usage * style(anthropic): format messages cache usage helper * fix(anthropic): accept integral float cache token counts Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(anthropic): accept integral float cache token counts * test(anthropic): cover cache usage edge cases * fix(gemini): preserve thoughtSignature for server-side tool responses When Gemini API returns toolCall and toolResponse parts, they might have different thoughtSignatures. Previously, LiteLLM merged them into a single dict, overwriting the response's thoughtSignature with the call's. This fix extracts them separately and re-injects them correctly. TAG=agy CONV=755b21d0-3200-40bc-bd1a-bb58a378a9a6 * fix(gemini): address PR comments on thoughtSignature handling - Fix orphan-response thoughtSignature regression by copying thought_signature to response_thought_signature - Add missing assertions in existing tests - Add new unit tests for orphan-response signature handling TAG=agy CONV=755b21d0-3200-40bc-bd1a-bb58a378a9a6 * feat(mcp): include server alias and server_id in mcp_info response - Add alias and server_id fields to mcp_info object in /mcp-rest/tools/list endpoint - Update rest_endpoints.py to surface alias from server config - Add test coverage in test_mcp_server.py and test_rest_endpoints.py Fixes #31015 * fix(proxy): reject non-finite spend via validate_finite_spend A NaN/-inf spend would bypass spend >= max_budget enforcement. Add a shared finite-value guard, defined above the litellm.proxy.* imports to avoid the module-level cyclic-import warning. * fix(proxy): require admin for any /key/update spend, reject non-finite Gate the admin check on the presence of `spend` (not a value diff): the DB spend lags the live cross-pod counter, so an "unchanged" spend on the non-admin path let a key owner / team member overwrite the live counter below real usage. Also reject NaN/+-inf spend before the DB write. * fix(proxy): invalidate spend counter on /user/update spend change A direct spend change on /user/update wrote the DB row but left the warm cross-pod counter at the stale value, so enforcement kept reading the old spend. Invalidate spend:user:{user_id} after the write (reseed-from-DB), and reject non-finite spend before the write. * fix(cache): route Bedrock semantic-cache sync embedding through the Router (#28244) The semantic cache's embedding model is a proxy Router alias whose AWS credentials (aws_role_name, aws_session_name) live only in the Router deployment's litellm_params. The sync embedding paths called litellm.embedding() directly, bypassing the Router, so they could neither resolve the alias nor assume the configured role; cross-account Bedrock semantic caching failed with "bedrock:InvokeModel is not authorized". On Redis this surfaced at proxy startup because redisvl's CustomTextVectorizer eagerly fires a dimension-probe embedding during cache construction, while llm_router is still None. Fix A: make the sync paths mirror the already-correct async paths. A shared, dependency-injected helper (litellm/caching/_embedding_router.py) decides whether to route through llm_router.embedding(...) when the model is a Router deployment, else fall back to direct litellm.embedding(...). Redis and qdrant sync set_cache/get_cache now precompute the embedding and pass vector= to the backend, exactly as the async astore/acheck already do. Both async _get_async_embedding methods are unified onto the same helper and now forward the caller's full metadata instead of a hand-picked subset. Fix B (Redis only): defer redisvl index construction from __init__ into a lazy, memoized llmcache property, so the dimension-probe embedding fires on first cache use, after llm_router is wired. A failed build is not memoized, so a transient outage recovers on the next request. Known limitation: resolve_embedding_router gates on an exact model-name match (same as the shipped async path); wildcard/alias/team-public routes still fall back to direct embedding. Tracked as a follow-up. * fix(cache): harden embedding-router and shrink Any surface (review) Address review feedback on the semantic-cache aws-role fix (#28244): - resolve_embedding_router now skips deployment entries missing model_name instead of raising KeyError on a malformed model_list (Greptile P2); add a regression test that fails on the old direct-key access. - Replace the `**kwargs: Any` passthrough on the four cache _get_embedding / _get_async_embedding helpers with an explicit, typed `metadata: Optional[Dict[str, Any]] = None` parameter. The helpers only ever consumed kwargs["metadata"], so this is behavior-preserving, makes the forwarded field obvious at the call site, and removes three bare-Any annotations (keeps the strict-rule ANN401 budget within ceiling). - Note in _build_llmcache that redisvl's dimension-probe embedding adds one extra billable embedding on the first cache request (Greptile P2). * fix(bedrock_mantle): correct responses routing for openai.gpt-5.x models Dashboard Test Connection for bedrock_mantle/openai.gpt-5.4 and openai.gpt-5.5 was failing with maximum recursion depth errors and "model does not exist" Route detection in the bedrock provider matched route tokens by plain substring, so the bedrock_mantle/ prefix was mistaken for the mantle/ invoke route and the body model was rewritten to bedrock_openai.gpt-5.5; route tokens now only match at a path-segment boundary so the bare model name is preserved A responses-mode model whose provider has no responses config bounced forever between the responses API and chat completions; the responses to completion fallback now tags its call so completion() does not bridge back, breaking the loop The Test Connection endpoint hardcoded the test mode to chat, which disabled mode auto-detection for responses-only models; the default is now None so the mode is detected from model capabilities acompletion() now drops a duplicate acompletion kwarg before building the partial and treats model_info=None as an empty dict to avoid a NoneType crash * test(bedrock_mantle): cover route guard and bridge flag; fix reportArgumentType regression Adds the regression coverage codecov flagged on the two responses to completion bridge guard lines and the bedrock route-prefix helper. The handler tests drive both the sync and async fallback paths with litellm.completion and litellm.acompletion mocked, and assert the forwarded kwargs carry _skip_responses_api_bridge=True, so dropping either flag line fails the suite. The common_utils tests assert that bedrock_mantle/openai.gpt-5.x no longer resolves to the mantle route while the genuine mantle/ and bedrock/mantle/ ids still do, exercising both branches of _model_has_route_prefix. Also aligns update_messages_with_model_file_ids model_id to Optional[str], matching its Responses API sibling, so the defensive model_info fallback no longer introduces a new reportArgumentType in completion(); the file-id lookup narrows model_id before the dict get * chore(ui): sync generated OpenAPI types for optional test_connection mode The test_model_connection mode body param default changed from chat to None so the mode is auto-detected from model capabilities, which makes the field optional in the proxy OpenAPI spec. Regenerate the committed schema so the dashboard types match: mode becomes optional and the description and default JSDoc follow the spec, keeping the Check UI API Types Sync gate green * refactor(bedrock): match all explicit route prefixes at path-segment boundary Migrates the remaining substring route checks to the existing _model_has_route_prefix helper so every explicit route token matches only as a leading path segment, consistent with get_bedrock_route and the mantle route. Covers _explicit_converse_route, _explicit_claude_platform_route, _explicit_invoke_route, _explicit_agent_route, _explicit_agentcore_route, _explicit_converse_like_route, _explicit_async_invoke_route and _explicit_openai_route. This also stops invoke/ from substring-matching async_invoke/. Route precedence and order are unchanged, and a note on the segment invariant is added to the helper docstring * test(bedrock): cover explicit route prefix segment matching Exercises all eight migrated _explicit_*_route helpers (converse, converse_like, invoke, async_invoke, agent, agentcore, claude_platform, openai) directly: each matches its token as a leading path segment and rejects the token glued to a preceding segment, so reverting any method to the old substring check fails the suite. Also asserts invoke/ no longer matches async_invoke/ models, the concrete improvement of the segment-boundary migration * test(proxy): assert negative spend is allowed (one-time grant use-case) Negative spend is intentionally permitted so admins can grant extra allowance for the current budget period only, without raising the recurring budget ceiling. Cover it explicitly in validate_finite_spend and via the /user/update invalidation test. * fix(google_genai): forward native generateContent top-level fields Google's native generateContent REST body carries safetySettings, toolConfig, cachedContent and labels at the top level as siblings of generationConfig. The proxy's :generateContent endpoint spread them into agenerate_content as loose kwargs and then dropped them, so callers had to wrap them in extra_body for them to take effect; safetySettings, for instance, was silently ignored The provider config now exposes the native top-level field names and setup_generate_content_call collects whichever are present, merging them into the outgoing request body through the existing extra_body merge so they reach Google verbatim. An explicit extra_body still wins on conflict. The sync generate_content_stream path now also forwards systemInstruction, matching the other three entry points Fixes #12671 Claude-Session: https://claude.ai/code/session_016MFtMXokCjT8u6mvyASudK * fix(proxy): resolve env refs for DB-stored models * fix(proxy): restrict DB env ref resolution * fix(proxy): block team DB env ref resolution * fix(lint): resolve ANN401/UP045/C901 strict-gate violations - Replace Optional[X] with X | None (UP045) in 8 files - Replace Any return/param types with concrete types or object (ANN401) - Extract _make_api_key_auth_header helper to reduce get_anthropic_headers complexity below C901 threshold (17 → 14) Co-Authored-By: Claude Sonnet 4.6 * fix(anthropic): preserve x-api-key for custom endpoints; opt-in Bearer via prefix Users who pass a key already prefixed with "Bearer " get Authorization: Bearer. All other keys continue to use x-api-key, preserving backward compatibility with custom api_base endpoints that expect x-api-key rather than Authorization. Also consolidates get_auth_header to reuse _make_api_key_auth_header helper, eliminating the duplicated custom-endpoint routing logic. Co-Authored-By: Claude Sonnet 4.6 * revert(anthropic): restore Bearer routing for non-sk-ant- keys on custom api_base The backwards-compat change broke existing tests that verify the intentional Bearer-for-custom-base behavior (Fixes #30926). Restore original logic while keeping the _make_api_key_auth_header helper for code deduplication. Co-Authored-By: Claude Sonnet 4.6 * fix(anthropic): gate Bearer-for-custom-base behind use_bearer_for_custom_base flag Previously the auth-header switch from x-api-key to Authorization: Bearer applied unconditionally for non-sk-ant- keys on a custom api_base, silently breaking existing deployments that proxied to gateways expecting x-api-key. Introduce use_bearer_for_custom_base: bool = False on _make_api_key_auth_header, get_anthropic_headers, and get_auth_header. validate_environment reads it from litellm_params so callers can opt in per-model without any API surface change. Tests updated to pass use_bearer_for_custom_base=True where Bearer behavior is asserted. Co-Authored-By: Claude Sonnet 4.6 * fix(redis): apply namespace prefix in delete_cache and async_delete_cache (#29981) DEL was the only Redis cache operation that skipped check_and_fix_namespace, so it targeted the raw SHA256 hash (e.g. 3997c4...) rather than the namespaced key (litellm:3997c4...). This caused two problems: a Redis NOPERM error on deployments with an ACL restricting DEL to the litellm:* pattern, and a silent no-op on all other deployments since the un-prefixed key was never stored. * style(anthropic): reformat common_utils.py with Black (--target-version py312) Co-Authored-By: Claude Sonnet 4.6 * fix: preserve cache metadata and spend counters * style: apply ruff format to streaming_iterator.py * refactor: reduce complexity of usage/spend helpers to satisfy strict ruff gate Extract Anthropic message_start cursor reset into _reset_anthropic_cursor_completion_tokens and the cross-pod spend-counter invalidation into _invalidate_user_spend_counter_if_changed, keeping both _calculate_usage_per_chunk and _update_single_user_helper under the max-complexity ceiling. Use builtin generics in the new signatures so no new UP006 violations are introduced. Behavior unchanged. --------- Co-authored-by: rupak-eng Co-authored-by: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com> Co-authored-by: songkuan-zheng Co-authored-by: Kannan Priyadharshan Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Marco Georgaklis Co-authored-by: Anjaiah Methuku Co-authored-by: Andrii Butko Co-authored-by: Kent Co-authored-by: kunal2002 Co-authored-by: Ali Khan Co-authored-by: jesco-absolut Co-authored-by: Claude Sonnet 4.6 Co-authored-by: Matt Hill Co-authored-by: Cursor Agent --- litellm/caching/_embedding_router.py | 45 +++ litellm/caching/caching.py | 5 +- litellm/caching/qdrant_semantic_cache.py | 82 +++-- litellm/caching/redis_semantic_cache.py | 165 +++++---- litellm/google_genai/main.py | 43 ++- .../prompt_templates/common_utils.py | 4 +- .../streaming_chunk_builder_utils.py | 59 ++++ .../llms/anthropic/batches/transformation.py | 4 +- litellm/llms/anthropic/chat/handler.py | 1 + litellm/llms/anthropic/common_utils.py | 45 ++- .../adapters/streaming_iterator.py | 29 +- .../adapters/transformation.py | 143 +++++--- litellm/llms/anthropic/files/handler.py | 2 +- .../llms/anthropic/files/transformation.py | 4 +- .../llms/anthropic/skills/transformation.py | 4 +- .../base_llm/google_genai/transformation.py | 13 + litellm/llms/bedrock/common_utils.py | 41 ++- .../llms/vertex_ai/gemini/transformation.py | 4 +- .../vertex_and_google_ai_studio_gemini.py | 9 +- litellm/main.py | 14 +- .../mcp_server/rest_endpoints.py | 23 +- .../health_endpoints/_health_endpoints.py | 5 +- .../management_endpoints/common_utils.py | 19 ++ .../internal_user_endpoints.py | 25 ++ .../key_management_endpoints.py | 39 ++- litellm/proxy/proxy_server.py | 61 +++- .../handler.py | 2 + tests/mcp_tests/test_mcp_server.py | 37 +- .../caching/test_embedding_router.py | 67 ++++ .../caching/test_qdrant_semantic_cache.py | 102 ++++++ .../test_litellm/caching/test_redis_cache.py | 37 +- .../caching/test_redis_semantic_cache.py | 271 ++++++++++++++- .../google_genai/test_google_genai_main.py | 227 ++++++++++++- .../test_streaming_chunk_builder_cursor.py | 321 ++++++++++++++++++ ...al_pass_through_adapters_transformation.py | 277 +++++++++++++++ .../test_streaming_iterator_combined_chunk.py | 54 +++ .../test_streaming_iterator_first_delta.py | 30 ++ .../anthropic/test_anthropic_common_utils.py | 107 ++++++ .../llms/bedrock/test_bedrock_common_utils.py | 120 +++++++ .../gemini/test_context_circulation.py | 125 ++++++- .../mcp_server/test_rest_endpoints.py | 68 ++++ .../management_endpoints/test_common_utils.py | 38 +++ .../test_internal_user_endpoints.py | 80 +++++ .../test_key_management_endpoints.py | 62 +++- .../proxy/proxy_server/test_proxy_config.py | 147 +++++++- .../test_handler.py | 73 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +- 47 files changed, 2865 insertions(+), 273 deletions(-) create mode 100644 litellm/caching/_embedding_router.py create mode 100644 tests/test_litellm/caching/test_embedding_router.py create mode 100644 tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py create mode 100644 tests/test_litellm/responses/litellm_completion_transformation/test_handler.py diff --git a/litellm/caching/_embedding_router.py b/litellm/caching/_embedding_router.py new file mode 100644 index 00000000000..1ec898012e9 --- /dev/null +++ b/litellm/caching/_embedding_router.py @@ -0,0 +1,45 @@ +"""Shared selection of the embedding path for semantic caches. + +Both the Redis and qdrant semantic caches need the same decision: when the +configured embedding model is a proxy Router deployment, embeddings must run +through the Router so per-deployment auth (e.g. Bedrock aws_role_name) is +applied. Otherwise fall back to a direct litellm embedding call. + +This module is dependency-injected: callers pass the proxy ``llm_router`` and +``llm_model_list`` in, so the decision logic is unit-testable without importing +``litellm.proxy.proxy_server``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from litellm.router import Router + + +def resolve_embedding_router( + embedding_model: str, + llm_router: Router | None, + llm_model_list: list[dict[str, Any]] | None, +) -> Router | None: + """Return ``llm_router`` iff it serves ``embedding_model`` as a deployment.""" + if llm_router is None: + return None + router_model_names: list[str] = ( + [m["model_name"] for m in llm_model_list if "model_name" in m] + if llm_model_list is not None + else [] + ) + if embedding_model in router_model_names: + return llm_router + return None + + +def build_router_embedding_metadata( + request_metadata: dict[str, Any] | None, +) -> dict[str, Any]: + """Forward the caller's full metadata, flagged as a semantic-cache embedding.""" + metadata: dict[str, Any] = dict(request_metadata or {}) + metadata["semantic-cache-embedding"] = True + return metadata diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index cb122e90102..5f2269d1945 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -574,8 +574,9 @@ class Cache: if prompt_kwarg in kwargs: cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg] - if isinstance(kwargs.get("metadata"), dict): - cache_lookup_kwargs["metadata"] = {} + metadata = kwargs.get("metadata") + if isinstance(metadata, dict): + cache_lookup_kwargs["metadata"] = dict(metadata) return cache_lookup_kwargs diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 68d3b8c20b3..504ef8a54eb 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.types.utils import EmbeddingResponse +from ._embedding_router import build_router_embedding_metadata, resolve_embedding_router from .base_cache import BaseCache @@ -219,37 +220,50 @@ class QdrantSemanticCache(BaseCache): cached_key = payload.get(self.CACHE_KEY_FIELD_NAME) return cached_key is not None and str(cached_key) == str(key) - async def _get_async_embedding(self, prompt: str, **kwargs) -> Any: - llm_model_list = None - llm_router = None - + def _get_embedding( + self, prompt: str, metadata: Dict[str, Any] | None = None + ) -> EmbeddingResponse: + """Embed via the proxy Router when it serves the model, else direct.""" try: - from litellm.proxy.proxy_server import ( - llm_model_list as proxy_llm_model_list, - llm_router as proxy_llm_router, - ) - - llm_model_list = proxy_llm_model_list - llm_router = proxy_llm_router + from litellm.proxy.proxy_server import llm_model_list, llm_router except ImportError: - pass + llm_model_list = None + llm_router = None - router_model_names = ( - [m["model_name"] for m in llm_model_list] - if llm_model_list is not None - else [] + router = resolve_embedding_router( + self.embedding_model, llm_router, llm_model_list ) - if llm_router is not None and self.embedding_model in router_model_names: - user_api_key = kwargs.get("metadata", {}).get("user_api_key", "") - return await llm_router.aembedding( + if router is not None: + return router.embedding( model=self.embedding_model, input=prompt, cache={"no-store": True, "no-cache": True}, - metadata={ - "user_api_key": user_api_key, - "semantic-cache-embedding": True, - "trace_id": kwargs.get("metadata", {}).get("trace_id", None), - }, + metadata=build_router_embedding_metadata(metadata), + ) + return litellm.embedding( + model=self.embedding_model, + input=prompt, + cache={"no-store": True, "no-cache": True}, + ) + + async def _get_async_embedding( + self, prompt: str, metadata: Dict[str, Any] | None = None + ) -> EmbeddingResponse: + try: + from litellm.proxy.proxy_server import llm_model_list, llm_router + except ImportError: + llm_model_list = None + llm_router = None + + router = resolve_embedding_router( + self.embedding_model, llm_router, llm_model_list + ) + if router is not None: + return await router.aembedding( + model=self.embedding_model, + input=prompt, + cache={"no-store": True, "no-cache": True}, + metadata=build_router_embedding_metadata(metadata), ) return await litellm.aembedding( @@ -269,11 +283,7 @@ class QdrantSemanticCache(BaseCache): # create an embedding for prompt embedding_response = cast( EmbeddingResponse, - litellm.embedding( - model=self.embedding_model, - input=prompt, - cache={"no-store": True, "no-cache": True}, - ), + self._get_embedding(prompt, metadata=kwargs.get("metadata")), ) # get the embedding @@ -312,11 +322,7 @@ class QdrantSemanticCache(BaseCache): # convert to embedding embedding_response = cast( EmbeddingResponse, - litellm.embedding( - model=self.embedding_model, - input=prompt, - cache={"no-store": True, "no-cache": True}, - ), + self._get_embedding(prompt, metadata=kwargs.get("metadata")), ) # get the embedding @@ -388,7 +394,9 @@ class QdrantSemanticCache(BaseCache): # get the prompt messages = kwargs["messages"] prompt = get_str_from_messages(messages) - embedding_response = await self._get_async_embedding(prompt, **kwargs) + embedding_response = await self._get_async_embedding( + prompt, metadata=kwargs.get("metadata") + ) # get the embedding embedding = embedding_response["data"][0]["embedding"] @@ -424,7 +432,9 @@ class QdrantSemanticCache(BaseCache): messages = kwargs["messages"] prompt = get_str_from_messages(messages) - embedding_response = await self._get_async_embedding(prompt, **kwargs) + embedding_response = await self._get_async_embedding( + prompt, metadata=kwargs.get("metadata") + ) # get the embedding embedding = embedding_response["data"][0]["embedding"] diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index cce4b75795f..e79392bb7f0 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -16,12 +16,13 @@ import os from typing import Any, Dict, List, Optional, Tuple, cast import litellm -from litellm._logging import print_verbose +from litellm._logging import print_verbose, verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) from litellm.types.utils import EmbeddingResponse +from ._embedding_router import build_router_embedding_metadata, resolve_embedding_router from .base_cache import BaseCache @@ -67,9 +68,6 @@ class RedisSemanticCache(BaseCache): Exception: If similarity_threshold is not provided or required Redis connection information is missing """ - from redisvl.extensions.llmcache import SemanticCache # type: ignore[import-not-found, import-untyped] - from redisvl.utils.vectorize import CustomTextVectorizer # type: ignore[import-not-found, import-untyped] - if index_name is None: index_name = self.DEFAULT_REDIS_INDEX_NAME @@ -107,15 +105,42 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Redis semantic-cache redis_url: {redis_url}") - # Initialize the Redis vectorizer and cache - cache_vectorizer = CustomTextVectorizer(self._get_embedding) + # Defer redisvl index construction until first use. redisvl's + # CustomTextVectorizer eagerly embeds a probe string at construction; + # building lazily ensures that probe runs after llm_router is wired so + # per-deployment auth (e.g. Bedrock aws_role_name) is applied. + self._index_name = index_name + self._redis_url = redis_url + self._llmcache = None - self.llmcache = self._init_semantic_cache( - semantic_cache_cls=SemanticCache, - index_name=index_name, - redis_url=redis_url, - cache_vectorizer=cache_vectorizer, - ) + @property + def llmcache(self) -> object: + if getattr(self, "_llmcache", None) is None: + self._llmcache = self._build_llmcache() + return self._llmcache + + @llmcache.setter + def llmcache(self, value: object) -> None: + self._llmcache = value + + def _build_llmcache(self) -> object: + # CustomTextVectorizer probes its embedding dimension at construction by + # embedding "dimension test", so the first cache request issues one extra + # billable embedding on top of the request's own. + from redisvl.extensions.llmcache import SemanticCache # type: ignore[import-not-found, import-untyped] + from redisvl.utils.vectorize import CustomTextVectorizer # type: ignore[import-not-found, import-untyped] + + try: + cache_vectorizer = CustomTextVectorizer(self._get_embedding) + return self._init_semantic_cache( + semantic_cache_cls=SemanticCache, + index_name=self._index_name, + redis_url=self._redis_url, + cache_vectorizer=cache_vectorizer, + ) + except Exception as e: + verbose_logger.error(f"Redis semantic-cache index build failed: {e}") + raise @classmethod def _cache_key_filterable_field(cls) -> Dict[str, str]: @@ -285,27 +310,43 @@ class RedisSemanticCache(BaseCache): return dict_method() return value - def _get_embedding(self, prompt: str) -> List[float]: + def _get_embedding( + self, prompt: str, metadata: Dict[str, Any] | None = None + ) -> List[float]: """ - Generate an embedding vector for the given prompt using the configured embedding model. - - Args: - prompt: The text to generate an embedding for - - Returns: - List[float]: The embedding vector + Routes through the proxy Router when the embedding model is a Router + deployment so per-deployment auth (e.g. Bedrock aws_role_name) applies, + mirroring ``_get_async_embedding``; otherwise embeds directly. """ - # Create an embedding from prompt - embedding_response = cast( - EmbeddingResponse, - litellm.embedding( - model=self.embedding_model, - input=prompt, - cache={"no-store": True, "no-cache": True}, - ), + try: + from litellm.proxy.proxy_server import llm_model_list, llm_router + except ImportError: + llm_model_list = None + llm_router = None + + router = resolve_embedding_router( + self.embedding_model, llm_router, llm_model_list ) - embedding = embedding_response["data"][0]["embedding"] - return embedding + if router is not None: + embedding_response = cast( + EmbeddingResponse, + router.embedding( + model=self.embedding_model, + input=prompt, + cache={"no-store": True, "no-cache": True}, + metadata=build_router_embedding_metadata(metadata), + ), + ) + else: + embedding_response = cast( + EmbeddingResponse, + litellm.embedding( + model=self.embedding_model, + input=prompt, + cache={"no-store": True, "no-cache": True}, + ), + ) + return embedding_response["data"][0]["embedding"] def _get_cache_logic(self, cached_response: Any) -> Any: """ @@ -357,7 +398,12 @@ class RedisSemanticCache(BaseCache): value_str = str(value) - store_kwargs: Dict[str, Any] = { + prompt_embedding = self._get_embedding( + prompt, metadata=kwargs.get("metadata") + ) + + store_kwargs: dict[str, Any] = { + "vector": prompt_embedding, "filters": self._get_cache_filters(key), } @@ -393,8 +439,12 @@ class RedisSemanticCache(BaseCache): # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. - check_kwargs: Dict[str, Any] = { + prompt_embedding = self._get_embedding( + prompt, metadata=kwargs.get("metadata") + ) + check_kwargs: dict[str, Any] = { "prompt": prompt, + "vector": prompt_embedding, "filter_expression": self._get_cache_key_filter_expression(key), } results = self.llmcache.check(**check_kwargs) @@ -435,49 +485,42 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Error retrieving from Redis semantic cache: {str(e)}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 - async def _get_async_embedding(self, prompt: str, **kwargs) -> List[float]: + async def _get_async_embedding( + self, prompt: str, metadata: Dict[str, Any] | None = None + ) -> List[float]: """ Asynchronously generate an embedding for the given prompt. Args: prompt: The text to generate an embedding for - **kwargs: Additional arguments that may contain metadata + metadata: Request metadata forwarded to the Router embedding call Returns: List[float]: The embedding vector """ - from litellm.proxy.proxy_server import llm_model_list, llm_router - - # Route the embedding request through the proxy if appropriate - router_model_names = ( - [m["model_name"] for m in llm_model_list] - if llm_model_list is not None - else [] - ) - try: - if llm_router is not None and self.embedding_model in router_model_names: - # Use the router for embedding generation - user_api_key = kwargs.get("metadata", {}).get("user_api_key", "") - embedding_response = await llm_router.aembedding( + from litellm.proxy.proxy_server import llm_model_list, llm_router + except ImportError: + llm_model_list = None + llm_router = None + + router = resolve_embedding_router( + self.embedding_model, llm_router, llm_model_list + ) + try: + if router is not None: + embedding_response = await router.aembedding( model=self.embedding_model, input=prompt, cache={"no-store": True, "no-cache": True}, - metadata={ - "user_api_key": user_api_key, - "semantic-cache-embedding": True, - "trace_id": kwargs.get("metadata", {}).get("trace_id", None), - }, + metadata=build_router_embedding_metadata(metadata), ) else: - # Generate embedding directly embedding_response = await litellm.aembedding( model=self.embedding_model, input=prompt, cache={"no-store": True, "no-cache": True}, ) - - # Extract and return the embedding vector return embedding_response["data"][0]["embedding"] except Exception as e: print_verbose(f"Error generating async embedding: {str(e)}") @@ -504,9 +547,11 @@ class RedisSemanticCache(BaseCache): value_str = str(value) # Generate embedding for the value (response) to cache - prompt_embedding = await self._get_async_embedding(prompt, **kwargs) + prompt_embedding = await self._get_async_embedding( + prompt, metadata=kwargs.get("metadata") + ) - store_kwargs: Dict[str, Any] = { + store_kwargs: dict[str, Any] = { "vector": prompt_embedding, "filters": self._get_cache_filters(key), } @@ -544,11 +589,13 @@ class RedisSemanticCache(BaseCache): return None # Generate embedding for the prompt - prompt_embedding = await self._get_async_embedding(prompt, **kwargs) + prompt_embedding = await self._get_async_embedding( + prompt, metadata=kwargs.get("metadata") + ) # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. - check_kwargs: Dict[str, Any] = { + check_kwargs: dict[str, Any] = { "prompt": prompt, "vector": prompt_embedding, "filter_expression": self._get_cache_key_filter_expression(key), diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index bdbb483dcf6..d35601c7e13 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -49,6 +49,7 @@ class GenerateContentSetupResult(BaseModel): custom_llm_provider: str generate_content_provider_config: Optional[BaseGoogleGenAIGenerateContentConfig] generate_content_config_dict: Dict[str, Any] + native_request_fields: dict[str, object] litellm_params: GenericLiteLLMParams litellm_logging_obj: LiteLLMLoggingObj litellm_call_id: Optional[str] @@ -152,6 +153,7 @@ class GenerateContentHelper: request_body={}, # Will be handled by adapter generate_content_provider_config=None, # type: ignore generate_content_config_dict=dict(config or {}), + native_request_fields={}, litellm_params=litellm_params, litellm_logging_obj=litellm_logging_obj, litellm_call_id=litellm_call_id, @@ -171,6 +173,12 @@ class GenerateContentHelper: system_instruction = kwargs.get("systemInstruction") or kwargs.get( "system_instruction" ) + # Native top-level REST fields arrive as loose kwargs and are otherwise dropped. + native_request_fields: dict[str, object] = { + field: kwargs[field] + for field in generate_content_provider_config.get_generate_content_request_top_level_fields() + if field in kwargs + } request_body = ( generate_content_provider_config.transform_generate_content_request( model=model, @@ -201,12 +209,29 @@ class GenerateContentHelper: request_body=request_body, generate_content_provider_config=generate_content_provider_config, generate_content_config_dict=generate_content_config_dict, + native_request_fields=native_request_fields, litellm_params=litellm_params, litellm_logging_obj=litellm_logging_obj, litellm_call_id=litellm_call_id, ) +def _merge_native_request_fields( + native_request_fields: dict[str, object], + extra_body: dict[str, object] | None, +) -> dict[str, object] | None: + """ + Merge native top-level request fields into ``extra_body`` so the HTTP handler + forwards them verbatim onto the outgoing request body. An explicit ``extra_body`` + value wins on conflict. Returns ``None`` only when there is genuinely nothing to + forward (no native fields and no caller-supplied ``extra_body``), preserving the + prior behavior without discarding an explicit ``extra_body={}``. + """ + if not native_request_fields and extra_body is None: + return None + return {**native_request_fields, **(extra_body or {})} + + @client async def agenerate_content( model: str, @@ -350,7 +375,9 @@ def generate_content( litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=_merge_native_request_fields( + setup_result.native_request_fields, extra_body + ), timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), @@ -447,7 +474,9 @@ async def agenerate_content_stream( litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=_merge_native_request_fields( + setup_result.native_request_fields, extra_body + ), timeout=timeout or request_timeout, _is_async=True, client=kwargs.get("client"), @@ -503,6 +532,11 @@ def generate_content_stream( **kwargs, ) + # Extract systemInstruction from kwargs to pass to handler + system_instruction = kwargs.get("systemInstruction") or kwargs.get( + "system_instruction" + ) + # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: if "stream" in kwargs: @@ -531,12 +565,15 @@ def generate_content_stream( litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=_merge_native_request_fields( + setup_result.native_request_fields, extra_body + ), timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), stream=True, litellm_metadata=kwargs.get("litellm_metadata", {}), + system_instruction=system_instruction, ) except Exception as e: diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index bf9ce3b0acb..4bdda6de2c8 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -466,7 +466,7 @@ def get_format_from_file_id(file_id: Optional[str]) -> Optional[str]: def update_messages_with_model_file_ids( messages: List[AllMessageValues], - model_id: str, + model_id: str | None, model_file_id_mapping: Dict[str, Dict[str, str]], ) -> List[AllMessageValues]: """ @@ -519,7 +519,7 @@ def update_messages_with_model_file_ids( if file_id: provider_file_id = ( model_file_id_mapping.get(file_id, {}).get(model_id) - if model_file_id_mapping + if model_file_id_mapping and model_id is not None else None ) if ( diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 66abb824610..1b1b652eaa0 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -585,6 +585,17 @@ class ChunkProcessor: # # Update usage information if needed prompt_tokens = 0 completion_tokens = 0 + # Anthropic's `message_start` SSE event carries usage.output_tokens=1 as a + # cursor/placeholder; the real value only arrives in `message_delta`. + # If a stream is cancelled before `message_delta` lands, the last-wins + # accumulator below leaves completion_tokens stuck at 1 — which then + # bypasses the `completion_tokens or token_counter(...)` fallback in + # calculate_usage() because 1 is truthy. Count the completion-bearing + # usage events so `_reset_anthropic_cursor_completion_tokens` can tell a + # legitimate single-token reply (Anthropic emits 1 in BOTH message_start + # AND message_delta, so >=2 events is positive evidence message_delta + # arrived) from a stale lone cursor. + completion_usage_updates = 0 ## anthropic prompt caching information ## cache_creation_input_tokens: Optional[int] = None cache_read_input_tokens: Optional[int] = None @@ -617,6 +628,7 @@ class ChunkProcessor: and usage_chunk_dict["completion_tokens"] > 0 ): completion_tokens = usage_chunk_dict["completion_tokens"] + completion_usage_updates += 1 if usage_chunk_dict["cache_creation_input_tokens"] is not None and ( usage_chunk_dict["cache_creation_input_tokens"] > 0 or cache_creation_input_tokens is None @@ -667,6 +679,12 @@ class ChunkProcessor: prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] + completion_tokens = self._reset_anthropic_cursor_completion_tokens( + chunks=chunks, + completion_tokens=completion_tokens, + completion_usage_updates=completion_usage_updates, + ) + return UsagePerChunk( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, @@ -678,6 +696,47 @@ class ChunkProcessor: prompt_tokens_details=prompt_tokens_details, ) + @staticmethod + def _reset_anthropic_cursor_completion_tokens( + chunks: list[dict[str, Any] | ModelResponse], + completion_tokens: int, + completion_usage_updates: int, + ) -> int: + """Reset a stale Anthropic ``message_start`` cursor placeholder to 0. + + See the ``completion_usage_updates`` comment in + ``_calculate_usage_per_chunk``. The accumulated value is NOT a stale + cursor when either it is > 1 (definitely not a placeholder) or we saw + >= 2 completion-bearing usage events (positive evidence ``message_delta`` + arrived). Otherwise — the only completion update we ever saw was the + Anthropic ``message_start`` cursor (=1) — reset to 0 so + ``calculate_usage()``'s ``or token_counter(text=...)`` fallback estimates + from the actually-received completion text instead of trusting the + placeholder. Gated on ``custom_llm_provider == "anthropic"`` so the + heuristic (which encodes Anthropic's specific message_start SSE shape) + does not silently affect other providers that may legitimately report + ``completion_tokens=1`` from a single usage event. + """ + saw_non_cursor_completion = ( + completion_tokens > 1 or completion_usage_updates >= 2 + ) + if saw_non_cursor_completion: + return completion_tokens + + custom_llm_provider: Optional[str] = None + if chunks: + first_chunk = chunks[0] + if isinstance(first_chunk, dict): + hp = first_chunk.get("_hidden_params") + else: + hp = getattr(first_chunk, "_hidden_params", None) + if isinstance(hp, dict): + custom_llm_provider = hp.get("custom_llm_provider") + + if custom_llm_provider == "anthropic" and completion_tokens == 1: + return 0 + return completion_tokens + def calculate_usage( self, chunks: List[Union[Dict[str, Any], ModelResponse]], diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index fd67a7fbaf1..7c4e9386d5f 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -43,7 +43,9 @@ class AnthropicBatchesConfig(BaseBatchesConfig): api_base: Optional[str] = None, ) -> dict: """Validate and prepare environment-specific headers and parameters.""" - auth_header = self.anthropic_model_info.get_auth_header(api_key) + if api_base is None and isinstance(litellm_params, dict): + api_base = litellm_params.get("api_base") + auth_header = self.anthropic_model_info.get_auth_header(api_key, api_base) if auth_header is None: raise ValueError( "Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params" diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index e7e0f12c455..b1212f93059 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -364,6 +364,7 @@ class AnthropicChatCompletion(BaseLLM): messages=messages, optional_params={**optional_params, "is_vertex_request": is_vertex_request}, litellm_params=litellm_params, + api_base=api_base, ) config = ProviderConfigManager.get_provider_chat_config( diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index eaa762204f2..d27187cf2e7 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -548,6 +548,19 @@ class AnthropicModelInfo(BaseLLMModelInfo): return list(set(betas)) + @staticmethod + def _make_api_key_auth_header( + api_key: str, api_base: str | None, use_bearer_for_custom_base: bool = False + ) -> dict: + if use_bearer_for_custom_base and ( + api_base + and "api.anthropic.com" not in api_base + and not api_key.startswith("sk-ant-") + ): + value = api_key if api_key.startswith("Bearer ") else f"Bearer {api_key}" + return {"authorization": value} + return {"x-api-key": api_key} + def get_anthropic_headers( self, api_key: Optional[str] = None, @@ -567,6 +580,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): user_anthropic_beta_headers: Optional[List[str]] = None, code_execution_tool_used: bool = False, container_with_skills_used: bool = False, + api_base: str | None = None, + use_bearer_for_custom_base: bool = False, ) -> dict: betas = set() # Anthropic no longer requires the prompt-caching beta header @@ -615,7 +630,11 @@ class AnthropicModelInfo(BaseLLMModelInfo): elif auth_token and not api_key: headers["authorization"] = f"Bearer {auth_token}" elif api_key: - headers["x-api-key"] = api_key + headers.update( + self._make_api_key_auth_header( + api_key, api_base, use_bearer_for_custom_base + ) + ) if user_anthropic_beta_headers is not None: betas.update(user_anthropic_beta_headers) @@ -644,6 +663,12 @@ class AnthropicModelInfo(BaseLLMModelInfo): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> Dict: + if api_base is None and isinstance(litellm_params, dict): + api_base = litellm_params.get("api_base") + use_bearer_for_custom_base: bool = bool( + isinstance(litellm_params, dict) + and litellm_params.get("use_bearer_for_custom_base", False) + ) # Check for Anthropic OAuth token in headers headers, api_key = optionally_handle_anthropic_oauth( headers=headers, api_key=api_key @@ -699,6 +724,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): effort_used=effort_used, code_execution_tool_used=code_execution_tool_used, container_with_skills_used=container_with_skills_used, + api_base=api_base, + use_bearer_for_custom_base=use_bearer_for_custom_base, ) headers = {**headers, **anthropic_headers} @@ -734,18 +761,24 @@ class AnthropicModelInfo(BaseLLMModelInfo): return auth_token or get_secret_str("ANTHROPIC_AUTH_TOKEN") @staticmethod - def get_auth_header(api_key: Optional[str] = None) -> Optional[dict]: + def get_auth_header( + api_key: str | None = None, + api_base: str | None = None, + use_bearer_for_custom_base: bool = False, + ) -> dict | None: """Resolve Anthropic credentials and return the appropriate auth header dict. - Checks ANTHROPIC_API_KEY first (-> x-api-key), then - ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer). + Checks ANTHROPIC_API_KEY first (-> x-api-key or Bearer depending on + use_bearer_for_custom_base), then ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer). Returns None if neither is available. """ resolved_key = AnthropicModelInfo.get_api_key(api_key) if resolved_key is not None: if is_anthropic_oauth_key(resolved_key): return {"authorization": f"Bearer {resolved_key}"} - return {"x-api-key": resolved_key} + return AnthropicModelInfo._make_api_key_auth_header( + resolved_key, api_base, use_bearer_for_custom_base + ) auth_token = AnthropicModelInfo.get_auth_token() if auth_token is not None: return {"authorization": f"Bearer {auth_token}"} @@ -759,7 +792,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): self, api_key: Optional[str] = None, api_base: Optional[str] = None ) -> List[str]: api_base = AnthropicModelInfo.get_api_base(api_base) - auth_header = AnthropicModelInfo.get_auth_header(api_key) + auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base) if api_base is None or auth_header is None: raise ValueError( "ANTHROPIC_API_BASE/ANTHROPIC_BASE_URL or ANTHROPIC_API_KEY/ANTHROPIC_AUTH_TOKEN is not set. Please set the environment variable, to query Anthropic's `/models` endpoint." diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index 9860997b1cd..a92834b5d71 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -205,32 +205,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if "delta" not in merged_chunk: merged_chunk["delta"] = {} - uncached_input_tokens = chunk.usage.prompt_tokens or 0 - if ( - hasattr(chunk.usage, "prompt_tokens_details") - and chunk.usage.prompt_tokens_details - ): - cached_tokens = ( - getattr(chunk.usage.prompt_tokens_details, "cached_tokens", 0) or 0 - ) - uncached_input_tokens -= cached_tokens + from .transformation import LiteLLMAnthropicMessagesAdapter - usage_dict: UsageDelta = { - "input_tokens": uncached_input_tokens, - "output_tokens": chunk.usage.completion_tokens or 0, - } - if ( - hasattr(chunk.usage, "_cache_creation_input_tokens") - and chunk.usage._cache_creation_input_tokens > 0 - ): - usage_dict["cache_creation_input_tokens"] = ( - chunk.usage._cache_creation_input_tokens - ) - if ( - hasattr(chunk.usage, "_cache_read_input_tokens") - and chunk.usage._cache_read_input_tokens > 0 - ): - usage_dict["cache_read_input_tokens"] = chunk.usage._cache_read_input_tokens + usage_dict: UsageDelta = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( + chunk.usage + ) merged_chunk["usage"] = usage_dict if self.applied_edits and "context_management" not in merged_chunk: merged_chunk["context_management"] = ContextManagementResponse( diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 46564a565ce..1fd593f63c8 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1399,6 +1399,95 @@ class LiteLLMAnthropicMessagesAdapter: return "tool_use" return "end_turn" + @staticmethod + def _positive_int(value: object) -> int: + if isinstance(value, bool): + return 0 + if isinstance(value, int) and value > 0: + return value + if isinstance(value, float) and value.is_integer() and value > 0: + return int(value) + return 0 + + @classmethod + def _first_positive_usage_value( + cls, usage: Usage, field_names: tuple[str, ...] + ) -> int: + for field_name in field_names: + value = cls._positive_int(getattr(usage, field_name, None)) + if value > 0: + return value + return 0 + + @classmethod + def _first_positive_prompt_tokens_detail_value( + cls, usage: Usage, field_names: tuple[str, ...] + ) -> int: + prompt_tokens_details = getattr(usage, "prompt_tokens_details", None) + if prompt_tokens_details is None: + return 0 + + for field_name in field_names: + if isinstance(prompt_tokens_details, dict): + value = cls._positive_int(prompt_tokens_details.get(field_name)) + else: + value = cls._positive_int( + getattr(prompt_tokens_details, field_name, None) + ) + if value > 0: + return value + return 0 + + @classmethod + def _get_cache_read_input_tokens(cls, usage: Usage) -> int: + explicit_value = cls._first_positive_usage_value( + usage, ("cache_read_input_tokens", "_cache_read_input_tokens") + ) + if explicit_value > 0: + return explicit_value + return cls._first_positive_prompt_tokens_detail_value(usage, ("cached_tokens",)) + + @classmethod + def _get_cache_creation_input_tokens(cls, usage: Usage) -> int: + explicit_value = cls._first_positive_usage_value( + usage, ("cache_creation_input_tokens", "_cache_creation_input_tokens") + ) + if explicit_value > 0: + return explicit_value + return cls._first_positive_prompt_tokens_detail_value( + usage, ("cache_creation_tokens", "cache_write_tokens") + ) + + @classmethod + def _translate_openai_usage_to_anthropic_usage_delta( + cls, usage: Usage + ) -> UsageDelta: + cache_read_input_tokens = cls._get_cache_read_input_tokens(usage) + cache_creation_input_tokens = cls._get_cache_creation_input_tokens(usage) + input_tokens = max( + (usage.prompt_tokens or 0) + - cache_read_input_tokens + - cache_creation_input_tokens, + 0, + ) + + usage_delta = UsageDelta( + input_tokens=input_tokens, + output_tokens=usage.completion_tokens or 0, + ) + if cache_creation_input_tokens > 0: + usage_delta["cache_creation_input_tokens"] = cache_creation_input_tokens + if cache_read_input_tokens > 0: + usage_delta["cache_read_input_tokens"] = cache_read_input_tokens + return usage_delta + + @classmethod + def _translate_openai_usage_to_anthropic_usage(cls, usage: Usage) -> AnthropicUsage: + return cast( + AnthropicUsage, + cls._translate_openai_usage_to_anthropic_usage_delta(usage), + ) + def translate_openai_response_to_anthropic( self, response: ModelResponse, @@ -1430,32 +1519,12 @@ class LiteLLMAnthropicMessagesAdapter: ) # extract usage usage: Usage = getattr(response, "usage") - uncached_input_tokens = usage.prompt_tokens or 0 - cached_tokens = 0 - if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: - cached_tokens = ( - getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 - ) - uncached_input_tokens -= cached_tokens - - anthropic_usage = AnthropicUsage( - input_tokens=uncached_input_tokens, - output_tokens=usage.completion_tokens or 0, - ) - if ( - hasattr(usage, "_cache_creation_input_tokens") - and usage._cache_creation_input_tokens > 0 - ): - anthropic_usage["cache_creation_input_tokens"] = ( - usage._cache_creation_input_tokens - ) - if cached_tokens > 0: - anthropic_usage["cache_read_input_tokens"] = cached_tokens + anthropic_usage = self._translate_openai_usage_to_anthropic_usage(usage) if polyfill_result is not None and polyfill_result.iterations_usage is not None: message_iteration: UsageIteration = { "type": "message", - "input_tokens": uncached_input_tokens, + "input_tokens": anthropic_usage["input_tokens"], "output_tokens": usage.completion_tokens or 0, } anthropic_usage["iterations"] = list(polyfill_result.iterations_usage) + [ @@ -1644,35 +1713,9 @@ class LiteLLMAnthropicMessagesAdapter: else: litellm_usage_chunk = None if litellm_usage_chunk is not None: - uncached_input_tokens = litellm_usage_chunk.prompt_tokens or 0 - cached_tokens = 0 - if ( - hasattr(litellm_usage_chunk, "prompt_tokens_details") - and litellm_usage_chunk.prompt_tokens_details - ): - cached_tokens = ( - getattr( - litellm_usage_chunk.prompt_tokens_details, - "cached_tokens", - 0, - ) - or 0 - ) - uncached_input_tokens -= cached_tokens - - usage_delta = UsageDelta( - input_tokens=uncached_input_tokens, - output_tokens=litellm_usage_chunk.completion_tokens or 0, + usage_delta = self._translate_openai_usage_to_anthropic_usage_delta( + litellm_usage_chunk ) - if ( - hasattr(litellm_usage_chunk, "_cache_creation_input_tokens") - and litellm_usage_chunk._cache_creation_input_tokens > 0 - ): - usage_delta["cache_creation_input_tokens"] = ( - litellm_usage_chunk._cache_creation_input_tokens - ) - if cached_tokens > 0: - usage_delta["cache_read_input_tokens"] = cached_tokens else: usage_delta = UsageDelta(input_tokens=0, output_tokens=0) message_block = MessageBlockDelta( diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index 56296df94a1..170cb086bb0 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -84,7 +84,7 @@ class AnthropicFilesHandler: # Get Anthropic API credentials api_base = self.anthropic_model_info.get_api_base(api_base) - auth_header = self.anthropic_model_info.get_auth_header(api_key) + auth_header = self.anthropic_model_info.get_auth_header(api_key, api_base) if auth_header is None: raise ValueError("Missing Anthropic API Key") diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py index ea9bf00f505..7ffb6beb4c7 100644 --- a/litellm/llms/anthropic/files/transformation.py +++ b/litellm/llms/anthropic/files/transformation.py @@ -95,7 +95,9 @@ class AnthropicFilesConfig(BaseFilesConfig): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: - auth_header = AnthropicModelInfo.get_auth_header(api_key) + if api_base is None and isinstance(litellm_params, dict): + api_base = litellm_params.get("api_base") + auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base) if auth_header is None: raise ValueError( "Anthropic API key is required. Set ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN environment variable or pass api_key parameter." diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index 4ea768b02af..2bfdf7ef4ba 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -38,10 +38,12 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): # Get API key from litellm_params if available api_key = None + api_base = None if litellm_params is not None: api_key = litellm_params.api_key + api_base = litellm_params.api_base - auth_header = AnthropicModelInfo.get_auth_header(api_key) + auth_header = AnthropicModelInfo.get_auth_header(api_key, api_base) if auth_header is None: raise ValueError( "ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN is required for Skills API" diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index e8b3bf1a576..8fb7eb9fde0 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -62,6 +62,19 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): "get_supported_generate_content_optional_params is not implemented" ) + def get_generate_content_request_top_level_fields(self) -> tuple[str, ...]: + """ + Native Google ``GenerateContentRequest`` fields that sit at the top level + (siblings of ``generationConfig``) rather than inside it. The proxy forwards + these verbatim from a native request so ``generateContent`` is a drop-in for + Google's REST API. + + Excludes ``contents``, ``model`` and ``tools`` (dedicated params), + ``systemInstruction`` (dedicated extraction) and ``generationConfig`` (mapped + to ``config``). + """ + return ("safetySettings", "toolConfig", "cachedContent", "labels") + @abstractmethod def map_generate_content_optional_params( self, diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 5e97394f459..69c0a8529b4 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -904,9 +904,12 @@ class BedrockModelInfo(BaseLLMModelInfo): "mantle/": "mantle", } - # Check explicit routes first + # Check explicit routes first. Match each prefix only as a leading path + # segment so the `bedrock_mantle/` provider prefix is never mistaken for + # the `mantle/` invoke route (which would mangle + # `bedrock_mantle/openai.gpt-5.5` into `bedrock_openai.gpt-5.5`). for prefix, route_type in route_mappings.items(): - if prefix in model: + if BedrockModelInfo._model_has_route_prefix(model, prefix): return route_type # Check for nova spec prefixes (nova/ and nova-2/) @@ -930,14 +933,14 @@ class BedrockModelInfo(BaseLLMModelInfo): """ Check if the model is an explicit converse route. """ - return "converse/" in model + return BedrockModelInfo._model_has_route_prefix(model, "converse/") @staticmethod def _explicit_claude_platform_route(model: str) -> bool: """ Check if the model is an explicit Claude Platform on AWS route. """ - return "claude_platform/" in model + return BedrockModelInfo._model_has_route_prefix(model, "claude_platform/") @staticmethod def get_claude_platform_model(model: str) -> str: @@ -967,42 +970,58 @@ class BedrockModelInfo(BaseLLMModelInfo): """ Check if the model is an explicit invoke route. """ - return "invoke/" in model + return BedrockModelInfo._model_has_route_prefix(model, "invoke/") @staticmethod def _explicit_agent_route(model: str) -> bool: """ Check if the model is an explicit agent route. """ - return "agent/" in model + return BedrockModelInfo._model_has_route_prefix(model, "agent/") @staticmethod def _explicit_agentcore_route(model: str) -> bool: """ Check if the model is an explicit agentcore route. """ - return "agentcore/" in model + return BedrockModelInfo._model_has_route_prefix(model, "agentcore/") + + @staticmethod + def _model_has_route_prefix(model: str, prefix: str) -> bool: + """Whether a route prefix (e.g. ``mantle/``) appears as a leading path segment. + + A route token is only valid at the start of the model id or immediately + after a ``/``. A plain substring check matches the ``bedrock_mantle/`` + provider prefix against the ``mantle/`` route, so the body model gets + mangled to ``bedrock_openai.gpt-5.5``; anchoring to a segment boundary + keeps the bare model id intact. + + ``f"/{prefix}" in model`` matches the token as a segment at any path + depth, not just the second segment; that is intentional and acceptable + for these short, unambiguous route tokens. + """ + return model.startswith(prefix) or f"/{prefix}" in model @staticmethod def _explicit_mantle_route(model: str) -> bool: """ Check if the model is an explicit mantle route (bedrock-mantle endpoint). """ - return "mantle/" in model + return BedrockModelInfo._model_has_route_prefix(model, "mantle/") @staticmethod def _explicit_converse_like_route(model: str) -> bool: """ Check if the model is an explicit converse like route. """ - return "converse_like/" in model + return BedrockModelInfo._model_has_route_prefix(model, "converse_like/") @staticmethod def _explicit_async_invoke_route(model: str) -> bool: """ Check if the model is an explicit async invoke route. """ - return "async_invoke/" in model + return BedrockModelInfo._model_has_route_prefix(model, "async_invoke/") @staticmethod def _explicit_openai_route(model: str) -> bool: @@ -1010,7 +1029,7 @@ class BedrockModelInfo(BaseLLMModelInfo): Check if the model is an explicit openai route. Used for Bedrock imported models that use OpenAI Chat Completions format. """ - return "openai/" in model + return BedrockModelInfo._model_has_route_prefix(model, "openai/") @staticmethod def get_bedrock_provider_config_for_messages_api( diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 79f006a83de..3bdcbd25949 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -1046,9 +1046,9 @@ def _gemini_convert_messages_with_history( if invocation.get("tool_type"): tr_dict["toolType"] = invocation["tool_type"] tr_part: Dict[str, Any] = {"toolResponse": tr_dict} - if "thought_signature" in invocation: + if "response_thought_signature" in invocation: tr_part["thoughtSignature"] = invocation[ - "thought_signature" + "response_thought_signature" ] assistant_content.append(tr_part) # type: ignore diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 0de34ec257c..423a5dd5d17 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1634,13 +1634,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): resp = tool_responses_by_id.pop(call_id, None) if resp is not None: merged["response"] = resp.get("response") - # Keep response signature if call didn't have one - if "thought_signature" not in merged and "thought_signature" in resp: - merged["thought_signature"] = resp["thought_signature"] + if "thought_signature" in resp: + merged["response_thought_signature"] = resp["thought_signature"] invocations.append(merged) # Any orphan responses (shouldn't happen, but be safe) for resp_id, resp_entry in tool_responses_by_id.items(): + if "thought_signature" in resp_entry: + resp_entry["response_thought_signature"] = resp_entry[ + "thought_signature" + ] invocations.append(resp_entry) return invocations if invocations else None diff --git a/litellm/main.py b/litellm/main.py index 1960c456f80..cdea7dbd684 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -633,6 +633,7 @@ async def acompletion( try: # Use a partial function to pass your keyword arguments + kwargs.pop("acompletion", None) func = partial(completion, **completion_kwargs, **kwargs) # Add the context to the function @@ -5063,6 +5064,12 @@ def completion( # type: ignore ######### unpacking kwargs ##################### args = locals() + # Set by the responses->completion fallback so completion() does not bridge + # back to the Responses API: that round-trip mutually recurses forever for a + # model whose model_cost mode is "responses" but whose provider has no + # Responses API config (get_provider_responses_api_config -> None). + skip_responses_api_bridge = kwargs.pop("_skip_responses_api_bridge", False) + skip_mcp_handler = kwargs.pop("_skip_mcp_handler", False) if not skip_mcp_handler and tools: from litellm.responses.mcp.chat_completions_handler import acompletion_with_mcp @@ -5358,7 +5365,7 @@ def completion( # type: ignore messages = update_messages_with_model_file_ids( messages=messages, - model_id=kwargs.get("model_info", {}).get("id", None), + model_id=(kwargs.get("model_info") or {}).get("id", None), model_file_id_mapping=cast( Dict[str, Dict[str, str]], kwargs.get("model_file_id_mapping") or {}, @@ -5560,7 +5567,10 @@ def completion( # type: ignore # detection when the deployment name differs from the model name. _azure_detection_model = base_model or model - if responses_api_model_info.get("mode") == "responses": + if ( + responses_api_model_info.get("mode") == "responses" + and not skip_responses_api_bridge + ): from litellm.completion_extras import responses_api_bridge optional_params, rs_val = ( diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 2149f079a3d..7b4f1e13a52 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -79,6 +79,7 @@ if MCP_AVAILABLE: ) from litellm.proxy._experimental.mcp_server.server import ( ListMCPToolsRestAPIResponseObject, + MCPInfo, MCPServer, _tool_name_matches, execute_mcp_tool, @@ -238,14 +239,24 @@ if MCP_AVAILABLE: ) return {} - def _create_tool_response_objects(tools, server_mcp_info): - """Helper function to create tool response objects.""" + def _create_tool_response_objects(tools, server: MCPServer): + """Helper function to create tool response objects. + + Enriches the server's ``mcp_info`` with ``server_id`` and ``alias`` so + REST clients can map the internal ``server_name`` to the user-facing + alias without needing access to the ``mcp_routes``-gated server listing. + """ + enriched_mcp_info: MCPInfo = { + **(server.mcp_info or {}), + "server_id": server.server_id, + "alias": server.alias, + } return [ ListMCPToolsRestAPIResponseObject( name=tool.name, description=tool.description, inputSchema=tool.inputSchema, - mcp_info=server_mcp_info, + mcp_info=enriched_mcp_info, ) for tool in tools ] @@ -405,7 +416,7 @@ if MCP_AVAILABLE: ) if not apply_tool_filters: - return _create_tool_response_objects(tools, server.mcp_info) + return _create_tool_response_objects(tools, server) # Always apply allowed_tools/disallowed_tools so the blacklist is # enforced even when no allowlist is set (matches the SSE/HTTP path). @@ -436,7 +447,7 @@ if MCP_AVAILABLE: if _tool_name_matches(tool.name, allowed_tools_for_server) ] - return _create_tool_response_objects(tools, server.mcp_info) + return _create_tool_response_objects(tools, server) async def _resolve_allowed_mcp_servers_for_tool_call( user_api_key_dict: UserAPIKeyAuth, @@ -587,6 +598,8 @@ if MCP_AVAILABLE: "mcp_info": { "server_name": "zapier", "logo_url": "https://www.zapier.com/logo.png", + "server_id": "a1b2c3d4-...", + "alias": "zapier_prod", } } ], diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 8a432eb2f42..928d00109fa 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1849,7 +1849,10 @@ async def test_model_connection( "responses", "ocr", ] - ] = fastapi.Body("chat", description="The mode to test the model with"), + ] = fastapi.Body( + None, + description="The mode to test the model with. If not provided, auto-detected from model capabilities.", + ), litellm_params: Dict = fastapi.Body( None, description="Parameters for litellm.completion, litellm.embedding for the health check", diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index f28bcc2bcd4..bc2da33672f 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,8 +1,27 @@ +import math from typing import TYPE_CHECKING, Any, Dict, Optional, Union from fastapi import HTTPException, status from pydantic import BaseModel + +# Defined above the `litellm.proxy.*` imports so the name is bound even when +# this module is imported first through the proxy import cycle (CodeQL: +# module-level cyclic import). Depends only on `math` + `HTTPException`. +def validate_finite_spend(spend: float | None) -> None: + """Reject NaN/±inf spend before it reaches the DB / spend counter. + + A non-finite spend would otherwise slip past `spend >= max_budget` + enforcement, since any comparison with NaN (and `-inf >= max_budget`) + is False, letting the entity keep spending past its configured budget. + """ + if spend is not None and not math.isfinite(spend): + raise HTTPException( + status_code=400, + detail={"error": f"spend must be a finite number. Received: {spend}"}, + ) + + from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.proxy._types import ( diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 8548d1c672c..99e276cfbb4 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_admin, _user_has_admin_view, require_caller_user_id_for_non_admin, + validate_finite_spend, ) from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, @@ -1256,6 +1257,25 @@ def _check_user_update_authz( ) +async def _invalidate_user_spend_counter_if_changed( + non_default_values: dict[str, Any], +) -> None: + """Invalidate the cross-pod spend counter after a direct ``spend`` change. + + A direct ``spend`` change must also invalidate the cross-pod spend counter + enforcement reads; the DB write alone leaves a warm counter at the stale + value. ``non_default_values["user_id"]`` is populated in every branch of the + caller (incl. the email-new-user insert path, whose response is a bare model + and not safely subscriptable). + """ + if non_default_values.get("spend") is not None: + from litellm.proxy.proxy_server import _invalidate_spend_counter + + await _invalidate_spend_counter( + counter_key=f"spend:user:{non_default_values['user_id']}" + ) + + async def _update_single_user_helper( user_request: UpdateUserRequest, user_api_key_dict: UserAPIKeyAuth, @@ -1336,6 +1356,9 @@ async def _update_single_user_helper( existing_metadata=existing_metadata or {}, ) + # Reject NaN/±inf spend before it can reach the DB / spend counter. + validate_finite_spend(non_default_values.get("spend")) + # Perform the update response: Optional[Dict[str, Any]] = None @@ -1384,6 +1407,8 @@ async def _update_single_user_helper( litellm_proxy_admin_name=litellm_proxy_admin_name, ) + await _invalidate_user_spend_counter_if_changed(non_default_values) + if response is None: raise HTTPException( status_code=400, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index db205803a47..9f316256ac9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -69,6 +69,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_admin, _set_object_metadata_field, _team_member_has_permission, + validate_finite_spend, ) from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, @@ -2210,6 +2211,9 @@ async def _validate_update_key_data( user_api_key_cache: Any, ) -> None: """Validate permissions and constraints for key update.""" + # Reject NaN/±inf spend before it can reach the DB / spend counter. + validate_finite_spend(data.spend) + _is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value _check_allowed_routes_caller_permission( @@ -2269,12 +2273,14 @@ async def _validate_update_key_data( # existing admin-only budget semantics). budget_limits uses # model_fields_set because an explicit null/[] clears the field # and must gate the same as setting or changing it. + # - spend gates on presence alone (not a value diff): the DB spend + # lags the live cross-pod counter, so letting an "unchanged" spend + # through the non-admin path would let a key owner / team member + # overwrite the live counter below real usage and silently weaken + # enforcement. _is_budget_change = ( (data.max_budget is not None and data.max_budget != existing_key_row.max_budget) - or ( - data.spend is not None - and data.spend != getattr(existing_key_row, "spend", None) - ) + or data.spend is not None or "budget_limits" in data.model_fields_set ) @@ -2609,15 +2615,24 @@ async def update_key_fn( ) if data.spend is not None: - try: - from litellm.proxy.proxy_server import _invalidate_spend_counter + from litellm.proxy.proxy_server import spend_counter_cache - token_to_invalidate = _hash_token_if_needed(key) - await _invalidate_spend_counter( - counter_key=f"spend:key:{token_to_invalidate}" - ) - except Exception: - pass + counter_key = f"spend:key:{_hash_token_if_needed(key)}" + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=data.spend, ttl=60 + ) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=data.spend, ttl=60 + ) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to update spend counter %s in Redis after key spend update: %s. " + "Budget checks may use stale value until counter expires.", + counter_key, + redis_err, + ) asyncio.create_task( KeyManagementEventHooks.async_key_updated_hook( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b0beaae02cf..3c91f1bc1d7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1114,6 +1114,33 @@ _OPENAPI_HTTP_METHODS = { # `_SSO_SENSITIVE_FIELDS` / `_CACHE_SENSITIVE_FIELDS` constants in the SSO # and cache endpoint files. _ALERTING_SENSITIVE_VARS: Set[str] = {"SLACK_WEBHOOK_URL", "SMTP_PASSWORD"} +_DB_LITELLM_PARAM_ENV_REF_KEYS = frozenset( + { + "api_key", + "client_secret", + "vertex_credentials", + "vertex_ai_credentials", + "aws_access_key_id", + "aws_secret_access_key", + } +) + + +def _db_model_is_team_scoped(model: object) -> bool: + model_info = getattr(model, "model_info", None) + if isinstance(model_info, BaseModel): + return getattr(model_info, "team_id", None) is not None + if isinstance(model_info, str): + try: + model_info = json.loads(model_info) + except (TypeError, ValueError): + model_info = None + if isinstance(model_info, dict) and model_info.get("team_id") is not None: + return True + if getattr(model_info, "team_id", None) is not None: + return True + model_name = getattr(model, "model_name", None) + return isinstance(model_name, str) and model_name.startswith("model_name_") def _strip_operation_id_method_suffix(operation_id: str) -> str: @@ -5255,6 +5282,24 @@ class ProxyConfig: deleted_deployments += 1 return deleted_deployments + def _resolve_db_litellm_param( + self, key: str, value: object, resolve_env_refs: bool = True + ) -> object: + if not isinstance(value, str): + return value + + decrypted_value = decrypt_value_helper( + value=value, key=key, return_original_value=True + ) + if ( + resolve_env_refs + and key in _DB_LITELLM_PARAM_ENV_REF_KEYS + and isinstance(decrypted_value, str) + and decrypted_value.startswith("os.environ/") + ): + return get_secret(decrypted_value) + return decrypted_value + def _add_deployment(self, db_models: list) -> int: """ Iterate through db models @@ -5272,15 +5317,13 @@ class ProxyConfig: ## ADD MODEL LOGIC for m in db_models: _litellm_params = m.litellm_params + resolve_env_refs = not _db_model_is_team_scoped(m) if isinstance(_litellm_params, dict): # decrypt values for k, v in _litellm_params.items(): - if isinstance(v, str): - # decrypt value - returns original value if decryption fails or no key is set - _value = decrypt_value_helper( - value=v, key=k, return_original_value=True - ) - _litellm_params[k] = _value + _litellm_params[k] = self._resolve_db_litellm_param( + key=k, value=v, resolve_env_refs=resolve_env_refs + ) _litellm_params = LiteLLM_Params(**_litellm_params) else: @@ -5308,15 +5351,15 @@ class ProxyConfig: _model_list: list = [] for m in new_models: _litellm_params = m.litellm_params + resolve_env_refs = not _db_model_is_team_scoped(m) if isinstance(_litellm_params, BaseModel): _litellm_params = _litellm_params.model_dump() if isinstance(_litellm_params, dict): # decrypt values for k, v in _litellm_params.items(): - decrypted_value = decrypt_value_helper( - value=v, key=k, return_original_value=True + _litellm_params[k] = self._resolve_db_litellm_param( + key=k, value=v, resolve_env_refs=resolve_env_refs ) - _litellm_params[k] = decrypted_value _litellm_params = LiteLLM_Params(**_litellm_params) else: verbose_proxy_logger.error( diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index a29f0a01b7a..1187640bf0c 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -59,6 +59,7 @@ class LiteLLMCompletionTransformationHandler: completion_args = {} completion_args.update(kwargs) completion_args.update(litellm_completion_request) + completion_args["_skip_responses_api_bridge"] = True litellm_completion_response: Union[ ModelResponse, litellm.CustomStreamWrapper @@ -107,6 +108,7 @@ class LiteLLMCompletionTransformationHandler: acompletion_args = {} acompletion_args.update(kwargs) acompletion_args.update(litellm_completion_request) + acompletion_args["_skip_responses_api_bridge"] = True litellm_completion_response: Union[ ModelResponse, litellm.CustomStreamWrapper diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 5f8fcbf835e..bb13a7ce8cc 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1883,10 +1883,11 @@ def test_get_server_auth_header_no_auth_headers(): def test_create_tool_response_objects(): - """Test _create_tool_response_objects function.""" + """Test _create_tool_response_objects enriches mcp_info with server_id and alias.""" from litellm.proxy._experimental.mcp_server.rest_endpoints import ( _create_tool_response_objects, ) + from litellm.types.mcp_server.mcp_server_manager import MCPServer from mcp.types import Tool as MCPTool # Create mock tools @@ -1903,20 +1904,32 @@ def test_create_tool_response_objects(): ), ] - server_mcp_info = { - "server_name": "zapier", + server = MCPServer( + server_id="a1b2c3d4", + name="zapier_internal", + alias="zapier", + transport="http", + mcp_info={ + "server_name": "zapier_internal", + "logo_url": "https://zapier.com/logo.png", + }, + ) + + result = _create_tool_response_objects(mock_tools, server) + + expected_mcp_info = { + "server_name": "zapier_internal", "logo_url": "https://zapier.com/logo.png", + "server_id": "a1b2c3d4", + "alias": "zapier", } - - result = _create_tool_response_objects(mock_tools, server_mcp_info) - assert len(result) == 2 assert result[0].name == "send_email" assert result[0].description == "Send an email" - assert result[0].mcp_info == server_mcp_info + assert result[0].mcp_info == expected_mcp_info assert result[1].name == "create_event" assert result[1].description == "Create a calendar event" - assert result[1].mcp_info == server_mcp_info + assert result[1].mcp_info == expected_mcp_info @pytest.mark.asyncio @@ -1930,6 +1943,8 @@ async def test_get_tools_for_single_server(): # Create a mock server (pin allowlist fields; MagicMock auto-attrs are truthy) mock_server = MagicMock() mock_server.mcp_info = {"server_name": "zapier"} + mock_server.server_id = "zapier_id" + mock_server.alias = "zapier_alias" mock_server.allowed_tools = None mock_server.disallowed_tools = None @@ -1963,7 +1978,11 @@ async def test_get_tools_for_single_server(): # Verify the result assert len(result) == 1 assert result[0].name == "send_email" - assert result[0].mcp_info == {"server_name": "zapier"} + assert result[0].mcp_info == { + "server_name": "zapier", + "server_id": "zapier_id", + "alias": "zapier_alias", + } @pytest.mark.asyncio diff --git a/tests/test_litellm/caching/test_embedding_router.py b/tests/test_litellm/caching/test_embedding_router.py new file mode 100644 index 00000000000..550095a112a --- /dev/null +++ b/tests/test_litellm/caching/test_embedding_router.py @@ -0,0 +1,67 @@ +import os +import sys +from unittest.mock import MagicMock + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.caching._embedding_router import ( + build_router_embedding_metadata, + resolve_embedding_router, +) + + +def test_resolve_returns_router_when_model_is_a_deployment(): + router = MagicMock() + assert ( + resolve_embedding_router("sem-embed", router, [{"model_name": "sem-embed"}]) + is router + ) + + +def test_resolve_returns_none_when_model_not_in_router(): + router = MagicMock() + assert ( + resolve_embedding_router("sem-embed", router, [{"model_name": "other"}]) is None + ) + + +def test_resolve_returns_none_when_router_is_none(): + assert ( + resolve_embedding_router("sem-embed", None, [{"model_name": "sem-embed"}]) + is None + ) + + +def test_resolve_returns_none_when_model_list_is_none(): + router = MagicMock() + assert resolve_embedding_router("sem-embed", router, None) is None + + +def test_resolve_skips_entries_missing_model_name(): + router = MagicMock() + model_list = [ + {"litellm_params": {"model": "bedrock/x"}}, + {"model_name": "sem-embed"}, + ] + assert resolve_embedding_router("sem-embed", router, model_list) is router + assert resolve_embedding_router("other", router, [{"litellm_params": {}}]) is None + + +def test_build_metadata_preserves_request_fields_and_adds_flag(): + md = build_router_embedding_metadata( + {"user_api_key": "sk-x", "user_api_key_team_id": "team-1", "trace_id": "t-1"} + ) + assert md == { + "user_api_key": "sk-x", + "user_api_key_team_id": "team-1", + "trace_id": "t-1", + "semantic-cache-embedding": True, + } + + +def test_build_metadata_handles_none_and_does_not_mutate_input(): + original = {"user_api_key": "sk-x"} + md = build_router_embedding_metadata(original) + assert md == {"user_api_key": "sk-x", "semantic-cache-embedding": True} + assert original == {"user_api_key": "sk-x"} + assert build_router_embedding_metadata(None) == {"semantic-cache-embedding": True} diff --git a/tests/test_litellm/caching/test_qdrant_semantic_cache.py b/tests/test_litellm/caching/test_qdrant_semantic_cache.py index 949e6ccc292..67d4e2d9892 100644 --- a/tests/test_litellm/caching/test_qdrant_semantic_cache.py +++ b/tests/test_litellm/caching/test_qdrant_semantic_cache.py @@ -1,5 +1,6 @@ import os import sys +import types from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -806,3 +807,104 @@ def test_qdrant_semantic_cache_large_vector_size(): ) create_payload = put_call.kwargs["json"] assert create_payload["vectors"]["size"] == 4096 + + +def _router_proxy_module(router, model_name): + mod = types.ModuleType("litellm.proxy.proxy_server") + mod.llm_router = router + mod.llm_model_list = [{"model_name": model_name}] + return mod + + +def test_qdrant_sync_get_cache_routes_through_router(monkeypatch): + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + cache.embedding_model = "sem-embed" + cache.qdrant_api_base = "http://test.qdrant.local" + cache.collection_name = "test_collection" + cache.headers = {"Content-Type": "application/json", "api-key": "test_key"} + cache.similarity_threshold = 0.8 + cache.sync_client = MagicMock() + search_response = MagicMock() + search_response.status_code = 200 + search_response.json.return_value = {"result": []} + cache.sync_client.post.return_value = search_response + + router = MagicMock() + router.embedding = MagicMock( + return_value={"data": [{"embedding": [0.3, 0.3, 0.3]}]} + ) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + with patch("litellm.embedding") as direct_embed: + result = cache.get_cache( + key="test_key", + messages=[{"content": "What is the capital of France?"}], + metadata={}, + ) + + assert result is None + router.embedding.assert_called_once() + assert router.embedding.call_args.kwargs["model"] == "sem-embed" + direct_embed.assert_not_called() + + +def test_qdrant_sync_set_cache_falls_back_to_direct(monkeypatch): + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + cache.embedding_model = "text-embedding-ada-002" + cache.qdrant_api_base = "http://test.qdrant.local" + cache.collection_name = "test_collection" + cache.headers = {"Content-Type": "application/json", "api-key": "test_key"} + cache.sync_client = MagicMock() + put_response = MagicMock() + put_response.status_code = 200 + cache.sync_client.put.return_value = put_response + + fake_proxy = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy.llm_router = None + fake_proxy.llm_model_list = None + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy) + + with patch( + "litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.1, 0.1]}]} + ) as direct_embed: + cache.set_cache( + key="test_key", + value={"content": "Paris"}, + messages=[{"content": "What is the capital of France?"}], + ) + + direct_embed.assert_called_once() + + +@pytest.mark.asyncio +async def test_qdrant_async_embedding_forwards_full_metadata(monkeypatch): + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + cache.embedding_model = "sem-embed" + + router = MagicMock() + router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]}) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + await cache._get_async_embedding( + "hello", + metadata={"user_api_key": "sk-x", "user_api_key_team_id": "team-1"}, + ) + + md = router.aembedding.call_args.kwargs["metadata"] + assert md["user_api_key"] == "sk-x" + assert md["user_api_key_team_id"] == "team-1" + assert md["semantic-cache-embedding"] is True diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 6aaadce93ce..37c938bb38d 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -3,7 +3,6 @@ import sys from unittest.mock import MagicMock, patch import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../..") @@ -94,6 +93,42 @@ async def test_redis_cache_async_increment_default_does_not_bump_existing_ttl( mock_redis_instance.expire.assert_not_awaited() +@pytest.mark.parametrize("namespace", [None, "litellm"]) +@pytest.mark.asyncio +async def test_async_delete_cache_applies_namespace( + namespace, monkeypatch, redis_no_ping +): + """async_delete_cache must prefix keys with the namespace, matching every + other cache operation. Without this, Redis NOPERM errors occur when an + ACL restricts DEL to the litellm:* pattern.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_instance = AsyncMock() + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_delete_cache(key="3997c4abcdef") + + expected_key = "litellm:3997c4abcdef" if namespace else "3997c4abcdef" + mock_redis_instance.delete.assert_awaited_once_with(expected_key) + + +@pytest.mark.parametrize("namespace", [None, "litellm"]) +def test_delete_cache_applies_namespace(namespace, monkeypatch, redis_no_ping): + """delete_cache must prefix keys with the namespace, matching every other + cache operation.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_client = MagicMock() + redis_cache.redis_client = mock_redis_client + + redis_cache.delete_cache(key="3997c4abcdef") + + expected_key = "litellm:3997c4abcdef" if namespace else "3997c4abcdef" + mock_redis_client.delete.assert_called_once_with(expected_key) + + @pytest.mark.asyncio async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "my-fake-host") diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index 13f9d00136d..1d3129d6467 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -104,6 +104,7 @@ def test_redis_semantic_cache_get_cache(monkeypatch): # Verify llmcache.check was called redis_semantic_cache.llmcache.check.assert_called_once_with( prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], filter_expression="cache-key-filter", ) @@ -138,10 +139,16 @@ def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch): ] ) - with patch.object( - redis_semantic_cache, - "_get_cache_key_filter_expression", - return_value="cache-key-filter", + with ( + patch( + "litellm.embedding", + return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}, + ), + patch.object( + redis_semantic_cache, + "_get_cache_key_filter_expression", + return_value="cache-key-filter", + ), ): metadata = {} result = redis_semantic_cache.get_cache( @@ -176,16 +183,21 @@ def test_redis_semantic_cache_set_cache_stores_cache_key_filter(monkeypatch): redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8) redis_semantic_cache.llmcache.store = MagicMock() - redis_semantic_cache.set_cache( - key="test_key", - value={"content": "Paris"}, - messages=[{"content": "What is the capital of France?"}], - ttl=60, - ) + with patch( + "litellm.embedding", + return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}, + ): + redis_semantic_cache.set_cache( + key="test_key", + value={"content": "Paris"}, + messages=[{"content": "What is the capital of France?"}], + ttl=60, + ) redis_semantic_cache.llmcache.store.assert_called_once_with( "What is the capital of France?", "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, ttl=60, ) @@ -299,10 +311,11 @@ def test_redis_semantic_cache_reraises_unexpected_isolated_index_error(monkeypat monkeypatch.setenv("REDIS_PASSWORD", "test_password") with pytest.raises(ValueError, match="connection failed"): - RedisSemanticCache( + cache = RedisSemanticCache( similarity_threshold=0.8, index_name="existing_index", ) + _ = cache.llmcache def test_redis_semantic_cache_reraises_unexpected_index_error(): @@ -534,6 +547,7 @@ def test_redis_semantic_cache_set_cache_uses_responses_string_input(): return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} ) redis_semantic_cache._get_ttl = MagicMock(return_value=None) + redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) redis_semantic_cache.set_cache( key="test_key", @@ -544,6 +558,7 @@ def test_redis_semantic_cache_set_cache_uses_responses_string_input(): redis_semantic_cache.llmcache.store.assert_called_once_with( "What is the capital of France?", "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, ) @@ -564,6 +579,7 @@ def test_redis_semantic_cache_get_cache_uses_responses_string_input(): } ] ) + redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) with patch.object( redis_semantic_cache, @@ -581,6 +597,7 @@ def test_redis_semantic_cache_get_cache_uses_responses_string_input(): assert metadata["semantic-similarity"] == pytest.approx(0.9) redis_semantic_cache.llmcache.check.assert_called_once_with( prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], filter_expression="cache-key-filter", ) @@ -594,6 +611,7 @@ def test_redis_semantic_cache_set_cache_flattens_structured_responses_input(): return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} ) redis_semantic_cache._get_ttl = MagicMock(return_value=None) + redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) redis_semantic_cache.set_cache( key="test_key", @@ -616,6 +634,7 @@ def test_redis_semantic_cache_set_cache_flattens_structured_responses_input(): redis_semantic_cache.llmcache.store.assert_called_once_with( "What is the capital of France?\nAnswer briefly.", "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, ) @@ -740,6 +759,7 @@ def test_redis_semantic_cache_get_cache_sets_similarity_when_no_results(): redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) redis_semantic_cache.llmcache = MagicMock() redis_semantic_cache.llmcache.check = MagicMock(return_value=[]) + redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) with patch.object( redis_semantic_cache, @@ -757,6 +777,7 @@ def test_redis_semantic_cache_get_cache_sets_similarity_when_no_results(): assert metadata["semantic-similarity"] == 0.0 redis_semantic_cache.llmcache.check.assert_called_once_with( prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], filter_expression="cache-key-filter", ) @@ -870,6 +891,63 @@ async def test_redis_semantic_cache_async_paths_set_similarity_on_misses(): ) +def test_redis_get_embedding_routes_through_router(monkeypatch): + import sys + import types + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "sem-embed" + + router = MagicMock() + router.embedding = MagicMock(return_value={"data": [{"embedding": [0.5, 0.6]}]}) + fake_proxy = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy.llm_router = router + fake_proxy.llm_model_list = [{"model_name": "sem-embed"}] + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy) + + with patch("litellm.embedding") as direct_embed: + vec = cache._get_embedding("hello", metadata={"user_api_key": "sk-x"}) + + assert vec == [0.5, 0.6] + router.embedding.assert_called_once() + assert router.embedding.call_args.kwargs["model"] == "sem-embed" + assert router.embedding.call_args.kwargs["input"] == "hello" + assert router.embedding.call_args.kwargs["cache"] == { + "no-store": True, + "no-cache": True, + } + assert router.embedding.call_args.kwargs["metadata"] == { + "user_api_key": "sk-x", + "semantic-cache-embedding": True, + } + direct_embed.assert_not_called() + + +def test_redis_get_embedding_falls_back_to_direct(monkeypatch): + import sys + import types + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "text-embedding-ada-002" + + fake_proxy = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy.llm_router = None + fake_proxy.llm_model_list = None + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy) + + with patch( + "litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2]}]} + ) as direct_embed: + vec = cache._get_embedding("hello") + + assert vec == [0.1, 0.2] + direct_embed.assert_called_once() + + def test_cache_get_cache_passes_responses_input_to_backend_cache(): from litellm.caching.caching import Cache @@ -893,7 +971,7 @@ def test_cache_get_cache_passes_responses_input_to_backend_cache(): ) -def test_cache_get_cache_filters_sensitive_kwargs_from_backend_cache(): +def test_cache_get_cache_filters_non_lookup_kwargs_from_backend_cache(): from litellm.caching.caching import Cache cache = Cache.__new__(Cache) @@ -927,7 +1005,11 @@ def test_cache_get_cache_filters_sensitive_kwargs_from_backend_cache(): forwarded_kwargs = cache.cache.get_cache.call_args.kwargs assert forwarded_kwargs == { "input": "What is the capital of France?", - "metadata": {"semantic-similarity": 0.7}, + "metadata": { + "user_api_key": "sk-secret", + "trace_id": "trace-id", + "semantic-similarity": 0.7, + }, } assert forwarded_kwargs["metadata"] is not metadata cache._get_cache_logic.assert_called_once_with( @@ -988,3 +1070,166 @@ def test_cache_get_cache_passes_responses_input_to_dynamic_cache(): cached_result={"content": "Paris"}, max_age=float("inf"), ) + + +def test_redis_sync_set_cache_passes_precomputed_vector(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.llmcache = MagicMock() + cache._get_cache_filters = MagicMock( + return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} + ) + cache._get_ttl = MagicMock(return_value=None) + cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) + + cache.set_cache( + key="test_key", + value={"content": "Paris"}, + messages=[{"content": "What is the capital of France?"}], + ) + + cache._get_embedding.assert_called_once() + cache.llmcache.store.assert_called_once_with( + "What is the capital of France?", + "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], + filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, + ) + + +def test_redis_sync_get_cache_passes_precomputed_vector(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.similarity_threshold = 0.8 + cache.llmcache = MagicMock() + cache.llmcache.check = MagicMock( + return_value=[ + { + "prompt": "What is the capital of France?", + "response": '{"content": "Paris"}', + "vector_distance": 0.1, + RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key", + } + ] + ) + cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) + + with patch.object( + cache, "_get_cache_key_filter_expression", return_value="cache-key-filter" + ): + result = cache.get_cache( + key="test_key", + messages=[{"content": "What is the capital of France?"}], + metadata={}, + ) + + assert result == {"content": "Paris"} + cache._get_embedding.assert_called_once() + cache.llmcache.check.assert_called_once_with( + prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], + filter_expression="cache-key-filter", + ) + + +@pytest.mark.asyncio +async def test_redis_async_embedding_forwards_full_metadata(monkeypatch): + import sys + import types + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "sem-embed" + + router = MagicMock() + router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]}) + fake_proxy = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy.llm_router = router + fake_proxy.llm_model_list = [{"model_name": "sem-embed"}] + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy) + + await cache._get_async_embedding( + "hello", + metadata={"user_api_key": "sk-x", "user_api_key_team_id": "team-1"}, + ) + + md = router.aembedding.call_args.kwargs["metadata"] + assert md["user_api_key"] == "sk-x" + assert md["user_api_key_team_id"] == "team-1" # FAILS today: team_id is dropped + assert md["semantic-cache-embedding"] is True + + +def test_redis_init_defers_redisvl_construction(monkeypatch): + semantic_cache_mock = MagicMock() + custom_vectorizer_mock = MagicMock() + + with patch.dict( + "sys.modules", + { + "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), + "redisvl.utils.vectorize": MagicMock( + CustomTextVectorizer=custom_vectorizer_mock + ), + }, + ): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + monkeypatch.setenv("REDIS_HOST", "localhost") + monkeypatch.setenv("REDIS_PORT", "6379") + monkeypatch.setenv("REDIS_PASSWORD", "test_password") + + cache = RedisSemanticCache(similarity_threshold=0.8) + + semantic_cache_mock.assert_not_called() + custom_vectorizer_mock.assert_not_called() + + first = cache.llmcache + semantic_cache_mock.assert_called_once() + custom_vectorizer_mock.assert_called_once() + + second = cache.llmcache + assert first is second + semantic_cache_mock.assert_called_once() + + +def test_redis_failed_llmcache_build_is_not_memoized(monkeypatch): + built_cache = MagicMock() + semantic_cache_mock = MagicMock( + side_effect=[ConnectionError("redis down"), built_cache] + ) + custom_vectorizer_mock = MagicMock() + + with patch.dict( + "sys.modules", + { + "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), + "redisvl.utils.vectorize": MagicMock( + CustomTextVectorizer=custom_vectorizer_mock + ), + }, + ): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + monkeypatch.setenv("REDIS_HOST", "localhost") + monkeypatch.setenv("REDIS_PORT", "6379") + monkeypatch.setenv("REDIS_PASSWORD", "test_password") + + cache = RedisSemanticCache(similarity_threshold=0.8) + + with pytest.raises(ConnectionError, match="redis down"): + _ = cache.llmcache + + assert cache.llmcache is built_cache + assert semantic_cache_mock.call_count == 2 + + +def test_redis_llmcache_setter_supported(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + sentinel = MagicMock() + cache.llmcache = sentinel + assert cache.llmcache is sentinel diff --git a/tests/test_litellm/google_genai/test_google_genai_main.py b/tests/test_litellm/google_genai/test_google_genai_main.py index 5854e4b55af..8f56b4e4bc0 100644 --- a/tests/test_litellm/google_genai/test_google_genai_main.py +++ b/tests/test_litellm/google_genai/test_google_genai_main.py @@ -2,6 +2,7 @@ """ Test to verify the Google GenAI generate_content adapter functionality """ + import json import os import sys @@ -42,4 +43,228 @@ async def test_agenerate_content_stream(): stream=True, ) mock_post.assert_called_once() - mock_post.call_args.kwargs["stream"] == True + assert mock_post.call_args.kwargs["stream"] is True + + +def _mock_gemini_post_response(): + """A minimal stand-in for a successful Gemini generateContent HTTP response.""" + from unittest.mock import MagicMock + + resp = MagicMock() + resp.status_code = 200 + resp.headers = {} + resp.json.return_value = { + "candidates": [ + { + "content": {"parts": [{"text": "hi"}], "role": "model"}, + "finishReason": "STOP", + } + ] + } + return resp + + +NATIVE_TOP_LEVEL_FIELD_CASES = [ + ( + "safetySettings", + [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}], + ), + ("toolConfig", {"functionCallingConfig": {"mode": "AUTO"}}), + ("cachedContent", "cachedContents/abc123"), + ("labels", {"team": "search"}), +] + + +@pytest.mark.parametrize("field_name, field_value", NATIVE_TOP_LEVEL_FIELD_CASES) +def test_native_top_level_field_forwarded_to_request_body(field_name, field_value): + """ + Regression for https://github.com/BerriAI/litellm/issues/12671 + + Google's native generateContent body carries fields like safetySettings at the + top level (siblings of generationConfig). The proxy spreads them as loose kwargs + into generate_content. They must reach Google's request body at the top level and + must NOT be silently dropped nor nested under generationConfig. + """ + from unittest.mock import patch + + from litellm.google_genai.main import generate_content + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + with patch.object( + HTTPHandler, "post", return_value=_mock_gemini_post_response() + ) as mock_post: + generate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + **{field_name: field_value}, + ) + + assert mock_post.called, "expected the request to reach the HTTP client" + body = mock_post.call_args.kwargs["json"] + assert ( + body[field_name] == field_value + ), f"{field_name} should be forwarded to Google at the top level" + assert field_name not in body.get("generationConfig", {}), ( + f"{field_name} must be a top-level sibling of generationConfig, " + "not nested inside it" + ) + + +@pytest.mark.asyncio +async def test_native_safety_settings_forwarded_async(): + """The async path (used by the proxy's :generateContent route) must also forward + native top-level fields.""" + from unittest.mock import AsyncMock, patch + + from litellm.google_genai.main import agenerate_content + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + safety_settings = [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"} + ] + + with patch.object( + AsyncHTTPHandler, + "post", + new_callable=AsyncMock, + return_value=_mock_gemini_post_response(), + ) as mock_post: + await agenerate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=safety_settings, + ) + + assert mock_post.called + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == safety_settings + assert "safetySettings" not in body.get("generationConfig", {}) + + +def test_native_fields_coexist_with_generation_config(): + """Forwarding native top-level fields must not regress the already-working + generationConfig path; both must land in their correct positions.""" + from unittest.mock import patch + + from litellm.google_genai.main import generate_content + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + safety_settings = [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"} + ] + + with patch.object( + HTTPHandler, "post", return_value=_mock_gemini_post_response() + ) as mock_post: + generate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=safety_settings, + generationConfig={"temperature": 0, "responseMimeType": "application/json"}, + ) + + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == safety_settings + generation_config = body["generationConfig"] + assert generation_config["temperature"] == 0 + assert generation_config["responseMimeType"] == "application/json" + + +def test_explicit_extra_body_overrides_native_top_level_field(): + """An explicit extra_body value takes precedence over the same top-level field.""" + from unittest.mock import patch + + from litellm.google_genai.main import generate_content + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + native = [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}] + override = [ + {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"} + ] + + with patch.object( + HTTPHandler, "post", return_value=_mock_gemini_post_response() + ) as mock_post: + generate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=native, + extra_body={"safetySettings": override}, + ) + + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == override + + +def test_native_fields_and_system_instruction_forwarded_on_sync_stream(): + """The sync streaming path (generate_content_stream) must forward native top-level + fields AND systemInstruction. The PR changed the merge here and newly added the + systemInstruction kwarg; without coverage a regression on either ships green.""" + from unittest.mock import patch + + from litellm.google_genai.main import generate_content_stream + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + safety_settings = [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"} + ] + system_instruction = {"parts": [{"text": "Be terse"}]} + + with patch.object( + HTTPHandler, "post", return_value=_mock_gemini_post_response() + ) as mock_post: + generate_content_stream( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=safety_settings, + systemInstruction=system_instruction, + ) + + assert mock_post.called + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == safety_settings + assert body["systemInstruction"] == system_instruction + assert "safetySettings" not in body.get("generationConfig", {}) + + +@pytest.mark.asyncio +async def test_native_fields_forwarded_on_async_stream(): + """The async streaming path (agenerate_content_stream) backs the proxy's + :streamGenerateContent route and must forward native top-level fields too.""" + from unittest.mock import AsyncMock, patch + + from litellm.google_genai.main import agenerate_content_stream + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + safety_settings = [ + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"} + ] + + with patch.object( + AsyncHTTPHandler, + "post", + new_callable=AsyncMock, + return_value=_mock_gemini_post_response(), + ) as mock_post: + await agenerate_content_stream( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "Say hi"}]}], + custom_llm_provider="gemini", + api_key="test-key", + safetySettings=safety_settings, + ) + + assert mock_post.called + body = mock_post.call_args.kwargs["json"] + assert body["safetySettings"] == safety_settings + assert "safetySettings" not in body.get("generationConfig", {}) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py new file mode 100644 index 00000000000..453c7490d98 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py @@ -0,0 +1,321 @@ +""" +Regression tests for the Anthropic message_start cursor=1 bug in +ChunkProcessor._calculate_usage_per_chunk. + +Background +---------- +Anthropic streams a `message_start` event that carries +`usage.output_tokens=1` as a placeholder ("cursor"). The real cumulative +output count only arrives in the final `message_delta` event. When a +stream is cancelled before `message_delta` lands (very common for +thinking models on long-tail prompts), the last-wins accumulator in +ChunkProcessor leaves completion_tokens stuck at 1. Because 1 is +truthy, the `completion_tokens or token_counter(text=...)` fallback in +calculate_usage() never fires, and the request is billed for 1 output +token even when several thousand tokens of text were actually streamed. + +These tests pin the post-fix behavior: completion_tokens should reset +to 0 when the only update we saw was the cursor, allowing the +text-based fallback to estimate from the real completion text. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor +from litellm.types.utils import ( + Delta, + ModelResponseStream, + StreamingChoices, + Usage, +) + + +def _make_chunk( + *, + content: str = "", + usage: Usage = None, + finish_reason: str = None, + custom_llm_provider: str = "anthropic", +) -> ModelResponseStream: + chunk = ModelResponseStream( + id="msg_test", + created=1738900000, + model="claude-sonnet-4-6", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=finish_reason, + index=0, + delta=Delta(content=content, role="assistant"), + ) + ], + usage=usage, + ) + # The cursor reset is now gated on provider; populate the same field the + # real streaming_handler sets (see litellm/litellm_core_utils/streaming_handler.py). + chunk._hidden_params = {"custom_llm_provider": custom_llm_provider} + return chunk + + +class TestAnthropicCursorBug: + """The core regression: completion_tokens=1 cursor must not leak through.""" + + def test_only_message_start_cursor_resets_completion_to_zero(self): + """ + Stream cancelled before message_delta — only the message_start cursor + (output_tokens=1) was seen. Per-chunk accumulator must reset to 0 so + token_counter fallback can estimate from completion text. + """ + # Anthropic message_start: input_tokens accurate, output_tokens=1 cursor + message_start = _make_chunk( + usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) + ) + # Several content_block_delta chunks (no usage attached) + text_chunks = [ + _make_chunk(content="Hello"), + _make_chunk(content=" world"), + _make_chunk(content=" this is partial."), + ] + chunks = [message_start, *text_chunks] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["prompt_tokens"] == 1024 + # The cursor value of 1 must NOT leak through — should be reset to 0 + # so the text-based fallback estimates the real completion length. + assert result["completion_tokens"] == 0, ( + "completion_tokens=1 from message_start cursor leaked through. " + "Should reset to 0 when only cursor was seen, so token_counter " + "fallback in calculate_usage() can estimate from completion text." + ) + + def test_message_start_plus_message_delta_uses_delta_value(self): + """ + Normal complete stream: message_start cursor=1, then message_delta=3847. + Last-wins must give 3847 (the real value). + """ + message_start = _make_chunk( + usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) + ) + text_chunks = [_make_chunk(content=t) for t in ["Hello", " world", "!"]] + # message_delta with the real cumulative output_tokens + message_delta = _make_chunk( + usage=Usage(prompt_tokens=1024, completion_tokens=3847, total_tokens=4871), + finish_reason="stop", + ) + chunks = [message_start, *text_chunks, message_delta] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["prompt_tokens"] == 1024 + assert result["completion_tokens"] == 3847 + + def test_calculate_usage_falls_back_to_token_counter_for_cursor_only(self): + """ + End-to-end via calculate_usage(): cursor-only stream + real completion + text should produce a token-counter estimate, NOT 1. + """ + message_start = _make_chunk( + usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) + ) + # ~50 visible chars ≈ ~12 tokens (anthropic-style tokenizer ballpark) + text_chunks = [ + _make_chunk(content="Based on your question, I think the answer is "), + _make_chunk(content="forty-two. Here is my reasoning: "), + ] + chunks = [message_start, *text_chunks] + completion_output = ( + "Based on your question, I think the answer is forty-two. " + "Here is my reasoning: " + ) + + processor = ChunkProcessor(chunks=chunks, messages=[]) + usage = processor.calculate_usage( + chunks=chunks, + model="claude-sonnet-4-6", + completion_output=completion_output, + messages=[], + ) + + # Should be a token_counter estimate of the text, not the cursor 1 + assert usage.completion_tokens > 1, ( + f"Expected token_counter estimate of completion text, got " + f"completion_tokens={usage.completion_tokens} (likely stuck at cursor)" + ) + + def test_cache_fields_preserved_from_message_start(self): + """cache_read / cache_creation come from message_start and must survive.""" + message_start_usage = Usage( + prompt_tokens=1024, completion_tokens=1, total_tokens=1025 + ) + # Anthropic puts these in message_start + message_start_usage.cache_read_input_tokens = 512 + message_start_usage.cache_creation_input_tokens = 128 + message_start = _make_chunk(usage=message_start_usage) + + chunks = [message_start, _make_chunk(content="hi")] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["cache_read_input_tokens"] == 512 + assert result["cache_creation_input_tokens"] == 128 + + def test_openai_streaming_unaffected(self): + """ + OpenAI's only usage chunk is the penultimate one (with + stream_options.include_usage=true), and it carries the real value + directly. Our cursor fix must not break this path — output > 1 + means saw_non_cursor_completion=True so no reset happens. + """ + # Simulate OpenAI: content chunks first, then ONE usage chunk at the end + text_chunks = [_make_chunk(content=t) for t in ["The", " answer", " is 42"]] + usage_chunk = _make_chunk( + usage=Usage(prompt_tokens=42, completion_tokens=15, total_tokens=57), + finish_reason="stop", + ) + chunks = [*text_chunks, usage_chunk] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["prompt_tokens"] == 42 + assert result["completion_tokens"] == 15 + + def test_single_token_completion_legitimate_case(self): + """ + Edge case: a stream that legitimately completes with output_tokens=1 + (e.g., model returns just "Yes."). Without saw_non_cursor_completion + we'd reset to 0 and fall through to token_counter — but token_counter + on a 1-token string also gives ~1, so billing is still approximately + correct. This test pins that the result is sane (1 or 0). + """ + message_start = _make_chunk( + usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21) + ) + text_chunk = _make_chunk(content="Yes.") + # Anthropic's message_delta also gives output_tokens=1 in this case + message_delta = _make_chunk( + usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21), + finish_reason="stop", + ) + chunks = [message_start, text_chunk, message_delta] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + usage = processor.calculate_usage( + chunks=chunks, + model="claude-sonnet-4-6", + completion_output="Yes.", + messages=[], + ) + + # Two completion-bearing usage events (message_start AND message_delta + # both with output_tokens=1) is positive evidence that message_delta + # arrived — saw_non_cursor_completion goes True via the count >= 2 + # branch and the reset is suppressed. Result: completion_tokens stays + # at the legitimate value of 1. + assert usage.completion_tokens == 1, ( + f"Legitimate single-token completion should bill exactly 1 token " + f"(message_start + message_delta both saw output_tokens=1, " + f"confirming message_delta arrived), got {usage.completion_tokens}" + ) + + def test_anthropic_cache_only_chunks_after_message_start_still_resets(self): + """ + Cache-only chunks (cache_read_input_tokens > 0 but completion_tokens=0) + following message_start should not be mistaken for completion progress. + The cursor=1 from message_start stays the only completion update; reset + must fire so token_counter estimates from completion text instead of + billing the placeholder. + """ + message_start_usage = Usage( + prompt_tokens=1024, completion_tokens=1, total_tokens=1025 + ) + message_start_usage.cache_read_input_tokens = 4096 + message_start = _make_chunk(usage=message_start_usage) + # Subsequent chunks with cache fields but no completion_tokens + cache_chunk_usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + cache_chunk_usage.cache_read_input_tokens = 4096 + cache_chunk = _make_chunk(content="partial", usage=cache_chunk_usage) + # No message_delta — stream was cancelled + chunks = [message_start, cache_chunk] + + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + + assert result["cache_read_input_tokens"] == 4096 + assert result["completion_tokens"] == 0, ( + "cache chunks alone don't count as completion progress — only " + "completion_tokens > 0 in a usage event proves real output happened. " + "Reset to 0 forces token_counter fallback." + ) + + +class TestProviderGuard: + """Class A: the cursor-reset heuristic must NOT silently affect non-Anthropic + providers, even if they happen to report completion_tokens=1.""" + + def test_non_anthropic_provider_completion_tokens_one_not_reset(self): + """ + Some non-Anthropic provider legitimately reports completion_tokens=1 + in its single usage chunk. Without the provider guard the cursor + heuristic would silently reset it to 0 and bill via token_counter, + producing a different (often inflated) number than what the provider + actually charged. + """ + chunks = [ + _make_chunk( + usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + finish_reason="stop", + custom_llm_provider="openai", + ), + ] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["completion_tokens"] == 1, ( + "Non-Anthropic providers must not be subject to the message_start " + "cursor reset — their completion_tokens=1 is the real value." + ) + + def test_unknown_provider_completion_tokens_one_not_reset(self): + """No custom_llm_provider on hidden_params (older path or custom + plugin) — heuristic must not fire.""" + chunk = _make_chunk( + usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + ) + # Explicitly clear hidden_params to simulate the unknown-provider case + chunk._hidden_params = {} + processor = ChunkProcessor(chunks=[chunk], messages=[]) + result = processor._calculate_usage_per_chunk(chunks=[chunk]) + assert result["completion_tokens"] == 1 + + +class TestNonAnthropicStreamingIntact: + """Make sure providers without cursor pattern still work.""" + + def test_completion_tokens_above_one_never_resets(self): + """Any chunk reporting completion_tokens > 1 sets saw_non_cursor + and prevents the reset.""" + chunks = [ + _make_chunk( + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + ), + ] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["completion_tokens"] == 5 + + def test_no_usage_chunks_leaves_zero(self): + """Stream with zero usage info → completion_tokens stays 0 + (token_counter fallback will handle it).""" + chunks = [_make_chunk(content="hi"), _make_chunk(content=" there")] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["prompt_tokens"] == 0 + assert result["completion_tokens"] == 0 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index f93f2404fa4..45820b9833f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -2175,6 +2175,283 @@ def test_translate_openai_response_to_anthropic_cache_tokens_from_prompt_tokens_ assert anthropic_response["usage"]["cache_read_input_tokens"] == 30 +def test_translate_openai_usage_to_anthropic_cache_tokens_from_dict_details_with_integral_floats(): + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + ) + usage.prompt_tokens_details = { + "cached_tokens": 30.0, + "cache_write_tokens": 20.0, + } + + anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( + usage + ) + + assert anthropic_usage["input_tokens"] == 70 + assert anthropic_usage["output_tokens"] == 50 + assert anthropic_usage["cache_read_input_tokens"] == 30 + assert anthropic_usage["cache_creation_input_tokens"] == 20 + + +def test_translate_openai_usage_to_anthropic_ignores_fractional_cache_tokens(): + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + ) + usage.prompt_tokens_details = { + "cached_tokens": 30.5, + "cache_creation_tokens": 20.25, + } + + anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( + usage + ) + + assert anthropic_usage["input_tokens"] == 120 + assert anthropic_usage["output_tokens"] == 50 + assert "cache_read_input_tokens" not in anthropic_usage + assert "cache_creation_input_tokens" not in anthropic_usage + + +def test_translate_openai_usage_to_anthropic_ignores_bool_cache_tokens(): + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + ) + usage.cache_read_input_tokens = True + usage.cache_creation_input_tokens = True + + anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta( + usage + ) + + assert anthropic_usage["input_tokens"] == 120 + assert anthropic_usage["output_tokens"] == 50 + assert "cache_read_input_tokens" not in anthropic_usage + assert "cache_creation_input_tokens" not in anthropic_usage + + +def test_translate_openai_response_to_anthropic_cache_creation_from_prompt_tokens_details(): + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message( + role="assistant", + content="Test response", + ), + ) + ], + model="gpt-4o-2024-08-06", + usage=usage, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_response = adapter.translate_openai_response_to_anthropic( + response=response, + tool_name_mapping=None, + ) + + assert anthropic_response["usage"]["input_tokens"] == 70 + assert anthropic_response["usage"]["output_tokens"] == 50 + assert anthropic_response["usage"]["cache_read_input_tokens"] == 30 + assert anthropic_response["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_openai_response_to_anthropic_cache_tokens_from_usage_fields(): + usage = Usage(prompt_tokens=120, completion_tokens=50, total_tokens=170) + usage.cache_read_input_tokens = 30 + usage.cache_creation_input_tokens = 20 + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message( + role="assistant", + content="Test response", + ), + ) + ], + model="claude-3-sonnet-20240229", + usage=usage, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_response = adapter.translate_openai_response_to_anthropic( + response=response, + tool_name_mapping=None, + ) + + assert anthropic_response["usage"]["input_tokens"] == 70 + assert anthropic_response["usage"]["output_tokens"] == 50 + assert anthropic_response["usage"]["cache_read_input_tokens"] == 30 + assert anthropic_response["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_openai_response_to_anthropic_cache_tokens_from_private_usage_fields(): + usage = Usage(prompt_tokens=120, completion_tokens=50, total_tokens=170) + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message( + role="assistant", + content="Test response", + ), + ) + ], + model="claude-3-sonnet-20240229", + usage=usage, + ) + response.usage._cache_read_input_tokens = 30 + response.usage._cache_creation_input_tokens = 20 + + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_response = adapter.translate_openai_response_to_anthropic( + response=response, + tool_name_mapping=None, + ) + + assert anthropic_response["usage"]["input_tokens"] == 70 + assert anthropic_response["usage"]["output_tokens"] == 50 + assert anthropic_response["usage"]["cache_read_input_tokens"] == 30 + assert anthropic_response["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_streaming_openai_response_to_anthropic_cache_tokens_from_prompt_tokens_details(): + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + response = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason="stop", + ) + ], + usage=usage, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + message_delta = adapter.translate_streaming_openai_response_to_anthropic( + response=response, + current_content_block_index=0, + ) + + assert message_delta["usage"]["input_tokens"] == 70 + assert message_delta["usage"]["output_tokens"] == 50 + assert message_delta["usage"]["cache_read_input_tokens"] == 30 + assert message_delta["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_streaming_openai_response_to_anthropic_cache_tokens_from_hidden_params_usage(): + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + response = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason="stop", + ) + ], + ) + response._hidden_params = {"usage": usage} + + adapter = LiteLLMAnthropicMessagesAdapter() + message_delta = adapter.translate_streaming_openai_response_to_anthropic( + response=response, + current_content_block_index=0, + ) + + assert message_delta["usage"]["input_tokens"] == 70 + assert message_delta["usage"]["output_tokens"] == 50 + assert message_delta["usage"]["cache_read_input_tokens"] == 30 + assert message_delta["usage"]["cache_creation_input_tokens"] == 20 + + +def test_translate_streaming_openai_response_to_anthropic_cache_tokens_with_applied_edits(): + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + response = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason="stop", + ) + ], + usage=usage, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + message_delta = adapter.translate_streaming_openai_response_to_anthropic( + response=response, + current_content_block_index=0, + applied_edits=[{"type": "compact_20260112"}], + ) + + assert message_delta["usage"]["input_tokens"] == 70 + assert message_delta["usage"]["output_tokens"] == 50 + assert message_delta["usage"]["cache_read_input_tokens"] == 30 + assert message_delta["usage"]["cache_creation_input_tokens"] == 20 + assert message_delta["context_management"]["applied_edits"][0]["type"] == ( + "compact_20260112" + ) + + # ===================================================================== # Web Search Tool Transformation Tests # ===================================================================== diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py index f74c5b61300..d67de0dcaf8 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -24,6 +24,7 @@ from litellm.types.utils import ( Message, ModelResponse, ModelResponseStream, + PromptTokensDetailsWrapper, StreamingChoices, Usage, ) @@ -88,6 +89,59 @@ def test_fake_stream_usage_preserved(): assert message_delta["usage"]["input_tokens"] == 10 +def test_delayed_usage_chunk_preserves_cache_tokens(): + usage = Usage( + prompt_tokens=120, + completion_tokens=5, + total_tokens=125, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + chunks = [ + ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Two."), + finish_reason=None, + ) + ], + ), + ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason="stop", + ) + ], + ), + ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(), + finish_reason=None, + ) + ], + usage=usage, + ), + ] + wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="gpt-4o") + events = list(wrapper) + + message_delta = next( + event for event in events if event.get("type") == "message_delta" + ) + + assert message_delta["usage"]["input_tokens"] == 70 + assert message_delta["usage"]["output_tokens"] == 5 + assert message_delta["usage"]["cache_read_input_tokens"] == 30 + assert message_delta["usage"]["cache_creation_input_tokens"] == 20 + + def test_splitter_passes_through_non_combined_chunks(): """A chunk with content but no finish_reason is not split.""" chunk = ModelResponseStream( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py index 8bc39a6d85e..d9d9474d6f3 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py @@ -30,7 +30,9 @@ from litellm.types.utils import ( ChatCompletionDeltaToolCall, Delta, Function, + PromptTokensDetailsWrapper, StreamingChoices, + Usage, ) @@ -107,6 +109,34 @@ def _input_json_deltas(events: List[dict]) -> List[str]: ] +def test_held_stop_reason_usage_merge_preserves_openai_cache_token_details(): + """OpenAI-compatible usage chunks carry cache reads in prompt_tokens_details.""" + wrapper = AnthropicStreamWrapper(completion_stream=iter([]), model="claude-x") + wrapper.holding_stop_reason_chunk = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"input_tokens": 0, "output_tokens": 0}, + } + + usage_chunk = MagicMock() + usage_chunk.usage = Usage( + prompt_tokens=120, + completion_tokens=50, + total_tokens=170, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30, + cache_creation_tokens=20, + ), + ) + + merged_chunk = wrapper._merge_usage_into_held_stop_reason_chunk(usage_chunk) + + assert merged_chunk["usage"]["input_tokens"] == 70 + assert merged_chunk["usage"]["output_tokens"] == 50 + assert merged_chunk["usage"]["cache_read_input_tokens"] == 30 + assert merged_chunk["usage"]["cache_creation_input_tokens"] == 20 + + def test_first_text_delta_after_tool_use_is_not_dropped_sync(): """A tool_use -> text transition (text resuming after a tool call) carries the resumed text's first token in the trigger chunk. Without the fix it was diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 03b72c504b8..b6c055f4390 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -159,6 +159,59 @@ class TestGetAnthropicHeaders: assert "authorization" not in headers assert "anthropic-dangerous-direct-browser-access" not in headers + def test_custom_api_base_uses_bearer_header(self): + """Custom api_base and non-standard API key should produce Authorization: Bearer header when opted in.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = config.get_anthropic_headers( + api_key="my-custom-ollama-token", + computer_tool_used=False, + prompt_caching_set=False, + pdf_used=False, + is_vertex_request=False, + api_base="https://ollama.com/", + use_bearer_for_custom_base=True, + ) + + assert headers["authorization"] == "Bearer my-custom-ollama-token" + assert "x-api-key" not in headers + + def test_custom_api_base_uses_bearer_header_already_starts_with_bearer(self): + """If the key already starts with Bearer and Bearer opt-in is enabled, use it directly.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = config.get_anthropic_headers( + api_key="Bearer my-custom-ollama-token", + computer_tool_used=False, + prompt_caching_set=False, + pdf_used=False, + is_vertex_request=False, + api_base="https://ollama.com/", + use_bearer_for_custom_base=True, + ) + + assert headers["authorization"] == "Bearer my-custom-ollama-token" + assert "x-api-key" not in headers + + def test_custom_api_base_uses_x_api_key_when_standard_key(self): + """If the key is standard sk-ant- key, use x-api-key even with custom api_base.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = config.get_anthropic_headers( + api_key=FAKE_REGULAR_KEY, + computer_tool_used=False, + prompt_caching_set=False, + pdf_used=False, + is_vertex_request=False, + api_base="https://ollama.com/", + ) + + assert headers["x-api-key"] == FAKE_REGULAR_KEY + assert "authorization" not in headers + def test_oauth_includes_standard_headers(self): """OAuth path should still include standard Anthropic headers.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -242,6 +295,46 @@ class TestValidateEnvironmentOAuth: assert updated_headers["x-api-key"] == FAKE_REGULAR_KEY assert "authorization" not in updated_headers + + def test_custom_api_base_via_param(self): + """validate_environment uses Bearer when use_bearer_for_custom_base is set in litellm_params.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = {} + + updated_headers = config.validate_environment( + headers=headers, + model="claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"use_bearer_for_custom_base": True}, + api_key="custom-api-key", + api_base="https://custom-gateway.com", + ) + + assert updated_headers["authorization"] == "Bearer custom-api-key" + assert "x-api-key" not in updated_headers + + def test_custom_api_base_via_litellm_params(self): + """validate_environment uses Bearer when api_base and use_bearer_for_custom_base are in litellm_params.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = {} + + updated_headers = config.validate_environment( + headers=headers, + model="claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"api_base": "https://custom-gateway.com", "use_bearer_for_custom_base": True}, + api_key="custom-api-key", + api_base=None, + ) + + assert updated_headers["authorization"] == "Bearer custom-api-key" + assert "x-api-key" not in updated_headers assert "anthropic-dangerous-direct-browser-access" not in updated_headers @@ -1004,6 +1097,20 @@ class TestGetAuthHeader: result = AnthropicModelInfo.get_auth_header() assert result == {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} + def test_custom_api_base_get_auth_header_uses_bearer(self): + """Non-standard API key and custom api_base returns Bearer when use_bearer_for_custom_base=True.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + result = AnthropicModelInfo.get_auth_header(api_key="my-custom-key", api_base="https://custom-gateway.com", use_bearer_for_custom_base=True) + assert result == {"authorization": "Bearer my-custom-key"} + + def test_custom_api_base_get_auth_header_uses_x_api_key_when_standard(self): + """Standard sk-ant- key with custom api_base should still return x-api-key.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + result = AnthropicModelInfo.get_auth_header(api_key=FAKE_REGULAR_KEY, api_base="https://custom-gateway.com") + assert result == {"x-api-key": FAKE_REGULAR_KEY} + class TestGetApiBaseFallbackChain: """Tests for AnthropicModelInfo.get_api_base() fallback to ANTHROPIC_BASE_URL.""" diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 6298eeb25e9..6b869076044 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -263,3 +263,123 @@ def test_bundled_bedrock_opus_model_info_declares_output_config_effort_ceiling( model_info = GetModelCostMap.load_local_model_cost_map()[model] assert model_info["bedrock_output_config_effort_ceiling"] == expected_ceiling + + +def test_route_prefix_matched_as_path_segment_not_substring(): + """Route tokens like ``mantle/`` must match only at a path-segment boundary. + + The ``bedrock_mantle/`` provider prefix contains the substring ``mantle/``; + a substring match misroutes ``bedrock_mantle/openai.gpt-5.5`` to the Claude + Mythos mantle config, whose request transform strips ``mantle/`` and mangles + the body model into ``bedrock_openai.gpt-5.5``. These assertions fail under + the old substring matching and pass once matching is anchored to ``startswith`` + or a ``/`` boundary. + """ + # The bedrock_mantle/ provider prefix must NOT be read as the mantle/ route. + assert ( + BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.5") != "mantle" + ) + assert ( + BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.4") == "invoke" + ) + assert ( + BedrockModelInfo._explicit_mantle_route("bedrock_mantle/openai.gpt-5.5") + is False + ) + + # A genuine mantle route still resolves, via the startswith branch... + assert ( + BedrockModelInfo.get_bedrock_route("mantle/anthropic.claude-mythos-preview") + == "mantle" + ) + # ...and via the mid-path "/mantle/" branch (after the bedrock/ provider prefix). + assert ( + BedrockModelInfo.get_bedrock_route( + "bedrock/mantle/anthropic.claude-mythos-preview" + ) + == "mantle" + ) + + +def test_model_has_route_prefix_exercises_both_branches(): + """``_model_has_route_prefix`` matches on ``startswith`` or a ``/`` boundary only.""" + # startswith branch + assert ( + BedrockModelInfo._model_has_route_prefix( + "mantle/anthropic.claude-mythos-preview", "mantle/" + ) + is True + ) + # f"/{prefix}" boundary branch + assert ( + BedrockModelInfo._model_has_route_prefix( + "bedrock/mantle/anthropic.claude-mythos-preview", "mantle/" + ) + is True + ) + # neither branch: the token only appears glued to another segment + assert ( + BedrockModelInfo._model_has_route_prefix( + "bedrock_mantle/openai.gpt-5.5", "mantle/" + ) + is False + ) + + +@pytest.mark.parametrize( + "route_method, token", + [ + (BedrockModelInfo._explicit_converse_route, "converse"), + (BedrockModelInfo._explicit_converse_like_route, "converse_like"), + (BedrockModelInfo._explicit_invoke_route, "invoke"), + (BedrockModelInfo._explicit_async_invoke_route, "async_invoke"), + (BedrockModelInfo._explicit_agent_route, "agent"), + (BedrockModelInfo._explicit_agentcore_route, "agentcore"), + (BedrockModelInfo._explicit_claude_platform_route, "claude_platform"), + (BedrockModelInfo._explicit_openai_route, "openai"), + ], + ids=[ + "converse", + "converse_like", + "invoke", + "async_invoke", + "agent", + "agentcore", + "claude_platform", + "openai", + ], +) +def test_explicit_route_helpers_match_token_only_as_path_segment(route_method, token): + """Each migrated ``_explicit_*_route`` matches its token only as a path segment. + + A leading segment (start of the id or right after a ``/``) matches; the token + glued onto a preceding segment does not. Reverting any method to the old + ``"/" in model`` substring check makes the non-segment case return True + and fails this test. + """ + # leading-segment forms match + assert route_method(f"{token}/some-model") is True + assert route_method(f"bedrock/{token}/some-model") is True + # the token only as a non-segment substring must not match + assert route_method(f"x{token}/y") is False + + +def test_explicit_invoke_route_does_not_match_async_invoke(): + """``invoke/`` must not substring-match ``async_invoke/`` models. + + This is the concrete improvement of the segment-boundary migration: the old + ``"invoke/" in model`` check wrongly classified async-invoke models as the + invoke route. + """ + async_invoke_model = "async_invoke/twelvelabs.marengo-embed-2-7-v1:0" + assert BedrockModelInfo._explicit_invoke_route(async_invoke_model) is False + assert ( + BedrockModelInfo._explicit_invoke_route(f"bedrock/{async_invoke_model}") + is False + ) + # ...while async_invoke/ is still detected as its own route. + assert BedrockModelInfo._explicit_async_invoke_route(async_invoke_model) is True + assert ( + BedrockModelInfo._explicit_async_invoke_route(f"bedrock/{async_invoke_model}") + is True + ) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py index 6d913ad5d1d..1422531edbf 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_context_circulation.py @@ -22,7 +22,6 @@ from litellm.llms.vertex_ai.gemini.transformation import ( ) from litellm.types.llms.vertex_ai import HttpxPartType - # --- Response extraction tests --- @@ -63,6 +62,7 @@ class TestExtractServerSideToolInvocations: assert result[0]["args"] == {"queries": ["weather Buenos Aires"]} assert result[0]["response"] == {"weather": "Sunny, 20°C"} assert result[0]["thought_signature"] == "sig_call_1" + assert result[0]["response_thought_signature"] == "sig_resp_1" def test_returns_none_when_no_server_side_tools(self): """No toolCall/toolResponse parts → returns None.""" @@ -145,6 +145,59 @@ class TestExtractServerSideToolInvocations: assert result[0]["id"] == "exec1" assert "response" not in result[0] + def test_extracts_tool_call_and_response_with_different_signatures(self): + """Case where toolCall and toolResponse have different signatures.""" + parts: List[HttpxPartType] = [ + { + "thoughtSignature": "sig_call_1", + "toolCall": { + "toolType": "GOOGLE_SEARCH_WEB", + "id": "abc123", + "args": {"queries": ["weather Buenos Aires"]}, + }, + }, + { + "thoughtSignature": "sig_resp_1", + "toolResponse": { + "toolType": "GOOGLE_SEARCH_WEB", + "id": "abc123", + "response": {"weather": "Sunny, 20°C"}, + }, + }, + ] + + result = VertexGeminiConfig._extract_server_side_tool_invocations(parts) + + assert result is not None + assert len(result) == 1 + assert result[0]["tool_type"] == "GOOGLE_SEARCH_WEB" + assert result[0]["id"] == "abc123" + assert result[0]["args"] == {"queries": ["weather Buenos Aires"]} + assert result[0]["response"] == {"weather": "Sunny, 20°C"} + assert result[0]["thought_signature"] == "sig_call_1" + assert result[0]["response_thought_signature"] == "sig_resp_1" + + def test_orphan_response_signature_extraction(self): + """Orphan toolResponse is captured and has response_thought_signature set.""" + parts: List[HttpxPartType] = [ + { + "thoughtSignature": "sig_resp_1", + "toolResponse": { + "toolType": "GOOGLE_SEARCH_WEB", + "id": "orphan123", + "response": {"result": "some response"}, + }, + }, + ] + + result = VertexGeminiConfig._extract_server_side_tool_invocations(parts) + + assert result is not None + assert len(result) == 1 + assert result[0]["id"] == "orphan123" + assert result[0]["response"] == {"result": "some response"} + assert result[0]["response_thought_signature"] == "sig_resp_1" + # --- Input re-injection tests --- @@ -200,6 +253,76 @@ class TestReInjectServerSideToolInvocations: "weather": "Sunny, 20°C" } + def test_roundtrip_single_invocation_with_different_signatures(self): + """Server-side invocations with different signatures for call and response.""" + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": "It's sunny in Buenos Aires.", + "provider_specific_fields": { + "server_side_tool_invocations": [ + { + "tool_type": "GOOGLE_SEARCH_WEB", + "id": "abc123", + "args": {"queries": ["weather Buenos Aires"]}, + "response": {"weather": "Sunny, 20°C"}, + "thought_signature": "sig_call", + "response_thought_signature": "sig_resp", + } + ] + }, + }, + {"role": "user", "content": "Thanks!"}, + ] + + contents = _gemini_convert_messages_with_history(messages) + + model_turn = [c for c in contents if c["role"] == "model"] + assert len(model_turn) == 1 + + parts = model_turn[0]["parts"] + tool_call_parts = [p for p in parts if "toolCall" in p] + tool_response_parts = [p for p in parts if "toolResponse" in p] + + assert len(tool_call_parts) == 1 + assert tool_call_parts[0]["thoughtSignature"] == "sig_call" + + assert len(tool_response_parts) == 1 + assert tool_response_parts[0]["thoughtSignature"] == "sig_resp" + + def test_roundtrip_orphan_response_signature(self): + """Orphan response signature is preserved and re-injected into toolResponse part.""" + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": "It's sunny in Buenos Aires.", + "provider_specific_fields": { + "server_side_tool_invocations": [ + { + "tool_type": "GOOGLE_SEARCH_WEB", + "id": "orphan123", + "response": {"result": "some response"}, + "response_thought_signature": "sig_orphan_resp", + } + ] + }, + }, + {"role": "user", "content": "Thanks!"}, + ] + + contents = _gemini_convert_messages_with_history(messages) + + model_turn = [c for c in contents if c["role"] == "model"] + assert len(model_turn) == 1 + + parts = model_turn[0]["parts"] + tool_response_parts = [p for p in parts if "toolResponse" in p] + + assert len(tool_response_parts) == 1 + assert tool_response_parts[0]["thoughtSignature"] == "sig_orphan_resp" + def test_no_invocations_no_extra_parts(self): """Without server_side_tool_invocations, no extra parts are added.""" messages = [ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 47c9396f121..27f9a311250 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1800,3 +1800,71 @@ class TestConnectionErrorMessage: message = rest_endpoints._connection_error_message(RuntimeError("weird")) assert "weird" not in message assert "proxy logs" in message.lower() + + +class TestToolResponseMcpInfoEnrichment: + """The REST tools/list response must expose the user-facing alias and the + server_id alongside the internal server_name so clients (agent builder UIs) + can map the internal config key to a friendly name without needing the + mcp_routes-gated server listing. + """ + + def test_enriches_mcp_info_with_alias_and_server_id(self): + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + server = MCPServer( + server_id="a1b2c3d4", + name="mcpAtlassian", + alias="atlassian", + server_name="mcpAtlassian", + transport=MCPTransport.http, + mcp_info={"server_name": "mcpAtlassian"}, + ) + tools = [ + MCPTool( + name="get_issue", + description="Fetch a Jira issue", + inputSchema={"type": "object"}, + ) + ] + + result = rest_endpoints._create_tool_response_objects(tools, server) + + assert result[0].mcp_info == { + "server_name": "mcpAtlassian", + "server_id": "a1b2c3d4", + "alias": "atlassian", + } + + def test_alias_none_is_explicit_in_mcp_info(self): + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + server = MCPServer( + server_id="server-uuid", + name="no_alias_server", + alias=None, + server_name="no_alias_server", + transport=MCPTransport.http, + mcp_info={"server_name": "no_alias_server"}, + ) + tools = [ + MCPTool( + name="ping", + description="Ping", + inputSchema={"type": "object"}, + ) + ] + + result = rest_endpoints._create_tool_response_objects(tools, server) + + assert result[0].mcp_info == { + "server_name": "no_alias_server", + "server_id": "server-uuid", + "alias": None, + } diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index 7a8d04507dc..dd9cbcf5232 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -570,3 +570,41 @@ class TestRequireCallerUserIdForNonAdmin: assert exc_info.value.status_code == 403 assert "Service-account keys" in str(exc_info.value.detail) + + +class TestValidateFiniteSpend: + """`validate_finite_spend` rejects NaN/±inf so a non-finite spend cannot + bypass `spend >= max_budget` enforcement (NaN/-inf compare false).""" + + def test_none_is_allowed(self): + from litellm.proxy.management_endpoints.common_utils import ( + validate_finite_spend, + ) + + assert validate_finite_spend(None) is None + + def test_finite_value_is_allowed(self): + from litellm.proxy.management_endpoints.common_utils import ( + validate_finite_spend, + ) + + assert validate_finite_spend(0.0) is None + assert validate_finite_spend(12.5) is None + # Negative spend is intentionally allowed. Admins may set a negative + # spend counter to grant an entity extra allowance for the current + # budget period only (e.g. a large one-time spend grant), effectively + # raising their headroom without raising the recurring budget ceiling. + # Future changes should continue to allow negative spend counters. + assert validate_finite_spend(-50.0) is None + + @pytest.mark.parametrize("bad", [float("nan"), float("inf"), float("-inf")]) + def test_non_finite_is_rejected(self, bad): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.common_utils import ( + validate_finite_spend, + ) + + with pytest.raises(HTTPException) as exc_info: + validate_finite_spend(bad) + assert exc_info.value.status_code == 400 diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index b4602e0ad8b..27e82df90c1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -3102,6 +3102,86 @@ async def test_ghsa_wvg4_proxy_admin_can_update_user_budget(mocker): assert result is not None +@pytest.mark.asyncio +async def test_admin_user_update_spend_invalidates_counter(mocker): + """A direct /user/update spend change must invalidate the cross-pod + spend counter so enforcement re-reads the new DB value.""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + mock_prisma_client = mocker.MagicMock() + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "target-user", "spend": 50.0} + existing_user.user_id = "target-user" + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( + return_value=existing_user + ) + mock_prisma_client.update_data = mocker.AsyncMock( + return_value={"user_id": "target-user", "spend": -25.0} + ) + mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + mock_invalidate = mocker.patch( + "litellm.proxy.proxy_server._invalidate_spend_counter", + new=mocker.AsyncMock(), + ) + + # Use a negative spend: this also implicitly validates that negative spend + # is allowed, which is desirable. Admins may grant an entity extra + # allowance for the current budget period only (a one-time spend grant) + # without raising the recurring budget ceiling. Future changes should + # continue allowing negative spend counters. + user_request = UpdateUserRequest(user_id="target-user", spend=-25) + admin_caller = UserAPIKeyAuth( + user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + await _update_single_user_helper( + user_request=user_request, user_api_key_dict=admin_caller + ) + mock_invalidate.assert_awaited_once_with(counter_key="spend:user:target-user") + + +@pytest.mark.asyncio +async def test_user_update_rejects_non_finite_spend(mocker): + """NaN/inf spend is rejected before any DB write or counter invalidation.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + mock_prisma_client = mocker.MagicMock() + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "target-user", "spend": 50.0} + existing_user.user_id = "target-user" + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( + return_value=existing_user + ) + mock_prisma_client.update_data = mocker.AsyncMock() + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + mock_invalidate = mocker.patch( + "litellm.proxy.proxy_server._invalidate_spend_counter", + new=mocker.AsyncMock(), + ) + + user_request = UpdateUserRequest(user_id="target-user", spend=float("nan")) + admin_caller = UserAPIKeyAuth( + user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + with pytest.raises(HTTPException) as exc: + await _update_single_user_helper( + user_request=user_request, user_api_key_dict=admin_caller + ) + assert exc.value.status_code == 400 + mock_prisma_client.update_data.assert_not_called() + mock_invalidate.assert_not_awaited() + + @pytest.mark.asyncio async def test_resolve_user_email_metadata_maps_page_user_ids_to_email(mocker): """Regression for LIT-3889. diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index b8ec8a8a388..97397fb06be 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6639,9 +6639,9 @@ async def test_reset_key_spend_success(monkeypatch): @pytest.mark.asyncio -async def test_update_key_spend_invalidates_counter(monkeypatch): +async def test_update_key_spend_updates_counter(monkeypatch): """ - Test that updating a key's spend via update_key_fn immediately invalidates the spend counter. + Test that updating a key's spend via update_key_fn immediately updates the spend counter. """ from litellm.proxy.management_endpoints.key_management_endpoints import ( update_key_fn, @@ -6676,15 +6676,17 @@ async def test_update_key_spend_invalidates_counter(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) monkeypatch.setattr("litellm.store_audit_logs", False) + mock_spend_counter_cache = MagicMock() + mock_spend_counter_cache.redis_cache = MagicMock() + mock_spend_counter_cache.redis_cache.async_set_cache = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.spend_counter_cache", + mock_spend_counter_cache, + ) - with ( - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" - ) as mock_delete_cache, - patch( - "litellm.proxy.proxy_server._invalidate_spend_counter" - ) as mock_invalidate, - ): + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" + ) as mock_delete_cache: mock_delete_cache.return_value = None user_api_key_dict = UserAPIKeyAuth( @@ -6704,7 +6706,12 @@ async def test_update_key_spend_invalidates_counter(monkeypatch): ) mock_delete_cache.assert_awaited_once() - mock_invalidate.assert_awaited_once_with(counter_key=f"spend:key:{hashed_key}") + mock_spend_counter_cache.in_memory_cache.set_cache.assert_called_once_with( + key=f"spend:key:{hashed_key}", value=0.0, ttl=60 + ) + mock_spend_counter_cache.redis_cache.async_set_cache.assert_awaited_once_with( + key=f"spend:key:{hashed_key}", value=0.0, ttl=60 + ) @pytest.mark.asyncio @@ -9718,6 +9725,39 @@ class TestKeyOwnerPrivilegeEscalation: assert exc_info.value.status_code == 403 mock_check.assert_called_once() + @pytest.mark.asyncio + async def test_creator_cannot_reset_own_spend_to_stale_value(self): + """Submitting `spend` equal to the stale DB value must still require + admin. The DB spend lags the live cross-pod counter, so an + "unchanged" spend on the non-admin path would let the creator + overwrite the live counter below real usage. Any explicit `spend` + is a budget change, regardless of value match.""" + existing = self._make_existing_key(created_by="creator-123") + existing.spend = 0.0 + # spend equals the stale DB value (0.0) — the old `!=` gate skipped + # the admin check here. + data = UpdateKeyRequest(key="sk-test", spend=0.0) + auth = self._make_auth(user_id="creator-123") + + mock_check = AsyncMock( + side_effect=HTTPException(status_code=403, detail="Not authorized") + ) + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access", + mock_check, + ): + with pytest.raises(HTTPException): + await _validate_update_key_data( + data=data, + existing_key_row=existing, + user_api_key_dict=auth, + llm_router=None, + premium_user=False, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + mock_check.assert_called_once() + @pytest.mark.asyncio async def test_assigned_user_blocked_from_model_escalation(self): data = UpdateKeyRequest(key="sk-test", models=["gpt-4", "claude-opus"]) diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 592232f45f5..196ff045208 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -10,8 +10,8 @@ from __future__ import annotations import os from types import SimpleNamespace -from typing import Any, Dict, List, Optional -from unittest.mock import AsyncMock, MagicMock, patch +from typing import Any, Dict +from unittest.mock import AsyncMock, MagicMock import pytest @@ -887,6 +887,84 @@ def test_ProxyConfig__add_deployment_invalid_litellm_params_skips(monkeypatch): assert pc._add_deployment(db_models=[bad]) == 0 +def test_ProxyConfig__add_deployment_resolves_env_refs_after_db_decrypt(monkeypatch): + monkeypatch.setenv("LITELLM_DB_MODEL_API_KEY", "resolved-secret") + monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret") + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + lambda value, key, return_original_value: value, + ) + fake_router = MagicMock() + fake_router.upsert_deployment = MagicMock(return_value=True) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router) + pc = ProxyConfig() + db_model = SimpleNamespace( + model_id="model-1", + model_name="env-model", + model_info={"id": "model-1"}, + litellm_params={ + "model": "openai/gpt-4o-mini", + "api_key": "os.environ/LITELLM_DB_MODEL_API_KEY", + "api_base": "os.environ/LITELLM_MASTER_KEY", + }, + blocked=False, + ) + + added = pc._add_deployment(db_models=[db_model]) + deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"] + + assert added == 1 + assert deployment.litellm_params.api_key == "resolved-secret" + assert deployment.litellm_params.api_base == "os.environ/LITELLM_MASTER_KEY" + + +def test_ProxyConfig__add_deployment_keeps_team_env_refs_literal(monkeypatch): + def fail_on_call(secret_name, *args, **kwargs): + raise AssertionError("team DB models should not resolve env refs") + + monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret") + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + lambda value, key, return_original_value: value, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.get_secret", fail_on_call) + fake_router = MagicMock() + fake_router.upsert_deployment = MagicMock(return_value=True) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router) + pc = ProxyConfig() + db_model = SimpleNamespace( + model_id="model-1", + model_name="model_name_team-1_abc", + model_info={"id": "model-1", "team_id": "team-1"}, + litellm_params={ + "model": "openai/gpt-4o-mini", + "api_key": "os.environ/LITELLM_MASTER_KEY", + "api_base": "https://attacker.example", + }, + blocked=False, + ) + + added = pc._add_deployment(db_models=[db_model]) + deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"] + + assert added == 1 + assert deployment.litellm_params.api_key == "os.environ/LITELLM_MASTER_KEY" + assert deployment.litellm_params.api_base == "https://attacker.example" + + +def test_ProxyConfig__resolve_db_litellm_param_skips_non_string_values(monkeypatch): + def fail_on_call(value, key, return_original_value): + raise AssertionError("decrypt_value_helper should only receive strings") + + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + fail_on_call, + ) + pc = ProxyConfig() + + assert pc._resolve_db_litellm_param(key="tpm", value=100) == 100 + + # --------------------------------------------------------------------------- # ProxyConfig.decrypt_model_list_from_db # --------------------------------------------------------------------------- @@ -919,6 +997,71 @@ def test_ProxyConfig_decrypt_model_list_from_db_returns_decrypted(monkeypatch): } +def test_ProxyConfig_decrypt_model_list_from_db_resolves_env_refs_after_db_decrypt( + monkeypatch, +): + monkeypatch.setenv("LITELLM_DB_MODEL_API_KEY", "resolved-secret") + monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret") + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + lambda value, key, return_original_value: ( + "os.environ/LITELLM_DB_MODEL_API_KEY" + if key == "api_key" + else "os.environ/LITELLM_MASTER_KEY" if key == "api_base" else value + ), + ) + pc = ProxyConfig() + m = SimpleNamespace( + model_id="model-1", + model_name="env-model", + model_info={"id": "model-1"}, + litellm_params={ + "api_key": "encrypted-env-ref", + "api_base": "encrypted-api-base-env-ref", + "model": "openai/gpt-4o-mini", + }, + blocked=False, + ) + + out = pc.decrypt_model_list_from_db(new_models=[m]) + + assert out[0]["litellm_params"]["api_key"] == "resolved-secret" + assert out[0]["litellm_params"]["api_base"] == "os.environ/LITELLM_MASTER_KEY" + + +def test_ProxyConfig_decrypt_model_list_from_db_keeps_team_env_refs_literal_after_db_decrypt( + monkeypatch, +): + def fail_on_call(secret_name, *args, **kwargs): + raise AssertionError("team DB models should not resolve env refs") + + monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret") + monkeypatch.setattr( + "litellm.proxy.proxy_server.decrypt_value_helper", + lambda value, key, return_original_value: ( + "os.environ/LITELLM_MASTER_KEY" if key == "api_key" else value + ), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.get_secret", fail_on_call) + pc = ProxyConfig() + m = SimpleNamespace( + model_id="model-1", + model_name="model_name_team-1_abc", + model_info={"id": "model-1", "team_id": "team-1"}, + litellm_params={ + "api_key": "encrypted-env-ref", + "api_base": "https://attacker.example", + "model": "openai/gpt-4o-mini", + }, + blocked=False, + ) + + out = pc.decrypt_model_list_from_db(new_models=[m]) + + assert out[0]["litellm_params"]["api_key"] == "os.environ/LITELLM_MASTER_KEY" + assert out[0]["litellm_params"]["api_base"] == "https://attacker.example" + + def test_ProxyConfig_decrypt_model_list_from_db_invalid_params_skips(): pc = ProxyConfig() bad = SimpleNamespace( diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py new file mode 100644 index 00000000000..04db7192364 --- /dev/null +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py @@ -0,0 +1,73 @@ +"""Regression tests for the responses -> completion fallback bridge guard. + +When the Responses API falls back to chat completions (no native responses +config), it must tag the forwarded ``litellm.completion`` / ``litellm.acompletion`` +call with ``_skip_responses_api_bridge=True`` so ``completion()`` does not bridge +the request straight back to the Responses API and mutually recurse forever. + +Both fallback paths are covered: the sync ``response_api_handler`` (``_is_async`` +False) and the async ``async_response_api_handler`` (``_is_async`` True). The +module-level ``litellm.completion`` / ``litellm.acompletion`` are patched to +capture the forwarded kwargs; if the flag-setting line is removed the captured +kwargs lack the flag and these tests fail. +""" + +import os +import sys +from unittest.mock import patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.responses.litellm_completion_transformation.handler import ( + LiteLLMCompletionTransformationHandler, +) + + +class _StopForwarding(Exception): + """Raised by the mocked (a)completion once the forwarded kwargs are captured.""" + + +def test_sync_fallback_tags_skip_responses_api_bridge(): + handler = LiteLLMCompletionTransformationHandler() + captured: dict = {} + + def fake_completion(**kwargs): + captured.update(kwargs) + raise _StopForwarding() + + with patch("litellm.completion", fake_completion): + with pytest.raises(_StopForwarding): + handler.response_api_handler( + model="gpt-4o", + input="hello", + responses_api_request={}, + custom_llm_provider="openai", + _is_async=False, + ) + + assert captured.get("_skip_responses_api_bridge") is True + + +@pytest.mark.asyncio +async def test_async_fallback_tags_skip_responses_api_bridge(): + handler = LiteLLMCompletionTransformationHandler() + captured: dict = {} + + async def fake_acompletion(**kwargs): + captured.update(kwargs) + raise _StopForwarding() + + with patch("litellm.acompletion", fake_acompletion): + coro = handler.response_api_handler( + model="gpt-4o", + input="hello", + responses_api_request={}, + custom_llm_provider="openai", + _is_async=True, + ) + with pytest.raises(_StopForwarding): + await coro + + assert captured.get("_skip_responses_api_bridge") is True diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index ee3d4e835ba..b53acf930f2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -21125,10 +21125,9 @@ export interface components { }; /** * Mode - * @description The mode to test the model with - * @default chat + * @description The mode to test the model with. If not provided, auto-detected from model capabilities. */ - mode: ("chat" | "completion" | "embedding" | "audio_speech" | "audio_transcription" | "image_generation" | "video_generation" | "batch" | "rerank" | "realtime" | "responses" | "ocr") | null; + mode?: ("chat" | "completion" | "embedding" | "audio_speech" | "audio_transcription" | "image_generation" | "video_generation" | "batch" | "rerank" | "realtime" | "responses" | "ocr") | null; /** * Model Info * @description Model info for the health check From 062d8ceeed152b8ca46bb40212c17b19f2107552 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 26 Jun 2026 09:31:02 +0530 Subject: [PATCH 12/16] fix(vertex_ai): prevent stale Vertex bearer token causing /v1/messages 401 after token expiry (#31276) * fix(vertex_ai): prevent stale Vertex bearer token causing /v1/messages 401 after token expiry Router shallow-copies litellm_params so extra_headers is a shared reference. The chat/completions path was calling headers.update() on that shared dict, persisting the Vertex OAuth bearer. After ~1 h the token expired and /v1/messages kept reusing it (skipping refresh due to Authorization-already-present guard). - Build a new headers dict in the Claude partner-models completion path instead of mutating the shared extra_headers object. - Always call _ensure_access_token() in validate_anthropic_messages_environment regardless of an existing Authorization header; the token cache makes this cheap. Co-authored-by: Cursor * fix(vertex_ai): copy headers in validate_anthropic_messages_environment to prevent shared-dict mutation Co-authored-by: Cursor --------- Co-authored-by: Cursor --- .../transformation.py | 28 ++--- .../vertex_ai_partner_models/main.py | 8 +- ...artner_models_anthropic_messages_config.py | 113 ++++++++++++++---- 3 files changed, 108 insertions(+), 41 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 8a92e7ec4a5..a633ef3298a 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -35,27 +35,21 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert Validate the environment for the request """ + # Work on a local copy — router shallow-copies litellm_params so the caller's + # headers dict may be the shared deployment extra_headers object. + headers = dict(headers) vertex_ai_project = VertexBase.safe_get_vertex_ai_project(litellm_params) vertex_ai_location = VertexBase.safe_get_vertex_ai_location(litellm_params) - project_id: Optional[str] = None - if "Authorization" not in headers: - vertex_credentials = VertexBase.safe_get_vertex_ai_credentials( - litellm_params - ) + vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(litellm_params) + access_token, project_id = self._ensure_access_token( + credentials=vertex_credentials, + project_id=vertex_ai_project, + custom_llm_provider="vertex_ai", + ) + headers["Authorization"] = f"Bearer {access_token}" - access_token, project_id = self._ensure_access_token( - credentials=vertex_credentials, - project_id=vertex_ai_project, - custom_llm_provider="vertex_ai", - ) - - headers["Authorization"] = f"Bearer {access_token}" - else: - # Authorization already in headers, but we still need project_id - project_id = vertex_ai_project - - # Always calculate api_base if not provided, regardless of Authorization header + # Calculate api_base if not provided if api_base is None: api_base = self.get_complete_vertex_url( custom_api_base=api_base, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index 960d3483848..78669f1e789 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -194,9 +194,11 @@ class VertexAIPartnerModels(VertexBase): encoding=encoding, ) elif "claude" in model: - if headers is None: - headers = {} - headers.update({"Authorization": "Bearer {}".format(access_token)}) + # Build a new dict so we never mutate the shared deployment extra_headers object. + headers = { + **(headers or {}), + "Authorization": "Bearer {}".format(access_token), + } optional_params.update( { diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 6f4bb4e59c2..d64b5c8d742 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -1,10 +1,11 @@ -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, ) +from litellm.llms.vertex_ai.vertex_ai_partner_models.main import VertexAIPartnerModels from litellm.types.router import GenericLiteLLMParams @@ -233,40 +234,36 @@ def test_both_compact_and_context_management_headers_added(): ), f"anthropic-beta should contain 'context-management-2025-06-27', got: {updated_headers['anthropic-beta']}" -def test_validate_environment_with_authorization_header_calculates_api_base(): - """Test that api_base is calculated even when Authorization header is already present""" +def test_validate_environment_always_refreshes_token_ignoring_stale_bearer(): + """Regression: stale Authorization in shared deployment extra_headers must not + skip token refresh on /v1/messages — _ensure_access_token is always called.""" config = VertexAIPartnerModelsAnthropicMessagesConfig() - # Simulate scenario where Authorization is already in headers (e.g., from cached extra_headers) - headers = {"Authorization": "Bearer existing-token"} + headers = {"Authorization": "Bearer EXPIRED"} litellm_params = { "vertex_project": "test-project", "vertex_location": "us-central1", - "extra_headers": {"anthropic-beta": "context-1m-2025-08-07"}, } - optional_params = {} - with patch.object( - config, "get_complete_vertex_url", return_value="https://mock-vertex-url" - ) as mock_get_url: + with ( + patch.object( + config, "_ensure_access_token", return_value=("fresh-token", "test-project") + ) as mock_ensure, + patch.object( + config, "get_complete_vertex_url", return_value="https://mock-vertex-url" + ), + ): updated_headers, api_base = config.validate_anthropic_messages_environment( headers=headers, model="claude-sonnet-4", messages=[], - optional_params=optional_params, + optional_params={}, litellm_params=litellm_params, api_base=None, ) - # Verify that api_base was calculated even though Authorization was already present - assert ( - api_base == "https://mock-vertex-url" - ), f"api_base should be calculated even with Authorization header. Got: {api_base}" - assert mock_get_url.called, "get_complete_vertex_url should be called" - - # Verify Authorization header is still present - assert ( - "Authorization" in updated_headers - ), "Authorization header should be preserved" + mock_ensure.assert_called_once() + assert updated_headers["Authorization"] == "Bearer fresh-token" + assert api_base == "https://mock-vertex-url" def test_transform_anthropic_messages_request_removes_scope_from_cache_control(): @@ -371,3 +368,77 @@ def test_provider_config_manager_reuses_vertex_anthropic_messages_config_instanc assert first_config is second_config finally: ProviderConfigManager._get_provider_anthropic_messages_config_cached.cache_clear() + + +def test_validate_environment_does_not_mutate_caller_headers(): + """Regression: beta headers (e.g. web-search) must not leak into the caller's + headers dict — which may be the shared deployment extra_headers object.""" + config = VertexAIPartnerModelsAnthropicMessagesConfig() + caller_headers: dict = {} + + with ( + patch.object( + config, "_ensure_access_token", return_value=("token", "test-project") + ), + patch.object( + config, "get_complete_vertex_url", return_value="https://mock-url" + ), + ): + config.validate_anthropic_messages_environment( + headers=caller_headers, + model="claude-sonnet-4", + messages=[], + optional_params={ + "tools": [{"type": "web_search_20250305", "name": "web_search"}] + }, + litellm_params={ + "vertex_ai_project": "p", + "vertex_ai_location": "us-central1", + }, + api_base=None, + ) + + assert ( + caller_headers == {} + ), "validate_anthropic_messages_environment must not mutate the caller's headers dict" + + +def test_vertex_claude_completion_does_not_mutate_shared_extra_headers(): + """Regression: router shallow-copies litellm_params so extra_headers is a shared + reference. Verify that the chat/completions path builds a new headers dict instead + of calling .update() on the shared object.""" + handler = VertexAIPartnerModels() + shared_extra_headers = {} # simulates deployment["litellm_params"]["extra_headers"] + + mock_response = MagicMock() + + with ( + patch.object( + handler, "_ensure_access_token", return_value=("ya29.fresh", "proj") + ), + patch.object( + handler, "get_complete_vertex_url", return_value="https://mock-url" + ), + patch( + "litellm.llms.anthropic.chat.AnthropicChatCompletion.completion", + return_value=mock_response, + ), + ): + handler.completion( + model="claude-haiku-4-5@20251001", + messages=[{"role": "user", "content": "hi"}], + model_response=MagicMock(), + print_verbose=lambda *a, **k: None, + encoding=None, + logging_obj=MagicMock(), + api_base=None, + optional_params={}, + custom_prompt_dict={}, + headers=shared_extra_headers, + timeout=30, + litellm_params={}, + ) + + assert ( + shared_extra_headers == {} + ), "extra_headers must not be mutated by completion()" From e5da5a3b6d57cf3c4f54f7e8d0449222b0d6ce61 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 26 Jun 2026 09:32:52 +0530 Subject: [PATCH 13/16] fix(proxy): skip model override when response has no model field (#31183) * fix(proxy): skip OpenAI model override for search responses Search responses omit a model field by spec but still set model on the request for routing, which caused noisy errors and dict injection. * fix(proxy): drop redundant search-specific model override skip The silent return for responses without a model field already covers SearchResponse objects; remove the extra search type check. Co-authored-by: Cursor * fix(proxy): skip model override for dict responses without model key Dict-shaped responses (e.g. search) must not get a spurious model field injected when they never had one; only override when model is present. Co-authored-by: Cursor * test(proxy): cover swallowed setattr failure in model override --------- Co-authored-by: Cursor Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 20 ++++---- .../proxy/test_common_request_processing.py | 46 +++++++++++++++++++ 2 files changed, 55 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 92f8dd8df98..c2ffd9b077d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -516,12 +516,13 @@ def _override_openai_response_model( LiteLLM internally prefixes some provider/deployment model identifiers (e.g. `hosted_vllm/...`). That internal identifier should not be returned to clients in the OpenAI `model` field. - Note: This is intentionally verbose. A model mismatch is a useful signal that an internal - model identifier is being stamped/preserved somewhere in the request/response pipeline. - We log mismatches as warnings (and then restamp to the client-requested value) so these - paths stay observable for maintainers/operators without breaking client compatibility. + Note: This is intentionally verbose at debug level. A model mismatch is a useful signal that an + internal model identifier is being stamped/preserved somewhere in the request/response pipeline. + We log mismatches as debug (and then restamp to the client-requested value) so these paths stay + observable for maintainers without breaking client compatibility or alarming operators. - Errors are reserved for cases where the proxy cannot read/override the response model field. + Responses that omit an OpenAI-style `model` field are left unchanged (silent return), + including dict responses with no `model` key. Exceptions: 1. If a fallback occurred (indicated by x-litellm-attempted-fallbacks header), @@ -577,6 +578,8 @@ def _override_openai_response_model( return if isinstance(response_obj, dict): + if "model" not in response_obj: + return downstream_model = response_obj.get("model") if downstream_model != requested_model: verbose_proxy_logger.debug( @@ -589,11 +592,6 @@ def _override_openai_response_model( return if not hasattr(response_obj, "model"): - verbose_proxy_logger.error( - "%s: cannot override response model; missing `model` attribute. response_type=%s", - log_context, - type(response_obj), - ) return downstream_model = getattr(response_obj, "model", None) @@ -608,7 +606,7 @@ def _override_openai_response_model( try: setattr(response_obj, "model", requested_model) except Exception as e: - verbose_proxy_logger.error( + verbose_proxy_logger.debug( "%s: failed to override response.model=%r on response_type=%s. error=%s", log_context, requested_model, diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 364861d6e31..3d3a2cc1013 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2022,6 +2022,52 @@ class TestOverrideOpenAIResponseModel: assert response_obj.model == requested_model + def test_skips_model_override_when_response_has_no_model_attribute(self): + from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult + + response_obj = SearchResponse( + results=[SearchResult(title="t", url="http://x.com", snippet="s")], + object="search", + ) + + _override_openai_response_model( + response_obj=response_obj, + requested_model="my-search-tool", + log_context="test_context", + ) + + assert not hasattr(response_obj, "model") + + def test_skips_model_override_for_dict_without_model_key(self): + response_obj = { + "object": "search", + "results": [{"title": "t", "url": "http://x.com", "snippet": "s"}], + } + + _override_openai_response_model( + response_obj=response_obj, + requested_model="my-search-tool", + log_context="test_context", + ) + + assert "model" not in response_obj + + def test_override_model_swallows_setattr_failure(self): + class ReadOnlyModelResponse: + @property + def model(self) -> str: + return "downstream-model" + + response_obj = ReadOnlyModelResponse() + + _override_openai_response_model( + response_obj=response_obj, + requested_model="my-model", + log_context="test_context", + ) + + assert response_obj.model == "downstream-model" + class TestIsAzureModelRouterRequest: """Tests for _is_azure_model_router_request helper""" From 29c254d3d38843ef4714c9c043e88f6b05f5b0e5 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 26 Jun 2026 09:04:29 +0300 Subject: [PATCH 14/16] fix(vertex): stop O(n^2) re-parse of accumulated Gemini stream JSON (#31297) handle_accumulated_json_chunk re-ran json.loads on the entire accumulated buffer after every fragment. For a streaming response fragmented across many chunks that is O(n^2) total work in a single GIL-holding C call, so a large enough Gemini response freezes the asyncio event loop for seconds, liveness probes time out, and the proxy pod gets killed and restarted. A complete Gemini stream value is a JSON object or array, so the buffer can only become parseable once its last non-whitespace byte can close one. Gate the json.loads attempt on that, which makes the common fragmented-response case parse roughly once instead of once per fragment. An 8MB payload drops from a 6.9s event-loop freeze to ~0.3s with identical parsed output. Resolves LIT-3503 Fixes #26181 --- .../vertex_and_google_ai_studio_gemini.py | 12 +++- ...test_vertex_and_google_ai_studio_gemini.py | 71 +++++++++++++++++++ 2 files changed, 80 insertions(+), 3 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 423a5dd5d17..32bdb17840d 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -3601,16 +3601,22 @@ class ModelResponseIterator: chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" message = chunk.replace("\n\n", "") - # Accumulate JSON data self.accumulated_json += message - # Try to parse the accumulated JSON + # json.loads on the whole buffer after every fragment is O(n^2) and + # holds the GIL, freezing the event loop for seconds on large responses + # (https://github.com/BerriAI/litellm/issues/26181). A complete Gemini + # chunk is a JSON object/array, so only attempt the parse once the + # buffer's last non-whitespace byte can close one. + stripped = self.accumulated_json.rstrip() + if not stripped or stripped[-1] not in "}]": + return None + try: _data = json.loads(self.accumulated_json) self.accumulated_json = "" # reset after successful parsing return self.chunk_parser(chunk=_data) except json.JSONDecodeError: - # If it's not valid JSON yet, continue to the next event return None def _common_chunk_parsing_logic( diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 1a2d0d86810..20e7aec377b 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2799,6 +2799,77 @@ def test_partial_json_chunk_on_first_chunk(): ), "Should switch to accumulated_json mode" +def test_accumulated_json_does_not_reparse_every_fragment(): + """ + Regression test for https://github.com/BerriAI/litellm/issues/26181 + + handle_accumulated_json_chunk used to call json.loads on the entire + accumulated buffer after EVERY fragment. For a large response fragmented + across many chunks that is O(n^2) work in a single GIL-holding C call, + which freezes the asyncio event loop for seconds and kills liveness probes. + + The buffer only becomes a complete JSON object on the final fragment, so a + correct implementation parses it ~once, not once per fragment. We assert the + full chunk still parses correctly AND that json.loads is not called on every + fragment (which is what made it quadratic). + """ + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + iterator = ModelResponseIterator( + streaming_response=MagicMock(), + sync_stream=True, + logging_obj=MagicMock(), + ) + iterator.chunk_type = "accumulated_json" + + text = "x" * 200_000 # no braces/brackets so only the final fragment closes + blob = json.dumps( + {"candidates": [{"content": {"role": "model", "parts": [{"text": text}]}}]} + ) + fragments = [blob[i : i + 4096] for i in range(0, len(blob), 4096)] + assert len(fragments) > 10, "need a multi-fragment payload to exercise the bug" + + parsed = None + with patch("json.loads", wraps=json.loads) as spy: + for fragment in fragments: + out = iterator.handle_accumulated_json_chunk(chunk=fragment) + if out is not None: + parsed = out + parse_calls = spy.call_count + + assert parsed is not None, "the reassembled chunk must still parse" + assert parsed.choices[0].delta.content == text, "content must be preserved intact" + + assert parse_calls <= 2, ( + f"json.loads was called {parse_calls} times for {len(fragments)} " + "fragments; the O(n^2) per-fragment re-parse has regressed" + ) + + +def test_accumulated_json_partial_fragment_returns_none_without_parsing(): + """A fragment that cannot complete the JSON must not trigger a json.loads + parse of the whole growing buffer (issue #26181).""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + iterator = ModelResponseIterator( + streaming_response=MagicMock(), + sync_stream=True, + logging_obj=MagicMock(), + ) + iterator.chunk_type = "accumulated_json" + + with patch("json.loads", wraps=json.loads) as spy: + result = iterator.handle_accumulated_json_chunk( + chunk='{"candidates": [{"content": {"parts": [{"text": "partial' + ) + assert result is None + assert spy.call_count == 0, "incomplete buffer should not be parsed" + + def test_google_ai_studio_presence_penalty_supported(): """ Test that presence_penalty is supported for Google AI Studio Gemini. From 248389c27651ef1b9476e161784c61fbf8eb3780 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 26 Jun 2026 09:37:29 +0300 Subject: [PATCH 15/16] fix(router): surface clean RateLimitError on mid-stream 429 with no fallbacks (#31298) When a streaming request hits a mid-stream 429 the streaming handler wraps it in the internal MidStreamFallbackError so the router can attempt fallbacks. With no fallbacks configured, async_function_with_fallbacks_common_utils falls through to re-raising that wrapper, which the streaming iterators caught and re-raised verbatim, so the client received MidStreamFallbackError (an internal type) rather than a clean RateLimitError (429). When the fallback path produces a MidStreamFallbackError that carries an original_exception (i.e. no fallback handled it), the iterators now raise that underlying provider exception instead of the wrapper, chained with from. Users with fallbacks are unaffected since their path never reaches this branch. Applied consistently to the chat async, chat sync, and responses streaming iterators. Resolves LIT-3503 Fixes #26015 --- litellm/router.py | 18 ++++ tests/test_litellm/test_router.py | 173 ++++++++++++++++++++++++++++++ 2 files changed, 191 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index 0008fcaa30b..1aba259a328 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2323,6 +2323,14 @@ class Router: verbose_router_logger.error( f"Fallback also failed: {fallback_error}" ) + # No fallback handled the mid-stream error, so surface the + # real provider exception (e.g. RateLimitError) instead of + # leaking the internal MidStreamFallbackError to the client + if ( + isinstance(fallback_error, MidStreamFallbackError) + and fallback_error.original_exception is not None + ): + raise fallback_error.original_exception from fallback_error raise fallback_error finally: # Close the underlying streams to release HTTP connections @@ -2754,6 +2762,11 @@ class Router: verbose_router_logger.error( f"Responses streaming fallback also failed: {fallback_error}" ) + if ( + isinstance(fallback_error, MidStreamFallbackError) + and fallback_error.original_exception is not None + ): + raise fallback_error.original_exception from fallback_error raise fallback_error finally: with anyio.CancelScope(shield=True): @@ -2890,6 +2903,11 @@ class Router: verbose_router_logger.error( f"Fallback also failed: {fallback_error}" ) + if ( + isinstance(fallback_error, MidStreamFallbackError) + and fallback_error.original_exception is not None + ): + raise fallback_error.original_exception from fallback_error raise fallback_error finally: if hasattr(model_response, "close"): diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 470d38caf10..7be176fffc7 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1,3 +1,4 @@ +import asyncio import copy import json import os @@ -12,6 +13,7 @@ sys.path.insert( import litellm +from litellm.exceptions import MidStreamFallbackError def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -2409,6 +2411,177 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): assert merged.total_tokens == 49 +def _midstream_rate_limit_error(): + rate_limit_error = litellm.RateLimitError( + message="vertex_ai_betaException - Resource exhausted.", + model="gemini", + llm_provider="vertex_ai_beta", + ) + midstream_error = MidStreamFallbackError( + message=str(rate_limit_error), + model="gemini", + llm_provider="vertex_ai_beta", + original_exception=rate_limit_error, + is_pre_first_chunk=True, + ) + return rate_limit_error, midstream_error + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks(): + """Regression for #26015: a mid-stream 429 with no fallbacks configured must + surface a clean RateLimitError, not leak the internal MidStreamFallbackError + wrapper to the client, and must terminate instead of hanging.""" + rate_limit_error, midstream_error = _midstream_rate_limit_error() + + router = litellm.Router( + model_list=[ + { + "model_name": "gemini", + "litellm_params": { + "model": "vertex_ai/gemini-2.0-flash", + "api_key": "fake-key", + }, + }, + ], + num_retries=0, + ) + + class _RaisingStream: + def __init__(self): + self.chunks = [] + + def __aiter__(self): + return self + + async def __anext__(self): + raise midstream_error + + stream = _RaisingStream() + setattr(stream, "model", "gemini") + setattr(stream, "custom_llm_provider", "vertex_ai_beta") + setattr(stream, "logging_obj", MagicMock()) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(side_effect=midstream_error), + ): + result = await router._acompletion_streaming_iterator( + model_response=stream, + messages=[{"role": "user", "content": "Hello"}], + initial_kwargs={"model": "gemini", "stream": True}, + ) + + async def _consume(): + async for _ in result: + pass + + with pytest.raises(litellm.RateLimitError) as exc_info: + await asyncio.wait_for(_consume(), timeout=10) + + assert not isinstance(exc_info.value, MidStreamFallbackError) + assert exc_info.value.status_code == 429 + assert exc_info.value is rate_limit_error + + +def test_completion_streaming_iterator_surfaces_rate_limit_without_fallbacks(): + """Sync counterpart of + test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks.""" + rate_limit_error, midstream_error = _midstream_rate_limit_error() + + router = litellm.Router( + model_list=[ + { + "model_name": "gemini", + "litellm_params": { + "model": "vertex_ai/gemini-2.0-flash", + "api_key": "fake-key", + }, + }, + ], + num_retries=0, + ) + + class _RaisingSyncStream: + def __init__(self): + self.model = "gemini" + self.custom_llm_provider = "vertex_ai_beta" + self.logging_obj = MagicMock() + self.chunks = [] + + def __iter__(self): + return self + + def __next__(self): + raise midstream_error + + with patch.object( + router, + "function_with_fallbacks", + side_effect=midstream_error, + ): + result = router._completion_streaming_iterator( + model_response=_RaisingSyncStream(), + messages=[{"role": "user", "content": "Hello"}], + initial_kwargs={"model": "gemini", "stream": True}, + ) + + with pytest.raises(litellm.RateLimitError) as exc_info: + list(result) + + assert not isinstance(exc_info.value, MidStreamFallbackError) + assert exc_info.value.status_code == 429 + assert exc_info.value is rate_limit_error + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_surfaces_rate_limit_without_fallbacks(): + """Responses-API counterpart of + test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks.""" + rate_limit_error, midstream_error = _midstream_rate_limit_error() + + router = litellm.Router( + model_list=[ + { + "model_name": "gemini", + "litellm_params": { + "model": "vertex_ai/gemini-2.0-flash", + "api_key": "fake-key", + }, + }, + ], + num_retries=0, + ) + src = _make_responses_iterator(error=midstream_error, model="gemini") + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(side_effect=midstream_error), + ): + wrapped = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={ + "model": "gemini", + "stream": True, + "input": "Hello", + "original_generic_function": litellm.aresponses, + }, + ) + + async def _consume(): + async for _ in wrapped: + pass + + with pytest.raises(litellm.RateLimitError) as exc_info: + await asyncio.wait_for(_consume(), timeout=10) + + assert not isinstance(exc_info.value, MidStreamFallbackError) + assert exc_info.value.status_code == 429 + assert exc_info.value is rate_limit_error + + @pytest.mark.asyncio async def test_async_function_with_fallbacks_common_utils(): """Test the async_function_with_fallbacks_common_utils method""" From 52e5b3ae98421763545c4a9430ae67432cdf9518 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Thu, 25 Jun 2026 23:41:08 -0700 Subject: [PATCH 16/16] build(docker): build the Admin UI from source in a build-platform-pinned stage (#31130) The monolith images shipped whatever UI bundle was committed to litellm/proxy/_experimental/out, so refreshing the UI for a release meant running build_ui.sh out of band and committing the regenerated bundle. Add a ui-builder stage to all three monolith Dockerfiles (root, database, non_root) that compiles the Next.js static export from this exact source and replaces the committed bundle before the final uv sync. The stage is pinned with FROM --platform=$BUILDPLATFORM so the architecture-independent static export compiles once on the native builder even in a multi-arch (linux/amd64,linux/arm64) build, rather than once per target arch under QEMU emulation. The destination is cleared before the COPY because COPY merges directories and would otherwise leave the committed bundle's content-hashed chunks behind alongside the fresh ones. build_admin_ui.sh still runs afterward so the enterprise custom-color override is preserved. The UI base image is pinned by digest to match LITELLM_BUILD_IMAGE, LITELLM_RUNTIME_IMAGE and UV_IMAGE, and .dockerignore now excludes the local .next/out so a developer's build artifacts never enter the context. --- .dockerignore | 2 ++ Dockerfile | 29 ++++++++++++++++++++++++++++- docker/Dockerfile.database | 29 ++++++++++++++++++++++++++++- docker/Dockerfile.non_root | 27 +++++++++++++++++++++++++++ 4 files changed, 85 insertions(+), 2 deletions(-) diff --git a/.dockerignore b/.dockerignore index 6b80caeaf9f..f3a80fee3e4 100644 --- a/.dockerignore +++ b/.dockerignore @@ -49,6 +49,8 @@ build/ *.egg-info/ .DS_Store **/node_modules +ui/litellm-dashboard/.next +ui/litellm-dashboard/out litellm-rust/target/ litellm/rust_bridge/_native*.so *.log diff --git a/Dockerfile b/Dockerfile index 681681f28cc..b6fef1a21fc 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,12 +1,33 @@ +# syntax=docker/dockerfile:1.7 + # Base image for building ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 # Runtime image ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a +# Pinned by digest like the other base images; bump explicitly on Node upgrades. +ARG UI_BUILD_IMAGE=node:20.18-alpine3.20@sha256:3488b10bf958af7125a176419d2d8a9937d895bf124012aae811651988d2ffe6 FROM $UV_IMAGE AS uvbin +# Admin UI builder. Pinned to the build platform so the architecture-independent +# Next.js static export compiles once natively even in a multi-arch build, +# instead of once per target arch under QEMU. +FROM --platform=$BUILDPLATFORM $UI_BUILD_IMAGE AS ui-builder + +ENV NEXT_TELEMETRY_DISABLED=1 \ + npm_config_fund=false \ + npm_config_audit=false + +WORKDIR /ui + +COPY ui/litellm-dashboard/package.json ui/litellm-dashboard/package-lock.json ./ +RUN --mount=type=cache,target=/root/.npm npm ci --prefer-offline + +COPY ui/litellm-dashboard/ ./ +RUN npm run build + # Builder stage FROM $LITELLM_BUILD_IMAGE AS builder @@ -48,7 +69,13 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr # Copy full source tree COPY . . -# Build Admin UI before final sync +# Replace the committed UI bundle with the one built from this exact source. +# Clearing first drops the committed bundle's content-hashed chunks that COPY +# would otherwise leave behind alongside the fresh ones. +RUN rm -rf litellm/proxy/_experimental/out +COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/ + +# Build Admin UI before final sync (applies the enterprise color override when present) RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh # Install project and workspace packages (fast - deps already cached) diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 50ef55e3261..b3af953511d 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -1,12 +1,33 @@ +# syntax=docker/dockerfile:1.7 + # Base image for building ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 # Runtime image ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a +# Pinned by digest like the other base images; bump explicitly on Node upgrades. +ARG UI_BUILD_IMAGE=node:20.18-alpine3.20@sha256:3488b10bf958af7125a176419d2d8a9937d895bf124012aae811651988d2ffe6 FROM $UV_IMAGE AS uvbin +# Admin UI builder. Pinned to the build platform so the architecture-independent +# Next.js static export compiles once natively even in a multi-arch build, +# instead of once per target arch under QEMU. +FROM --platform=$BUILDPLATFORM $UI_BUILD_IMAGE AS ui-builder + +ENV NEXT_TELEMETRY_DISABLED=1 \ + npm_config_fund=false \ + npm_config_audit=false + +WORKDIR /ui + +COPY ui/litellm-dashboard/package.json ui/litellm-dashboard/package-lock.json ./ +RUN --mount=type=cache,target=/root/.npm npm ci --prefer-offline + +COPY ui/litellm-dashboard/ ./ +RUN npm run build + FROM $LITELLM_BUILD_IMAGE AS builder WORKDIR /app @@ -46,7 +67,13 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr # Copy full source tree COPY . . -# Build Admin UI before final sync +# Replace the committed UI bundle with the one built from this exact source. +# Clearing first drops the committed bundle's content-hashed chunks that COPY +# would otherwise leave behind alongside the fresh ones. +RUN rm -rf litellm/proxy/_experimental/out +COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/ + +# Build Admin UI before final sync (applies the enterprise color override when present) RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh # Install project and workspace packages (fast - deps already cached) diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 6bb925aa723..c24cb9008f0 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -1,11 +1,32 @@ +# syntax=docker/dockerfile:1.7 + # Base images ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370 ARG PROXY_EXTRAS_SOURCE=published ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a +# Pinned by digest like the other base images; bump explicitly on Node upgrades. +ARG UI_BUILD_IMAGE=node:20.18-alpine3.20@sha256:3488b10bf958af7125a176419d2d8a9937d895bf124012aae811651988d2ffe6 FROM $UV_IMAGE AS uvbin +# Admin UI builder. Pinned to the build platform so the architecture-independent +# Next.js static export compiles once natively even in a multi-arch build, +# instead of once per target arch under QEMU. +FROM --platform=$BUILDPLATFORM $UI_BUILD_IMAGE AS ui-builder + +ENV NEXT_TELEMETRY_DISABLED=1 \ + npm_config_fund=false \ + npm_config_audit=false + +WORKDIR /ui + +COPY ui/litellm-dashboard/package.json ui/litellm-dashboard/package-lock.json ./ +RUN --mount=type=cache,target=/root/.npm npm ci --prefer-offline + +COPY ui/litellm-dashboard/ ./ +RUN npm run build + FROM $LITELLM_BUILD_IMAGE AS builder ARG PROXY_EXTRAS_SOURCE WORKDIR /app @@ -53,6 +74,12 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ # Copy full source tree COPY . . +# Replace the committed UI bundle with the one built from this exact source. +# Clearing first drops the committed bundle's content-hashed chunks that COPY +# would otherwise leave behind alongside the fresh ones. +RUN rm -rf litellm/proxy/_experimental/out +COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/ + # Set non-root flag for build time consistency ENV LITELLM_NON_ROOT=true