Merge branch 'main' into litellm_rag_team_provider_credentials

Keeps main's Path import next to the branch's Final import in the batches endpoint tests, and takes the integration test's closed port from allocate_port so it follows main's TQ012 rule

.circleci/config.yml

@@ -1293,6 +1293,19 @@ jobs:
                   uv venv --python 3.12 .venv-mcp-peer
                   uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
                   echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
+      - when:
+          condition:
+            equal: [sdk, << parameters.suite >>]
+          steps:
+            - run:
+                name: Install redis-server 7.2.16
+                command: |
+                  mkdir -p /tmp/redis-bin
+                  cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
+                  docker cp "${cid}:/usr/local/bin/redis-server" /tmp/redis-bin/redis-server
+                  docker rm "$cid"
+                  /tmp/redis-bin/redis-server --version | grep -q 'v=7.2.16'
+                  echo 'export PATH="/tmp/redis-bin:$PATH"' >> "$BASH_ENV"
       - start_postgres:
           image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
           server_args: "-c max_connections=300 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000"

.github/scripts/assert_ci_coverage.py

@@ -32,6 +32,7 @@ IGNORE_ARG_RE: Final = re.compile(r"--ignore(?:-glob)?[= ](\S+)")
 DOCKERFILE_TOKEN_RE = re.compile(r"[A-Za-z0-9_./-]*Dockerfile[A-Za-z0-9_.-]*")
 COMMENT_RE = re.compile(r"^\s*#.*$", re.MULTILINE)
 GLOB_CHARS = frozenset("*?")
+DEDICATED_INTEGRATION_WORKFLOWS: Final = MappingProxyType({"image-scan.yml": "tests/integration/database/image/"})
 
 # Trees whose jobs are sharded with no catch-all bucket, so every child that holds
 # tests has to be named by some shard or it runs nowhere. A child listed here is
@@ -560,6 +561,17 @@ def _integration_github_files(runner: pathlib.Path) -> frozenset[str]:
     return frozenset(ast.literal_eval(values))
 
 
+def _owns(workflow: str, relative_path: str) -> bool:
+    owned_root: Final = DEDICATED_INTEGRATION_WORKFLOWS.get(workflow)
+    return owned_root is not None and relative_path.startswith(owned_root)
+
+
+def _tokens_outside_owner(relative_path: str, workflow_tokens: Mapping[str, frozenset[str]]) -> frozenset[str]:
+    return frozenset[str]().union(
+        *(tokens for workflow, tokens in workflow_tokens.items() if not _owns(workflow, relative_path))
+    )
+
+
 def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozenset[str], tuple[Finding, ...]]:
     runner: Final = repo_root / "tests/integration/run.py"
     if not runner.exists():
@@ -611,15 +623,20 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
         )
         != 1
     )
-    gha_tokens: Final = _invoked_test_tokens(
-        scalar
-        for path in (repo_root / ".github/workflows").glob("*.y*ml")
-        for scalar in _scalars(yaml.safe_load(path.read_text()), path.name)
+    workflow_tokens: Final = MappingProxyType(
+        {
+            path.name: _invoked_test_tokens(_scalars(yaml.safe_load(path.read_text()), path.name))
+            for path in (repo_root / ".github/workflows").glob("*.y*ml")
+        }
+    )
+    gha_tokens: Final = frozenset[str]().union(*workflow_tokens.values())
+    circle_owned: Final = (
+        frozenset(str(path.relative_to(repo_root)) for path in integration_root.rglob("test_*.py")) - github_files
     )
     findings: Final = (
         tuple(
-            Finding(path, "integration contract is also selected by GitHub Actions")
-            for path in paths
+            Finding(path, "integration file outside GITHUB_FILES is selected by GitHub Actions")
+            for path in sorted(circle_owned)
             if any(_token_covers(token, path) for token in gha_tokens)
         )
         + tuple(
@@ -632,6 +649,11 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
             for path in sorted(github_files)
             if not (repo_root / path).is_file()
         )
+        + tuple(
+            Finding(path, "GitHub-owned integration contract runs outside a dedicated integration workflow")
+            for path in sorted(github_files)
+            if any(_token_covers(token, path) for token in _tokens_outside_owner(path, workflow_tokens))
+        )
     )
     browser_commands: Final = tuple(
         scalar.value
@@ -688,13 +710,9 @@ def main() -> int:
     scalars = _all_scalars()
 
     integration_paths, ownership_findings = _integration_ownership()
-    test_findings = (
-        _uncovered_tests(
-            allowlist,
-            _invoked_selections(scalars)
-            + (Selection(included=integration_paths, ignored=frozenset()),),
-        )
-        + ownership_findings
+    test_findings = _uncovered_tests(
+        allowlist,
+        _invoked_selections(scalars) + (Selection(included=integration_paths, ignored=frozenset()),),
     )
     dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars))
     stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles())
@@ -711,13 +729,25 @@ def main() -> int:
             test_findings,
             "Add each to a job's test path, or list it in .github/ci-coverage-allowlist.yml with a reason.",
         )
+    if ownership_findings:
+        browser_remedy = (
+            " CircleCI's browser group owns the integrationCritical Playwright contracts, so no GitHub workflow may "
+            "select them and the default Playwright config must exclude them."
+            if any("browser" in finding.detail or "Playwright" in finding.detail for finding in ownership_findings)
+            else ""
+        )
+        _report(
+            "integration contracts with the wrong CI owner",
+            ownership_findings,
+            "Follow tests/integration/AGENTS.md, which says which CI job owns each integration file." + browser_remedy,
+        )
     if dockerfile_findings:
         _report(
             "Dockerfiles that no CI job builds",
             dockerfile_findings,
             "Build each in a workflow, or list it in .github/ci-coverage-allowlist.yml with a reason.",
         )
-    if stale_findings or test_findings or dockerfile_findings:
+    if stale_findings or test_findings or ownership_findings or dockerfile_findings:
         return 1
 
     _write(

.github/workflows/test-linting.yml

@@ -21,7 +21,7 @@ jobs:
       pull-requests: read
       actions: read
     runs-on: ubuntu-latest
-    timeout-minutes: 15
+    timeout-minutes: 20
     services:
       redis:
         image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f

.github/workflows/test-rust.yml

@@ -16,6 +16,10 @@ on:
       - "litellm/proxy/guardrails/**"
       - "litellm/utils.py"
       - "litellm/ocr/**"
+      - "litellm/chat_completions/dispatch.py"
+      - "litellm/embeddings/dispatch.py"
+      - "litellm/messages/dispatch.py"
+      - "litellm/responses/dispatch.py"
       - "litellm/llms/base_llm/ocr/**"
       - "litellm/llms/custom_httpx/llm_http_handler.py"
       - "Makefile"
@@ -46,6 +50,10 @@ on:
       - "litellm/proxy/guardrails/**"
       - "litellm/utils.py"
       - "litellm/ocr/**"
+      - "litellm/chat_completions/dispatch.py"
+      - "litellm/embeddings/dispatch.py"
+      - "litellm/messages/dispatch.py"
+      - "litellm/responses/dispatch.py"
       - "litellm/llms/base_llm/ocr/**"
       - "litellm/llms/custom_httpx/llm_http_handler.py"
       - "Makefile"

.github/workflows/test-unit.yml

@@ -346,6 +346,7 @@ jobs:
               tests/unit/llms/tinyfish
               tests/unit/llms/together_ai
               tests/unit/llms/tool_loop
+              tests/unit/llms/topaz
               tests/unit/llms/triton
               tests/unit/llms/v0
               tests/unit/llms/valkey
@@ -417,6 +418,7 @@ jobs:
               tests/unit/rag
               tests/unit/realtime_api
               tests/unit/rerank_api
+              tests/unit/search
               tests/unit/vector_stores
               tests/unit/videos
             python-version: "3.12"
@@ -454,11 +456,11 @@ jobs:
               tests/unit/telemetry tests/unit/proxy/telemetry
               tests/unit/tracing
             python-version: "3.12"
-            workers: 2
+            workers: 4
             reruns: 0
             timeout-minutes: 20
             job-timeout-minutes: 60
-            dist: loadscope
+            dist: load
 
           - shard: proxy-auth
             test-path: |-
@@ -563,7 +565,6 @@ jobs:
 
           - shard: mcp-elicitation
             test-path: |-
-              --override-ini=pythonpath=tests
               tests/unit/experimental_mcp_client/test_mcp_client.py
               tests/unit/proxy/_experimental/mcp_server/test_capabilities.py
               tests/unit/proxy/_experimental/mcp_server/test_interactions.py
@@ -571,9 +572,8 @@ jobs:
               tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
               tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
               tests/unit/proxy/_experimental/mcp_server/test_operations.py
-              tests/integration/mcp/test_interactions.py
             python-version: "3.12"
-            workers: 2
+            workers: 4
             reruns: 0
             timeout-minutes: 20
             job-timeout-minutes: 60
@@ -696,7 +696,6 @@ jobs:
               tests/unit/responses/litellm_completion_transformation
               tests/unit/responses/test_additional_tools.py
               tests/unit/responses/test_custom_tool_call.py
-              tests/unit/responses/test_dispatch.py
               tests/unit/responses/test_metadata_codex_callback.py
               tests/unit/responses/test_no_duplicate_spend_logs.py
               tests/unit/responses/test_null_test_fix.py
@@ -1451,7 +1450,6 @@ jobs:
             tests/unit/caching/test_redis_connection_pool.py \
             tests/unit/caching/test_redis_cluster_cache.py \
             tests/unit/caching/test_evicted_client_closer.py \
-            tests/integration/sdk/test_redis_cluster_iam_auth.py \
             --tb=short -vv -rs -p no:rerunfailures -p no:pytest-retry --timeout=120 --durations=20 \
             --cov=./litellm --cov-report=xml:coverage.xml
 

AGENTS.md

@@ -29,7 +29,7 @@ A test must only fail when litellm code changes. Never pin facts we don't own (a
 
 `tests/unit/` mirrors `litellm/` in a parallel path (see `tests/unit/AGENTS.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
 
-End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `AGENTS.md`
+End-to-end tests belong in `tests/e2e/`. A rule about one directory's files, like its harness or which CI runs it, goes in that directory's `AGENTS.md`, never only in a README or here. Read a directory's `AGENTS.md` before editing files in it
 
 When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
 
@@ -72,7 +72,7 @@ When referencing or running models (coding, QA'ing, writing docs, writing tests,
 
 Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
 
-If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
+Internal contributors branch a new PR off the repository's current default branch, with a branch name prefixed with litellm_
 
 Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
 

Makefile

@@ -290,7 +290,7 @@ test-rust-extension:
 	"$$temporary/venv/bin/python" -I -m mypy.stubtest \
 		--mypy-config-file tests/unit/rust_bridge/stubtest.ini \
 		litellm.rust_bridge._native && \
-	LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
+	LITELLM_LOCAL_MODEL_COST_MAP=True \
 	"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
 
 rust-sqlx-prepare:

ci_cd/generate_model_prices_schema.py

@@ -189,6 +189,10 @@ INTEGER_KEYS: dict[str, JsonSchema] = {
         **NONNEG_INTEGER,
         "description": "Smallest prefix the provider will actually cache; absent means the provider default applies.",
     },
+    "token_count_framing": {
+        **NONNEG_INTEGER,
+        "description": "Tokens added once per request to a local count estimate, for the provider's chat framing.",
+    },
     "tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."},
     "rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."},
 }
@@ -283,6 +287,11 @@ def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
             "enum": ["none", "minimal", "low", "medium", "high", "xhigh"],
         },
         "comment": STRING,
+        "token_count_fallback_model": {
+            "type": "string",
+            "minLength": 1,
+            "description": "Model that counts tokens for this one when the provider cannot count it.",
+        },
         "audio_transcription_config": STRING,
         "vertex_ai_audio_api": {
             "type": "string",

cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json

@@ -6563,6 +6563,87 @@
       ],
       "title": "litellm_project_model_rate_limit_used_metric",
       "type": "timeseries"
+    },
+    {
+      "collapsed": false,
+      "gridPos": {
+        "h": 1,
+        "w": 24,
+        "x": 0,
+        "y": 447
+      },
+      "id": 115,
+      "panels": [],
+      "title": "Agent tracing export (needs LITELLM_LENS_URL)",
+      "type": "row"
+    },
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "${DS_PROMETHEUS}"
+      },
+      "description": "Request records the Lens exporter delivered to the Lens service against the records it dropped before the service accepted them",
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "palette-classic"
+          },
+          "custom": {
+            "drawStyle": "line",
+            "fillOpacity": 10,
+            "lineWidth": 1,
+            "showPoints": "never",
+            "spanNulls": false
+          },
+          "unit": "short"
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 8,
+        "w": 24,
+        "x": 0,
+        "y": 448
+      },
+      "id": 116,
+      "options": {
+        "legend": {
+          "calcs": [],
+          "displayMode": "list",
+          "placement": "bottom",
+          "showLegend": true
+        },
+        "tooltip": {
+          "mode": "multi",
+          "sort": "desc"
+        }
+      },
+      "targets": [
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "${DS_PROMETHEUS}"
+          },
+          "editorMode": "code",
+          "expr": "sum(rate(litellm_lens_export_rows_written_total[$__rate_interval]))",
+          "legendFormat": "written",
+          "range": true,
+          "refId": "A"
+        },
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "${DS_PROMETHEUS}"
+          },
+          "editorMode": "code",
+          "expr": "sum(rate(litellm_lens_export_rows_dropped_total[$__rate_interval]))",
+          "legendFormat": "dropped",
+          "range": true,
+          "refId": "B"
+        }
+      ],
+      "title": "litellm_lens_export_rows_written and litellm_lens_export_rows_dropped rate",
+      "type": "timeseries"
     }
   ],
   "preload": false,

cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md

@@ -1,10 +1,10 @@
 # LiteLLM All Prometheus Metrics dashboard
 
-Every `litellm_*` metric family the proxy can expose on `/metrics` (141 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about
+Every `litellm_*` metric family the proxy can expose on `/metrics` (143 families across 98 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes), project model rate limits, and the agent tracing export to Lens. Panel titles are the metric names so you can grep the JSON for the metric you care about
 
 Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/unit/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard
 
-The first eleven rows need only `callbacks: ["prometheus"]`. The last three rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected
+The first eleven rows and the project model rate limits row need only `callbacks: ["prometheus"]`. The other rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, the agent tracing export row needs `LITELLM_LENS_URL`, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected
 
 ## Pre-requisites
 

docker/docker-compose.quickstart.yml

@@ -11,7 +11,7 @@
 # beyond local evaluation, pin the image to a specific release tag.
 services:
   litellm:
-    image: docker.litellm.ai/berriai/litellm:main-stable
+    image: docker.litellm.ai/berriai/litellm:latest
     ports:
       # LITELLM_BIND is empty by default, so this stays "4000:4000". The quickstart
       # script sets it to "127.0.0.1:" so new installs listen on this machine only.

enterprise/litellm_enterprise/proxy/auth/user_api_key_auth.py

@@ -1,14 +1,16 @@
-from typing import Any, Optional
+from collections.abc import Awaitable, Callable
 
 from fastapi import Request
 
 from litellm._logging import verbose_proxy_logger
-from litellm.proxy._types import ProxyException, UserAPIKeyAuth
+from litellm.proxy._types import ProxyException
 
 
 async def enterprise_custom_auth(
-    request: Request, api_key: str, user_custom_auth: Optional[Any]
-) -> Optional[UserAPIKeyAuth]:
+    request: Request,
+    api_key: str,
+    user_custom_auth: Callable[[Request, str], Awaitable[object]] | None,
+) -> object:
     from litellm_enterprise.proxy.proxy_server import custom_auth_settings
 
     if user_custom_auth is None:

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -2,27 +2,34 @@
 Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
 """
 
-from collections.abc import Sequence
+import asyncio
+from collections.abc import Mapping, Sequence
 from dataclasses import replace as dataclasses_replace
 from datetime import datetime, timedelta, timezone
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
 
+from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
+    ManagedObjectCreatorAttribution,
+    ManagedObjectRow,
+)
+
+import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
 from litellm.constants import (
-    CLI_SESSION_KEY_PREFIX,
+    BATCH_LINE_ITEMS_MAX_PART_BYTES,
     MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
     MAX_OBJECTS_PER_POLL_CYCLE,
 )
 from litellm.repositories.table_repositories import ManagedObjectRepository
 from litellm.repositories.team_repository import TeamRepository
-from litellm.repositories.user_repository import UserRepository
 from litellm.repositories.verification_token_repository import VerificationTokenRepository
 
 if TYPE_CHECKING:
     from prisma import models as prisma_models
 
+    from litellm.batches.batch_line_items import BatchFileProvider
     from litellm.integrations.prometheus import PrometheusLogger
     from litellm.proxy.utils import PrismaClient, ProxyLogging
     from litellm.repositories.prisma_protocols import TableActions
@@ -47,29 +54,6 @@ TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
 )
 
 
-class _ManagedObjectRow(Protocol):
-    @property
-    def id(self) -> str: ...
-
-    @property
-    def unified_object_id(self) -> str: ...
-
-    @property
-    def created_by(self) -> str | None: ...
-
-    @property
-    def file_object(self) -> object: ...
-
-    @property
-    def org_id(self) -> str | None: ...
-
-    @property
-    def api_key(self) -> str | None: ...
-
-    @property
-    def team_id(self) -> str | None: ...
-
-
 class _ReadableFileContent(Protocol):
     async def read(self) -> bytes: ...
 
@@ -87,14 +71,10 @@ async def _file_content_bytes(file_content: object) -> bytes:
     return cast(bytes, file_content)
 
 
-def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]":
+def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[ManagedObjectRow]":
     return ManagedObjectRepository(prisma_client).table
 
 
-def _user_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_UserTable]":
-    return UserRepository(prisma_client).table
-
-
 def _token_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_VerificationToken]":
     return VerificationTokenRepository(prisma_client).table
 
@@ -103,6 +83,11 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li
     return TeamRepository(prisma_client).table
 
 
+def _row_metadata(row: object) -> dict[str, object]:
+    metadata: Final = getattr(row, "metadata", None)
+    return cast(dict[str, object], metadata) if isinstance(metadata, dict) else {}  # cast-ok: prisma Json object column
+
+
 class CheckBatchCost:
     def __init__(
         self,
@@ -118,10 +103,15 @@ class CheckBatchCost:
         self.prisma_client: PrismaClient = prisma_client
         self.llm_router: Router = llm_router
         self._track_unmanaged_batch_cost = track_unmanaged_batch_cost
+        self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckBatchCost")
         # Cached after the first poll cycle. Once we know the column is absent we skip
         # the guaranteed-failing primary query on every subsequent cycle.
         self._has_batch_processed_column: bool = True
         self.batch_processed_support_confirmed: bool = False
+        self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
+        self._line_item_store_turn: Final = asyncio.Lock()
+        self._poll_idle: Final = asyncio.Event()
+        self._poll_idle.set()
 
     @staticmethod
     def _is_missing_batch_processed_column_error(err: Exception) -> bool:
@@ -149,133 +139,6 @@ class CheckBatchCost:
             return
         self.batch_processed_support_confirmed = True
 
-    async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> dict[str, str | None]:
-        """
-        Look up user email and key alias by user_id for enriching the S3 callback metadata.
-        Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None).
-        Returns an empty dict when user_id is None: batches created by a team or service
-        account key carry no user id, and find_unique(where={"user_id": None}) raises.
-        """
-        if not user_id:
-            return {}
-        try:
-            user_row: prisma_models.LiteLLM_UserTable | None = await _user_table(self.prisma_client).find_unique(
-                where={"user_id": user_id}
-            )
-            if user_row is None:
-                return {}
-            return {
-                "user_api_key_user_email": getattr(user_row, "user_email", None),
-                "user_api_key_alias": getattr(user_row, "user_alias", None),
-            }
-        except Exception as e:
-            verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
-            return {}
-
-    async def _get_key_alias(self, batch_id: str, api_key: str | None, created_by: str | None) -> str | None:
-        """Resolve the creating virtual key's alias from its hashed token."""
-        if not api_key:
-            return None
-        if created_by and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}":
-            return api_key
-        try:
-            key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table(
-                self.prisma_client
-            ).find_unique(where={"token": api_key})
-            return getattr(key_row, "key_alias", None) if key_row is not None else None
-        except Exception as e:
-            verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}")
-            return None
-
-    async def _get_team_alias(self, team_id: str | None) -> str | None:
-        """Resolve a team's alias from its id."""
-        if not team_id:
-            return None
-        try:
-            team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique(
-                where={"team_id": team_id}
-            )
-            return getattr(team_row, "team_alias", None) if team_row is not None else None
-        except Exception as e:
-            verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
-            return None
-
-    async def _get_org_id(self, job: "_ManagedObjectRow", batch_id: str) -> str | None:
-        org_id: Final = getattr(job, "org_id", None)
-        if org_id:
-            return org_id
-        api_key: Final = getattr(job, "api_key", None)
-        team_id: Final = getattr(job, "team_id", None)
-        if api_key:
-            try:
-                key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table(
-                    self.prisma_client
-                ).find_unique(where={"token": api_key})
-                key_org_id: Final = (
-                    cast(str | None, getattr(key_row, "organization_id", None))
-                    if key_row is not None
-                    else None
-                )
-                if key_org_id:
-                    return key_org_id
-            except Exception as e:
-                verbose_proxy_logger.error(
-                    f"CheckBatchCost: could not resolve the key's org for batch {batch_id}, "
-                    f"still trying the team's: {e}"
-                )
-        if not team_id:
-            return None
-        try:
-            team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique(
-                where={"team_id": team_id}
-            )
-            return cast(str | None, getattr(team_row, "organization_id", None)) if team_row is not None else None
-        except Exception as e:
-            verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}")
-            return None
-
-    async def _build_creator_attribution_metadata(
-        self, job: "_ManagedObjectRow", batch_id: str
-    ) -> dict[str, object]:
-        """
-        Rebuild the spend-tracking metadata for the key, team, and tags that created the
-        batch so the batch-cost spend log is attributed the same way a non-batch request
-        is. Rows created before api_key and request_tags were persisted carry only
-        created_by and team_id, and fall back to those. A named creating key owns
-        user_api_key_alias; when it has no alias, or the key has since been rotated or
-        deleted, the field keeps the creating user's alias that _get_user_info filled in,
-        because a resolvable name is more useful on the spend row than a null.
-
-        user_api_key_org_id must be resolved here too: the spend update writer reads it
-        off this metadata to increment organization spend, so leaving it out silently
-        drops batch cost from org accounting for keys and teams that belong to one.
-        """
-        api_key = getattr(job, "api_key", None)
-        team_id = getattr(job, "team_id", None)
-        request_tags = getattr(job, "request_tags", None)
-
-        metadata: dict[str, object] = {
-            "user_api_key_user_id": job.created_by,
-            "user_api_key": api_key,
-            "user_api_key_hash": api_key,
-            "user_api_key_team_id": team_id,
-            **(await self._get_user_info(batch_id, job.created_by)),
-        }
-
-        key_alias = await self._get_key_alias(batch_id, api_key, job.created_by)
-        if key_alias is not None:
-            metadata["user_api_key_alias"] = key_alias
-        team_alias = await self._get_team_alias(team_id)
-        if team_alias is not None:
-            metadata["user_api_key_team_alias"] = team_alias
-        org_id: Final = await self._get_org_id(job, batch_id)
-        if org_id is not None:
-            metadata["user_api_key_org_id"] = org_id
-        if isinstance(request_tags, list) and request_tags:
-            metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)]
-
-        return metadata
-
     async def _cleanup_stale_managed_objects(self) -> None:
         """
         Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
@@ -317,7 +180,7 @@ class CheckBatchCost:
                 f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
             )
 
-    async def _fallback_find_jobs(self) -> "Sequence[_ManagedObjectRow]":
+    async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
         """Query batch jobs without the batch_processed filter (for older schemas)."""
         return await _managed_object_table(self.prisma_client).find_many(
             where={
@@ -337,7 +200,7 @@ class CheckBatchCost:
             order={"created_at": "asc"},
         )
 
-    async def _retire_job(self, job: "_ManagedObjectRow", reason: str) -> None:
+    async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
         """
         Take a row that can never be costed out of the poll page. Leaving it selectable
         would burn one of the MAX_OBJECTS_PER_POLL_CYCLE slots on every future cycle, and
@@ -364,7 +227,7 @@ class CheckBatchCost:
             "so it will no longer be polled"
         )
 
-    async def _claim_job_for_costing(self, job: "_ManagedObjectRow") -> bool:
+    async def _claim_job_for_costing(self, job: "ManagedObjectRow") -> bool:
         """
         Atomically flip batch_processed from false to true, returning whether this pod won
         the row. Every pod and uvicorn worker schedules its own poller against the shared
@@ -390,7 +253,7 @@ class CheckBatchCost:
             return False
         return claimed > 0
 
-    async def _release_job_claim(self, job: "_ManagedObjectRow") -> None:
+    async def _release_job_claim(self, job: "ManagedObjectRow") -> None:
         """Give a claimed row back once billing it failed, so a later poll cycle retries it.
 
         Safe to match on batch_processed=True: while this poller is active the retrieve
@@ -411,7 +274,7 @@ class CheckBatchCost:
             )
 
     @staticmethod
-    def _has_unified_id_without_model(job: "_ManagedObjectRow") -> bool:
+    def _has_unified_id_without_model(job: "ManagedObjectRow") -> bool:
         """A unified id that decodes but carries no model_id can never be routed."""
         from litellm.proxy.openai_files_endpoints.common_utils import (
             convert_b64_uid_to_unified_uid,
@@ -460,10 +323,10 @@ class CheckBatchCost:
         return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
 
     async def _finalize_unbilled_terminal_job(
-        self, job: "_ManagedObjectRow", response: "LiteLLMBatch"
-    ) -> None:
+        self, job: "ManagedObjectRow", response: "LiteLLMBatch"
+    ) -> bool:
         """Persist a terminal batch that has nothing billable, converting any raw
-        provider file ids to managed ids, and take it out of the poll page."""
+        provider file ids to managed ids, and take it out of the poll page. Returns whether it was persisted."""
         try:
             from litellm.proxy.openai_files_endpoints.common_utils import (
                 _is_base64_encoded_unified_file_id,
@@ -491,10 +354,12 @@ class CheckBatchCost:
             verbose_proxy_logger.info(
                 f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
             )
+            return True
         except Exception as db_err:
             verbose_proxy_logger.error(
                 f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}"
             )
+            return False
 
     @staticmethod
     def _record_error(
@@ -505,7 +370,7 @@ class CheckBatchCost:
 
     def _resolve_job_routing(
         self,
-        job: "_ManagedObjectRow",
+        job: "ManagedObjectRow",
         prom_logger: Optional["PrometheusLogger"],
     ) -> Optional[Tuple[str, str]]:
         """
@@ -582,7 +447,7 @@ class CheckBatchCost:
 
     def _resolve_unmanaged_provider_routing(
         self,
-        job: "_ManagedObjectRow",
+        job: "ManagedObjectRow",
         prom_logger: Optional["PrometheusLogger"],
         llm_provider: str,
         bare_model_name: str,
@@ -678,7 +543,7 @@ class CheckBatchCost:
     @classmethod
     def _get_managed_file_model_name(
         cls,
-        job: "_ManagedObjectRow",
+        job: "ManagedObjectRow",
         deployment_info: "Deployment",
     ) -> Optional[str]:
         """
@@ -698,7 +563,7 @@ class CheckBatchCost:
         )
 
     @staticmethod
-    def _get_input_file_id(job: "_ManagedObjectRow") -> Optional[str]:
+    def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
         import json
 
         from litellm.types.utils import LiteLLMBatch
@@ -719,7 +584,7 @@ class CheckBatchCost:
 
     async def _track_completed_batch_cost(
         self,
-        job: "_ManagedObjectRow",
+        job: "ManagedObjectRow",
         response: "LiteLLMBatch",
         model_id: str,
         batch_id: str,
@@ -735,9 +600,9 @@ class CheckBatchCost:
 
         """
         from litellm.batches.batch_utils import (
+            calculate_batch_cost_and_usage,
             count_error_file_failed_requests,
             get_file_content_as_dictionary,
-            calculate_batch_cost_and_usage,
         )
         from litellm.files.main import afile_content
         from litellm.files.types import FileContentCallOptions
@@ -911,7 +776,7 @@ class CheckBatchCost:
                 },
                 **({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}),
                 "metadata": {
-                    **(await self._build_creator_attribution_metadata(job, batch_id)),
+                    **(await self._creator_attribution.batch_metadata(job, batch_id)),
                     # spend logs read the deployment identity off these metadata keys, so
                     # without them the batch cost row carries no model_id or model_group
                     "model_info": {"id": model_id},
@@ -956,7 +821,191 @@ class CheckBatchCost:
 
         return model_name, str(llm_provider) if llm_provider else None
 
-    async def check_batch_cost(self):
+    async def _retire_unbilled_terminal_job(
+        self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str
+    ) -> None:
+        """Persist a terminal batch that has nothing billable, then store its lines, reading its input file by the id
+        the provider reported rather than the managed id that was persisted."""
+        provider_input_file_id: Final = response.input_file_id
+        if await self._finalize_unbilled_terminal_job(job, response):
+            response.input_file_id = provider_input_file_id
+            self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=False)
+
+    def _store_batch_line_items_in_background(
+        self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
+    ) -> None:
+        """Store a retired batch's lines off the poll loop, one batch at a time and only between polls, so neither
+        pricing the lines nor a slow or hung bucket holds up billing or retiring the other batches. While a poll
+        cycle's worth of batches is still waiting, the lines of the next one are dropped with an error."""
+        if not litellm.batch_line_item_callbacks:
+            return
+        if len(self._line_item_stores) >= MAX_OBJECTS_PER_POLL_CYCLE:
+            verbose_proxy_logger.error(
+                f"CheckBatchCost: {len(self._line_item_stores)} batches are still waiting to store their line items, "
+                f"so the line items of batch {batch_id} (job {job.id}) were not stored"
+            )
+            return
+        store: Final = asyncio.create_task(
+            self._store_batch_line_items_in_turn(job, response, model_id, batch_id, billed=billed)
+        )
+        self._line_item_stores.add(store)
+        store.add_done_callback(self._line_item_stores.discard)
+
+    async def _store_batch_line_items_in_turn(
+        self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
+    ) -> None:
+        async with self._line_item_store_turn:
+            await self._poll_idle.wait()
+            await self._store_batch_line_items(job, response, model_id, batch_id, billed=billed)
+
+    async def _store_batch_line_items(
+        self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
+    ) -> None:
+        """Send every line of a finished batch to batch_line_item_callbacks, one part at a time. A billed batch's
+        output file is read again here, so nothing waiting for its turn holds a batch's files. Failures, including an
+        output file that can no longer be read, are logged and never raised, so the batch stays billed and retired
+        whether or not its lines were stored."""
+        if not litellm.batch_line_item_callbacks:
+            return
+        from litellm.batches.batch_line_items import (
+            BatchFileProvider,
+            BatchLineItemSource,
+            batch_line_item_parts,
+            dispatch_batch_line_items,
+            line_item_metadata,
+        )
+        from litellm.batches.batch_utils import fetch_batch_managed_file_content
+        from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
+        from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info
+
+        try:
+            deployment_info: Final = self.llm_router.get_deployment(model_id=model_id)
+            if deployment_info is None:
+                verbose_proxy_logger.error(
+                    f"CheckBatchCost: deployment {model_id} is no longer in the router, so the line items of batch "
+                    f"{batch_id} (job {job.id}) were not stored"
+                )
+                return
+            litellm_model_name: Final = deployment_info.litellm_params.model
+            model_name, llm_provider, _, _ = get_llm_provider(
+                model=litellm_model_name,
+                custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
+            )
+            cost_provider: Final = cast(BatchFileProvider, llm_provider)
+            credentials: Final = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
+            file_access: Final = {
+                **credentials,
+                "_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
+            }
+            output_content: Final = (
+                await fetch_batch_managed_file_content(
+                    response.output_file_id, custom_llm_provider=cost_provider, litellm_params=file_access
+                )
+                if billed and response.output_file_id
+                else b""
+            )
+            error_content: Final = await self._batch_file_content(response.error_file_id, cost_provider, file_access)
+            if not output_content and not error_content:
+                return
+            redact_bodies: Final = await self._creator_turned_off_message_logging(job)
+            source: Final = BatchLineItemSource(
+                batch_id=job.unified_object_id,
+                finished_at=response.completed_at
+                or response.failed_at
+                or response.expired_at
+                or response.cancelled_at
+                or response.created_at,
+                model=model_name,
+                custom_llm_provider=llm_provider,
+                cost_provider=cost_provider,
+                model_info=deployment_pricing_model_info(model_id=model_id, deployment_model=litellm_model_name),
+                metadata=line_item_metadata(await self._creator_attribution.batch_metadata(job, batch_id)),
+            )
+            input_content: Final = (
+                b""
+                if redact_bodies
+                else await self._batch_file_content(response.input_file_id, cost_provider, file_access)
+            )
+            parts: Final = batch_line_item_parts(
+                source, input_content, (output_content, error_content), BATCH_LINE_ITEMS_MAX_PART_BYTES
+            )
+            handed = 0
+            # pricing a part's lines is CPU-bound, so each part is built off the event loop
+            while (payload := await asyncio.to_thread(next, parts, None)) is not None:
+                await dispatch_batch_line_items(payload, redact_bodies=redact_bodies)
+                handed += len(payload["lines"])
+            if handed:
+                verbose_proxy_logger.info(
+                    f"CheckBatchCost: handed the {handed} line items of batch {batch_id} to batch_line_item_callbacks"
+                )
+        except Exception as error:  # noqa: BLE001  # storing the lines is best effort, the batch is already billed
+            verbose_proxy_logger.exception(
+                f"CheckBatchCost: the line items of batch {batch_id} (job {job.id}) were not stored: {error}"
+            )
+
+    @staticmethod
+    async def _batch_file_content(
+        file_id: str | None, custom_llm_provider: "BatchFileProvider", file_access: Mapping[str, object]
+    ) -> bytes:
+        """A batch file's raw JSONL, or nothing when the batch has no such file or the provider will not serve it."""
+        from litellm.batches.batch_utils import fetch_batch_managed_file_content
+
+        if not file_id:
+            return b""
+        try:
+            return await fetch_batch_managed_file_content(
+                file_id, custom_llm_provider=custom_llm_provider, litellm_params=file_access
+            )
+        except Exception as error:  # noqa: BLE001  # a missing input or error file leaves those lines out, not the rest
+            verbose_proxy_logger.warning(f"CheckBatchCost: could not read batch file {file_id}: {error}")
+            return b""
+
+    async def _creator_turned_off_message_logging(self, job: "ManagedObjectRow") -> bool:
+        """Whether the creating key's or team's logging settings turn off message logging, falling back to the
+        proxy-wide setting. A batch whose saved credential is not a virtual key's hash, such as one created under JWT
+        or OAuth2 auth or with the master key, has only its team's and the proxy's settings. A key or team that can no
+        longer be read counts as turned off."""
+        from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
+        from litellm.proxy._types import UserAPIKeyAuth
+        from litellm.proxy.litellm_pre_call_utils import get_dynamic_logging_metadata
+        from litellm.proxy.proxy_server import proxy_config
+        from litellm.secret_managers.main import str_to_bool
+
+        api_key: Final = getattr(job, "api_key", None)
+        team_id: Final = getattr(job, "team_id", None)
+        has_key_row: Final = api_key is not None and is_valid_sha256_hash(api_key)
+        try:
+            key_row: Final = (
+                await _token_table(self.prisma_client).find_unique(where={"token": api_key}) if has_key_row else None
+            )
+            team_row: Final = (
+                await _team_table(self.prisma_client).find_unique(where={"team_id": team_id}) if team_id else None
+            )
+        except Exception as error:
+            verbose_proxy_logger.error(
+                f"CheckBatchCost: could not read the logging settings of batch job {job.id}'s creator, "
+                f"so its line items are stored without request and response: {error}"
+            )
+            return True
+        if (has_key_row and key_row is None) or (team_id and team_row is None):
+            return True
+        settings: Final = get_dynamic_logging_metadata(
+            UserAPIKeyAuth(metadata=_row_metadata(key_row), team_metadata=_row_metadata(team_row), team_id=team_id),
+            proxy_config,
+        )
+        callback_vars: Final = (settings.callback_vars if settings is not None else None) or {}
+        creator_setting: Final = str_to_bool(callback_vars.get("turn_off_message_logging"))
+        return creator_setting if creator_setting is not None else litellm.turn_off_message_logging is True
+
+    async def check_batch_cost(self) -> None:
+        """Run one poll cycle. Batch line items wait until it has billed and retired every batch it found."""
+        self._poll_idle.clear()
+        try:
+            await self._poll_batches()
+        finally:
+            self._poll_idle.set()
+
+    async def _poll_batches(self):
         """
         Check if the batch JOB has been tracked.
         - get all status="validating" and file_purpose="batch" jobs
@@ -1070,7 +1119,7 @@ class CheckBatchCost:
                             f"CheckBatchCost: output file {response.output_file_id} of batch {batch_id} "
                             f"does not exist at the provider; retiring job {job.id} unbilled"
                         )
-                        await self._finalize_unbilled_terminal_job(job, response)
+                        await self._retire_unbilled_terminal_job(job, response, model_id, batch_id)
                         continue
                     verbose_proxy_logger.error(
                         f"CheckBatchCost: failed to track cost for batch {batch_id} "
@@ -1100,6 +1149,10 @@ class CheckBatchCost:
                     verbose_proxy_logger.error(
                         f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
                     )
+                    if not self._has_batch_processed_column:
+                        # nothing claimed the row, so it stays in the poll page and its lines wait for that retry
+                        continue
+                self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True)
 
             elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
                 from litellm.proxy.openai_files_endpoints.common_utils import (
@@ -1112,7 +1165,7 @@ class CheckBatchCost:
                         f"has not appeared yet; leaving job {job.id} for the next poll cycle"
                     )
                     continue
-                await self._finalize_unbilled_terminal_job(job, response)
+                await self._retire_unbilled_terminal_job(job, response, model_id, batch_id)
 
         # Record polling run metrics (always, even if nothing was processed)
         if prom_logger:

enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py

@@ -1,12 +1,36 @@
 """
 Polls LiteLLM_ManagedObjectTable to check if the response is complete.
-Cost tracking is handled by the get-responses call, which prices normally only because the
-poll stamps itself with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN; user-facing reads of the
-same route are non-inference and free.
+A background response is billed exactly once: whichever client read or poll first sees it
+finished moves its row out of a pending status, then logs the body that read already fetched as
+one spend event stamped with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN and the creating key's
+attribution, priced against the deployment that served it. The poll's status checks are stamped
+with BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN and client reads are marked managed by the
+managed files hook, so both stay free. The spend event carries the id the client holds, the row's
+unified_object_id, so its spend log is found under that id with the _background_cost suffix.
+A read that finds no row, because the create could not store one, stores it under the reading key
+and bills it the same way, so a response the create left untracked is still billed once. A streamed
+background create is free on its own; its stream stores the row under the creating key and bills
+the finished response through it before the client sees response.completed.
+A read the caller retargets with client-side credentials, and a delete of a still-pending response,
+settle the bill from a fetch through the configured deployment instead: the client controls what a
+retargeted read returns, and a deleted response leaves the poll nothing to fetch.
+A row stored before this billing model keeps the encrypted create id as its model_object_id. The
+first lookup by that id rekeys it by the provider id, and the poll retires such a row without
+billing it once a read has already tracked the response under the provider id.
 """
 
+import uuid
+from collections.abc import Awaitable, Callable, Mapping
+from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Dict, Final, Protocol, cast
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Final, cast
+
+from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
+    ManagedObjectCreatorAttribution,
+    ManagedObjectRow,
+)
+from pydantic import BaseModel, ConfigDict, TypeAdapter
 
 import litellm
 from litellm._logging import verbose_proxy_logger
@@ -14,36 +38,124 @@ from litellm.constants import (
     INTERNAL_CALL_ORIGIN_METADATA_KEY,
     MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
     MAX_OBJECTS_PER_POLL_CYCLE,
+    PENDING_RESPONSE_STATUSES,
     STALE_OBJECT_CLEANUP_BATCH_SIZE,
+    TERMINAL_RESPONSE_STATUSES,
 )
+from litellm.proxy.hooks.responses_id_security import decrypted_response_id
 from litellm.repositories.table_repositories import ManagedObjectRepository
 from litellm.responses.utils import ResponsesAPIRequestUtils
 from litellm.types.llms.openai import ResponsesAPIResponse
-from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
+from litellm.types.router import Deployment
+from litellm.types.utils import (
+    BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
+    BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
+)
 
 if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
     from litellm.proxy.utils import PrismaClient, ProxyLogging
     from litellm.repositories.prisma_protocols import TableActions
     from litellm.router import Router
 
-TERMINAL_RESPONSE_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete"})
+_STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
+_RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses"
 
 
-class _ManagedObjectRow(Protocol):
-    @property
-    def id(self) -> str: ...
-
-    @property
-    def unified_object_id(self) -> str: ...
-
-    @property
-    def created_by(self) -> str | None: ...
-
-    @property
-    def file_object(self) -> object: ...
+@dataclass(frozen=True, slots=True)
+class _PollTarget:
+    response_id: str
+    model_metadata: Mapping[str, object]
+    deployment: Deployment | None
 
 
-def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]":
+@dataclass(frozen=True, slots=True)
+class BilledBackgroundResponse:
+    response: ResponsesAPIResponse
+    model: str
+    custom_llm_provider: str | None
+    litellm_params: Mapping[str, object]
+    metadata: Mapping[str, object]
+
+
+@dataclass(frozen=True, slots=True)
+class BilledCharge:
+    cost: float | None
+    logging_obj: "LiteLLMLogging"
+
+
+BilledResponseLogger = Callable[[BilledBackgroundResponse], Awaitable[BilledCharge]]
+UntrackedResponseTracker = Callable[[ResponsesAPIResponse], Awaitable[bool]]
+
+
+def _without_read_price(response: ResponsesAPIResponse) -> ResponsesAPIResponse:
+    unpriced: Final = response.model_copy(deep=True)
+    unpriced.hidden_params = {key: value for key, value in response.hidden_params.items() if key != "response_cost"}
+    return unpriced
+
+
+async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
+    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+
+    logging_obj: Final = LiteLLMLogging(
+        model=billed.model,
+        messages=[{"role": "user", "content": "<retrieve_response>"}],
+        stream=False,
+        call_type=_RETRIEVE_RESPONSE_CALL_TYPE,
+        start_time=datetime.now(),
+        litellm_call_id=str(uuid.uuid4()),
+        function_id=str(uuid.uuid4()),
+    )
+    logging_obj.update_environment_variables(
+        litellm_params={**billed.litellm_params, "metadata": dict(billed.metadata)},
+        optional_params={},
+        custom_llm_provider=billed.custom_llm_provider,
+    )
+    unpriced: Final = _without_read_price(billed.response)
+    cost: Final = logging_obj.response_cost_calculator(result=unpriced)
+    end_time: Final = datetime.now()
+    GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
+        async_coroutine=logging_obj.async_success_handler(
+            result=unpriced, start_time=logging_obj.start_time, end_time=end_time
+        )
+    )
+    logging_obj.handle_sync_success_callbacks_for_async_calls(
+        result=unpriced, start_time=logging_obj.start_time, end_time=end_time
+    )
+    return BilledCharge(cost=cost, logging_obj=logging_obj)
+
+
+class _DeploymentBaseModel(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    base_model: str | None = None
+
+
+def _deployment_litellm_params(deployment: Deployment) -> Mapping[str, object]:
+    from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials
+
+    api_base: Final = deployment.litellm_params.api_base
+    base_model: Final = _DeploymentBaseModel.model_validate(deployment.litellm_params, from_attributes=True).base_model
+    return MappingProxyType(
+        {
+            **({"api_base": mask_api_base_credentials(api_base)} if api_base else {}),
+            **({"base_model": base_model} if base_model else {}),
+        }
+    )
+
+
+def _deployment_metadata(deployment: Deployment) -> Mapping[str, object]:
+    return MappingProxyType(
+        {
+            "model_info": deployment.model_info.model_dump(exclude_none=True),
+            "model_group": deployment.model_name,
+            "deployment": deployment.litellm_params.model,
+        }
+    )
+
+
+def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[ManagedObjectRow]":
     return ManagedObjectRepository(prisma_client).table
 
 
@@ -51,12 +163,17 @@ def _is_response_gone_at_provider(error: Exception, provider_response_id: str) -
     return getattr(error, "status_code", None) == 404 and provider_response_id in str(error)
 
 
+def _keyed_by_issued_id(job: ManagedObjectRow) -> bool:
+    return decrypted_response_id(job.model_object_id) != job.model_object_id
+
+
 class CheckResponsesCost:
     def __init__(
         self,
         proxy_logging_obj: "ProxyLogging",
         prisma_client: "PrismaClient",
         llm_router: "Router",
+        log_billed_response: BilledResponseLogger = log_billed_background_response,
     ):
         from litellm.proxy.utils import PrismaClient, ProxyLogging
         from litellm.router import Router
@@ -64,15 +181,17 @@ class CheckResponsesCost:
         self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
         self.prisma_client: PrismaClient = prisma_client
         self.llm_router: Router = llm_router
+        self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckResponsesCost")
+        self._log_billed_response: Final = log_billed_response
 
-    def _resolve_deployment(self, response_id: str) -> bool:
+    def _resolve_deployment(self, response_id: str) -> Deployment | None:
         model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
-        return model_id is not None and self.llm_router.get_deployment(model_id=model_id) is not None
+        return self.llm_router.get_deployment(model_id=model_id) if model_id is not None else None
 
     async def _get_response(
         self,
         response_id: str,
-        litellm_metadata: Dict[str, str],
+        litellm_metadata: dict[str, object],
         via_router: bool,
     ) -> ResponsesAPIResponse:
         """Fetch the upstream response, using deployment credentials when available.
@@ -91,9 +210,17 @@ class CheckResponsesCost:
         )
         return cast(ResponsesAPIResponse, router_response)
 
-    async def _expire_stale_rows(
-        self, cutoff: datetime, batch_size: int
-    ) -> int:
+    async def _fetch_unbilled(self, target: _PollTarget) -> ResponsesAPIResponse:
+        return await self._get_response(
+            response_id=target.response_id,
+            litellm_metadata={
+                **target.model_metadata,
+                INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
+            },
+            via_router=target.deployment is not None,
+        )
+
+    async def _expire_stale_rows(self, cutoff: datetime, batch_size: int) -> int:
         """Execute the bounded UPDATE that marks stale rows as 'stale_expired'.
 
         Isolated so it can be swapped / mocked in tests without touching the
@@ -140,14 +267,154 @@ class CheckResponsesCost:
                 f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
             )
 
+    def _target(self, response_id: str, model_name: str | None) -> _PollTarget:
+        return _PollTarget(
+            response_id=response_id,
+            model_metadata=MappingProxyType({"model": model_name, "model_group": model_name} if model_name else {}),
+            deployment=self._resolve_deployment(response_id),
+        )
+
+    def _poll_target(self, job: ManagedObjectRow) -> _PollTarget:
+        stored_model: Final = _STORED_RESPONSE_ADAPTER.validate_python(job.file_object).get("model")
+        return self._target(
+            decrypted_response_id(job.unified_object_id), stored_model if isinstance(stored_model, str) else None
+        )
+
+    async def _claim_for_billing(self, job: ManagedObjectRow) -> bool:
+        claimed: Final = await _managed_object_table(self.prisma_client).update_many(
+            # bounded-ok: PENDING_RESPONSE_STATUSES is a two-value constant tuple
+            where={"id": job.id, "status": {"in": list(PENDING_RESPONSE_STATUSES)}},
+            data={"status": "completed"},
+        )
+        return claimed > 0
+
+    async def _release_billing_claim(self, job: ManagedObjectRow) -> None:
+        await _managed_object_table(self.prisma_client).update_many(
+            where={"id": job.id, "status": "completed"},
+            data={"status": job.status or PENDING_RESPONSE_STATUSES[0]},
+        )
+
+    async def _billed_response(
+        self,
+        job: ManagedObjectRow,
+        target: _PollTarget,
+        response: ResponsesAPIResponse,
+    ) -> BilledBackgroundResponse:
+        deployment: Final = target.deployment
+        return BilledBackgroundResponse(
+            response=response.model_copy(update={"id": job.unified_object_id}),
+            model=deployment.litellm_params.model if deployment is not None else (response.model or ""),
+            custom_llm_provider=ResponsesAPIRequestUtils.decode_responses_api_response_id(target.response_id).get(
+                "custom_llm_provider"
+            ),
+            litellm_params=_deployment_litellm_params(deployment) if deployment is not None else MappingProxyType({}),
+            metadata=MappingProxyType(
+                {
+                    **(await self._creator_attribution.response_metadata(job, job.unified_object_id)),
+                    INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
+                    **target.model_metadata,
+                    **(_deployment_metadata(deployment) if deployment is not None else {}),
+                }
+            ),
+        )
+
+    async def _bill_once(
+        self,
+        job: ManagedObjectRow,
+        target: _PollTarget,
+        response: ResponsesAPIResponse,
+    ) -> BilledCharge | None:
+        if not await self._claim_for_billing(job):
+            verbose_proxy_logger.info(f"Response {job.unified_object_id} was already billed elsewhere, skipping")
+            return None
+        try:
+            charge: Final = await self._log_billed_response(await self._billed_response(job, target, response))
+        except BaseException:
+            await self._release_billing_claim(job)
+            raise
+        verbose_proxy_logger.info(f"Response {job.unified_object_id} reached a terminal status and was billed")
+        return charge
+
+    async def _job_for_provider_id(self, provider_response_id: str) -> ManagedObjectRow | None:
+        return await _managed_object_table(self.prisma_client).find_first(
+            where={"model_object_id": provider_response_id, "file_purpose": "response"}
+        )
+
+    async def _rekey_by_provider_id(self, job: ManagedObjectRow) -> None:
+        await _managed_object_table(self.prisma_client).update_many(
+            where={"id": job.id}, data={"model_object_id": decrypted_response_id(job.model_object_id)}
+        )
+
+    async def _tracked_job(self, response_id: str) -> ManagedObjectRow | None:
+        if tracked := await self._job_for_provider_id(decrypted_response_id(response_id)):
+            return tracked
+        stored_by_issued_id: Final = await _managed_object_table(self.prisma_client).find_first(
+            where={"unified_object_id": response_id, "file_purpose": "response"}
+        )
+        if stored_by_issued_id is not None and _keyed_by_issued_id(stored_by_issued_id):
+            await self._rekey_by_provider_id(stored_by_issued_id)
+        return stored_by_issued_id
+
+    async def _superseded_by_a_provider_id_row(self, job: ManagedObjectRow) -> bool:
+        if not _keyed_by_issued_id(job):
+            return False
+        if await self._job_for_provider_id(decrypted_response_id(job.model_object_id)) is None:
+            await self._rekey_by_provider_id(job)
+            return False
+        await self._claim_for_billing(job)
+        return True
+
+    async def _track_from_deployment(
+        self, response_id: str, track_untracked: UntrackedResponseTracker | None
+    ) -> BilledCharge | None:
+        target: Final = self._target(decrypted_response_id(response_id), None)
+        if track_untracked is None or target.deployment is None:
+            return None
+        return await self.bill_if_finished(response_id, await self._fetch_unbilled(target), track_untracked)
+
+    async def _start_tracking(
+        self, response_id: str, response: ResponsesAPIResponse, track_untracked: UntrackedResponseTracker | None
+    ) -> ManagedObjectRow | None:
+        if track_untracked is None or not await track_untracked(response):
+            return None
+        return await self._tracked_job(response_id)
+
+    async def bill_if_finished(
+        self,
+        response_id: str,
+        response: ResponsesAPIResponse,
+        track_untracked: UntrackedResponseTracker | None = None,
+    ) -> BilledCharge | None:
+        if response.status not in TERMINAL_RESPONSE_STATUSES:
+            return None
+        job: Final = await self._tracked_job(response_id) or await self._start_tracking(
+            response_id, response, track_untracked
+        )
+        if job is None or job.status not in PENDING_RESPONSE_STATUSES:
+            return None
+        return await self._bill_once(job, self._poll_target(job), response)
+
+    async def bill_from_deployment(
+        self, response_id: str, track_untracked: UntrackedResponseTracker | None = None
+    ) -> BilledCharge | None:
+        job: Final = await self._tracked_job(response_id)
+        if job is None:
+            return await self._track_from_deployment(response_id, track_untracked)
+        if job.status not in PENDING_RESPONSE_STATUSES:
+            return None
+        target: Final = self._poll_target(job)
+        response: Final = await self._fetch_unbilled(target)
+        if response.status not in TERMINAL_RESPONSE_STATUSES:
+            return None
+        return await self._bill_once(job, target, response)
+
     async def check_responses_cost(self):
         """
         Check if background responses are complete and track their cost.
         - Get all status="queued" or "in_progress" and file_purpose="response" jobs
-        - Query the provider to check if response is complete
-        - Cost is tracked by the get-responses call, billed because the poll is stamped
-          with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
-        - Mark responses in a terminal state as complete in the database
+        - Query the provider, unbilled, to check if the response is complete
+        - Claim each terminal job by moving its row to completed, then log the body that check fetched
+          as its one billed event, so two pollers racing over the same job bill it once
         - Mark responses the provider no longer has (404 through a resolved
           deployment) as stale_expired
         """
@@ -160,99 +427,53 @@ class CheckResponsesCost:
 
         jobs = await _managed_object_table(self.prisma_client).find_many(
             where={
-                "status": {"in": ["queued", "in_progress"]},
+                # bounded-ok: PENDING_RESPONSE_STATUSES is a two-value constant tuple
+                "status": {"in": list(PENDING_RESPONSE_STATUSES)},
                 "file_purpose": "response",
             },
             take=MAX_OBJECTS_PER_POLL_CYCLE,
             order={"created_at": "asc"},
         )
-        
+
         verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check")
-        completed_jobs: Final[list[_ManagedObjectRow]] = []
-        expired_jobs: Final[list[_ManagedObjectRow]] = []
+        expired_jobs: Final[list[ManagedObjectRow]] = []
 
         for job in jobs:
             unified_object_id = job.unified_object_id
 
             try:
-                from litellm.proxy.hooks.responses_id_security import (
-                    ResponsesIDSecurity,
-                )
-
-                # Get the stored response object to extract model information
-                stored_response = job.file_object
-                model_name = stored_response.get("model", None)
-
-                # Decrypt the response ID
-                responses_id_security, _, _ = ResponsesIDSecurity()._decrypt_response_id(unified_object_id)
-
-                # Prepare metadata with model information for cost tracking
-                litellm_metadata = {
-                    "user_api_key_user_id": job.created_by or "default-user-id",
-                    INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
-                }
-
-                # Add model information if available
-                if model_name:
-                    litellm_metadata["model"] = model_name
-                    litellm_metadata["model_group"] = model_name  # Use same value for model_group
-                via_router = self._resolve_deployment(responses_id_security)
+                if await self._superseded_by_a_provider_id_row(job):
+                    verbose_proxy_logger.info(f"Response {unified_object_id} is billed under its provider id row")
+                    continue
+                target = self._poll_target(job)
             except Exception as e:
-                verbose_proxy_logger.warning(
-                    f"Skipping job {unified_object_id} due to error: {e}"
-                )
+                verbose_proxy_logger.warning(f"Skipping job {unified_object_id} due to error: {e}")
                 continue
 
-            provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id(
-                responses_id_security
-            ).get("response_id", responses_id_security)
+            provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id(target.response_id).get(
+                "response_id", target.response_id
+            )
             try:
-                response = await self._get_response(
-                    response_id=responses_id_security,
-                    litellm_metadata=litellm_metadata,
-                    via_router=via_router,
-                )
+                response = await self._fetch_unbilled(target)
 
-                verbose_proxy_logger.debug(
-                    f"Response {unified_object_id} status: {response.status}, model: {model_name}"
-                )
+                verbose_proxy_logger.debug(f"Response {unified_object_id} status: {response.status}")
+                if response.status in TERMINAL_RESPONSE_STATUSES:
+                    await self._bill_once(job, target, response)
 
             except Exception as e:
-                if via_router and _is_response_gone_at_provider(e, provider_response_id):
+                if target.deployment is not None and _is_response_gone_at_provider(e, provider_response_id):
                     verbose_proxy_logger.info(
                         f"Response {unified_object_id} no longer available at provider (404), marking stale_expired: {e}"
                     )
                     expired_jobs.append(job)
                     continue
-                verbose_proxy_logger.warning(
-                    f"Skipping job {unified_object_id} due to error: {e}"
-                )
+                verbose_proxy_logger.warning(f"Skipping job {unified_object_id} due to error: {e}")
                 continue
 
-            if response.status in TERMINAL_RESPONSE_STATUSES:
-                verbose_proxy_logger.info(
-                    f"Response {unified_object_id} has terminal status {response.status}, marking as complete"
-                )
-                completed_jobs.append(job)
-
-        # Mark completed jobs in the database
-        if len(completed_jobs) > 0:
-            await _managed_object_table(self.prisma_client).update_many(
-                # bounded-ok: at most MAX_OBJECTS_PER_POLL_CYCLE rows per cycle, the find_many take above
-                where={"id": {"in": [job.id for job in completed_jobs]}},
-                data={"status": "completed"},
-            )
-            verbose_proxy_logger.info(
-                f"Marked {len(completed_jobs)} response jobs as completed"
-            )
-
         if len(expired_jobs) > 0:
             await _managed_object_table(self.prisma_client).update_many(
                 # bounded-ok: at most MAX_OBJECTS_PER_POLL_CYCLE rows per cycle, the find_many take above
                 where={"id": {"in": [job.id for job in expired_jobs]}},
                 data={"status": "stale_expired"},
             )
-            verbose_proxy_logger.info(
-                f"Marked {len(expired_jobs)} response jobs as stale_expired"
-            )
-
+            verbose_proxy_logger.info(f"Marked {len(expired_jobs)} response jobs as stale_expired")

enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py

@@ -0,0 +1,218 @@
+"""Creator attribution for the spend row a managed-object cost poll writes.
+
+A batch or a background response is billed by a poll job that runs long after the creating
+request returned, so the managed object row's creator columns are the only record of which
+key, team, organization, user, end user, and tags carry the spend.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Final, Protocol, cast
+
+from pydantic import TypeAdapter
+
+from litellm._logging import verbose_proxy_logger
+from litellm.constants import CLI_SESSION_KEY_PREFIX
+from litellm.repositories.team_repository import TeamRepository
+from litellm.repositories.user_repository import UserRepository
+from litellm.repositories.verification_token_repository import VerificationTokenRepository
+
+if TYPE_CHECKING:
+    from prisma import models as prisma_models
+
+    from litellm.proxy.utils import PrismaClient
+    from litellm.repositories.prisma_protocols import TableActions
+
+
+class ManagedObjectRow(Protocol):
+    @property
+    def id(self) -> str: ...
+
+    @property
+    def unified_object_id(self) -> str: ...
+
+    @property
+    def model_object_id(self) -> str: ...
+
+    @property
+    def created_by(self) -> str | None: ...
+
+    @property
+    def status(self) -> str | None: ...
+
+    @property
+    def file_object(self) -> object: ...
+
+    @property
+    def org_id(self) -> str | None: ...
+
+    @property
+    def api_key(self) -> str | None: ...
+
+    @property
+    def team_id(self) -> str | None: ...
+
+    @property
+    def request_tags(self) -> object: ...
+
+    @property
+    def end_user_id(self) -> str | None: ...
+
+
+_TAG_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object])
+_KEY_OWNER_SCOPE_PREFIX: Final = "key:"
+
+
+def _string_tags(request_tags: object) -> list[str] | None:
+    if not isinstance(request_tags, list):
+        return None
+    return [tag for tag in _TAG_LIST_ADAPTER.validate_python(request_tags) if isinstance(tag, str)] or None
+
+
+def _creator_user_id(created_by: str | None) -> str | None:
+    if created_by is None or created_by.startswith(_KEY_OWNER_SCOPE_PREFIX):
+        return None
+    return created_by
+
+
+def _user_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_UserTable]:
+    return UserRepository(prisma_client).table
+
+
+def _token_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_VerificationToken]:
+    return VerificationTokenRepository(prisma_client).table
+
+
+def _team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_TeamTable]:
+    return TeamRepository(prisma_client).table
+
+
+@dataclass(frozen=True, slots=True)
+class _CreatorRows:
+    key: prisma_models.LiteLLM_VerificationToken | None
+    team: prisma_models.LiteLLM_TeamTable | None
+    user: prisma_models.LiteLLM_UserTable | None
+
+
+class ManagedObjectCreatorAttribution:
+    """Rebuilds the spend-tracking metadata of whoever created a managed object, so the poll
+    job's spend row is attributed like the creating request: key hash, team, organization,
+    user, aliases, and tags.
+
+    Rows created before api_key and request_tags were persisted carry only created_by and
+    team_id, and fall back to those. A named creating key owns user_api_key_alias; when it
+    has no alias, or the key has since been rotated or deleted, the field keeps the creating
+    user's alias, because a resolvable name is more useful on the spend row than a null.
+    user_api_key_org_id is resolved too: the spend update writer reads it off this metadata
+    to increment organization spend, so leaving it out drops the cost from org accounting.
+    A background response also carries its end user and the key, team, and user budgets, the
+    way the read that billed it before did, and never names a key-scoped creator as its user.
+    """
+
+    def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None:
+        self._prisma_client: Final = prisma_client
+        self._poller_name: Final = poller_name
+
+    async def batch_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
+        rows: Final = await self._rows(job, object_id)
+        return {"user_api_key_user_id": job.created_by, **_shared_metadata(job, rows)}
+
+    async def response_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
+        rows: Final = await self._rows(job, object_id)
+        end_user_id: Final = getattr(job, "end_user_id", None)
+        return {
+            "user_api_key_user_id": _creator_user_id(job.created_by),
+            **_shared_metadata(job, rows),
+            **({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
+            **_budget_fields(rows.key, "user_api_key_spend", "user_api_key_max_budget"),
+            **_budget_fields(rows.team, "user_api_key_team_spend", "user_api_key_team_max_budget"),
+            **_budget_fields(rows.user, "user_api_key_user_spend", "user_api_key_user_max_budget"),
+        }
+
+    async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows:
+        return _CreatorRows(
+            key=await self._key_row(object_id, getattr(job, "api_key", None), job.created_by),
+            team=await self._team_row(object_id, getattr(job, "team_id", None)),
+            user=await self._user_row(object_id, _creator_user_id(job.created_by)),
+        )
+
+    async def _user_row(self, object_id: str, user_id: str | None) -> prisma_models.LiteLLM_UserTable | None:
+        """Objects created by a team or service account key carry no user id, and
+        find_unique(where={"user_id": None}) raises, so those resolve no row."""
+        if not user_id:
+            return None
+        try:
+            return await _user_table(self._prisma_client).find_unique(where={"user_id": user_id})
+        except Exception as e:
+            verbose_proxy_logger.error(f"{self._poller_name}: could not look up user {user_id} for {object_id}: {e}")
+            return None
+
+    async def _key_row(
+        self, object_id: str, api_key: str | None, created_by: str | None
+    ) -> prisma_models.LiteLLM_VerificationToken | None:
+        if not api_key or _is_cli_session_key(api_key, created_by):
+            return None
+        try:
+            return await _token_table(self._prisma_client).find_unique(where={"token": api_key})
+        except Exception as e:
+            verbose_proxy_logger.error(f"{self._poller_name}: could not look up the creating key for {object_id}: {e}")
+            return None
+
+    async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
+        if not team_id:
+            return None
+        try:
+            return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
+        except Exception as e:
+            verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}")
+            return None
+
+
+def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
+    api_key: Final = getattr(job, "api_key", None)
+    team_id: Final = getattr(job, "team_id", None)
+    key_alias: Final = _key_alias(api_key, job.created_by, rows.key)
+    team_alias: Final = getattr(rows.team, "team_alias", None)
+    org_id: Final = getattr(job, "org_id", None) or _organization_id(rows.key) or _organization_id(rows.team)
+    tags: Final = _string_tags(getattr(job, "request_tags", None))
+    return {
+        "user_api_key": api_key,
+        "user_api_key_hash": api_key,
+        "user_api_key_team_id": team_id,
+        **(
+            {
+                "user_api_key_user_email": getattr(rows.user, "user_email", None),
+                "user_api_key_alias": getattr(rows.user, "user_alias", None),
+            }
+            if rows.user is not None
+            else {}
+        ),
+        **({"user_api_key_alias": key_alias} if key_alias is not None else {}),
+        **({"user_api_key_team_alias": team_alias} if team_alias is not None else {}),
+        **({"user_api_key_org_id": org_id} if org_id is not None else {}),
+        **({"tags": tags} if tags else {}),
+    }
+
+
+def _is_cli_session_key(api_key: str, created_by: str | None) -> bool:
+    return bool(created_by) and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}"
+
+
+def _key_alias(
+    api_key: str | None, created_by: str | None, key_row: prisma_models.LiteLLM_VerificationToken | None
+) -> str | None:
+    if api_key and _is_cli_session_key(api_key, created_by):
+        return api_key
+    return getattr(key_row, "key_alias", None) if key_row is not None else None
+
+
+def _organization_id(row: object) -> str | None:
+    return cast("str | None", getattr(row, "organization_id", None)) if row is not None else None
+
+
+def _budget_fields(row: object, spend_key: str, max_budget_key: str) -> Mapping[str, object]:
+    if row is None:
+        return {}
+    return {spend_key: getattr(row, "spend", None), max_budget_key: getattr(row, "max_budget", None)}

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -24,7 +24,9 @@ from uuid import NAMESPACE_URL, uuid5
 
 import httpx
 from fastapi import HTTPException
-from pydantic import ValidationError
+from openai import APIConnectionError, AsyncOpenAI
+from openai.types.file_deleted import FileDeleted
+from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, Unpack
 
 import litellm
@@ -35,17 +37,17 @@ from litellm.caching.caching import DualCache
 from litellm.constants import (
     BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS,
     BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
+    MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
     MAX_FILE_LIST_LIMIT,
+    PENDING_RESPONSE_STATUSES,
 )
 from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
+from litellm.litellm_core_utils.internal_call_metadata import requests_background
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     extract_file_metadata,
 )
-from openai import APIConnectionError, AsyncOpenAI
-from openai.types.file_deleted import FileDeleted
-
 from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
 from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
 from litellm.llms.base_llm.managed_resources.isolation import (
@@ -54,6 +56,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
     can_access_resource,
     resolve_resource_owner_id,
 )
+from litellm.llms.xai.batches.transformation import is_xai_batch_results_id
 from litellm.proxy._types import (
     CallTypes,
     LiteLLM_ManagedFileTable,
@@ -87,6 +90,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio
     request_tags_from_metadata,
 )
 from litellm.repositories.managed_file_repository import ManagedFileRepository
+from litellm.responses.utils import ResponsesAPIRequestUtils
 from litellm.types.llms.openai import (  # pyright: ignore[reportAttributeAccessIssue]
     AllMessageValues,
     AsyncCursorPage,
@@ -99,9 +103,11 @@ from litellm.types.llms.openai import (  # pyright: ignore[reportAttributeAccess
     ResponsesAPIResponse,
 )
 from litellm.types.utils import (
+    ADDRESSED_RESPONSE_ID_FIELD,
     CallTypesLiteral,
     LiteLLMBatch,
     LiteLLMFineTuningJob,
+    LlmProviders,
     LLMResponseTypes,
     SpecialEnums,
 )
@@ -194,6 +200,32 @@ def _is_transient_file_retrieve_error(error: Exception) -> bool:
     return isinstance(error, (httpx.TransportError, APIConnectionError, asyncio.TimeoutError))
 
 
+_MEASURED_SIZE_SAVE_ATTEMPTS: Final = 3
+
+
+def _stored_json_text(raw_file_object: object) -> str:
+    """JSON text of a stored file object exactly as read, for a write that must not race a newer one."""
+    if isinstance(raw_file_object, BaseModel):
+        return raw_file_object.model_dump_json()
+    return json.dumps(raw_file_object)
+
+
+def _with_measured_size(file_object: OpenAIFileObject, size_bytes: int | None) -> OpenAIFileObject:
+    """Fill a file object whose provider reported no size (0 bytes) with a size measured elsewhere."""
+    if not size_bytes or file_object.bytes:
+        return file_object
+    return file_object.model_copy(update={"bytes": size_bytes})
+
+
+def _over_stored_object(file_object: OpenAIFileObject, stored_object: OpenAIFileObject | None) -> OpenAIFileObject:
+    """Refreshed details for a stored row, keeping its size when none is reported and its stored created_at."""
+    if stored_object is None:
+        return file_object
+    return _with_measured_size(file_object, stored_object.bytes).model_copy(
+        update={"created_at": stored_object.created_at}
+    )
+
+
 def _proxy_llm_router() -> Router | None:
     import litellm.proxy.proxy_server as proxy_server_module
 
@@ -252,6 +284,23 @@ def _has_provider_file_retrieve_route(
     return bool(router_credentials) or _model_name_file_retrieve_provider(model_name) is not None
 
 
+def _is_xai_batch_output(
+    *,
+    llm_router: Router | None,
+    model_id: str | None,
+    model_name: str | None,
+    provider_file_id: str,
+) -> bool:
+    """Whether a batch output row holds an xAI batch's results, from its batch id and its deployment's provider."""
+    if not is_xai_batch_results_id(provider_file_id):
+        return False
+    router_credentials: Final = _provider_file_retrieve_credentials(llm_router=llm_router, model_id=model_id)
+    provider: Final = (router_credentials or {}).get("custom_llm_provider") or _model_name_file_retrieve_provider(
+        model_name
+    )
+    return provider == LlmProviders.XAI.value
+
+
 class _ManagedFileRow(Protocol):
     unified_file_id: str
     file_object: OpenAIFileObject
@@ -301,6 +350,8 @@ class _ManagedObjectTableActions(Protocol):
 
     async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
 
+    async def create(self, data: Mapping[str, object]) -> "PrismaManagedObjectRow": ...
+
 
 class _ManagedResourceDatabase(Protocol):
     @property
@@ -347,6 +398,39 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag
     return prisma_client.db.litellm_managedobjecttable
 
 
+_JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
+_POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00"
+
+
+def _storable_text(text: str) -> str:
+    return text.replace(_POSTGRES_UNSTORABLE_CHARACTER, "")
+
+
+def _storable_json_value(value: JsonValue) -> JsonValue:
+    if isinstance(value, str):
+        return _storable_text(value)
+    if isinstance(value, list):
+        return [_storable_json_value(item) for item in value]
+    if isinstance(value, dict):
+        return {_storable_text(key): _storable_json_value(item) for key, item in value.items()}
+    return value
+
+
+def _postgres_storable_json(file_object: BaseModel) -> str:
+    dumped: Final = file_object.model_dump_json()
+    if "\\u0000" not in dumped:
+        return dumped
+    return _JSON_VALUE_ADAPTER.dump_json(_storable_json_value(_JSON_VALUE_ADAPTER.validate_json(dumped))).decode()
+
+
+def _without_model_output(
+    file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse,
+) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse:
+    return (
+        file_object.model_copy(update={"output": []}) if isinstance(file_object, ResponsesAPIResponse) else file_object
+    )
+
+
 def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
     hidden_params: Final = cast(  # cast-ok: _hidden_params is an untyped attribute the upload path sets
         "Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
@@ -360,6 +444,40 @@ def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, s
     )
 
 
+_MANAGED_RESPONSE_READ_CALL_TYPES: Final = frozenset({"aget_responses", "acancel_responses"})
+_REQUEST_FIELDS: Final = TypeAdapter(dict[str, object])
+
+
+def _addresses_a_proxy_issued_response(data: Mapping[str, object]) -> bool:
+    from litellm.proxy.hooks.responses_id_security import decrypted_response_id
+
+    addressed_id: Final = data.get(ADDRESSED_RESPONSE_ID_FIELD) or data.get("response_id")
+    return (
+        isinstance(addressed_id, str)
+        and ResponsesAPIRequestUtils.get_model_id_from_response_id(decrypted_response_id(addressed_id)) is not None
+    )
+
+
+def _is_streamed_background_create(request: Mapping[str, object]) -> bool:
+    return (
+        request.get("stream") is True
+        and isinstance(request.get("litellm_metadata"), dict)
+        and requests_background(request)
+    )
+
+
+def _with_managed_background_marker(request: Mapping[str, object]) -> dict[str, object]:
+    metadata_field: Final = "litellm_metadata" if isinstance(request.get("litellm_metadata"), dict) else "metadata"
+    existing_metadata: Final = request.get(metadata_field)
+    return {
+        **request,
+        metadata_field: {
+            **(_REQUEST_FIELDS.validate_python(existing_metadata) if isinstance(existing_metadata, dict) else {}),
+            MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: True,
+        },
+    }
+
+
 _MANAGED_FILES_TARGET: Final = "managed_files"
 _PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0)
 
@@ -484,11 +602,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
     ) -> None:
         """Persist a managed object row, caching it and upserting it in the DB.
 
-        persist_attribution is set only by the batch create, which is the one caller
-        that can speak for the creator; it gates the api_key and request_tags columns
-        that CheckBatchCost bills against, so a later poll or retrieve of the same
-        batch cannot record itself as the paying key. Like created_by and team_id,
-        both are written only in the upsert create branch, never on update.
+        persist_attribution is set only by the batch and background response creates,
+        the callers that can speak for the creator; it gates the api_key, request_tags, and,
+        for a response, end_user_id columns that the cost poll jobs bill against, so a later poll or
+        retrieve of the same object cannot record itself as the paying key. Like created_by and team_id,
+        they are written only in the upsert create branch, never on update.
+
+        A stored response keeps no model output: billing reads only its status and model from
+        the row, so the row never holds generated content.
 
         create_if_missing is cleared by callers that observe a batch they did not
         create, such as a poll. They still refresh status and file_object, but a
@@ -501,11 +622,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         It is written only in the upsert create branch.
         """
         verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
+        stored_object: Final = _without_model_output(file_object)
         litellm_managed_object = LiteLLM_ManagedObjectTable(
             unified_object_id=unified_object_id,
             model_object_id=model_object_id,
             file_purpose=file_purpose,
-            file_object=file_object,
+            file_object=stored_object,
         )
         await self.internal_usage_cache.async_set_cache(
             key=unified_object_id,
@@ -516,32 +638,35 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         from prisma import Json
 
         api_key = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None
+        end_user_id: Final = user_api_key_dict.end_user_id if file_purpose == "response" else None
         attribution_columns = (
             {
                 **({"api_key": api_key} if api_key is not None else {}),
-                **({"request_tags": Json(list(request_tags))} if request_tags else {}),
+                **({"request_tags": Json([_storable_text(tag) for tag in request_tags])} if request_tags else {}),
+                **({"end_user_id": end_user_id} if end_user_id else {}),
             }
             if persist_attribution
             else {}
         )
         # FIX: Update status and file_object on every operation to keep state in sync
         update_columns: Final = {
-            "file_object": file_object.model_dump_json(),
+            "file_object": _postgres_storable_json(stored_object),
             "status": file_object.status,
             "updated_by": user_api_key_dict.user_id,
         }
+        managed_object_table: Final = _managed_object_table(self.prisma_client)
         if not create_if_missing:
-            await _managed_object_table(self.prisma_client).update_many(
+            await managed_object_table.update_many(
                 where={"unified_object_id": unified_object_id},
                 data=update_columns,
             )
             return
-        await _managed_object_table(self.prisma_client).upsert(
+        await managed_object_table.upsert(
             where={"unified_object_id": unified_object_id},
             data={
                 "create": {
                     "unified_object_id": unified_object_id,
-                    "file_object": file_object.model_dump_json(),
+                    "file_object": _postgres_storable_json(stored_object),
                     "model_object_id": model_object_id,
                     "file_purpose": file_purpose,
                     "created_by": resolve_resource_owner_id(user_api_key_dict),
@@ -556,6 +681,42 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             },
         )
 
+    async def track_untracked_background_response(
+        self,
+        response_id: str,
+        response: ResponsesAPIResponse,
+        user_api_key_dict: UserAPIKeyAuth,
+        request_tags: Sequence[str] | None = None,
+    ) -> bool:
+        if not _addresses_a_proxy_issued_response({"response_id": response_id}):
+            return False
+        from prisma import Json
+        from prisma.errors import UniqueViolationError
+
+        from litellm.proxy.hooks.responses_id_security import decrypted_response_id
+
+        api_key: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None
+        try:
+            await _managed_object_table(self.prisma_client).create(
+                data={
+                    "unified_object_id": response_id,
+                    "model_object_id": decrypted_response_id(response_id),
+                    "file_purpose": "response",
+                    "file_object": _postgres_storable_json(_without_model_output(response)),
+                    "status": PENDING_RESPONSE_STATUSES[0],
+                    "created_by": resolve_resource_owner_id(user_api_key_dict),
+                    "team_id": user_api_key_dict.team_id,
+                    "org_id": await self._resolve_creator_org_id(user_api_key_dict),
+                    "updated_by": user_api_key_dict.user_id,
+                    **({"api_key": api_key} if api_key is not None else {}),
+                    **({"end_user_id": user_api_key_dict.end_user_id} if user_api_key_dict.end_user_id else {}),
+                    **({"request_tags": Json([_storable_text(tag) for tag in request_tags])} if request_tags else {}),
+                }
+            )
+        except UniqueViolationError:
+            return False
+        return True
+
     @with_service_target(_MANAGED_FILES_TARGET)
     async def get_unified_file_id(
         self, file_id: str, litellm_parent_otel_span: Optional[Span] = None
@@ -845,6 +1006,57 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             litellm_parent_otel_span=None,
         )
 
+    async def _save_xai_refreshed_file_object(
+        self,
+        stored: LiteLLM_ManagedFileTable,
+        file_object: OpenAIFileObject,
+    ) -> OpenAIFileObject | None:
+        """Save an xAI batch output's refreshed details, keeping a known size, only over the row as read; drop the cached copy if a newer write won."""
+        stored_object: Final = _parse_managed_file_object(stored.file_object, stored.unified_file_id)
+        sized_file_object: Final = _over_stored_object(file_object, stored_object)
+        if not await ManagedFileRepository(self.prisma_client).update_file_object(
+            stored.unified_file_id,
+            sized_file_object,
+            if_file_object=_stored_json_text(stored.file_object) if stored.file_object is not None else None,
+        ):
+            await self.internal_usage_cache.async_set_cache(
+                key=stored.unified_file_id, value=None, litellm_parent_otel_span=None
+            )
+            return None
+        await self.internal_usage_cache.async_set_cache(
+            key=stored.unified_file_id,
+            value=stored.model_copy(update={"file_object": sized_file_object}).model_dump(),
+            litellm_parent_otel_span=None,
+        )
+        return sized_file_object
+
+    async def _save_batch_output_file_object(
+        self,
+        stored: LiteLLM_ManagedFileTable,
+        file_object: OpenAIFileObject,
+        size_bytes: int | None,
+    ) -> None:
+        """Save an xAI batch output's details; after losing to a refresh, size the refreshed details instead of replacing them."""
+        current_row: LiteLLM_ManagedFileTable | None = stored  # rebind-ok: the row as last read
+        details: OpenAIFileObject = file_object  # rebind-ok: a concurrent refresh's details replace the job's copy
+        for _ in range(_MEASURED_SIZE_SAVE_ATTEMPTS):
+            if current_row is None:
+                return
+            saved = await self._save_xai_refreshed_file_object(  # rebind-ok: per attempt
+                current_row, _with_measured_size(details, size_bytes)
+            )
+            if saved is not None or not size_bytes:
+                return
+            latest_row = await ManagedFileRepository(  # rebind-ok: per attempt
+                self.prisma_client, use_writer=True
+            ).table.find_first(where={"unified_file_id": stored.unified_file_id})
+            current_row = (
+                None if latest_row is None else LiteLLM_ManagedFileTable.model_validate(latest_row.model_dump())
+            )
+            latest_object = None if current_row is None else current_row.file_object  # rebind-ok: per attempt
+            if latest_object is not None and not latest_object.litellm_details_fallback:
+                details = latest_object
+
     async def store_batch_output_file(
         self,
         *,
@@ -885,6 +1097,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             )
             return
         if stored_object is not None and not stored_object.litellm_details_fallback:
+            if (
+                stored_file is not None
+                and size_bytes
+                and not stored_object.bytes
+                and _is_xai_batch_output(
+                    llm_router=_proxy_llm_router(),
+                    model_id=model_id,
+                    model_name=model_name,
+                    provider_file_id=provider_file_id,
+                )
+            ):
+                await self._save_batch_output_file_object(stored_file, stored_object, size_bytes)
             return
 
         fallback_written_recently: Final = (
@@ -924,13 +1148,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             )
         )
 
+        xai_batch_output: Final = _is_xai_batch_output(
+            llm_router=_proxy_llm_router(),
+            model_id=model_id,
+            model_name=model_name,
+            provider_file_id=provider_file_id,
+        )
+        if stored_file is not None and xai_batch_output:
+            await self._save_batch_output_file_object(stored_file, file_object, size_bytes)
+            return
         if stored_file is not None:
             await self._save_refreshed_file_object(stored_file, file_object)
             return
 
         await self.store_unified_file_id(
             file_id=unified_file_id,
-            file_object=file_object,
+            file_object=_with_measured_size(file_object, size_bytes) if xai_batch_output else file_object,
             litellm_parent_otel_span=litellm_parent_otel_span,
             model_mappings={model_id: provider_file_id} if model_id else {},
             user_api_key_dict=owner,
@@ -1258,6 +1491,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 unified_vs_ids = self.get_vector_store_ids_from_file_search_tools(tools)
                 if unified_vs_ids:
                     await self.check_vector_store_ids_access(unified_vs_ids, user_api_key_dict)
+            create_request: Final = _REQUEST_FIELDS.validate_python(data)
+            if _is_streamed_background_create(create_request):
+                return _with_managed_background_marker(create_request)
         elif call_type == CallTypes.afile_content.value:
             retrieve_file_id = cast(Optional[str], data.get("file_id"))
             potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False
@@ -1315,6 +1551,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 data[accessor_key] = get_batch_id_from_unified_batch_id(potential_llm_object_id)
             elif retrieve_object_id and accessor_key == "batch_id":
                 await self.enforce_batch_object_access(retrieve_object_id, user_api_key_dict)
+        elif call_type in _MANAGED_RESPONSE_READ_CALL_TYPES:
+            request: Final = _REQUEST_FIELDS.validate_python(data)
+            if _addresses_a_proxy_issued_response(request):
+                return _with_managed_background_marker(request)
         elif call_type == CallTypes.acreate_fine_tuning_job.value:
             fine_tuning_file_ids: Final = tuple(
                 file_id
@@ -1943,8 +2183,20 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                     )
                     if refreshed_file_object is None:
                         return _public_file_object(file_object, file_id)
-                    await self._save_refreshed_file_object(stored_file_object, refreshed_file_object)
-                    return _public_file_object(refreshed_file_object, file_id)
+                    if not _is_xai_batch_output(
+                        llm_router=llm_router if llm_router is not None else _proxy_llm_router(),
+                        model_id=model_id,
+                        model_name=model_id,
+                        provider_file_id=provider_file_id,
+                    ):
+                        await self._save_refreshed_file_object(stored_file_object, refreshed_file_object)
+                        return _public_file_object(refreshed_file_object, file_id)
+                    saved_file_object: Final = await self._save_xai_refreshed_file_object(
+                        stored_file_object, refreshed_file_object
+                    )
+                    return _public_file_object(
+                        saved_file_object or _over_stored_object(refreshed_file_object, file_object), file_id
+                    )
                 except Exception as error:
                     verbose_logger.warning(
                         "Failed to refresh batch file object for "

helm/litellm-helm/README.md

@@ -126,6 +126,11 @@ Set `billingMetrics.caSecretName` only when the collector is a private or test o
 | `db.secret.name`          | If `db.useExisting` is `true`, the name of the Kubernetes Secret that contains credentials.                                                                                                                                                                                               | `postgres`                                                                                 |
 | `db.secret.usernameKey`   | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the username for authenticating with the Postgres instance.                                                                                                                                    | `username`                                                                                 |
 | `db.secret.passwordKey`   | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the password associates with the above user.                                                                                                                                                   | `password`                                                                                 |
+| `db.username`             | If `db.useExisting` is `true`, the database user as a plain value. When set, `db.secret.usernameKey` is not read.                                                                                                                                                                        | `""`                                                                                       |
+| `db.port`                 | If `db.useExisting` is `true` with token auth, the Postgres port used to build the connection URL. Password auth takes the port from `db.url`.                                                                                                                                           | `5432`                                                                                     |
+| `db.schema`               | With token auth, the Postgres schema holding the LiteLLM tables, rendered as `DATABASE_SCHEMA` when the tables are not in `public`. Ignored with password auth, where the schema stays in `db.url` as `?schema=`.                                                                        | `""`                                                                                       |
+| `db.useIAMAuth`           | If `db.useExisting` is `true` and `db.deployStandalone` is `false`, authenticate with an AWS RDS IAM token instead of a password. See [Token auth](#token-auth-to-an-existing-postgres-rds-iam-or-microsoft-entra-id).                                                                  | `false`                                                                                    |
+| `db.useAzureEntraAuth`    | If `db.useExisting` is `true` and `db.deployStandalone` is `false`, authenticate with a Microsoft Entra ID token instead of a password. Mutually exclusive with `db.useIAMAuth`.                                                                                                        | `false`                                                                                    |
 | `db.useStackgresOperator` | Not yet implemented.                                                                                                                                                                                                                                                                      | `false`                                                                                    |
 | `db.deployStandalone`     | Deploy a standalone, single instance deployment of Postgres, using the Bitnami postgresql chart. This is useful for getting started but doesn't provide HA or (by default) data backups.                                                                                                  | `true`                                                                                     |
 | `postgresql.*`            | If `db.deployStandalone` is `true`, configuration passed to the Bitnami postgresql chart. See the [Bitnami Documentation](https://github.com/bitnami/charts/tree/main/bitnami/postgresql) for full configuration details. See [values.yaml](./values.yaml) for the default configuration. | See [values.yaml](./values.yaml)                                                           |
@@ -156,6 +161,22 @@ data:
 type: Opaque
 ```
 
+#### Token auth to an existing Postgres (RDS IAM or Microsoft Entra ID)
+
+When the database accepts short-lived tokens instead of a password (AWS RDS or Aurora with IAM database authentication, or Azure Database for PostgreSQL with Microsoft Entra ID), there is no password to put in a Secret. Turn on `db.useIAMAuth` or `db.useAzureEntraAuth` and the chart stops referencing `db.secret.passwordKey` and `db.url`. It sets `IAM_TOKEN_DB_AUTH` (or `AZURE_POSTGRESQL_AUTH`) and emits `DATABASE_HOST`, `DATABASE_PORT`, `DATABASE_USERNAME` and `DATABASE_NAME` for the proxy, the migrations job and the collector sidecar. Each of them mints its own token at startup and builds `DATABASE_URL` from those parts, and the in-container pgbouncer (`db.connectionPool`) refreshes the token before it expires.
+
+```yaml
+db:
+  useExisting: true
+  deployStandalone: false
+  useIAMAuth: true
+  endpoint: litellm.cluster-abc123.us-west-1.rds.amazonaws.com
+  database: litellm
+  username: litellm_iam   # the role granted rds_iam; no db.secret needed
+```
+
+The pod needs cloud credentials for the token: on EKS use IRSA or Pod Identity on the chart's `serviceAccount` (and `migrationJob.serviceAccountName` when the Helm hook runs before the ServiceAccount exists), on AKS use workload identity. Set `AWS_RDS_REGION` under `envVars` when the RDS hostname is not a standard regional endpoint (the token must be signed for the region the database lives in). TLS options that would otherwise go in `db.url` are set as `DATABASE_SSLMODE` and `DATABASE_SSLROOTCERT` under `envVars`. Leaving `db.username` empty reads the user from `db.secret.usernameKey` as before.
+
 #### Examples for `environmentSecrets` and `environemntConfigMaps`
 
 ```yaml

helm/litellm-helm/templates/_helpers.tpl

@@ -162,6 +162,77 @@ taken before the change, which by that point no longer exists.
 {{- end -}}
 {{- end -}}
 
+{{/*
+Whether the existing Postgres is reached with a short-lived token (AWS RDS IAM
+or Microsoft Entra ID) instead of a password.
+*/}}
+{{- define "litellm.db.tokenAuth" -}}
+{{- if or .Values.db.useIAMAuth .Values.db.useAzureEntraAuth -}}true{{- end -}}
+{{- end -}}
+
+{{- define "litellm.validateDbTokenAuth" -}}
+{{- if and .Values.db.useIAMAuth .Values.db.useAzureEntraAuth -}}
+{{- fail "db.useIAMAuth and db.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source" -}}
+{{- end -}}
+{{- if and (include "litellm.db.tokenAuth" .) (or (not .Values.db.useExisting) .Values.db.deployStandalone) -}}
+{{- fail "db.useIAMAuth and db.useAzureEntraAuth need db.useExisting: true and db.deployStandalone: false (token auth targets an existing RDS or Azure Database for PostgreSQL; the bundled postgresql subchart only supports password auth)" -}}
+{{- end -}}
+{{- end -}}
+
+{{/*
+Database env for an existing Postgres, shared by the proxy container, the
+collector sidecar and the migrations Job. Under token auth every entrypoint
+mints its own token and assembles DATABASE_URL itself, so db.url is left out.
+*/}}
+{{- define "litellm.existingDbEnv" -}}
+{{- $tokenAuth := include "litellm.db.tokenAuth" . -}}
+- name: DATABASE_USERNAME
+  {{- if .Values.db.username }}
+  value: {{ .Values.db.username | quote }}
+  {{- else }}
+  valueFrom:
+    secretKeyRef:
+      name: {{ .Values.db.secret.name }}
+      key: {{ .Values.db.secret.usernameKey }}
+  {{- end }}
+{{- if .Values.db.useIAMAuth }}
+- name: IAM_TOKEN_DB_AUTH
+  value: "true"
+{{- else if .Values.db.useAzureEntraAuth }}
+- name: AZURE_POSTGRESQL_AUTH
+  value: "true"
+{{- else }}
+- name: DATABASE_PASSWORD
+  valueFrom:
+    secretKeyRef:
+      name: {{ .Values.db.secret.name }}
+      key: {{ .Values.db.secret.passwordKey }}
+{{- end }}
+- name: DATABASE_HOST
+  {{- if .Values.db.secret.endpointKey }}
+  valueFrom:
+    secretKeyRef:
+      name: {{ .Values.db.secret.name }}
+      key: {{ .Values.db.secret.endpointKey }}
+  {{- else }}
+  value: {{ .Values.db.endpoint }}
+  {{- end }}
+{{- if $tokenAuth }}
+- name: DATABASE_PORT
+  value: {{ .Values.db.port | quote }}
+{{- end }}
+- name: DATABASE_NAME
+  value: {{ .Values.db.database }}
+{{- if and $tokenAuth .Values.db.schema }}
+- name: DATABASE_SCHEMA
+  value: {{ .Values.db.schema | quote }}
+{{- end }}
+{{- if not $tokenAuth }}
+- name: DATABASE_URL
+  value: {{ .Values.db.url | quote }}
+{{- end }}
+{{- end -}}
+
 {{/*
 Environment shared by the proxy container and the opt-in collector sidecar:
 database, pgbouncer, master key, redis, user envVars. Both containers must see
@@ -169,6 +240,7 @@ the same DATABASE_URL and REDIS_* so the sidecar reaches the pod's pgbouncer
 and the same spend transaction buffer.
 */}}
 {{- define "litellm.proxyEnv" -}}
+{{- include "litellm.validateDbTokenAuth" . -}}
 - name: HOST
   value: "{{ .Values.listen | default "0.0.0.0" }}"
 - name: PORT
@@ -189,29 +261,7 @@ and the same spend transaction buffer.
 - name: DATABASE_NAME
   value: litellm
 {{- else if .Values.db.useExisting }}
-- name: DATABASE_USERNAME
-  valueFrom:
-    secretKeyRef:
-      name: {{ .Values.db.secret.name }}
-      key: {{ .Values.db.secret.usernameKey }}
-- name: DATABASE_PASSWORD
-  valueFrom:
-    secretKeyRef:
-      name: {{ .Values.db.secret.name }}
-      key: {{ .Values.db.secret.passwordKey }}
-- name: DATABASE_HOST
-  {{- if .Values.db.secret.endpointKey }}
-  valueFrom:
-    secretKeyRef:
-      name: {{ .Values.db.secret.name }}
-      key: {{ .Values.db.secret.endpointKey }}
-  {{- else }}
-  value: {{ .Values.db.endpoint }}
-  {{- end }}
-- name: DATABASE_NAME
-  value: {{ .Values.db.database }}
-- name: DATABASE_URL
-  value: {{ .Values.db.url | quote }}
+{{ include "litellm.existingDbEnv" . }}
 {{- end }}
 {{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }}
 - name: DATABASE_READER_HOST

helm/litellm-helm/templates/migrations-job.yaml

@@ -68,30 +68,9 @@ spec:
           {{- end }}
           workingDir: "/app"
           env:
+            {{- include "litellm.validateDbTokenAuth" . -}}
             {{- if .Values.db.useExisting }}
-            - name: DATABASE_USERNAME
-              valueFrom:
-                secretKeyRef:
-                  name: {{ .Values.db.secret.name }}
-                  key: {{ .Values.db.secret.usernameKey }}
-            - name: DATABASE_PASSWORD
-              valueFrom:
-                secretKeyRef:
-                  name: {{ .Values.db.secret.name }}
-                  key: {{ .Values.db.secret.passwordKey }}
-            - name: DATABASE_HOST
-              {{- if .Values.db.secret.endpointKey }}
-              valueFrom:
-                secretKeyRef:
-                  name: {{ .Values.db.secret.name }}
-                  key: {{ .Values.db.secret.endpointKey }}
-              {{- else }}
-              value: {{ .Values.db.endpoint }}
-              {{- end }}
-            - name: DATABASE_NAME
-              value: {{ .Values.db.database }}
-            - name: DATABASE_URL
-              value: {{ .Values.db.url | quote }}
+            {{- include "litellm.existingDbEnv" . | nindent 12 }}
             {{- else if .Values.db.deployStandalone }}
             - name: DATABASE_URL
               value: postgresql://{{ .Values.postgresql.auth.username }}:{{ .Values.postgresql.auth.password }}@{{ .Release.Name }}-postgresql/{{ .Values.postgresql.auth.database }}

helm/litellm-helm/tests/database_auth_tests.yaml

@@ -0,0 +1,288 @@
+suite: test existing database token auth
+templates:
+  - deployment.yaml
+  - migrations-job.yaml
+  - configmap-litellm.yaml
+tests:
+  - it: should keep password auth from the secret when no token auth is set
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USERNAME
+            valueFrom:
+              secretKeyRef:
+                name: postgres
+                key: username
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+            valueFrom:
+              secretKeyRef:
+                name: postgres
+                key: password
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL
+            value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: AZURE_POSTGRESQL_AUTH
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PORT
+          any: true
+
+  - it: should render RDS IAM auth without a password secret on the proxy and the collector sidecar
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useIAMAuth: true
+      db.username: litellm_iam
+      db.endpoint: writer.cluster-abc.us-west-1.rds.amazonaws.com
+      db.port: 5433
+      collector.enabled: true
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USERNAME
+            value: litellm_iam
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_HOST
+            value: writer.cluster-abc.us-west-1.rds.amazonaws.com
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PORT
+            value: "5433"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_NAME
+            value: litellm
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: AZURE_POSTGRESQL_AUTH
+          any: true
+      - contains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+            value: "true"
+      - notContains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+
+  - it: should give the migrations job the same token auth env as the proxy
+    template: migrations-job.yaml
+    set:
+      migrationJob.enabled: true
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useIAMAuth: true
+      db.username: litellm_iam
+      db.endpoint: writer.cluster-abc.us-west-1.rds.amazonaws.com
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USERNAME
+            value: litellm_iam
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_HOST
+            value: writer.cluster-abc.us-west-1.rds.amazonaws.com
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PORT
+            value: "5432"
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL
+          any: true
+
+  - it: should carry db.schema as DATABASE_SCHEMA under token auth and omit it when empty
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useIAMAuth: true
+      db.username: litellm_iam
+      db.schema: tenant1
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SCHEMA
+            value: tenant1
+
+  - it: should leave the schema to db.url under password auth even when db.schema is set
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+      db.schema: tenant1
+      db.url: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)?schema=tenant1
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SCHEMA
+          any: true
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL
+            value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)?schema=tenant1
+
+  - it: should not emit DATABASE_SCHEMA when db.schema is empty
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useIAMAuth: true
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SCHEMA
+          any: true
+
+  - it: should read the user from the secret when db.username is empty under token auth
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useIAMAuth: true
+      db.secret.name: rds-identity
+      db.secret.usernameKey: user
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USERNAME
+            valueFrom:
+              secretKeyRef:
+                name: rds-identity
+                key: user
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+
+  - it: should render Entra auth with AZURE_POSTGRESQL_AUTH and no password
+    template: deployment.yaml
+    set:
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useAzureEntraAuth: true
+      db.username: litellm@example.onmicrosoft.com
+      db.endpoint: litellm.postgres.database.azure.com
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: AZURE_POSTGRESQL_AUTH
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_USERNAME
+            value: litellm@example.onmicrosoft.com
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_PASSWORD
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_URL
+          any: true
+
+  - it: should fail when both token sources are enabled
+    template: deployment.yaml
+    set:
+      migrationJob.enabled: false
+      db.useExisting: true
+      db.deployStandalone: false
+      db.useIAMAuth: true
+      db.useAzureEntraAuth: true
+    asserts:
+      - failedTemplate:
+          errorMessage: "db.useIAMAuth and db.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source"
+
+  - it: should fail when token auth is set while the bundled postgres is still deployed
+    template: deployment.yaml
+    set:
+      migrationJob.enabled: false
+      db.useExisting: true
+      db.useIAMAuth: true
+    asserts:
+      - failedTemplate:
+          errorPattern: "need db.useExisting: true and db.deployStandalone: false"
+
+  - it: should fail when token auth is set without useExisting
+    template: migrations-job.yaml
+    set:
+      migrationJob.enabled: true
+      db.deployStandalone: false
+      db.useAzureEntraAuth: true
+    asserts:
+      - failedTemplate:
+          errorPattern: "need db.useExisting: true and db.deployStandalone: false"

helm/litellm-helm/values.yaml

@@ -376,6 +376,30 @@ db:
   endpoint: localhost
   database: litellm
   url: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)
+  # Authenticate with a short-lived token instead of a password. useIAMAuth is
+  # AWS RDS / Aurora IAM database authentication (the pod needs AWS credentials,
+  # e.g. IRSA or EKS Pod Identity, and the database user needs GRANT rds_iam);
+  # useAzureEntraAuth is Azure Database for PostgreSQL with a Microsoft Entra ID
+  # identity (workload identity). Mutually exclusive, and both need
+  # useExisting: true with deployStandalone: false. The chart then sets
+  # IAM_TOKEN_DB_AUTH or AZURE_POSTGRESQL_AUTH, emits DATABASE_HOST, DATABASE_PORT,
+  # DATABASE_USERNAME and DATABASE_NAME and skips DATABASE_PASSWORD and db.url:
+  # the proxy, the migrations job and the collector sidecar each mint their own
+  # token and build DATABASE_URL from those parts. Set DATABASE_SSLMODE /
+  # DATABASE_SSLROOTCERT under envVars for TLS options you would otherwise put
+  # in db.url.
+  useIAMAuth: false
+  useAzureEntraAuth: false
+  # Postgres port, used to build the token-auth URL (password auth takes the
+  # port from db.url)
+  port: 5432
+  # Postgres schema for the LiteLLM tables, rendered as DATABASE_SCHEMA under
+  # token auth only (password auth carries it as ?schema= in db.url)
+  schema: ""
+  # Database user as a plain value. Empty means read it from
+  # db.secret.usernameKey. With token auth the user name is the only credential
+  # the chart needs, so setting it here removes the need for db.secret entirely.
+  username: ""
   secret:
     name: postgres
     usernameKey: username

helm/litellm/templates/_helpers.tpl

@@ -475,6 +475,27 @@ harmless no-op for the Job and authoritative for the app pods.
   {{- end }}
 {{- end }}
 {{- end }}
+{{- $sslCertReqs := $root.Values.redis.sslCertReqs }}
+{{- if and $sslCertReqs (not (has $sslCertReqs (list "none" "optional" "required"))) }}
+{{- fail (printf "redis.sslCertReqs must be one of none, optional, required (got %q)" $sslCertReqs) }}
+{{- end }}
+{{- if and $urlKey (or $root.Values.redis.ssl $sslCertReqs $root.Values.redis.sslCaCerts) }}
+{{- fail "redis.ssl, redis.sslCertReqs and redis.sslCaCerts cannot be combined with redis.passwordSecret.urlKey, use a rediss:// URL with ssl_cert_reqs / ssl_ca_certs query parameters instead" }}
+{{- end }}
+{{- if or $root.Values.redis.host $hostKey }}
+{{- if $root.Values.redis.ssl }}
+- name: REDIS_SSL
+  value: "true"
+{{- end }}
+{{- if $sslCertReqs }}
+- name: REDIS_SSL_CERT_REQS
+  value: {{ $sslCertReqs | quote }}
+{{- end }}
+{{- if $root.Values.redis.sslCaCerts }}
+- name: REDIS_SSL_CA_CERTS
+  value: {{ $root.Values.redis.sslCaCerts | quote }}
+{{- end }}
+{{- end }}
 {{- with $component.extraEnv }}
 {{ toYaml . }}
 {{- end }}

helm/litellm/tests/redis_env_tests.yaml

@@ -3,6 +3,7 @@ templates:
   - gateway/deployment.yaml
   - gateway/configmap.yaml
   - backend/deployment.yaml
+  - migrations-job.yaml
 values:
   - ./values/required.yaml
 tests:
@@ -320,3 +321,168 @@ tests:
     asserts:
       - failedTemplate:
           errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set"
+
+  - it: redis omits TLS env vars by default
+    template: gateway/deployment.yaml
+    set:
+      redis.host: redis.example.com
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CERT_REQS
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CA_CERTS
+          any: true
+
+  - it: redis ssl false omits REDIS_SSL
+    template: gateway/deployment.yaml
+    set:
+      redis.host: redis.example.com
+      redis.ssl: false
+      redis.sslCertReqs: required
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+          any: true
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CERT_REQS
+            value: "required"
+
+  - it: gateway renders redis TLS env vars
+    template: gateway/deployment.yaml
+    set:
+      redis.host: redis.example.com
+      redis.ssl: true
+      redis.sslCertReqs: none
+      redis.sslCaCerts: /etc/redis-tls/ca.crt
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CERT_REQS
+            value: "none"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CA_CERTS
+            value: "/etc/redis-tls/ca.crt"
+
+  - it: backend renders redis TLS env vars
+    template: backend/deployment.yaml
+    set:
+      redis.host: redis.example.com
+      redis.ssl: true
+      redis.sslCertReqs: required
+      redis.sslCaCerts: /etc/redis-tls/ca.crt
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CERT_REQS
+            value: "required"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CA_CERTS
+            value: "/etc/redis-tls/ca.crt"
+
+  - it: migration job renders redis TLS env vars
+    template: migrations-job.yaml
+    set:
+      db.deployStandalone: false
+      db.useExisting: true
+      db.endpoint: postgres.example.com
+      db.secret.name: pg-secret
+      redis.host: redis.example.com
+      redis.ssl: true
+      redis.sslCertReqs: optional
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CERT_REQS
+            value: "optional"
+
+  - it: redis TLS values reject a secret-backed URL
+    template: gateway/deployment.yaml
+    set:
+      redis.passwordSecret.name: redis-secret
+      redis.passwordSecret.urlKey: connection-url
+      redis.sslCaCerts: /etc/redis-tls/ca.crt
+    asserts:
+      - failedTemplate:
+          errorMessage: "redis.ssl, redis.sslCertReqs and redis.sslCaCerts cannot be combined with redis.passwordSecret.urlKey, use a rediss:// URL with ssl_cert_reqs / ssl_ca_certs query parameters instead"
+
+  - it: hostKey mode renders redis TLS env vars
+    template: gateway/deployment.yaml
+    set:
+      redis.passwordSecret.name: redis-secret
+      redis.passwordSecret.hostKey: hostname
+      redis.ssl: true
+      redis.sslCaCerts: /etc/redis-tls/ca.crt
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CA_CERTS
+            value: "/etc/redis-tls/ca.crt"
+
+  - it: redis TLS env vars are omitted when redis is not configured
+    template: gateway/deployment.yaml
+    set:
+      redis.ssl: true
+      redis.sslCertReqs: none
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: REDIS_SSL_CERT_REQS
+          any: true
+
+  - it: redis sslCertReqs rejects unknown values
+    template: gateway/deployment.yaml
+    set:
+      redis.host: redis.example.com
+      redis.ssl: true
+      redis.sslCertReqs: strict
+    asserts:
+      - failedTemplate:
+          errorMessage: 'redis.sslCertReqs must be one of none, optional, required (got "strict")'

helm/litellm/values.yaml

@@ -275,10 +275,21 @@ database:
 # self-hosted Redis Cluster). The chart emits REDIS_CLUSTER_NODES from
 # `host` / `port` as the single seed; the cluster client discovers the
 # remaining nodes from CLUSTER SLOTS at startup.
+#
+# For a TLS Redis endpoint set `ssl: true`, which the chart emits as REDIS_SSL.
+# `sslCertReqs` (none / optional / required) becomes REDIS_SSL_CERT_REQS and
+# `sslCaCerts` (path to a CA bundle mounted in the pod, e.g. via
+# `extraVolumeMounts`) becomes REDIS_SSL_CA_CERTS. All three are omitted when
+# unset. They apply to `host` / `hostKey` mode only. With `passwordSecret.urlKey`
+# the proxy reads TLS settings from the URL itself, so use `rediss://` and the
+# `ssl_cert_reqs` / `ssl_ca_certs` query parameters there instead.
 redis:
   cluster: false
   host: ""
   port: 6379
+  ssl: false
+  sslCertReqs: ""
+  sslCaCerts: ""
   passwordSecret:
     name: "" # Leave empty for auth-less Redis
     passwordKey: password

litellm-proxy-extras/litellm_proxy_extras/_redaction.py

@@ -0,0 +1,50 @@
+import functools
+import os
+import re
+from collections.abc import Callable
+from typing import Final
+from urllib.parse import unquote, urlsplit
+
+_REDACTED: Final = "REDACTED"
+_PASSWORD_QUERY_KEYS: Final = frozenset(("password", "sslpassword"))
+
+
+@functools.cache
+def _secret_shape_redactor() -> Callable[[str], str]:
+    try:
+        from litellm._logging import redact_secrets
+    except ImportError:
+        return lambda text: text
+    return redact_secrets
+
+
+def _url_passwords(url: str) -> frozenset[str]:
+    try:
+        parts: Final = urlsplit(url)
+    except ValueError:
+        return frozenset()
+    query_pairs: Final = tuple(pair.partition("=") for pair in parts.query.split("&"))
+    raw_query_passwords: Final = tuple(
+        value for key, separator, value in query_pairs if separator and key.lower() in _PASSWORD_QUERY_KEYS
+    )
+    raw_passwords: Final = ((parts.password,) if parts.password else ()) + raw_query_passwords
+    return frozenset(password for password in raw_passwords + tuple(map(unquote, raw_passwords)) if password)
+
+
+def _configured_database_passwords() -> frozenset[str]:
+    database_url: Final = os.getenv("DATABASE_URL")
+    direct_url: Final = os.getenv("DIRECT_URL")
+    database_passwords: Final = _url_passwords(database_url) if database_url else frozenset()
+    direct_passwords: Final = _url_passwords(direct_url) if direct_url else frozenset()
+    return database_passwords | direct_passwords
+
+
+def _redact_credentials(text: str) -> str:
+    """Mask configured database passwords before passing the text to LiteLLM redaction."""
+    passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True)
+    alternation: Final = "|".join(re.escape(password) for password in passwords)
+    password_pattern: Final = (
+        re.compile(rf"(?<![A-Za-z0-9_])(?:{alternation})(?![A-Za-z0-9_])", re.IGNORECASE) if passwords else None
+    )
+    result: Final = password_pattern.sub(_REDACTED, text) if password_pattern is not None else text
+    return _secret_shape_redactor()(result)

litellm-proxy-extras/litellm_proxy_extras/migration_lock.py

@@ -1,4 +1,5 @@
 import random
+import re
 import time
 from collections.abc import Generator, Mapping
 from contextlib import contextmanager
@@ -7,6 +8,7 @@ from typing import TYPE_CHECKING, Final
 from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
 
 from litellm_proxy_extras._logging import logger
+from litellm_proxy_extras._redaction import _redact_credentials
 from litellm_proxy_extras.prisma_toolchain import (
     MIGRATION_LOCK_TIMEOUT_ENV_VAR,
     migration_lock_timeout,
@@ -14,6 +16,26 @@ from litellm_proxy_extras.prisma_toolchain import (
 
 MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
 PRISMA_ADVISORY_LOCK_KEY: Final = 72707369
+_SERVER_ANSWERED_MARKERS: Final = ("FATAL:", "SSL", "certificate")
+
+
+@dataclass(frozen=True, slots=True)
+class AdvisoryLock:
+    key: int
+    name: str
+    holder: str
+
+
+COORDINATOR_LOCK: Final = AdvisoryLock(
+    MIGRATION_LOCK_KEY,
+    "v2 migration coordinator lock",
+    "another v2 migration resolver",
+)
+PRISMA_LOCK: Final = AdvisoryLock(
+    PRISMA_ADVISORY_LOCK_KEY,
+    "Prisma migration lock",
+    "another Prisma migration",
+)
 
 if TYPE_CHECKING:
     import psycopg
@@ -49,6 +71,28 @@ def _try_lock(connection: "psycopg.Connection[tuple[object, ...]]", key: int = M
     return row is not None and row.acquired
 
 
+def _coordination_failure(database_url: str, exc: "psycopg.Error", *, connecting: bool) -> str:
+    import psycopg
+
+    prefix: Final = "Lost or could not establish v2 migration coordination with the database: "
+    raw: Final = str(exc)
+    unquoted: Final = re.sub(r"\"[^\"]*\"|'[^']*'", "", raw)
+    detail: Final = _redact_credentials(raw).strip()
+    if (
+        not connecting
+        or not isinstance(exc, psycopg.OperationalError)
+        or any(marker in unquoted for marker in _SERVER_ANSWERED_MARKERS)
+    ):
+        return prefix + detail
+    parsed_url: Final = urlsplit(database_url)
+    try:
+        port: Final = parsed_url.port
+    except ValueError:
+        return prefix + f"P1001: Can't reach database server. {detail}"
+    host: Final = parsed_url.hostname or "localhost"
+    return prefix + f"P1001: Can't reach database server at `{host}:{port or 5432}`. {detail}"
+
+
 @dataclass(frozen=True, slots=True)
 class MigrationCoordinator:
     connection: "psycopg.Connection[tuple[object, ...]]"
@@ -75,39 +119,45 @@ class MigrationCoordinator:
 
 
 @contextmanager
-def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, None]:
+def migration_lock(
+    database_url: str,
+    lock: AdvisoryLock = COORDINATOR_LOCK,
+) -> Generator[MigrationCoordinator, None, None]:
     import psycopg
 
     wait_seconds: Final = migration_lock_timeout()
     deadline: Final = time.monotonic() + wait_seconds
     try:
-        with psycopg.connect(database_url, connect_timeout=10, autocommit=True) as connection:
+        connection: Final = psycopg.connect(database_url, connect_timeout=10, autocommit=True)
+    except psycopg.Error as exc:
+        raise RuntimeError(_coordination_failure(database_url, exc, connecting=True)) from None
+    try:
+        with connection:
             coordinator: Final = MigrationCoordinator(connection)
-            logger.info("Waiting for the v2 migration coordinator lock (up to %ss)", wait_seconds)
+            logger.info("Waiting for the %s (up to %ss)", lock.name, wait_seconds)
             while time.monotonic() < deadline:
                 with connection.transaction():
-                    if _try_lock(connection):
-                        logger.info("Acquired the v2 migration coordinator lock")
+                    if _try_lock(connection, lock.key):
+                        logger.info("Acquired the %s", lock.name)
 
                         yield coordinator
                         coordinator.check_connection()
                         return
                 time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
     except psycopg.Error as exc:
-        raise RuntimeError(f"Lost or could not establish v2 migration coordination with the database: {exc}") from exc
+        raise RuntimeError(_coordination_failure(database_url, exc, connecting=False)) from None
     raise RuntimeError(
-        f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. "
+        f"Timed out waiting for {lock.holder} after {wait_seconds}s. "
         f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}."
     )
 
 
 @contextmanager
 def prisma_migration_lock(database_url: str) -> Generator[None, None, None]:
-    """Hold the coordinator lock and Prisma's own lock key as transaction-level locks on one
-    pinned connection while a Prisma command runs with its session advisory lock disabled, so a
-    transaction pooler can never strand that lock on an idle backend after the command exits."""
-    with migration_lock(database_url) as coordinator:
-        coordinator.acquire_prisma_lock()
+    """Hold only Prisma's key as an xact lock on a pinned connection with Prisma's session lock disabled.
+
+    This keeps deploy independent of the coordinator key held by request-log index builds."""
+    with migration_lock(database_url, PRISMA_LOCK) as coordinator:
         coordinator.release_statement_snapshot()
         yield
 

litellm-proxy-extras/litellm_proxy_extras/migrations/20261015000000_managed_object_end_user_id/migration.sql

@@ -0,0 +1,4 @@
+SET LOCAL lock_timeout = '10s';
+
+-- AlterTable
+ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "end_user_id" TEXT;

litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql

@@ -0,0 +1,2 @@
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "idx_background_interaction_settlement_settled_at"
+ON "LiteLLM_BackgroundInteractionSettlement" ("settled_at");

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -1179,6 +1179,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
   org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
   api_key String?
   request_tags Json? @default("[]")
+  end_user_id String?
   updated_at DateTime @updatedAt
   updated_by String?
 
@@ -1980,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement {
   outcome             String?
 
   @@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at")
+  @@index([settled_at], map: "idx_background_interaction_settlement_settled_at")
 }
 
 model LiteLLM_Lens {

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -1,4 +1,3 @@
-import functools
 import glob
 import os
 import random
@@ -13,10 +12,10 @@ from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
-from urllib.parse import unquote, urlsplit
 
 from litellm_proxy_extras import prisma_toolchain
 from litellm_proxy_extras._logging import logger
+from litellm_proxy_extras._redaction import _redact_credentials
 from litellm_proxy_extras.migration_lock import held_migration_lock
 from litellm_proxy_extras.prisma_toolchain import (
     PRISMA_COMMAND_TIMEOUT_ENV_VAR,
@@ -212,51 +211,6 @@ def _max_migration_timestamp(names) -> int:
     return max(_migration_timestamp(n) for n in names)
 
 
-_REDACTED: Final = "REDACTED"
-_PASSWORD_QUERY_KEYS: Final = frozenset(("password", "sslpassword"))
-
-
-@functools.cache
-def _secret_shape_redactor() -> Callable[[str], str]:
-    try:
-        from litellm._logging import redact_secrets
-    except ImportError:
-        return lambda text: text
-    return redact_secrets
-
-
-def _url_passwords(url: str) -> frozenset[str]:
-    try:
-        parts: Final = urlsplit(url)
-    except ValueError:
-        return frozenset()
-    query_pairs: Final = tuple(pair.partition("=") for pair in parts.query.split("&"))
-    raw_query_passwords: Final = tuple(
-        value for key, separator, value in query_pairs if separator and key.lower() in _PASSWORD_QUERY_KEYS
-    )
-    raw_passwords: Final = ((parts.password,) if parts.password else ()) + raw_query_passwords
-    return frozenset(password for password in raw_passwords + tuple(map(unquote, raw_passwords)) if password)
-
-
-def _configured_database_passwords() -> frozenset[str]:
-    database_url: Final = os.getenv("DATABASE_URL")
-    direct_url: Final = os.getenv("DIRECT_URL")
-    database_passwords: Final = _url_passwords(database_url) if database_url else frozenset()
-    direct_passwords: Final = _url_passwords(direct_url) if direct_url else frozenset()
-    return database_passwords | direct_passwords
-
-
-def _redact_credentials(text: str) -> str:
-    """Mask configured database passwords before passing the text to LiteLLM redaction."""
-    passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True)
-    alternation: Final = "|".join(re.escape(password) for password in passwords)
-    password_pattern: Final = (
-        re.compile(rf"(?<![A-Za-z0-9_])(?:{alternation})(?![A-Za-z0-9_])", re.IGNORECASE) if passwords else None
-    )
-    result: Final = password_pattern.sub(_REDACTED, text) if password_pattern is not None else text
-    return _secret_shape_redactor()(result)
-
-
 def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]:
     if isinstance(command, tuple):
         return tuple(_redact_credentials(str(argument)) for argument in command)
@@ -789,16 +743,23 @@ class ProxyExtrasDBManager:
         (connection_limit, pool_timeout, schema, pgbouncer, sslaccept, ...) and
         translate Prisma's TLS params back, since libpq reads ``sslcert`` as a
         client certificate where Prisma reads it as the CA."""
-        from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
+        from urllib.parse import parse_qsl, quote, unquote, urlencode, urlparse, urlunparse
 
         parsed: Final = urlparse(url)
+        netloc: Final = parsed.netloc
+        userinfo, separator, hostport = netloc.rpartition("@")
+        user, password_separator, password = userinfo.partition(":")
+        canonical_userinfo: Final = quote(unquote(user), safe="") + (
+            password_separator + quote(unquote(password), safe="") if password_separator else ""
+        )
+        canonical_netloc: Final = f"{canonical_userinfo}{separator}{hostport}" if separator else netloc
         if not parsed.query:
-            return url
+            return url if canonical_netloc == netloc else urlunparse(parsed._replace(netloc=canonical_netloc))
         pairs: Final = tuple(parse_qsl(parsed.query))
         kept: Final = tuple((k, v) for k, v in pairs if k in LIBPQ_URL_PARAMS)
         sslaccept: Final = next((v for k, v in pairs if k == "sslaccept"), None)
         libpq_pairs: Final = ProxyExtrasDBManager._libpq_tls_params(kept, sslaccept)
-        return urlunparse(parsed._replace(query=urlencode(libpq_pairs, quote_via=quote)))
+        return urlunparse(parsed._replace(netloc=canonical_netloc, query=urlencode(libpq_pairs, quote_via=quote)))
 
     @staticmethod
     def _libpq_tls_params(

litellm-rust/AGENTS.md

@@ -15,6 +15,7 @@ Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `st
 - `crates/<crate>/tests/` is only for tests that drive several modules together end to end, such as a full request through a route. A test of one type or function never goes there
 - A helper shared by several inline test modules goes in a `#[cfg(test)] mod test_support;` at the crate root, and only once a second module needs it
 - A test for another crate's item belongs in that crate, not in a downstream one
+- Tests that call a real provider go in `crates/<crate>/tests/live/<provider>.rs` behind a `tests/live/main.rs`. Each reads the provider's key from the environment through an `Option` fixture and returns early when it is unset, never `#[ignore]`, so `cargo test` stays offline until the key is exported
 - Never set `autotests = false` or hand-list `[[test]]` targets. Every file directly under `tests/` is discovered by cargo, and a shared helper goes in `tests/<name>/mod.rs` or `tests/<subject>/support.rs` so it is not picked up as a test crate of its own
 
 ## Test fixtures and cases

litellm-rust/Cargo.lock

@@ -3406,6 +3406,7 @@ dependencies = [
 name = "litellm-auth-gcp"
 version = "0.1.0"
 dependencies = [
+ "base64 0.22.1",
  "gcp_auth",
  "google-cloud-auth",
  "http 1.4.2",
@@ -4088,6 +4089,7 @@ dependencies = [
  "litellm-cache",
  "litellm-cache-memory",
  "litellm-cache-response",
+ "litellm-core-utils",
  "litellm-host",
  "litellm-host-native",
  "litellm-http",

litellm-rust/crates/auth-gcp/Cargo.toml

@@ -11,6 +11,7 @@ google-sdk = ["dep:google-cloud-auth", "dep:http"]
 [dependencies]
 litellm-auth-types.workspace = true
 
+base64.workspace = true
 moka.workspace = true
 serde.workspace = true
 serde_json.workspace = true

litellm-rust/crates/auth-gcp/src/auth.rs

@@ -1,4 +1,4 @@
-use std::{future::Future, path::Path, pin::Pin, sync::Arc};
+use std::{future::Future, pin::Pin, sync::Arc};
 
 use gcp_auth::{CustomServiceAccount, TokenProvider};
 use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential};
@@ -8,7 +8,8 @@ use sha2::{Digest, Sha256};
 
 use crate::{
     config::{
-        CredentialSource, VertexConfig, credential_source, get_vertex_ai_project, non_empty_env,
+        CredentialSource, VertexConfig, credential_source, get_vertex_ai_project,
+        is_base64_encoded_json, names_a_credentials_file, non_empty_env,
     },
     constants::{
         CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT, VERTEX_AI_API_KEY_ENV,
@@ -176,14 +177,7 @@ impl VertexProviderLoader for GcpProviderLoader {
                     .map_err(auth_acquisition_error)?,
                 ),
                 CredentialSource::Trusted(configured) => {
-                    let configured = configured.expose();
-                    let service_account = if Path::new(configured).is_file() {
-                        CustomServiceAccount::from_file(configured)
-                    } else {
-                        CustomServiceAccount::from_json(configured)
-                    }
-                    .map_err(auth_acquisition_error)?;
-                    Arc::new(service_account)
+                    Arc::new(load_trusted_service_account(configured.expose())?)
                 }
                 CredentialSource::ApplicationCredentials(path) => {
                     Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?)
@@ -240,6 +234,23 @@ enum CredentialCacheKey {
     Adc,
 }
 
+const BASE64_CREDENTIALS: &str = "The `vertex_credentials` value is base64-encoded JSON, not the JSON itself. \
+    Set it to a readable file path, or to the credentials JSON itself.";
+
+fn load_trusted_service_account(configured: &str) -> Result<CustomServiceAccount, Error> {
+    if is_base64_encoded_json(configured) {
+        return Err(Error::InvalidConfiguration(ErrorDetail::Message(
+            BASE64_CREDENTIALS.into(),
+        )));
+    }
+    if names_a_credentials_file(configured) {
+        CustomServiceAccount::from_file(configured)
+    } else {
+        CustomServiceAccount::from_json(configured)
+    }
+    .map_err(auth_acquisition_error)
+}
+
 fn auth_acquisition_error(error: gcp_auth::Error) -> Error {
     Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed(
         "Vertex AI credentials",
@@ -264,6 +275,73 @@ mod tests {
         assert_eq!(validate_request_credentials(credentials).is_ok(), accepted);
     }
 
+    #[rstest::rstest]
+    #[case::missing_file(
+        "/nonexistent/vertex-key.json",
+        "failed to open application credentials file"
+    )]
+    #[case::malformed_json(
+        " {\"type\": \"service_account\",",
+        "failed to deserialize ApplicationCredentials"
+    )]
+    #[case::quoted_json(
+        "'{\"type\": \"service_account\"}'",
+        "failed to deserialize ApplicationCredentials"
+    )]
+    #[tokio::test]
+    async fn trusted_credentials_are_read_as_a_file_unless_they_are_json(
+        #[case] configured: &str,
+        #[case] expected: &str,
+    ) {
+        let Err(error) = GcpProviderLoader
+            .load(CredentialSource::Trusted(SecretValue::new(configured)))
+            .await
+        else {
+            panic!("{configured:?} should not load");
+        };
+        assert!(error.to_string().contains(expected), "{error}");
+    }
+
+    #[rstest::rstest]
+    #[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")]
+    #[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")]
+    #[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")]
+    #[tokio::test]
+    async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) {
+        assert!(
+            trusted_load_error(configured)
+                .await
+                .contains("is base64-encoded JSON, not the JSON itself")
+        );
+    }
+
+    async fn trusted_load_error(configured: &str) -> String {
+        match GcpProviderLoader
+            .load(CredentialSource::Trusted(SecretValue::new(configured)))
+            .await
+        {
+            Ok(_) => panic!("{configured:?} should not load"),
+            Err(error) => error.to_string(),
+        }
+    }
+
+    #[rstest::rstest]
+    #[tokio::test]
+    async fn a_relative_key_file_named_like_json_is_read_as_a_file() {
+        let key = tempfile::Builder::new()
+            .prefix("{vertex")
+            .suffix(".json")
+            .tempfile_in(".")
+            .unwrap();
+        std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
+        let relative_name = key.path().file_name().unwrap().to_str().unwrap();
+        let absolute_path = std::fs::canonicalize(key.path()).unwrap();
+        assert_eq!(
+            trusted_load_error(relative_name).await,
+            trusted_load_error(absolute_path.to_str().unwrap()).await
+        );
+    }
+
     #[rstest::rstest]
     fn inline_and_trusted_credentials_never_share_a_cache_entry() {
         assert_ne!(

litellm-rust/crates/auth-gcp/src/config.rs

@@ -1,5 +1,9 @@
 use std::{collections::BTreeMap, path::Path};
 
+use base64::{
+    Engine, alphabet,
+    engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig},
+};
 use litellm_auth_types::{Error, ErrorDetail, InputSource, SecretValue, Sourced};
 use serde_json::{Map, Value};
 
@@ -80,9 +84,10 @@ pub fn get_vertex_ai_project_from_credentials(
         CredentialSource::Inline(configured) => configured.expose().to_string(),
         CredentialSource::Trusted(configured) => {
             let configured = configured.expose();
-            match std::fs::read_to_string(configured) {
-                Ok(contents) if Path::new(configured).is_file() => contents,
-                _ => configured.to_string(),
+            if names_a_credentials_file(configured) {
+                std::fs::read_to_string(configured).ok()?
+            } else {
+                configured.to_string()
             }
         }
         CredentialSource::ApplicationCredentials(path) => std::fs::read_to_string(path).ok()?,
@@ -105,6 +110,32 @@ pub fn get_vertex_ai_location(
         .or_else(|| VertexParams::LOCATION.resolve(&|_| None, env_lookup))
 }
 
+const INLINE_JSON_LEADING_NOISE: [char; 7] = [' ', '\t', '\r', '\n', '\u{feff}', '"', '\''];
+
+fn is_inline_json(configured: &str) -> bool {
+    configured
+        .trim_start_matches(INLINE_JSON_LEADING_NOISE)
+        .starts_with('{')
+}
+
+pub(crate) fn names_a_credentials_file(configured: &str) -> bool {
+    !is_inline_json(configured) || Path::new(configured).is_file()
+}
+
+const BASE64_ANY_PADDING: GeneralPurpose = GeneralPurpose::new(
+    &alphabet::STANDARD,
+    GeneralPurposeConfig::new().with_decode_padding_mode(DecodePaddingMode::Indifferent),
+);
+
+pub(crate) fn is_base64_encoded_json(configured: &str) -> bool {
+    let compact: String = configured.split_whitespace().collect();
+    let quoted = compact.trim_matches(['"', '\'']);
+    let encoded = quoted.strip_prefix("base64:").unwrap_or(quoted);
+    BASE64_ANY_PADDING
+        .decode(encoded.replace('-', "+").replace('_', "/"))
+        .is_ok_and(|decoded| serde_json::from_slice::<Map<String, Value>>(&decoded).is_ok())
+}
+
 #[derive(Clone, Debug)]
 pub enum CredentialSource {
     Inline(SecretValue),

litellm-rust/crates/auth-gcp/tests/vertex_config.rs

@@ -93,3 +93,22 @@ fn project_is_read_from_the_credentials_that_would_authenticate() {
         None
     );
 }
+
+#[rstest]
+fn project_is_read_from_a_relative_key_file_named_like_json() {
+    let key = tempfile::Builder::new()
+        .prefix("{vertex")
+        .suffix(".json")
+        .tempfile_in(".")
+        .unwrap();
+    std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
+    let relative_name = key.path().file_name().unwrap().to_str().unwrap();
+    let from_file = VertexConfig::from_params(&VertexParams {
+        vertex_credentials: Some(relative_name.into()),
+        ..VertexParams::default()
+    });
+    assert_eq!(
+        get_vertex_ai_project_from_credentials(&from_file, &|_| None).as_deref(),
+        Some("file-project")
+    );
+}

litellm-rust/crates/gateway-inference/src/messages.rs

@@ -24,7 +24,7 @@ const OPENROUTER_API_HEADERS: [&str; 2] = [
     "x-openrouter-experimental-metadata",
 ];
 const ANTHROPIC_API_HEADER_PROVIDERS: &str =
-    "anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter";
+    "anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter,deepseek";
 
 pub async fn create(
     State(gateway): State<Arc<Gateway>>,

litellm-rust/crates/gateway-inference/tests/messages.rs

@@ -16,7 +16,7 @@ use wiremock::{
 #[rstest]
 #[case::anthropic("anthropic/test-model", "/v1/messages", true)]
 #[case::azure("azure_ai/test-model", "/anthropic/v1/messages", false)]
-#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", false)]
+#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", true)]
 #[tokio::test]
 async fn messages_reaches_the_provider_and_preserves_json_or_sse(
     #[case] model: &str,

litellm-rust/crates/inference-messages/src/prepare.rs

@@ -124,7 +124,7 @@ fn prepare_provider_request(
         url,
         body: transformed,
         environment,
-        timeout,
+        timeout: timeout.or_else(|| shaping.settings.request_timeout()),
         api_key: api_key.map(SecretValue::new),
     })
 }
@@ -253,6 +253,41 @@ mod tests {
         );
     }
 
+    #[rstest]
+    #[case::nothing_configured(None, None, None)]
+    #[case::configured_setting(None, Some(1800.0), Some(Duration::from_secs(1800)))]
+    #[case::call_timeout_wins(
+        Some(Duration::from_secs(30)),
+        Some(1800.0),
+        Some(Duration::from_secs(30))
+    )]
+    #[case::call_timeout_only(Some(Duration::from_secs(30)), None, Some(Duration::from_secs(30)))]
+    #[case::unusable_setting_ignored(None, Some(-1.0), None)]
+    #[case::zero_setting_ignored(None, Some(0.0), None)]
+    fn timeout_falls_back_to_the_configured_request_timeout(
+        shaping: MessagesShaping,
+        #[case] call_timeout: Option<Duration>,
+        #[case] request_timeout: Option<f64>,
+        #[case] expected: Option<Duration>,
+    ) {
+        let prepared = prepare(MessagesCall {
+            body: body(json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})),
+            api_key: Some("sk-test".into()),
+            api_base: Some("https://anthropic.test".into()),
+            custom_llm_provider: Some("anthropic".into()),
+            litellm_params: Default::default(),
+            extra_headers: None,
+            provider_specific_header: None,
+            timeout: call_timeout,
+            shaping: MessagesShaping {
+                settings: MessagesSettings { request_timeout, ..shaping.settings },
+                ..shaping
+            },
+        })
+        .unwrap();
+        assert_eq!(prepared.timeout, expected);
+    }
+
     fn prepared_body(fields: Value, shaping: MessagesShaping) -> Result<Value, Error> {
         prepare(MessagesCall {
             body: body(fields),

litellm-rust/crates/inference-messages/src/types.rs

@@ -60,6 +60,16 @@ pub struct MessagesSettings {
     pub reasoning_auto_summary: bool,
     #[serde(default)]
     pub additional_drop_params: Vec<String>,
+    #[serde(default)]
+    pub request_timeout: Option<f64>,
+}
+
+impl MessagesSettings {
+    pub fn request_timeout(&self) -> Option<Duration> {
+        self.request_timeout
+            .and_then(|secs| Duration::try_from_secs_f64(secs).ok())
+            .filter(|timeout| !timeout.is_zero())
+    }
 }
 
 #[cfg(test)]
@@ -120,13 +130,15 @@ mod tests {
             },
             "drop_params": true,
             "reasoning_auto_summary": true,
-            "additional_drop_params": ["metadata.user_id", "thinking"]
+            "additional_drop_params": ["metadata.user_id", "thinking"],
+            "request_timeout": 1800.0
         }),
         MessagesShaping {
             settings: MessagesSettings {
                 drop_params: true,
                 reasoning_auto_summary: true,
                 additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()],
+                request_timeout: Some(1800.0),
             },
             capabilities: MessagesModelCapabilities {
                 supports_reasoning: true,

litellm-rust/crates/inference-responses/Cargo.toml

@@ -10,6 +10,7 @@ bytes.workspace = true
 futures-util.workspace = true
 litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] }
 litellm-cache-response.workspace = true
+litellm-core-utils.workspace = true
 litellm-host.workspace = true
 litellm-http.workspace = true
 litellm-inference.workspace = true

litellm-rust/crates/inference-responses/src/prepare.rs

@@ -1,7 +1,11 @@
-use litellm_host::interceptors::RequestContext;
+use litellm_core_utils::get_llm_provider_logic::{
+    CustomLlmProvider, LlmProviders, get_custom_llm_provider,
+};
+use litellm_host::interceptors::{ProviderIdentity, RequestContext};
 use litellm_llms::{
     base_llm::responses::transformation::BaseResponsesApiConfig,
     openai::responses::transformation::OpenAiResponsesApiConfig,
+    openrouter::responses::transformation::OPENROUTER_RESPONSES_CONFIG,
 };
 use litellm_secrets::source::SecretSource;
 use serde_json::Value;
@@ -16,10 +20,10 @@ pub(super) async fn prepare(
     call: ResponsesCall,
     secrets: &dyn SecretSource,
 ) -> Result<ProviderResponsesRequest, Error> {
-    let identity = resolve_provider(&call.model, call.custom_llm_provider.as_deref())?;
+    let ResponsesProvider { identity, config } =
+        resolve_provider(&call.model, call.custom_llm_provider.as_deref())?;
     let provider = identity.provider.as_str();
     let model = identity.model.as_str();
-    let config: &'static dyn BaseResponsesApiConfig = &OpenAiResponsesApiConfig;
     let snapshot = secrets
         .resolve(config.secret_names(call.api_key.as_deref(), call.api_base.as_deref()))
         .await?;
@@ -52,20 +56,107 @@ pub(super) async fn prepare(
     })
 }
 
+pub(super) struct ResponsesProvider {
+    pub identity: ProviderIdentity,
+    pub config: &'static dyn BaseResponsesApiConfig,
+}
+
+/// Python's `get_llm_provider`: a bare model is OpenAI's, and with an explicit OpenRouter
+/// provider a model without a vendor segment keeps the `openrouter/` prefix as its wire id.
 pub(super) fn resolve_provider(
     model: &str,
     custom_llm_provider: Option<&str>,
-) -> Result<litellm_host::interceptors::ProviderIdentity, Error> {
-    let provider = custom_llm_provider.unwrap_or("openai");
-    if provider != "openai" {
-        return Err(Error::Unsupported("native HTTP responses provider"));
-    }
-    let resolved = model.strip_prefix("openai/").unwrap_or(model);
-    if resolved.is_empty() || resolved.contains('/') {
+) -> Result<ResponsesProvider, Error> {
+    let explicit = custom_llm_provider.filter(|provider| !provider.is_empty());
+    let (provider, resolved) = match get_custom_llm_provider(model, explicit) {
+        Some(CustomLlmProvider {
+            model,
+            custom_llm_provider,
+        }) => (custom_llm_provider, model),
+        None => ("openai", model),
+    };
+    if resolved.is_empty() {
         return Err(Error::InvalidProvider(model.into()));
     }
-    Ok(litellm_host::interceptors::ProviderIdentity {
-        model: resolved.into(),
-        provider: provider.into(),
+    let (config, wire_model): (&'static dyn BaseResponsesApiConfig, String) =
+        match provider.parse::<LlmProviders>() {
+            Ok(LlmProviders::Openai) if resolved.contains('/') => {
+                return Err(Error::InvalidProvider(model.into()));
+            }
+            Ok(LlmProviders::Openai) => (&OpenAiResponsesApiConfig, resolved.into()),
+            Ok(LlmProviders::Openrouter) if explicit.is_some() && !resolved.contains('/') => (
+                &OPENROUTER_RESPONSES_CONFIG,
+                format!("openrouter/{resolved}"),
+            ),
+            Ok(LlmProviders::Openrouter) => (&OPENROUTER_RESPONSES_CONFIG, resolved.into()),
+            _ => return Err(Error::Unsupported("native HTTP responses provider")),
+        };
+    Ok(ResponsesProvider {
+        identity: ProviderIdentity {
+            model: wire_model,
+            provider: provider.into(),
+        },
+        config,
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use super::*;
+
+    /// Pinned against Python's `get_llm_provider` on 2026-10-11.
+    #[rstest]
+    #[case::bare("gpt-5", None, "openai", "gpt-5")]
+    #[case::openai_prefix("openai/gpt-5", None, "openai", "gpt-5")]
+    #[case::openai_explicit("gpt-5", Some("openai"), "openai", "gpt-5")]
+    #[case::openrouter_prefix(
+        "openrouter/anthropic/claude",
+        None,
+        "openrouter",
+        "anthropic/claude"
+    )]
+    #[case::openrouter_explicit(
+        "openrouter/anthropic/claude",
+        Some("openrouter"),
+        "openrouter",
+        "anthropic/claude"
+    )]
+    #[case::openrouter_explicit_unprefixed(
+        "anthropic/claude",
+        Some("openrouter"),
+        "openrouter",
+        "anthropic/claude"
+    )]
+    #[case::openrouter_native_id(
+        "openrouter/auto",
+        Some("openrouter"),
+        "openrouter",
+        "openrouter/auto"
+    )]
+    #[case::openrouter_native_id_implicit("openrouter/auto", None, "openrouter", "auto")]
+    fn providers_and_wire_models_match_python(
+        #[case] model: &str,
+        #[case] custom_llm_provider: Option<&str>,
+        #[case] provider: &str,
+        #[case] wire_model: &str,
+    ) {
+        let resolved = resolve_provider(model, custom_llm_provider).unwrap();
+        assert_eq!(resolved.identity.provider, provider);
+        assert_eq!(resolved.identity.model, wire_model);
+    }
+
+    #[rstest]
+    #[case::other_prefix("other/model", None)]
+    #[case::other_explicit("model", Some("other"))]
+    #[case::openai_with_vendor_segment("openai/other/model", None)]
+    #[case::empty_model("openai/", None)]
+    #[case::conflicting_prefix("openrouter/anthropic/claude", Some("openai"))]
+    fn everything_else_is_rejected(#[case] model: &str, #[case] custom_llm_provider: Option<&str>) {
+        assert!(matches!(
+            resolve_provider(model, custom_llm_provider),
+            Err(Error::InvalidProvider(_) | Error::Unsupported(_))
+        ));
+    }
+}

litellm-rust/crates/inference-responses/tests/live/main.rs

@@ -0,0 +1,7 @@
+//! Real provider calls, one module per provider. Each test skips when its provider's key is not
+//! in the environment, so `cargo test` stays offline by default.
+
+#[path = "../support/mod.rs"]
+mod support;
+
+mod openrouter;

litellm-rust/crates/inference-responses/tests/live/openrouter.rs

@@ -0,0 +1,98 @@
+use std::{sync::Arc, time::Duration};
+
+use futures_util::TryStreamExt;
+use litellm_inference_responses::types::{ResponsesCall, ResponsesOutput};
+use litellm_inference_testing::RecordingSecrets;
+use litellm_llms::base_llm::responses::streaming::decode;
+use litellm_llms_types::formats::responses::streaming::ResponsesEventType;
+use rstest::{fixture, rstest};
+use serde_json::json;
+
+use crate::support::responses_route;
+
+const DEFAULT_MODEL: &str = "openrouter/openai/gpt-5.5";
+
+/// `OPENROUTER_API_KEY` as a secret source; tests skip when it is unset.
+#[fixture]
+fn secrets() -> Option<Arc<RecordingSecrets>> {
+    let key = std::env::var("OPENROUTER_API_KEY").ok()?;
+    Some(Arc::new(RecordingSecrets::new([(
+        "OPENROUTER_API_KEY",
+        key.as_str(),
+    )])))
+}
+
+fn call(stream: bool) -> ResponsesCall {
+    ResponsesCall {
+        model: std::env::var("OPENROUTER_LIVE_MODEL").unwrap_or_else(|_| DEFAULT_MODEL.into()),
+        input: json!("Reply with the single word: pong"),
+        optional_params: json!({"max_output_tokens": 32, "stream": stream})
+            .as_object()
+            .unwrap()
+            .clone(),
+        api_key: None,
+        api_base: None,
+        custom_llm_provider: None,
+        extra_headers: None,
+        timeout: Some(Duration::from_secs(60)),
+    }
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_complete_response_carries_output_and_usage(secrets: Option<Arc<RecordingSecrets>>) {
+    let Some(secrets) = secrets else { return };
+    let output = responses_route(secrets)
+        .execute(call(false), &(), None)
+        .await
+        .unwrap();
+    let ResponsesOutput::Complete(response) = output else {
+        panic!("expected a complete response");
+    };
+    let body = serde_json::to_value(response).unwrap();
+    let text = body["output"]
+        .as_array()
+        .into_iter()
+        .flatten()
+        .filter_map(|item| item["content"].as_array())
+        .flatten()
+        .filter_map(|part| part["text"].as_str())
+        .collect::<String>();
+    assert!(text.to_ascii_lowercase().contains("pong"), "{body}");
+    assert!(
+        body["usage"]["total_tokens"].as_u64().unwrap() > 0,
+        "{body}"
+    );
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_stream_ends_with_response_completed(secrets: Option<Arc<RecordingSecrets>>) {
+    let Some(secrets) = secrets else { return };
+    let output = responses_route(secrets)
+        .execute(call(true), &(), None)
+        .await
+        .unwrap();
+    let ResponsesOutput::Stream { chunks, .. } = output else {
+        panic!("expected SSE");
+    };
+    let wire = chunks.try_collect::<Vec<_>>().await.unwrap().concat();
+    let events = decode(&wire).unwrap();
+    assert_eq!(
+        events.first().map(|event| &event.event_type),
+        Some(&ResponsesEventType::ResponseCreated)
+    );
+    assert!(
+        events
+            .iter()
+            .any(|event| event.event_type == ResponsesEventType::OutputTextDelta)
+    );
+    let last = events.last().unwrap();
+    assert_eq!(last.event_type, ResponsesEventType::ResponseCompleted);
+    assert!(
+        last.data["response"]["usage"]["total_tokens"]
+            .as_u64()
+            .unwrap()
+            > 0
+    );
+}

litellm-rust/crates/inference-responses/tests/openrouter.rs

@@ -0,0 +1,150 @@
+use std::sync::Arc;
+
+use futures_util::TryStreamExt;
+use litellm_inference_responses::{
+    Error,
+    types::{ResponsesCall, ResponsesOutput},
+};
+use litellm_inference_testing::{RecordingSecrets, no_secrets};
+use rstest::{fixture, rstest};
+use serde_json::json;
+use wiremock::ResponseTemplate;
+
+mod support;
+use support::*;
+
+#[fixture]
+fn call() -> ResponsesCall {
+    ResponsesCall {
+        model: "openrouter/anthropic/test-model".into(),
+        input: json!("hello"),
+        optional_params: Default::default(),
+        api_key: None,
+        api_base: None,
+        custom_llm_provider: None,
+        extra_headers: None,
+        timeout: None,
+    }
+}
+
+fn response_body() -> serde_json::Value {
+    json!({"id": "resp_1", "model": "anthropic/test-model", "output": [{"type": "message", "content": []}], "usage": {"total_tokens": 7, "cost": 0.00017}})
+}
+
+#[rstest]
+#[case::explicit_key(Some("explicit"), &[], "Bearer explicit")]
+#[case::managed_key(None, &[("OPENROUTER_API_KEY", "managed")], "Bearer managed")]
+#[case::legacy_name(None, &[("OR_API_KEY", "legacy")], "Bearer legacy")]
+#[tokio::test]
+async fn the_payload_goes_to_the_openrouter_responses_endpoint_with_a_bearer(
+    call: ResponsesCall,
+    #[case] api_key: Option<&str>,
+    #[case] secrets: &[(&str, &str)],
+    #[case] expected: &str,
+) {
+    let upstream = upstream([json_response(response_body())]).await;
+    let output = responses_route(Arc::new(RecordingSecrets::new(secrets.iter().copied())))
+        .execute(
+            ResponsesCall {
+                api_key: api_key.map(str::to_owned),
+                api_base: Some(format!("{}/api/v1", upstream.uri())),
+                optional_params: json!({"max_output_tokens": 16})
+                    .as_object()
+                    .unwrap()
+                    .clone(),
+                ..call
+            },
+            &(),
+            None,
+        )
+        .await
+        .unwrap();
+    let ResponsesOutput::Complete(response) = output else {
+        panic!("expected a complete response");
+    };
+    assert_eq!(serde_json::to_value(response).unwrap(), response_body());
+    let request = only_request(&upstream).await;
+    assert_eq!(request.url.path(), "/api/v1/responses");
+    assert_eq!(request.header("authorization"), Some(expected));
+    assert_eq!(
+        request.json(),
+        json!({"max_output_tokens": 16, "model": "anthropic/test-model", "input": "hello"})
+    );
+}
+
+#[rstest]
+#[tokio::test]
+async fn the_secret_manager_supplies_the_base_and_key(call: ResponsesCall) {
+    let upstream = upstream([json_response(response_body())]).await;
+    let base = format!("{}/api/v1/", upstream.uri());
+    let secrets = Arc::new(RecordingSecrets::new([
+        ("OPENROUTER_API_KEY", "managed"),
+        ("OPENROUTER_API_BASE", base.as_str()),
+    ]));
+    responses_route(secrets.clone())
+        .execute(call, &(), None)
+        .await
+        .unwrap();
+    assert_eq!(
+        secrets.requested(),
+        ["OPENROUTER_API_KEY", "OR_API_KEY", "OPENROUTER_API_BASE"]
+    );
+    let request = only_request(&upstream).await;
+    assert_eq!(request.url.path(), "/api/v1/responses");
+    assert_eq!(request.header("authorization"), Some("Bearer managed"));
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_missing_key_fails_before_transport(call: ResponsesCall) {
+    let upstream = upstream([json_response(response_body())]).await;
+    let result = responses_route(no_secrets())
+        .execute(
+            ResponsesCall {
+                api_base: Some(upstream.uri()),
+                ..call
+            },
+            &(),
+            None,
+        )
+        .await;
+    assert!(matches!(
+        result,
+        Err(Error::Auth(litellm_auth::Error::MissingApiKey {
+            provider: "OpenRouter",
+            ..
+        }))
+    ));
+    assert!(received(&upstream).await.is_empty());
+}
+
+#[rstest]
+#[tokio::test]
+async fn streaming_relays_every_openrouter_event_untouched(call: ResponsesCall) {
+    let sse = "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\"}}\n\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\"}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"usage\":{\"cost\":0.0002}}}\n\n";
+    let upstream =
+        upstream([ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")]).await;
+    let output = responses_route(no_secrets())
+        .execute(
+            ResponsesCall {
+                api_key: Some("key".into()),
+                api_base: Some(upstream.uri()),
+                optional_params: json!({"stream": true}).as_object().unwrap().clone(),
+                ..call
+            },
+            &(),
+            None,
+        )
+        .await
+        .unwrap();
+    let ResponsesOutput::Stream { chunks, .. } = output else {
+        panic!("expected SSE");
+    };
+    assert_eq!(
+        chunks.try_collect::<Vec<_>>().await.unwrap().concat(),
+        sse.as_bytes()
+    );
+    let request = only_request(&upstream).await;
+    assert_eq!(request.json()["stream"], true);
+    assert_eq!(request.json()["model"], "anthropic/test-model");
+}

litellm-rust/crates/llms/AGENTS.md

@@ -58,6 +58,7 @@
 
 - `AnthropicConfig`, `AmazonConverseConfig` and the non-OCR base traits are partial. Preserve their acceptance gates, passthrough and host fallback contracts
 - `OpenAiResponsesApiConfig` implements WebSocket transformations and direct HTTP Responses, without Python's model-specific param rewriting or Responses-to-Chat emulation
+- OpenAI-compatible Responses hosts (`openrouter`) reuse the OpenAI request and response transforms as free functions and own only their credentials and base URL
 - Chat `supported_openai_param_mappings` lists OpenAI/provider name pairs, unlike Python's name-only list. Audio `map_transcription_params` is a Rust filtering helper
 - Auth resolution is split between configs and route preparation in `inference-<format>`
 

litellm-rust/crates/llms/src/openai/responses/transformation.rs

@@ -165,34 +165,46 @@ impl BaseResponsesApiConfig for OpenAiResponsesApiConfig {
         input: Value,
         params: Map<String, Value>,
     ) -> Result<Value, Error> {
-        if !input.is_string() && !input.is_array() {
-            return Err(Error::InvalidRequest(
-                "responses input must be a string or an array".into(),
-            ));
-        }
-        if params
-            .get("stream")
-            .is_some_and(|stream| !stream.is_boolean())
-        {
-            return Err(Error::InvalidRequest("stream must be a boolean".into()));
-        }
-        Ok(Value::Object(
-            params
-                .into_iter()
-                .chain([
-                    ("model".into(), Value::String(model.into())),
-                    ("input".into(), input),
-                ])
-                .collect(),
-        ))
+        responses_request(model, input, params)
     }
 
     fn transform_response_api_response(&self, body: Value) -> Result<ResponsesApiResponse, Error> {
-        serde_json::from_value(body)
-            .map_err(|error| Error::InvalidResponse(error.to_string().into()))
+        responses_response(body)
     }
 }
 
+/// The OpenAI Responses request body, shared by every OpenAI-compatible host.
+pub fn responses_request(
+    model: &str,
+    input: Value,
+    params: Map<String, Value>,
+) -> Result<Value, Error> {
+    if !input.is_string() && !input.is_array() {
+        return Err(Error::InvalidRequest(
+            "responses input must be a string or an array".into(),
+        ));
+    }
+    if params
+        .get("stream")
+        .is_some_and(|stream| !stream.is_boolean())
+    {
+        return Err(Error::InvalidRequest("stream must be a boolean".into()));
+    }
+    Ok(Value::Object(
+        params
+            .into_iter()
+            .chain([
+                ("model".into(), Value::String(model.into())),
+                ("input".into(), input),
+            ])
+            .collect(),
+    ))
+}
+
+pub fn responses_response(body: Value) -> Result<ResponsesApiResponse, Error> {
+    serde_json::from_value(body).map_err(|error| Error::InvalidResponse(error.to_string().into()))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

litellm-rust/crates/llms/src/openrouter/mod.rs

@@ -1 +1,2 @@
 pub mod messages;
+pub mod responses;

litellm-rust/crates/llms/src/openrouter/responses/AGENTS.md

@@ -0,0 +1,10 @@
+# rules
+
+- Send the OpenAI Responses payload to `/api/v1/responses` and relay the OpenAI SSE events
+- Authenticate with a bearer token from `OPENROUTER_API_KEY`, falling back to `OR_API_KEY`
+- Only `OPENROUTER_API_BASE` overrides the default base, never the OpenAI base variables
+- No native WebSocket mode
+
+# references
+
+- https://openrouter.ai/docs/api_reference/responses/overview

litellm-rust/crates/llms/src/openrouter/responses/mod.rs

@@ -0,0 +1 @@
+pub mod transformation;

litellm-rust/crates/llms/src/openrouter/responses/transformation.rs

@@ -0,0 +1,174 @@
+use litellm_auth::{CredentialPlacement, SecretValue};
+use litellm_core_utils::settings::resolve_non_empty;
+use litellm_llms_types::formats::responses::ResponsesApiResponse;
+use serde_json::{Map, Value};
+
+use crate::{
+    Error,
+    base_llm::{
+        auth::{AuthScheme, ValidatedEnvironment},
+        responses::transformation::BaseResponsesApiConfig,
+    },
+    openai::responses::transformation::{responses_request, responses_response},
+};
+
+const API_KEY_ENVS: [&str; 2] = ["OPENROUTER_API_KEY", "OR_API_KEY"];
+const API_BASE_ENV: &str = "OPENROUTER_API_BASE";
+const DEFAULT_BASE: &str = "https://openrouter.ai/api/v1";
+const RESPONSES_PATH: &str = "/responses";
+
+pub struct OpenRouterResponsesApiConfig;
+pub const OPENROUTER_RESPONSES_CONFIG: OpenRouterResponsesApiConfig = OpenRouterResponsesApiConfig;
+
+impl BaseResponsesApiConfig for OpenRouterResponsesApiConfig {
+    fn secret_names(
+        &self,
+        api_key: Option<&str>,
+        api_base: Option<&str>,
+    ) -> &'static [&'static str] {
+        match (
+            api_key.is_some_and(|key| !key.trim().is_empty()),
+            api_base.is_some_and(|base| !base.trim().is_empty()),
+        ) {
+            (true, true) => &[],
+            (true, false) => &[API_BASE_ENV],
+            (false, true) => &API_KEY_ENVS,
+            (false, false) => &["OPENROUTER_API_KEY", "OR_API_KEY", API_BASE_ENV],
+        }
+    }
+
+    fn validate_environment(
+        &self,
+        headers: Vec<(String, String)>,
+        api_key: Option<&str>,
+        lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> Result<ValidatedEnvironment, Error> {
+        let key = resolve_non_empty(api_key, lookup, &API_KEY_ENVS).ok_or(
+            litellm_auth::Error::MissingApiKey {
+                provider: "OpenRouter",
+                environment_variable: API_KEY_ENVS[0],
+            },
+        )?;
+        Ok(ValidatedEnvironment {
+            headers: litellm_http::request::with_default_headers(
+                headers,
+                &[("content-type", "application/json")],
+            ),
+            auth: AuthScheme::Credential {
+                placement: CredentialPlacement::Bearer,
+                secret: SecretValue::new(key),
+            },
+        })
+    }
+
+    fn get_complete_url(
+        &self,
+        api_base: Option<&str>,
+        lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> String {
+        let base = resolve_non_empty(api_base, lookup, &[API_BASE_ENV])
+            .unwrap_or_else(|| DEFAULT_BASE.into());
+        format!("{}{RESPONSES_PATH}", base.trim_end_matches('/'))
+    }
+
+    fn transform_responses_api_request(
+        &self,
+        model: &str,
+        input: Value,
+        params: Map<String, Value>,
+    ) -> Result<Value, Error> {
+        responses_request(model, input, params)
+    }
+
+    fn transform_response_api_response(&self, body: Value) -> Result<ResponsesApiResponse, Error> {
+        responses_response(body)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use super::*;
+
+    fn env(values: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option<String> {
+        move |name| {
+            values
+                .iter()
+                .find(|(key, _)| *key == name)
+                .map(|(_, value)| (*value).to_owned())
+        }
+    }
+
+    #[rstest]
+    #[case::default(None, &[], "https://openrouter.ai/api/v1/responses")]
+    #[case::explicit(Some("https://proxy.test/api/v1"), &[], "https://proxy.test/api/v1/responses")]
+    #[case::trailing_slash(Some("https://proxy.test/api/v1/"), &[], "https://proxy.test/api/v1/responses")]
+    #[case::blank_falls_to_secret(Some(" "), &[("OPENROUTER_API_BASE", "https://secret.test/v1")], "https://secret.test/v1/responses")]
+    #[case::openai_base_is_ignored(None, &[("OPENAI_BASE_URL", "https://openai.test/v1")], "https://openrouter.ai/api/v1/responses")]
+    fn url_is_the_base_plus_responses(
+        #[case] api_base: Option<&str>,
+        #[case] secrets: &'static [(&'static str, &'static str)],
+        #[case] expected: &str,
+    ) {
+        assert_eq!(
+            OPENROUTER_RESPONSES_CONFIG.get_complete_url(api_base, &env(secrets)),
+            expected
+        );
+    }
+
+    #[rstest]
+    #[case::explicit(Some("explicit"), &[("OPENROUTER_API_KEY", "managed")], "explicit")]
+    #[case::managed(None, &[("OPENROUTER_API_KEY", "managed"), ("OR_API_KEY", "legacy")], "managed")]
+    #[case::legacy_name(None, &[("OR_API_KEY", "legacy")], "legacy")]
+    #[case::blank_explicit(Some(" "), &[("OR_API_KEY", "legacy")], "legacy")]
+    fn bearer_comes_from_the_call_then_either_secret_name(
+        #[case] api_key: Option<&str>,
+        #[case] secrets: &'static [(&'static str, &'static str)],
+        #[case] expected: &str,
+    ) {
+        let environment = OPENROUTER_RESPONSES_CONFIG
+            .validate_environment(vec![], api_key, &env(secrets))
+            .unwrap();
+        let AuthScheme::Credential { placement, secret } = environment.auth else {
+            panic!("expected a credential");
+        };
+        assert_eq!(placement, CredentialPlacement::Bearer);
+        assert_eq!(secret.expose(), expected);
+        assert_eq!(
+            environment.headers,
+            vec![("content-type".to_owned(), "application/json".to_owned())]
+        );
+    }
+
+    #[rstest]
+    #[case::nothing(&[])]
+    #[case::openai_key_only(&[("OPENAI_API_KEY", "openai")])]
+    fn a_missing_key_names_the_primary_variable(
+        #[case] secrets: &'static [(&'static str, &'static str)],
+    ) {
+        assert!(matches!(
+            OPENROUTER_RESPONSES_CONFIG.validate_environment(vec![], None, &env(secrets)),
+            Err(Error::Auth(litellm_auth::Error::MissingApiKey {
+                provider: "OpenRouter",
+                environment_variable: "OPENROUTER_API_KEY",
+            }))
+        ));
+    }
+
+    #[rstest]
+    #[case::both_given(Some("key"), Some("base"), &[])]
+    #[case::key_given(Some("key"), None, &["OPENROUTER_API_BASE"])]
+    #[case::base_given(None, Some("base"), &["OPENROUTER_API_KEY", "OR_API_KEY"])]
+    #[case::neither(None, None, &["OPENROUTER_API_KEY", "OR_API_KEY", "OPENROUTER_API_BASE"])]
+    fn only_missing_settings_are_looked_up(
+        #[case] api_key: Option<&str>,
+        #[case] api_base: Option<&str>,
+        #[case] expected: &[&str],
+    ) {
+        assert_eq!(
+            OPENROUTER_RESPONSES_CONFIG.secret_names(api_key, api_base),
+            expected
+        );
+    }
+}

litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs

@@ -23,7 +23,7 @@ fn run_chat_completions(
 ) -> PyResult<Py<PyAny>> {
     let host = InferenceHost::new(
         call.view()?,
-        "litellm.rust_bridge.chat_completions.route_host",
+        "litellm.rust_bridge.host.routes.chat_completions",
     );
     run_inference::<ChatCompletionsRoute, _>(
         py,

litellm-rust/crates/python-bridge/src/routes/inference.rs

@@ -21,10 +21,7 @@ use serde_json::{Map, Value};
 use super::{NativeCall, RequestView};
 use crate::{
     errors::{RustUpstreamError, route_error_to_pyerr},
-    marshal::{
-        RouteOptions, optional_timeout, project_optional_fields, public_response,
-        python_timeout_seconds,
-    },
+    marshal::{RouteOptions, optional_timeout, public_response, python_timeout_seconds},
 };
 
 pub(super) struct InferenceHost {
@@ -115,10 +112,7 @@ impl InferenceHost {
         py: Python<'_>,
         arguments: &Bound<'_, PyDict>,
     ) -> PyResult<Map<String, Value>> {
-        let names: Vec<String> = py.import(self.module)?.getattr("PARAMETERS")?.extract()?;
-        project_optional_fields(names.iter().map(String::as_str), |name| {
-            self.argument(py, arguments, name)
-        })
+        self.view.body_params(py, arguments, self.module)
     }
 
     pub fn response(&self, py: Python<'_>, response: &impl Serialize) -> PyResult<Py<PyAny>> {

litellm-rust/crates/python-bridge/src/routes/messages/host.rs

@@ -26,34 +26,9 @@ use crate::{
     routes::RequestView,
 };
 
-const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.messages.route_host";
+const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages";
 const REQUEST_ERROR_MARKER: &str = "messages_request_error";
 
-const BODY_FIELDS: [&str; 22] = [
-    "max_tokens",
-    "metadata",
-    "stop_sequences",
-    "stream",
-    "system",
-    "temperature",
-    "thinking",
-    "tool_choice",
-    "tools",
-    "top_k",
-    "inference_geo",
-    "top_p",
-    "mcp_servers",
-    "context_management",
-    "compaction",
-    "container",
-    "output_format",
-    "speed",
-    "output_config",
-    "cache_control",
-    "reasoning_effort",
-    "safeguards",
-];
-
 fn merge_headers(
     forwarded: Option<Map<String, Value>>,
     extra_headers: Option<Map<String, Value>>,
@@ -143,7 +118,7 @@ impl MessagesPythonHost {
         let model = string("model")?.ok_or_else(|| PyValueError::new_err("model is required"))?;
         let messages =
             argument("messages")?.ok_or_else(|| PyValueError::new_err("messages is required"))?;
-        let fields = project_optional_fields(BODY_FIELDS, argument)?;
+        let fields = self.view.body_params(py, arguments, ROUTE_HOST_MODULE)?;
         let body = [
             ("model".to_string(), Value::String(model.clone())),
             ("messages".to_string(), from_py(&messages)?),

litellm-rust/crates/python-bridge/src/routes/mod.rs

@@ -13,6 +13,7 @@ use litellm_callbacks_legacy_python::{LegacyLogging, LoggingOperation, PublicCal
 use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol};
 use litellm_host_python::{
     CallOptions, HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args,
+    from_py,
 };
 use litellm_host_python::{Error as HostError, present};
 use pyo3::{
@@ -20,6 +21,7 @@ use pyo3::{
     prelude::*,
     types::{PyDict, PyMapping, PyTuple},
 };
+use serde_json::{Map, Value};
 
 /// The public call as Python bound it: `base` holds the positional arguments by name plus
 /// the signature defaults, `kwargs` the caller's keyword dict.
@@ -92,6 +94,22 @@ impl RequestView {
         present(arguments, self.request(py)?, name)
     }
 
+    /// The body params `module.body_params` keeps from the keywords the call would send, so
+    /// each route's Python host decides which params reach the provider.
+    pub(crate) fn body_params(
+        &self,
+        py: Python<'_>,
+        arguments: &Bound<'_, PyDict>,
+        module: &str,
+    ) -> PyResult<Map<String, Value>> {
+        let request = effective_py_args(self.request(py)?, arguments)?;
+        from_py(
+            &py.import(module)?
+                .getattr("body_params")?
+                .call1((request,))?,
+        )
+    }
+
     pub(crate) fn close(&mut self) {
         self.0 = None;
     }

litellm-rust/crates/python-bridge/src/routes/ocr/host.rs

@@ -77,7 +77,7 @@ impl OcrPythonHost {
             _ => "",
         };
         let mapped = py
-            .import("litellm.rust_bridge.ocr.route_host")
+            .import("litellm.rust_bridge.host.routes.ocr")
             .and_then(|module| module.getattr("map_failure"))
             .and_then(|map| map.call1((error.value(py), self.view.request(py)?, provider)))
             .and_then(|mapped| mapped.extract::<Py<PyBaseException>>().map_err(PyErr::from));
@@ -115,7 +115,7 @@ impl PythonBinding for OcrPythonHost {
         py: Python<'_>,
         response: LiteLLMOcrResponse,
     ) -> PyResult<Py<PyAny>> {
-        public_response(py, "litellm.rust_bridge.ocr.route_host", &response)
+        public_response(py, "litellm.rust_bridge.host.routes.ocr", &response)
     }
 
     fn encode_stream_head(

litellm-rust/crates/python-bridge/src/routes/responses/mod.rs

@@ -21,7 +21,7 @@ use super::{
 };
 use crate::errors::RustBridgeDeclined;
 
-const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.responses.route_host";
+const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.responses";
 
 fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
     let resolved = call.resolved()?;

litellm/__init__.py

@@ -118,6 +118,7 @@ success_callback: List[CALLBACK_TYPES] = []
 failure_callback: List[CALLBACK_TYPES] = []
 service_callback: List[CALLBACK_TYPES] = []
 audit_log_callbacks: List[CALLBACK_TYPES] = []
+batch_line_item_callbacks: Sequence[CALLBACK_TYPES] = ()
 # logging_callback_manager is lazy-loaded via __getattr__
 _custom_logger_compatible_callbacks_literal = Literal[
     "lago",
@@ -206,7 +207,7 @@ _async_failure_callback: List[Union[str, Callable[..., object], "CustomLogger"]]
     []
 )  # internal variable - async custom callbacks are routed here.
 pre_call_rules: List[Callable] = []
-post_call_rules: List[Callable] = []
+post_call_rules: List[Callable[..., object]] = []
 turn_off_message_logging: Optional[bool] = False
 standard_logging_payload_excluded_fields: Optional[List[str]] = (
     None  # Fields to exclude from StandardLoggingPayload before callbacks receive it
@@ -392,6 +393,7 @@ key_alias_pattern: str | None = None
 enable_gemini_default_thinking_level_low: bool = (
     False  # opt-in: force thinkingLevel low/minimal for Gemini 3 thinking param mapping
 )
+gemini_include_thoughts_default: bool = True
 ####################
 logging: bool = True
 enable_loadbalancing_on_batch_endpoints: Optional[bool] = None
@@ -1175,7 +1177,9 @@ def _build_models_by_provider() -> dict:
         "edenai": edenai_models,
         "datarobot": datarobot_models,
         "vertex_ai": vertex_chat_models
+        | vertex_code_chat_models
         | vertex_text_models
+        | vertex_code_text_models
         | vertex_anthropic_models
         | vertex_vision_models
         | vertex_language_models
@@ -1183,7 +1187,7 @@ def _build_models_by_provider() -> dict:
         | vertex_minimax_models
         | vertex_moonshot_models
         | vertex_zai_models,
-        "ai21": ai21_models,
+        "ai21": ai21_models | ai21_chat_models,
         "bedrock": bedrock_models | bedrock_converse_models,
         "petals": petals_models,
         "ollama": ollama_models,
@@ -1414,6 +1418,7 @@ from .exceptions import (
     InvalidRequestError,
     BadRequestError,
     ImageFetchError,
+    MalformedRequestFieldError as MalformedRequestFieldError,
     VectorStoreSearchError,
     NotFoundError,
     PermissionDeniedError,
@@ -2287,12 +2292,14 @@ if TYPE_CHECKING:
     get_response_string: Callable[..., str]
     supports_function_calling: Callable[..., bool]
     supports_web_search: Callable[..., bool]
+    supports_tool_search: Callable[..., bool]
     supports_url_context: Callable[..., bool]
     supports_response_schema: Callable[..., bool]
     supports_parallel_function_calling: Callable[..., bool]
     supports_vision: Callable[..., bool]
     supports_audio_input: Callable[..., bool]
     supports_video_input: Callable[..., bool]
+    supports_box_questions: Callable[..., bool]
     supports_audio_output: Callable[..., bool]
     supports_system_messages: Callable[..., bool]
     supports_reasoning: Callable[..., bool]

litellm/_internal_context.py

@@ -37,6 +37,16 @@ REDIS_FAMILIES_METADATA_KEY: Final = "families"
 _service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", default=None)
 
 
+@contextmanager
+def internal_sub_call() -> Generator[None]:
+    """A nested call the parent call logs and bills, so the parent must fold this call's cost into its own."""
+    token: Final = is_internal_call.set(True)
+    try:
+        yield
+    finally:
+        is_internal_call.reset(token)
+
+
 @contextmanager
 def post_response_phase() -> Generator[None]:
     """Work the caller no longer waits for (success callbacks, response-cache writes), including tasks it spawns."""

litellm/_lazy_imports_registry.py

@@ -31,12 +31,14 @@ UTILS_NAMES: Final = (
     "create_tokenizer",
     "supports_function_calling",
     "supports_web_search",
+    "supports_tool_search",
     "supports_url_context",
     "supports_response_schema",
     "supports_parallel_function_calling",
     "supports_vision",
     "supports_audio_input",
     "supports_video_input",
+    "supports_box_questions",
     "supports_audio_output",
     "supports_system_messages",
     "supports_reasoning",
@@ -477,6 +479,7 @@ _UTILS_IMPORT_MAP: Final = {
     "create_tokenizer": (".utils", "create_tokenizer"),
     "supports_function_calling": (".utils", "supports_function_calling"),
     "supports_web_search": (".utils", "supports_web_search"),
+    "supports_tool_search": (".utils", "supports_tool_search"),
     "supports_url_context": (".utils", "supports_url_context"),
     "supports_response_schema": (".utils", "supports_response_schema"),
     "supports_parallel_function_calling": (
@@ -486,6 +489,7 @@ _UTILS_IMPORT_MAP: Final = {
     "supports_vision": (".utils", "supports_vision"),
     "supports_audio_input": (".utils", "supports_audio_input"),
     "supports_video_input": (".utils", "supports_video_input"),
+    "supports_box_questions": (".utils", "supports_box_questions"),
     "supports_audio_output": (".utils", "supports_audio_output"),
     "supports_system_messages": (".utils", "supports_system_messages"),
     "supports_reasoning": (".utils", "supports_reasoning"),

litellm/_logging.py

@@ -6,12 +6,14 @@ import logging
 import os
 import re
 import sys
-from collections.abc import Iterator
+from collections.abc import Iterator, Mapping
 from datetime import datetime
 from logging import Formatter
 from typing import Final, TextIO
 from urllib.parse import unquote
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm.constants import (
     LITELLM_TRUNCATED_PAYLOAD_FIELD,
@@ -84,6 +86,7 @@ _redact_string = redact_string
 
 _REDACTED_RECORD_ATTR: Final = "litellm_redacted"
 _REDACTED_STAMP: Final = object()
+_DECODED_LOG_PAYLOAD: Final = TypeAdapter(object)
 _UNREDACTED_SCALAR_TYPES: Final = (bool, int, float, type(None))
 
 
@@ -91,6 +94,10 @@ def _is_redacted(record: logging.LogRecord) -> bool:
     return getattr(record, _REDACTED_RECORD_ATTR, None) is _REDACTED_STAMP
 
 
+def _record_attributes(record: logging.LogRecord) -> Mapping[str, object]:
+    return record.__dict__
+
+
 def _scrubbing_changed_nothing(scrubbed: object, original: object) -> bool:
     try:
         return bool(scrubbed == original)
@@ -438,7 +445,7 @@ def _process_record(record: logging.LogRecord, *, base64_limit: int, text_limit:
                     key, substituted_color if key == "color_message" and substituted_color is not None else value
                 ),
             )
-            for key, value in record.__dict__.items()
+            for key, value in _record_attributes(record).items()
             if key not in _STANDARD_RECORD_ATTRS
             and key != _REDACTED_RECORD_ATTR
             and not isinstance(value, _UNREDACTED_SCALAR_TYPES)
@@ -669,7 +676,7 @@ def _try_parse_json_message(message: str) -> dict[str, object] | None:
     msg_stripped: Final = message.strip()
     if not (msg_stripped.startswith("{") or msg_stripped.startswith("[")):
         return None
-    parsed: Final[object] = safe_json_loads(message, default=None)
+    parsed: Final = _DECODED_LOG_PAYLOAD.validate_python(safe_json_loads(message, default=None))
     if parsed is None or not isinstance(parsed, dict):
         return None
     return parsed
@@ -699,7 +706,7 @@ def _try_parse_embedded_python_dict(message: str) -> dict[str, object] | None:
                 if depth == 0:
                     substr = message[start : j + 1]
                     try:
-                        result: object = ast.literal_eval(substr)
+                        result: object = _DECODED_LOG_PAYLOAD.validate_python(ast.literal_eval(substr))
                         if isinstance(result, dict) and len(result) > 0:
                             return result
                     except (ValueError, SyntaxError, TypeError):
@@ -757,7 +764,7 @@ class JsonFormatter(Formatter):
                     json_record[key] = value
 
         # Include extra attributes passed via logger.debug("msg", extra={...})
-        for key, value in record.__dict__.items():
+        for key, value in _record_attributes(record).items():
             if key not in _NON_EXTRA_RECORD_ATTRS and key not in json_record:
                 json_record[key] = value
 

litellm/_redis.py

@@ -19,6 +19,7 @@ from urllib.parse import urlsplit, urlunsplit
 
 import redis
 import redis.asyncio as async_redis
+from pydantic import TypeAdapter
 from redis.credentials import CredentialProvider
 
 from litellm import get_secret, get_secret_str
@@ -40,6 +41,7 @@ from ._logging import verbose_logger
 _generate_gcp_iam_access_token = generate_gcp_iam_access_token
 
 AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
+_UNWRAPPED_INIT: Final = TypeAdapter(object)
 
 _AWS_IAM_KWARG_NAMES: Final = (
     "aws_iam_auth",
@@ -64,7 +66,7 @@ def _unwrapped_init_args(cls: type) -> frozenset[str]:
     ``@deprecated_args`` too, which the same class of bug as ``_init_arg_names``
     would otherwise silently empty this allowlist through (see its docstring).
     """
-    spec: Final = inspect.getfullargspec(inspect.unwrap(cls.__init__))
+    spec: Final = inspect.getfullargspec(_UNWRAPPED_INIT.validate_python(inspect.unwrap(cls.__init__)))
     return frozenset(spec.args + spec.kwonlyargs)
 
 
@@ -112,7 +114,7 @@ def _init_arg_names(cls: type) -> frozenset[str]:
         name
         for klass in inspect.getmro(cls)
         if klass is not object
-        for spec in (inspect.getfullargspec(inspect.unwrap(klass.__init__)),)
+        for spec in (inspect.getfullargspec(_UNWRAPPED_INIT.validate_python(inspect.unwrap(klass.__init__))),)
         for name in spec.args + spec.kwonlyargs
     )
 
@@ -935,6 +937,10 @@ def get_redis_connection_pool(
     return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs)
 
 
+def _mask_for_display(masker: SensitiveDataMasker, redis_kwargs: Mapping[str, object]) -> Mapping[str, object]:
+    return masker.mask_dict(redis_kwargs)
+
+
 def _redis_kwargs_for_logging(redis_kwargs: Mapping[str, object]) -> Mapping[str, object]:
     return {
         key: "<credential provider>"
@@ -966,7 +972,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
         masker = SensitiveDataMasker()
 
         # Mask sensitive data in redis_kwargs
-        masked_redis_kwargs = masker.mask_dict(redis_kwargs_for_logging)
+        masked_redis_kwargs = _mask_for_display(masker, redis_kwargs_for_logging)
 
         # Create main panel title
         title: Final = Text("Redis Configuration", style="bold blue")
@@ -1029,7 +1035,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
     except ImportError:
         # Fallback to simple logging if rich is not available
         masker = SensitiveDataMasker()
-        masked_redis_kwargs = masker.mask_dict(redis_kwargs_for_logging)
+        masked_redis_kwargs = _mask_for_display(masker, redis_kwargs_for_logging)
         verbose_logger.info("Redis configuration: %s", masked_redis_kwargs)
     except Exception as e:
         verbose_logger.error("Error pretty printing Redis configuration: %s", e)

litellm/a2a_protocol/litellm_completion_bridge/handler.py

@@ -11,6 +11,7 @@ A2A Streaming Events (in order):
 """
 
 from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
+from types import MappingProxyType
 from typing import Any, Final
 
 import litellm
@@ -21,6 +22,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
     A2AStreamingContext,
 )
 from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
+from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS
 from litellm.interactions.agents.utils import merge_agent_headers
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.types.utils import ModelResponse
@@ -44,6 +46,32 @@ _AGENT_ONLY_PARAMS: Final = frozenset(
 )
 
 
+def bridge_model_name(litellm_params: Mapping[str, object]) -> str:
+    """The litellm model string a bridge agent completes with: ``<custom_llm_provider>/<model>``.
+
+    A provider bridged on chat names a real upstream endpoint with ``model``, so a missing one stays empty for that
+    provider to reject or ignore instead of becoming a made-up ``agent`` endpoint.
+    """
+    custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
+    default_model: Final = "" if custom_llm_provider in A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS else "agent"
+    configured_model: Final = litellm_params.get("model")
+    model: Final = configured_model if isinstance(configured_model, str) else default_model
+    if not custom_llm_provider or model.startswith(f"{custom_llm_provider}/"):
+        return model
+    return f"{custom_llm_provider}/{model}"
+
+
+def agent_completion_kwargs(litellm_params: Mapping[str, object]) -> Mapping[str, object]:
+    """The agent's litellm_params that are valid litellm.acompletion kwargs (api_key, api_base, ...)."""
+    return MappingProxyType(
+        {
+            key: value
+            for key, value in litellm_params.items()
+            if key not in ("model", "custom_llm_provider") and key not in _AGENT_ONLY_PARAMS
+        }
+    )
+
+
 class A2ACompletionBridgeHandler:
     """
     Static methods for handling A2A requests via LiteLLM completion.
@@ -64,16 +92,7 @@ class A2ACompletionBridgeHandler:
         # Transform A2A message to OpenAI format
         openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
 
-        # Get completion params
-        custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
-        model: Final[str] = litellm_params.get("model", "agent")
-
-        # Build full model string if provider specified
-        # Skip prepending if model already starts with the provider prefix
-        if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
-            full_model = f"{custom_llm_provider}/{model}"
-        else:
-            full_model = model
+        full_model: Final = bridge_model_name(litellm_params)
 
         if stream:
             verbose_logger.info("A2A completion bridge streaming: model=%s, api_base=%s", full_model, api_base)
@@ -87,13 +106,7 @@ class A2ACompletionBridgeHandler:
             "api_base": api_base,
             "stream": stream,
         }
-        # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
-        litellm_params_to_add: Final = {
-            k: v
-            for k, v in litellm_params.items()
-            if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
-        }
-        completion_params.update(litellm_params_to_add)
+        completion_params.update(agent_completion_kwargs(litellm_params))
         # Apply forward metadata AFTER the litellm_params merge so the helper
         # sees any agent-owner-configured ``extra_body.metadata`` and can keep
         # those keys authoritative over the client-supplied A2A metadata.

litellm/a2a_protocol/main.py

@@ -14,7 +14,9 @@ import datetime
 import uuid
 from collections.abc import AsyncIterator, Coroutine, Mapping
 from types import MappingProxyType, ModuleType
-from typing import TYPE_CHECKING, Any, Final, Optional, cast
+from typing import TYPE_CHECKING, Final, Optional, cast
+
+from typing_extensions import TypedDict, Unpack
 
 import litellm
 from litellm._logging import verbose_logger, verbose_proxy_logger
@@ -89,6 +91,10 @@ from litellm.a2a_protocol.exceptions import A2ALocalhostURLError
 A2ACardResolver: Final = LiteLLMA2ACardResolver
 
 
+class _SendMessageOptions(TypedDict, total=False, extra_items=object):
+    pass
+
+
 def _set_usage_on_logging_obj(
     kwargs: Mapping[str, object],
     prompt_tokens: int,
@@ -563,7 +569,7 @@ async def asend_message(
 def send_message(
     a2a_client: "A2AClientType",
     request: "SendMessageRequest",
-    **kwargs: Any,
+    **kwargs: Unpack[_SendMessageOptions],
 ) -> LiteLLMSendMessageResponse | Coroutine[object, object, LiteLLMSendMessageResponse]:
     """
     Sync: Send a message to an A2A agent.

litellm/a2a_protocol/streaming_iterator.py

@@ -7,16 +7,21 @@ from collections.abc import AsyncIterator
 from datetime import datetime
 from typing import TYPE_CHECKING, Final
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm._logging import verbose_logger
 from litellm.a2a_protocol.cost_calculator import A2ACostCalculator
 from litellm.a2a_protocol.utils import A2ARequestUtils
 from litellm.litellm_core_utils.asyncify import asyncify
+from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 
 if TYPE_CHECKING:
     from a2a.compat.v0_3.types import SendStreamingMessageRequest, SendStreamingMessageResponse
 
+_OPAQUE_PAYLOAD: Final = TypeAdapter(object)
+
 
 class A2AStreamingIterator:
     """
@@ -160,8 +165,11 @@ class A2AStreamingIterator:
         # Add final chunk result if available
         if self.final_chunk:
             try:
-                chunk_dict: Final = self.final_chunk.model_dump(mode="json", exclude_none=True)
-                result["result"] = chunk_dict.get("result", {})
+                chunk_dict: Final = as_str_mapping(
+                    _OPAQUE_PAYLOAD.validate_python(self.final_chunk.model_dump(mode="json", exclude_none=True))
+                )
+                if chunk_dict is not None:
+                    result["result"] = chunk_dict.get("result", {})
             except Exception:
                 pass
 

litellm/anthropic_interface/messages/__init__.py

@@ -11,7 +11,6 @@ This is an __init__.py file to allow the following interface
 """
 
 from collections.abc import AsyncIterator, Coroutine, Iterator
-from typing import Any
 
 from litellm.messages import (
     anthropic_messages as _async_anthropic_messages,
@@ -102,8 +101,8 @@ def create(
 ) -> (
     AnthropicMessagesResponse
     | Iterator[bytes]
-    | AsyncIterator[Any]
-    | Coroutine[Any, Any, AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]]
+    | AsyncIterator[object]
+    | Coroutine[object, object, AnthropicMessagesResponse | AsyncIterator[object] | Iterator[bytes]]
 ):
     """
     Async wrapper for Anthropic's messages API

litellm/assistants/main.py

@@ -129,6 +129,7 @@ def get_assistants(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             aget_assistants=aget_assistants,
         )
@@ -312,6 +313,7 @@ def create_assistants(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             create_assistant_data=create_assistant_data,
             client=client,
             async_create_assistants=async_create_assistants,
@@ -461,6 +463,7 @@ def delete_assistant(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             assistant_id=assistant_id,
             client=client,
             async_delete_assistants=async_delete_assistants,
@@ -637,6 +640,7 @@ def create_thread(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             acreate_thread=acreate_thread,
         )
@@ -783,6 +787,7 @@ def get_thread(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             aget_thread=aget_thread,
         )
@@ -960,6 +965,7 @@ def add_message(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             a_add_message=a_add_message,
         )
@@ -1111,6 +1117,7 @@ def get_messages(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             aget_messages=aget_messages,
         )
@@ -1308,6 +1315,7 @@ def run_thread(
             timeout=timeout,
             max_retries=optional_params.max_retries,
             organization=organization,
+            litellm_params=litellm_params_dict,
             client=client,
             arun_thread=arun_thread,
             event_handler=event_handler,

litellm/batches/batch_line_items.py

@@ -0,0 +1,250 @@
+"""Per-line records of a finished batch for the loggers listed in batch_line_item_callbacks."""
+
+import asyncio
+import hashlib
+import re
+from collections import Counter
+from collections.abc import Iterable, Iterator, Mapping
+from dataclasses import dataclass
+from itertools import accumulate, chain, count
+from typing import Final, Literal, TypeAlias
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.batches.batch_utils import batch_cost_and_usage, iter_batch_input_lines, parse_batch_output_line
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.litellm_logging import init_custom_logger_compatible_class, is_valid_sha256_hash
+from litellm.llms.vertex_ai.files.transformation import vertex_batch_row_custom_id
+from litellm.types.utils import (
+    ModelInfo,
+    StandardBatchLineItem,
+    StandardBatchLineItemMetadata,
+    StandardBatchLineItemsPayload,
+    StandardBatchLineItemUsage,
+)
+
+BatchFileProvider: TypeAlias = Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"]
+
+_VERTEX_AI_UNKNOWN_CUSTOM_ID: Final = "unknown"  # what a Vertex AI batch row that carries no custom_id is read as
+_LINE: Final = re.compile(rb"[^\n]+")
+
+
+@dataclass(frozen=True, slots=True)
+class BatchLineItemSource:
+    batch_id: str
+    finished_at: int
+    model: str
+    custom_llm_provider: str | None
+    cost_provider: BatchFileProvider
+    model_info: ModelInfo | None
+    metadata: StandardBatchLineItemMetadata
+
+
+def line_item_metadata(attribution: Mapping[str, object]) -> StandardBatchLineItemMetadata:
+    return StandardBatchLineItemMetadata(
+        user_api_key_hash=_creator_key_hash(attribution),
+        user_api_key_alias=_text(attribution, "user_api_key_alias"),
+        user_api_key_team_id=_text(attribution, "user_api_key_team_id"),
+        user_api_key_team_alias=_text(attribution, "user_api_key_team_alias"),
+        user_api_key_user_id=_text(attribution, "user_api_key_user_id"),
+        user_api_key_org_id=_text(attribution, "user_api_key_org_id"),
+    )
+
+
+def _text(attribution: Mapping[str, object], name: str) -> str | None:
+    value: Final = attribution.get(name)
+    return value if isinstance(value, str) else None
+
+
+def _creator_key_hash(attribution: Mapping[str, object]) -> str | None:
+    """The creating credential as a SHA-256 hash. A virtual key is saved hashed already, while a credential saved as
+    is, such as an OAuth2 access token, is hashed here so the object never holds it."""
+    credential: Final = _text(attribution, "user_api_key_hash")
+    if credential is None or is_valid_sha256_hash(credential):
+        return credential
+    return hashlib.sha256(credential.encode()).hexdigest()
+
+
+def _line_key(entry: Mapping[str, object]) -> str | None:
+    """A line's custom_id. Bedrock calls it recordId, and native Vertex rows carry it in the row key or request labels."""
+    key: Final = entry.get("custom_id", entry.get("recordId"))
+    return key if isinstance(key, str) else vertex_batch_row_custom_id(entry)
+
+
+def batch_line_item_parts(
+    source: BatchLineItemSource, input_content: bytes, result_contents: Iterable[bytes], max_part_bytes: int
+) -> Iterator[StandardBatchLineItemsPayload]:
+    """Pair every output and error line with the input line that shares its custom_id, priced one line at a time, in
+    parts that each close once their raw request and result lines reach max_part_bytes. Input lines stay unparsed until
+    their result line is reached, so only one part's records are held at a time."""
+    sized: Final = _sized_line_items(source, input_content, result_contents)
+    for part in count():
+        lines = tuple(_until_full(sized, max_part_bytes))
+        if not lines:
+            return
+        yield StandardBatchLineItemsPayload(
+            batch_id=source.batch_id, finished_at=source.finished_at, part=part, lines=lines
+        )
+
+
+def _sized_line_items(
+    source: BatchLineItemSource, input_content: bytes, result_contents: Iterable[bytes]
+) -> Iterator[tuple[StandardBatchLineItem, int]]:
+    """Each output and error line's record, with the raw bytes of that line and of the input line it pairs with."""
+    request_spans: Final = _request_spans_by_custom_id(input_content)
+    for raw_result in chain.from_iterable(map(iter_batch_input_lines, result_contents)):
+        entry = parse_batch_output_line(raw_result)
+        if entry is None:
+            continue
+        custom_id = _line_key(entry)
+        span = _request_span_for(source, request_spans, custom_id)
+        request = None if span is None else parse_batch_output_line(input_content[span[0] : span[1]])
+        yield (
+            _line_item(source, custom_id, request, entry),
+            len(raw_result) + (0 if span is None else span[1] - span[0]),
+        )
+
+
+def _until_full(
+    sized: Iterator[tuple[StandardBatchLineItem, int]], max_part_bytes: int
+) -> Iterator[StandardBatchLineItem]:
+    """Take records until their raw bytes reach max_part_bytes, leaving the rest in sized for the next part."""
+    for line, part_bytes in accumulate(sized, _with_bytes_so_far):
+        yield line
+        if part_bytes >= max_part_bytes:
+            return
+
+
+def _with_bytes_so_far(
+    so_far: tuple[StandardBatchLineItem, int], line: tuple[StandardBatchLineItem, int]
+) -> tuple[StandardBatchLineItem, int]:
+    return line[0], so_far[1] + line[1]
+
+
+def _request_spans_by_custom_id(input_content: bytes) -> Mapping[str | None, tuple[int, int]]:
+    """Where each input line starts and ends, by its custom_id. A custom_id that several input lines share, as when a
+    Vertex AI embeddings entry is split into one row per input, is paired with none of them."""
+    keyed: Final = tuple(_keyed_spans(input_content))
+    counts: Final = Counter(key for key, _ in keyed)
+    return {key: span for key, span in keyed if counts[key] == 1}
+
+
+def _keyed_spans(input_content: bytes) -> Iterator[tuple[str | None, tuple[int, int]]]:
+    for line in _LINE.finditer(input_content):
+        entry = parse_batch_output_line(line[0]) if line[0].strip() else None
+        if entry is not None:
+            yield _line_key(entry), line.span()
+
+
+def _line_item(
+    source: BatchLineItemSource, custom_id: str | None, request: Mapping[str, object] | None, entry: dict[str, object]
+) -> StandardBatchLineItem:
+    priced: Final = batch_cost_and_usage(
+        (entry,), source.cost_provider, model_name=source.model, model_info=source.model_info
+    )
+    succeeded: Final = priced.successful_requests > 0
+    return StandardBatchLineItem(
+        batch_id=source.batch_id,
+        custom_id=custom_id,
+        status="success" if succeeded else "failure",
+        request=request,
+        response=entry,
+        usage=StandardBatchLineItemUsage(
+            prompt_tokens=priced.usage.prompt_tokens,
+            completion_tokens=priced.usage.completion_tokens,
+            total_tokens=priced.usage.total_tokens,
+        )
+        if succeeded
+        else None,
+        response_cost=priced.cost,
+        model=source.model,
+        custom_llm_provider=source.custom_llm_provider,
+        metadata=source.metadata,
+    )
+
+
+def _request_span_for(
+    source: BatchLineItemSource, request_spans: Mapping[str | None, tuple[int, int]], custom_id: str | None
+) -> tuple[int, int] | None:
+    """Where the input line that shares a result line's custom_id is. A Vertex AI line read as "unknown" pairs with
+    none."""
+    if custom_id is None or (source.cost_provider == "vertex_ai" and custom_id == _VERTEX_AI_UNKNOWN_CUSTOM_ID):
+        return None
+    return request_spans.get(custom_id)
+
+
+def without_request_and_response(payload: StandardBatchLineItemsPayload) -> StandardBatchLineItemsPayload:
+    return StandardBatchLineItemsPayload(
+        batch_id=payload["batch_id"],
+        finished_at=payload["finished_at"],
+        part=payload["part"],
+        lines=tuple(
+            StandardBatchLineItem(
+                batch_id=line["batch_id"],
+                custom_id=line["custom_id"],
+                status=line["status"],
+                request=None,
+                response=None,
+                usage=line["usage"],
+                response_cost=line["response_cost"],
+                model=line["model"],
+                custom_llm_provider=line["custom_llm_provider"],
+                metadata=line["metadata"],
+            )
+            for line in payload["lines"]
+        ),
+    )
+
+
+async def dispatch_batch_line_items(payload: StandardBatchLineItemsPayload, redact_bodies: bool) -> None:
+    """Hand one part's lines to every batch_line_item_callbacks logger. A failing logger is logged and never stops the
+    rest."""
+    await asyncio.gather(
+        *(_deliver(callback, payload, redact_bodies) for callback in litellm.batch_line_item_callbacks)
+    )
+
+
+async def _deliver(
+    callback: litellm.CALLBACK_TYPES, payload: StandardBatchLineItemsPayload, redact_bodies: bool
+) -> None:
+    try:
+        logger: Final = _resolve_logger(callback)
+        if logger is None:
+            verbose_logger.error(
+                "batch_line_item_callbacks: %r cannot store batch line items, part %d of batch %s was not stored",
+                callback,
+                payload["part"],
+                payload["batch_id"],
+            )
+            return
+        await logger.async_log_batch_line_items(
+            without_request_and_response(payload)
+            if redact_bodies or _logger_turned_off_message_logging(logger)
+            else payload
+        )
+    except Exception as error:  # noqa: BLE001  # one failing sink must not stop the other sinks or the cost poller
+        verbose_logger.exception(
+            "batch_line_item_callbacks: %r failed to store the %d lines of part %d of batch %s: %s",
+            callback,
+            len(payload["lines"]),
+            payload["part"],
+            payload["batch_id"],
+            error,
+        )
+
+
+def _resolve_logger(callback: litellm.CALLBACK_TYPES) -> CustomLogger | None:
+    match callback:
+        case CustomLogger():
+            return callback
+        case "s3_v2" | "gcs_bucket":
+            return init_custom_logger_compatible_class(callback, internal_usage_cache=None, llm_router=None)
+        case _:
+            return None
+
+
+def _logger_turned_off_message_logging(logger: CustomLogger) -> bool:
+    return (
+        getattr(logger, "turn_off_message_logging", False) is True
+        or getattr(logger, "message_logging", True) is not True
+    )

litellm/batches/batch_utils.py

@@ -1,10 +1,12 @@
 import json
-from collections.abc import Iterable, Iterator, Mapping
+from collections.abc import Iterable, Iterator, Mapping, Sequence
 from dataclasses import dataclass
 from dataclasses import replace as dataclasses_replace
 from enum import Enum
 from typing import Any, Final, Literal, cast
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
@@ -17,7 +19,7 @@ from litellm.llms.vertex_ai.batches.transformation import (
 )
 from litellm.types.llms.openai import Batch
 from litellm.types.utils import ModelInfo, Usage
-from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS
+from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS
 from litellm.utils import token_counter
 
 
@@ -35,6 +37,7 @@ class BatchCostUsageResult:
 
 
 _COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"})
+_DECODED_BATCH_LINE: Final = TypeAdapter(object)
 
 
 def _uses_native_vertex_output(
@@ -70,7 +73,17 @@ def batch_cost_is_final(batch: Batch) -> bool:
 
 
 async def calculate_batch_cost_and_usage(
-    file_content_dictionary: list[dict[str, object]],
+    file_content_dictionary: Sequence[dict[str, object]],
+    custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
+    model_name: str | None = None,
+    model_info: ModelInfo | None = None,
+) -> BatchCostUsageResult:
+    """Calculate the cost and usage of a batch, see batch_cost_and_usage."""
+    return batch_cost_and_usage(file_content_dictionary, custom_llm_provider, model_name, model_info)
+
+
+def batch_cost_and_usage(
+    file_content_dictionary: Sequence[dict[str, object]],  # mutable-ok: the rows calculate_batch_cost_and_usage takes
     custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
     model_name: str | None = None,
     model_info: ModelInfo | None = None,
@@ -436,7 +449,7 @@ def _provider_output_file_id(output_file_id: str) -> str:
     return extracted
 
 
-async def _fetch_batch_managed_file_content(
+async def fetch_batch_managed_file_content(
     file_id: str,
     custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai",
     litellm_params: Mapping[str, object] | None = None,
@@ -467,6 +480,9 @@ async def _fetch_batch_managed_file_content(
     return _file_content.content
 
 
+_fetch_batch_managed_file_content: Final = fetch_batch_managed_file_content
+
+
 async def _fetch_batch_output_file_content(
     batch: Batch,
     custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai",
@@ -553,6 +569,7 @@ def extract_file_access_credentials(
             # A federated deployment holds no api_key, so without these the fetch that reads a
             # finished batch's output has nothing to authenticate with and its cost is never billed.
             *sorted(ANTHROPIC_WIF_KWARGS_KEYS),
+            *sorted(OPENAI_WIF_KWARGS_KEYS),
         )
         for key in credential_keys:
             if key in litellm_params:
@@ -604,14 +621,14 @@ def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict[str, object
     whole batch's cost accounting.
     """
     for line in iter_batch_input_lines(file_content):
-        entry = _parse_batch_output_line(line)
+        entry = parse_batch_output_line(line)
         if entry is not None:
             yield entry
 
 
-def _parse_batch_output_line(line: bytes) -> dict[str, object] | None:
+def parse_batch_output_line(line: bytes) -> dict[str, object] | None:
     try:
-        parsed: Final[object] = json.loads(line)
+        parsed: Final = _DECODED_BATCH_LINE.validate_python(json.loads(line))
     except ValueError as e:
         verbose_logger.warning("skipping malformed batch output line: %s", str(e))
         return None

litellm/batches/main.py

@@ -127,7 +127,7 @@ async def acreate_batch(
     metadata: dict[str, str] | None = None,
     extra_headers: dict[str, str] | None = None,
     extra_body: dict[str, str] | None = None,
-    output_expires_after: dict[str, Any] | None = None,
+    output_expires_after: FileExpiresAfter | None = None,
     **kwargs,
 ) -> LiteLLMBatch:
     """
@@ -191,9 +191,9 @@ def create_batch(
     metadata: dict[str, str] | None = None,
     extra_headers: dict[str, str] | None = None,
     extra_body: dict[str, str] | None = None,
-    output_expires_after: dict[str, Any] | None = None,
+    output_expires_after: FileExpiresAfter | None = None,
     **kwargs,
-) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
+) -> LiteLLMBatch | Coroutine[object, None, LiteLLMBatch]:
     """
     Creates and executes a batch from an uploaded file of request
 
@@ -246,7 +246,7 @@ def create_batch(
             extra_body=extra_body,
         )
         if output_expires_after is not None:
-            _create_batch_request["output_expires_after"] = cast(FileExpiresAfter, output_expires_after)
+            _create_batch_request["output_expires_after"] = output_expires_after
         if model is not None:
             provider_config = ProviderConfigManager.get_provider_batches_config(
                 model=model,
@@ -309,6 +309,7 @@ def create_batch(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 _is_async=_is_async,
+                litellm_params=litellm_params,
             )
         elif custom_llm_provider == "azure":
             api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
@@ -471,6 +472,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
             organization=organization,
             timeout=timeout,
             max_retries=optional_params.max_retries,
+            litellm_params=litellm_params,
         )
     elif custom_llm_provider == "azure":
         api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
@@ -568,7 +570,7 @@ def retrieve_batch(
     extra_headers: dict[str, str] | None = None,
     extra_body: dict[str, str] | None = None,
     **kwargs,
-) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
+) -> LiteLLMBatch | Coroutine[object, None, LiteLLMBatch]:
     """
     Retrieves a batch.
 
@@ -819,6 +821,7 @@ def list_batches(
                 organization=organization,
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
+                litellm_params=litellm_params,
             )
         elif custom_llm_provider == "azure":
             api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
@@ -941,7 +944,7 @@ def cancel_batch(
     extra_headers: dict[str, str] | None = None,
     extra_body: dict[str, str] | None = None,
     **kwargs,
-) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
+) -> LiteLLMBatch | Coroutine[object, None, LiteLLMBatch]:
     """
     Cancels a batch.
 
@@ -1016,6 +1019,7 @@ def cancel_batch(
                 organization=organization,
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
+                litellm_params=litellm_params,
             )
         elif custom_llm_provider == "azure":
             api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")

litellm/budget_manager.py

@@ -13,6 +13,8 @@ import threading
 import time
 from typing import Final, Literal
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm.constants import (
     DAYS_IN_A_MONTH,
@@ -22,6 +24,8 @@ from litellm.constants import (
 )
 from litellm.utils import ModelResponse
 
+_BUDGET_SERVER_REPLY: Final = TypeAdapter(object)
+
 
 class BudgetManager:
     def __init__(
@@ -211,6 +215,5 @@ class BudgetManager:
         elif self.client_type == "hosted":
             url: Final = self.api_base + "/set_budget"
             data: Final = {"project_name": self.project_name, "user_dict": self.user_dict}
-            response = litellm.module_level_client.post(url, headers=self.headers, json=data)
-            response = response.json()
-            return response
+            response: Final = litellm.module_level_client.post(url, headers=self.headers, json=data)
+            return _BUDGET_SERVER_REPLY.validate_python(response.json())

litellm/caching/azure_blob_cache.py

@@ -65,7 +65,7 @@ class AzureBlobCache(BaseCache):
             print_verbose(f"Get Azure Blob Cache: key: {key}")
             as_bytes: Final = self.container_client.download_blob(key).readall()
             as_str: Final = as_bytes.decode("utf-8")
-            cached_response: Final = json.loads(as_str)
+            cached_response: Final[object] = json.loads(as_str)
 
             verbose_logger.debug(
                 "Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s",
@@ -86,7 +86,7 @@ class AzureBlobCache(BaseCache):
             blob: Final = await self.async_container_client.download_blob(key)
             as_bytes: Final = await blob.readall()
             as_str: Final = as_bytes.decode("utf-8")
-            cached_response: Final = json.loads(as_str)
+            cached_response: Final[object] = json.loads(as_str)
             verbose_logger.debug(
                 "Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s",
                 key,

litellm/caching/caching.py

@@ -18,14 +18,14 @@ from contextlib import contextmanager
 from enum import Enum
 from typing import Any, Final, Literal, cast
 
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 
 import litellm
 from litellm._internal_context import current_service_target, service_target
 from litellm._logging import verbose_logger
 from litellm.constants import CACHED_STREAMING_CHUNK_DELAY
 from litellm.integrations.otel.runtime import phase_span
-from litellm.litellm_core_utils.model_param_helper import ModelParamHelper
+from litellm.litellm_core_utils.model_param_helper import ModelParamHelper, peek_reasoning_summary_aliases
 from litellm.types.caching import *
 from litellm.types.utils import EmbeddingResponse, is_litellm_owned_kwarg
 
@@ -42,6 +42,8 @@ from .redis_semantic_cache import RedisSemanticCache
 from .s3_cache import S3Cache
 
 RESPONSE_CACHE_TARGET: Final = "llm_response"
+_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"})
+_DECODED_LITERAL: Final = TypeAdapter(object)
 
 
 @contextmanager
@@ -374,7 +376,7 @@ class Cache:
         )
         return "".join(f"{field}: {value}" for field, value in scope_values if value is not None)
 
-    def get_cache_key(self, **kwargs) -> str:
+    def get_cache_key(self, **kwargs: object) -> str:
         """
         Get the cache key for the given arguments.
 
@@ -398,21 +400,26 @@ class Cache:
         for param in kwargs:
             if param in scope_excluded_params:
                 continue
-            if param in combined_kwargs:
+            if param in combined_kwargs or param in _ALWAYS_KEYED_PROVIDER_PARAMS:
                 param_value: str | None = self._get_param_value(param, kwargs)
+                if param in _ALWAYS_KEYED_PROVIDER_PARAMS and not param_value:
+                    continue
                 if param_value is not None:
                     cache_key += f"{param}: {param_value}"
             elif not is_litellm_owned_kwarg(param):
                 if litellm.enable_caching_on_provider_specific_optional_params is True:  # feature flagged for now
                     if kwargs[param] is None:
                         continue  # ignore None params
-                    param_value = kwargs[param]
-                    cache_key += f"{param}: {param_value}"
+                    cache_key += f"{param}: {kwargs[param]}"
 
         if is_semantic_cache:
             cache_key += self._get_semantic_cache_tenant_scope(kwargs)
 
-        hashed_cache_key = Cache._get_hashed_cache_key(cache_key)
+        reasoning_summary: Final = peek_reasoning_summary_aliases(kwargs)
+        reasoning_summary_material: Final = (
+            "" if reasoning_summary is None else f"reasoning_summary: {reasoning_summary}"
+        )
+        hashed_cache_key = Cache._get_hashed_cache_key(cache_key + reasoning_summary_material)
         hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs)
         verbose_logger.debug(
             "\nCreated cache key: %s (source material length: %d)",
@@ -599,9 +606,11 @@ class Cache:
                 if isinstance(cached_response, dict):
                     pass
                 else:
-                    cached_response = json.loads(cached_response)  # Convert string to dictionary
+                    cached_response = _DECODED_LITERAL.validate_python(
+                        json.loads(cached_response)
+                    )  # Convert string to dictionary
             except Exception:
-                cached_response = ast.literal_eval(cached_response)
+                cached_response = _DECODED_LITERAL.validate_python(ast.literal_eval(cached_response))
             return cached_response
         return cached_result
 

litellm/caching/caching_handler.py

@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
 from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError
 
 import litellm
-from litellm._internal_context import post_response_phase
+from litellm._internal_context import is_internal_call, post_response_phase
 from litellm._logging import print_verbose, verbose_logger
 from litellm.caching import InMemoryCache
 from litellm.caching.caching import S3Cache, response_cache_phase
@@ -511,7 +511,9 @@ class LLMCachingHandler:
                         custom_llm_provider=custom_llm_provider,
                     )
 
-                    if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result):
+                    if not is_internal_call.get() and not _should_defer_streaming_cache_hit_callbacks(
+                        cached_result=cached_result
+                    ):
                         logging_obj.handle_sync_success_callbacks_for_async_calls(
                             result=cached_result,
                             start_time=start_time,
@@ -820,7 +822,12 @@ class LLMCachingHandler:
             start_time (datetime): The start time of the operation.
             end_time (datetime): The end time of the operation.
             cache_hit (bool): Whether it was a cache hit.
+
+        An internal sub-call's hit fires nothing, since the parent call that folds it in logs it.
         """
+        if is_internal_call.get():
+            return
+
         from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 
         GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
@@ -1088,6 +1095,7 @@ class LLMCachingHandler:
             kwargs=kwargs,
             start_time=self.start_time,
             end_time=datetime.datetime.now(),
+            cache_hit=True,
         )
         return cached_result
 

litellm/caching/gcs_cache.py

@@ -7,6 +7,8 @@ import json
 from typing import Final
 from urllib.parse import quote
 
+from pydantic import TypeAdapter
+
 from litellm._logging import print_verbose, verbose_logger
 from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
 from litellm.llms.custom_httpx.http_handler import (
@@ -17,6 +19,8 @@ from litellm.llms.custom_httpx.http_handler import (
 
 from .base_cache import BaseCache
 
+_DECODED_JSON: Final = TypeAdapter(object)
+
 
 class GCSCache(BaseCache):
     def __init__(
@@ -70,7 +74,7 @@ class GCSCache(BaseCache):
             url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{quote(object_name, safe='')}?alt=media"
             response: Final = self.sync_client.get(url=url, headers=headers)
             if response.status_code == 200:
-                cached_response: Final = json.loads(response.text)
+                cached_response: Final = _DECODED_JSON.validate_python(json.loads(response.text))
                 verbose_logger.debug(
                     "Got GCS Cache: key: %s, cached_response %s. Type Response %s",
                     key,
@@ -90,7 +94,7 @@ class GCSCache(BaseCache):
             url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{quote(object_name, safe='')}?alt=media"
             response: Final = await self.async_client.get(url=url, headers=headers)
             if response.status_code == 200:
-                return json.loads(response.text)
+                return _DECODED_JSON.validate_python(json.loads(response.text))
             return None
         except Exception as e:
             verbose_logger.error("GCS Caching: async_get_cache() - Got exception from GCS: %s", e)

litellm/caching/in_memory_cache.py

@@ -51,6 +51,7 @@ class InMemoryCache(BaseCache):
         self.ttl_dict: dict = {}
         self.expiration_heap: list[tuple[float, str]] = []
         self._increment_lock = threading.Lock()
+        self._write_lock = threading.Lock()
         self._clock = clock if clock is not None else lambda: time.time()
 
     def check_value_size(self, value: Any):
@@ -96,6 +97,9 @@ class InMemoryCache(BaseCache):
         """
         return key in self.ttl_dict and self._clock() > self.ttl_dict[key]
 
+    def _cache_keys(self) -> tuple[str, ...]:
+        return tuple(self.cache_dict)  # pyright: ignore[reportUnknownVariableType]  # cache_dict is untyped legacy state
+
     def _remove_key(self, key: str) -> None:
         """
         Remove a key from both cache_dict and ttl_dict
@@ -133,13 +137,23 @@ class InMemoryCache(BaseCache):
                 # Case 3: Entry is valid and not expired
                 break
 
+        if len(self.expiration_heap) > 2 * self.max_size_in_memory:
+            self.expiration_heap = [entry for entry in self.expiration_heap if self.ttl_dict.get(entry[1]) == entry[0]]
+            heapq.heapify(self.expiration_heap)
+
         # Step 2: Evict if cache is still full
-        while len(self.cache_dict) >= self.max_size_in_memory:
+        while self.expiration_heap and len(self.cache_dict) >= self.max_size_in_memory:
             expiration_time, key = heapq.heappop(self.expiration_heap)
             # Skip if key was removed or updated
             if self.ttl_dict.get(key) == expiration_time:
                 self._remove_key(key)
 
+        if not self.expiration_heap:
+            untracked_keys: Final = self._cache_keys()
+            overflow: Final = len(untracked_keys) - self.max_size_in_memory + 1
+            for untracked_key in untracked_keys[: max(overflow, 0)]:
+                self._remove_key(untracked_key)
+
         # de-reference the removed item
         # https://www.geeksforgeeks.org/diagnosing-and-fixing-memory-leaks-in-python/
         # One of the most common causes of memory leaks in Python is the retention of objects that are no longer being used.
@@ -150,7 +164,7 @@ class InMemoryCache(BaseCache):
         Check if ttl is set for a key
         """
         ttl_time: Final = self.ttl_dict.get(key)
-        if ttl_time is None or float(ttl_time) < self._clock():  # if ttl is not set, allow override
+        if ttl_time is None or ttl_time < self._clock():  # if ttl is not set, allow override
             return True
         else:
             return False
@@ -160,21 +174,22 @@ class InMemoryCache(BaseCache):
         if self.max_size_in_memory == 0:
             return  # Don't cache anything if max size is 0
 
-        # Always prune expired/outdated heap roots before inserting.
-        # This keeps expiration_heap bounded even when the live cache stays
-        # below max_size_in_memory and keys are reinserted after TTL expiry.
-        self.evict_cache()
-        if not self.check_value_size(value):
-            return
+        with self._write_lock:
+            # Always prune expired/outdated heap roots before inserting.
+            # This keeps expiration_heap bounded even when the live cache stays
+            # below max_size_in_memory and keys are reinserted after TTL expiry.
+            self.evict_cache()
+            if not self.check_value_size(value):
+                return
 
-        self.cache_dict[key] = value
-        if self.allow_ttl_override(key):  # if ttl is not set, set it to default ttl
-            if "ttl" in kwargs and kwargs["ttl"] is not None:
-                self.ttl_dict[key] = self._clock() + float(kwargs["ttl"])
-                heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
-            else:
-                self.ttl_dict[key] = self._clock() + self.default_ttl
-                heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
+            self.cache_dict[key] = value
+            if self.allow_ttl_override(key):  # if ttl is not set, set it to default ttl
+                if "ttl" in kwargs and kwargs["ttl"] is not None:
+                    self.ttl_dict[key] = self._clock() + float(kwargs["ttl"])
+                    heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
+                else:
+                    self.ttl_dict[key] = self._clock() + self.default_ttl
+                    heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
 
     async def async_set_cache(self, key, value, **kwargs):
         self.set_cache(key=key, value=value, **kwargs)
@@ -258,14 +273,15 @@ class InMemoryCache(BaseCache):
         return results
 
     def flush_cache(self):
-        self.cache_dict.clear()
-        self.ttl_dict.clear()
-        self.expiration_heap.clear()
+        with self._write_lock:
+            self.cache_dict.clear()
+            self.ttl_dict.clear()
+            self.expiration_heap.clear()
 
     async def disconnect(self):
         pass
 
-    def delete_cache(self, key):
+    def delete_cache(self, key: str) -> None:
         self._remove_key(key)
 
     async def async_get_ttl(self, key: str) -> int | None:

litellm/caching/qdrant_semantic_cache.py

@@ -15,6 +15,8 @@ import os
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Protocol, cast
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm._logging import print_verbose
 from litellm.constants import (
@@ -38,6 +40,7 @@ from ._embedding_router import (
 from .base_cache import BaseCache
 
 _WAIT_FOR_INDEXING: Final = MappingProxyType({"wait": "true"})
+_EMBEDDING_VECTOR: Final = TypeAdapter(object)
 
 if TYPE_CHECKING:
     from litellm.router import Router
@@ -295,7 +298,7 @@ class QdrantSemanticCache(BaseCache):
         )
 
         # get the embedding
-        embedding: Final = embedding_response["data"][0]["embedding"]
+        embedding: Final = _EMBEDDING_VECTOR.validate_python(embedding_response["data"][0]["embedding"])
 
         value = str(value)
         assert isinstance(value, str)
@@ -334,7 +337,7 @@ class QdrantSemanticCache(BaseCache):
         )
 
         # get the embedding
-        embedding: Final = embedding_response["data"][0]["embedding"]
+        embedding: Final = _EMBEDDING_VECTOR.validate_python(embedding_response["data"][0]["embedding"])
 
         data: Final = {
             "vector": embedding,
@@ -404,7 +407,7 @@ class QdrantSemanticCache(BaseCache):
         embedding_response: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
 
         # get the embedding
-        embedding: Final = embedding_response["data"][0]["embedding"]
+        embedding: Final = _EMBEDDING_VECTOR.validate_python(embedding_response["data"][0]["embedding"])
 
         value = str(value)
         assert isinstance(value, str)
@@ -440,7 +443,7 @@ class QdrantSemanticCache(BaseCache):
         embedding_response: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
 
         # get the embedding
-        embedding: Final = embedding_response["data"][0]["embedding"]
+        embedding: Final = _EMBEDDING_VECTOR.validate_python(embedding_response["data"][0]["embedding"])
 
         data: Final = {
             "vector": embedding,

litellm/caching/redis_cluster_cache.py

@@ -9,7 +9,7 @@ import asyncio
 import contextvars
 import os
 import weakref
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar
 
@@ -24,7 +24,7 @@ if TYPE_CHECKING:
 
     pipeline = Pipeline
     async_redis_client = Redis
-    Span = _Span | Any
+    Span = _Span
 else:
     pipeline = Any
     async_redis_client = Any
@@ -126,13 +126,13 @@ class RedisClusterCache(RedisCache):
             return call()
         return self._connection_gate.run(call, self._get_async_client_cache_key())
 
-    def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
+    def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]:
         """
         Overrides `_run_redis_mget_operation` in redis_cache.py
         """
         return self.redis_client.mget_nonatomic(keys=keys)
 
-    async def _async_run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
+    async def _async_run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]:
         """
         Overrides `_async_run_redis_mget_operation` in redis_cache.py
         """

litellm/caching/redis_semantic_cache.py

@@ -117,7 +117,8 @@ class RedisSemanticCache(BaseCache):
                 password = password or os.environ["REDIS_PASSWORD"]
             except KeyError as e:
                 # Raise a more informative exception if any of the required keys are missing
-                missing_var: Final[object] = e.args[0]
+                key_error_args: Final[tuple[object, ...]] = e.args
+                missing_var: Final = key_error_args[0]
                 raise ValueError(
                     f"Missing required Redis configuration: {missing_var}. Provide {missing_var} or redis_url."
                 ) from e
@@ -310,9 +311,9 @@ class RedisSemanticCache(BaseCache):
                         return
             return
 
-        content = getattr(value, "content", None)
-        if content is not None:
-            cls._collect_responses_input_text(content, prompt_parts)
+        attribute_content: Final[object] = getattr(value, "content", None)
+        if attribute_content is not None:
+            cls._collect_responses_input_text(attribute_content, prompt_parts)
             return
 
         for text_key in ("text", "output", "input_text", "output_text"):

litellm/caching/valkey_semantic_cache.py

@@ -21,6 +21,7 @@ from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from typing import Any, Final
 
+from pydantic import TypeAdapter
 from redis import Redis
 from redis.asyncio import Redis as AsyncRedis
 from redis.commands.search.field import TagField, VectorField
@@ -34,6 +35,8 @@ from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector
 from ._embedding_router import resolve_embedding_timeout
 from .redis_semantic_cache import RedisSemanticCache
 
+_OPAQUE_SEARCH_FIELD: Final = TypeAdapter(object)
+
 
 @dataclass(frozen=True, slots=True)
 class _ValkeyCacheHit:
@@ -226,7 +229,7 @@ class ValkeySemanticCache(RedisSemanticCache):
         if not docs:
             return None
         doc: Final = docs[0]
-        response_field: Final[object] = getattr(doc, cls.RESPONSE_FIELD_NAME)
+        response_field: Final = _OPAQUE_SEARCH_FIELD.validate_python(getattr(doc, cls.RESPONSE_FIELD_NAME))
         distance_field: Final[str | bytes | float] = getattr(doc, cls.DISTANCE_FIELD_NAME)
         return _ValkeyCacheHit(
             response=str(response_field),

litellm/completion_extras/litellm_responses_transformation/transformation.py

@@ -35,6 +35,7 @@ from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
 from litellm.llms.base_llm.bridges.completion_transformation import (
     CompletionTransformationBridge,
 )
+from litellm.responses.litellm_completion_transformation.tool_search import is_anthropic_tool_search_chat_tool
 from litellm.responses.sse_output_recovery import (
     parse_sse_json_chunk,
     record_output_item_chunk,
@@ -196,7 +197,7 @@ def _drops_audio_input(model: str, litellm_params: Mapping[str, object]) -> bool
     return not _supports_audio_input(model, litellm_params)
 
 
-def _provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]:
+def responses_provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]:
     return MappingProxyType(
         {
             key: value
@@ -304,6 +305,136 @@ def _reasoning_items_from_output_items(output_items: Sequence[object]) -> tuple[
     )
 
 
+_TOOL_SEARCH_ITEM_TYPES: Final = frozenset({"tool_search_call", "tool_search_output"})
+_TOOL_SEARCH_REPLAY_ITEM_TYPES: Final = frozenset({"reasoning", *_TOOL_SEARCH_ITEM_TYPES})
+_TOOL_SEARCH_ITEMS_KEY: Final = "tool_search_items"
+
+
+def _json_field(value: object, key: str) -> object:
+    return cast(Mapping[str, object], value).get(key) if isinstance(value, dict) else None  # cast-ok: untyped JSON
+
+
+def _tool_search_replay_item(item: object) -> Mapping[str, object] | None:
+    raw: Final[object] = item.model_dump(exclude_unset=True) if isinstance(item, BaseModel) else item
+    if not isinstance(raw, dict):
+        return None
+    fields: Final = cast(dict[str, object], raw)  # cast-ok: provider JSON item, keys forwarded verbatim
+    if fields.get("type") not in _TOOL_SEARCH_REPLAY_ITEM_TYPES:
+        return None
+    return {key: value for key, value in fields.items() if value is not None}
+
+
+def _tool_search_items_from_output_items(output_items: object) -> tuple[Mapping[str, object], ...]:
+    if not isinstance(output_items, list):
+        return ()
+    entries: Final = cast(Sequence[object], output_items)  # cast-ok: each entry is re-validated
+    items: Final = tuple(item for item in (_tool_search_replay_item(entry) for entry in entries) if item is not None)
+    last_search: Final = max(
+        (index for index, item in enumerate(items) if item.get("type") in _TOOL_SEARCH_ITEM_TYPES), default=-1
+    )
+    return items[: last_search + 1]
+
+
+def _tool_search_provider_fields(
+    items: Sequence[Mapping[str, object]],
+) -> dict[str, object] | None:  # mutable-ok: Message.provider_specific_fields is typed as a dict
+    return {_TOOL_SEARCH_ITEMS_KEY: [dict(item) for item in items]} if items else None
+
+
+def _tool_search_input_items(msg: object) -> tuple[Mapping[str, object], ...]:
+    return _tool_search_items_from_output_items(
+        _json_field(_json_field(msg, "provider_specific_fields"), _TOOL_SEARCH_ITEMS_KEY)
+    )
+
+
+def carries_tool_search_items(provider_specific_fields: object) -> bool:
+    return bool(_json_field(provider_specific_fields, _TOOL_SEARCH_ITEMS_KEY))
+
+
+def _tool_definition_name(tool: object) -> tuple[str, ...]:
+    name: Final = _json_field(tool, "server_label" if _json_field(tool, "type") == "mcp" else "name")
+    return (name,) if isinstance(name, str) else ()
+
+
+def _namespace_members(tool: object) -> Sequence[object]:
+    members: Final = _json_field(tool, "tools")
+    if _json_field(tool, "type") != "namespace" or not isinstance(members, list):
+        return (tool,)
+    return cast(Sequence[object], members)  # cast-ok: untyped JSON, each member re-checked
+
+
+def _loaded_tools(item: Mapping[str, object]) -> Sequence[object]:
+    tools: Final = item.get("tools")
+    if item.get("type") != "tool_search_output" or not isinstance(tools, list):
+        return ()
+    return cast(Sequence[object], tools)  # cast-ok: untyped JSON, each tool re-checked
+
+
+def _json_list(value: object) -> Sequence[object] | None:
+    return cast(Sequence[object], value) if isinstance(value, list) else None  # cast-ok: untyped JSON
+
+
+def _message_loaded_tools(message: object) -> Iterator[object]:
+    loaded: Final = chain.from_iterable(_loaded_tools(item) for item in _tool_search_input_items(message))
+    return chain.from_iterable(_namespace_members(tool) for tool in loaded)
+
+
+def replayed_tool_search_tools(request: object) -> tuple[object, ...]:
+    messages: Final = _json_list(_json_field(request, "messages")) or ()
+    return tuple(chain.from_iterable(_message_loaded_tools(message) for message in messages))
+
+
+def replayed_tool_search_tool_names(request: object) -> tuple[str, ...]:
+    return tuple(chain.from_iterable(_tool_definition_name(tool) for tool in replayed_tool_search_tools(request)))
+
+
+def _loads_any_tool(message: object, names: frozenset[str]) -> bool:
+    loaded_names: Final = chain.from_iterable(_tool_definition_name(tool) for tool in _message_loaded_tools(message))
+    return any(name in names for name in loaded_names)
+
+
+def _without_tool_search_items(message: object) -> object:
+    fields: Final = _json_field(message, "provider_specific_fields")
+    if not isinstance(fields, dict):
+        return message
+    kept: Final = {
+        key: value
+        for key, value in cast(Mapping[str, object], fields).items()  # cast-ok: untyped JSON
+        if key != _TOOL_SEARCH_ITEMS_KEY
+    }
+    return {**cast(Mapping[str, object], message), "provider_specific_fields": kept}  # cast-ok: dict checked above
+
+
+def without_tool_search_items_loading(
+    request: object, names: frozenset[str]
+) -> list[object] | None:  # mutable-ok: request messages stay a list for every downstream handler
+    messages: Final = _json_list(_json_field(request, "messages"))
+    if messages is None:
+        return None
+    return [_without_tool_search_items(message) if _loads_any_tool(message, names) else message for message in messages]
+
+
+def _effective_tools(request: object) -> Sequence[object]:
+    override: Final = _json_list(_json_field(_json_field(request, "extra_body"), "tools"))
+    return override if override is not None else (_json_list(_json_field(request, "tools")) or ())
+
+
+def _requests_tool_search(request: object) -> bool:
+    return any(_json_field(tool, "type") == "tool_search" for tool in _effective_tools(request))
+
+
+def _reasoning_ids(items: Sequence[Mapping[str, object]]) -> frozenset[str]:
+    ids: Final = (item.get("id") for item in items if item.get("type") == "reasoning")
+    return frozenset(item_id for item_id in ids if isinstance(item_id, str))
+
+
+def _assistant_replay_items(msg: "AllMessageValues", replay_tool_search: bool) -> tuple[Mapping[str, object], ...]:
+    searched: Final = _tool_search_input_items(msg) if replay_tool_search else ()
+    searched_reasoning_ids: Final = _reasoning_ids(searched)
+    reasoning: Final = (item for item in _reasoning_input_items(msg) if item.get("id") not in searched_reasoning_ids)
+    return (*searched, *reasoning)
+
+
 def _as_chat_reasoning_items(
     reasoning_items: Sequence[_BuiltReasoningItem | ChatCompletionReasoningItem],
 ) -> list[ChatCompletionReasoningItem] | None:
@@ -346,6 +477,7 @@ def _fold_choices_into_one(
     tool_calls: Sequence[Mapping[str, object]],
     trailing_reasoning_content: str | None,
     trailing_reasoning_item: _BuiltReasoningItem | None,
+    tool_search_items: Sequence[Mapping[str, object]],
     structured_output: bool = False,
 ) -> "Choices":
     """Clients that only read choices[0] must still see every text and tool call.
@@ -373,6 +505,7 @@ def _fold_choices_into_one(
             tool_calls=list(tool_calls) or None,
             reasoning_content=" ".join(value for value in reasoning_contents if value) or None,
             reasoning_items=_as_chat_reasoning_items(reasoning_items),
+            provider_specific_fields=_tool_search_provider_fields(tool_search_items),
         ),
         finish_reason="tool_calls" if tool_calls else "stop",
         index=0,
@@ -457,7 +590,9 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False):
     provider_specific_fields: Mapping[str, object]
 
 
-def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
+def tool_call_dict_from_output_item(
+    item: Mapping[str, Any], index: int, *, include_namespace: bool = False
+) -> _ChatToolCallDict:
     """Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat
     completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw
     string payload in ``input`` rather than ``arguments``; both map to
@@ -486,12 +621,24 @@ def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _Cha
         provider_specific_fields = vars(raw_provider_fields)
     else:
         provider_specific_fields = None
-    if provider_specific_fields:
-        tool_call_dict["provider_specific_fields"] = provider_specific_fields
-        function_chunk["provider_specific_fields"] = provider_specific_fields
+    namespace: Final = item.get("namespace") if include_namespace else None
+    merged_provider_fields: Final = {
+        **(provider_specific_fields or {}),
+        **({"namespace": namespace} if isinstance(namespace, str) else {}),
+    }
+    if merged_provider_fields:
+        tool_call_dict["provider_specific_fields"] = merged_provider_fields
+        function_chunk["provider_specific_fields"] = merged_provider_fields
     return tool_call_dict
 
 
+def _tool_call_namespace(tool_call: object) -> str | None:
+    top_level: Final = _json_field(_json_field(tool_call, "provider_specific_fields"), "namespace")
+    function_fields: Final = _json_field(_json_field(tool_call, "function"), "provider_specific_fields")
+    nested: Final = _json_field(function_fields, "namespace")
+    return top_level if isinstance(top_level, str) else nested if isinstance(nested, str) else None
+
+
 def _flat_responses_tool_choice(choice_type: str, name: str) -> ToolChoiceFunctionParam | ToolChoiceCustomParam:
     if choice_type == "custom":
         return ToolChoiceCustomParam(type="custom", name=name)
@@ -517,6 +664,50 @@ def _reasoning_item_to_response_input(
     }
 
 
+def _requested_defer_loading(tool: Mapping[str, object], function_tool: Mapping[str, object]) -> object:
+    tool_level: Final = tool.get("defer_loading")
+    return tool_level if tool_level is not None else function_tool.get("defer_loading")
+
+
+def _defer_loading_field(
+    tool: Mapping[str, object],
+    function_tool: Mapping[str, object],
+    *,
+    tool_search: Literal["hosted", "client", "absent"],
+    drop_params: bool,
+) -> Mapping[str, bool]:
+    defer_loading: Final = _requested_defer_loading(tool, function_tool)
+    if defer_loading is None or (defer_loading is False and tool_search == "absent"):
+        return {}
+    if tool_search == "client":
+        verbose_logger.debug(
+            "Chat provider: dropping defer_loading=%r beside a client-executed tool_search", defer_loading
+        )
+        return {}
+    if isinstance(defer_loading, bool) and tool_search == "hosted":
+        return {"defer_loading": defer_loading}
+    if drop_params:
+        verbose_logger.debug("Chat provider: dropping defer_loading=%r under drop_params", defer_loading)
+        return {}
+    problem: Final = (
+        'defer_loading needs a {"type": "tool_search"} tool in the same request'
+        if isinstance(defer_loading, bool)
+        else f"defer_loading must be a boolean, got {defer_loading!r}"
+    )
+    raise litellm.UnsupportedParamsError(
+        message=f"{problem}. To drop it, set `litellm.drop_params=True` or pass `drop_params=True`."
+    )
+
+
+def _without_dropped_defer_loading(
+    tool: Mapping[str, object], *, tool_search: Literal["hosted", "client", "absent"], drop_params: bool
+) -> Mapping[str, object]:
+    if tool.get("type") != "function" or "defer_loading" not in tool:
+        return tool
+    kept: Final = _defer_loading_field(tool, {}, tool_search=tool_search, drop_params=drop_params)
+    return {**{key: value for key, value in tool.items() if key != "defer_loading"}, **kept}
+
+
 class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
     """
     Handler for transforming /chat/completions api requests to litellm.responses requests
@@ -597,9 +788,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         *,
         drop_params: bool = False,
         keep_prompt_cache_breakpoints: bool = False,
+        replay_tool_search: bool = False,
+        include_namespace: bool = False,
     ) -> tuple[list[object], str | None]:
         converted_input_items, instructions = self._convert_chat_completion_messages_to_responses_input(
-            messages, drop_params=drop_params
+            messages,
+            drop_params=drop_params,
+            replay_tool_search=replay_tool_search,
+            include_namespace=include_namespace,
         )
         return (
             converted_input_items
@@ -609,7 +805,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         )
 
     def _convert_chat_completion_messages_to_responses_input(
-        self, messages: list["AllMessageValues"], *, drop_params: bool = False
+        self,
+        messages: list["AllMessageValues"],
+        *,
+        drop_params: bool = False,
+        replay_tool_search: bool = False,
+        include_namespace: bool = False,
     ) -> tuple[list[object], str | None]:
         input_items: Final[list[object]] = []
         instructions: str | None = None
@@ -691,7 +892,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                         )
                     )
             elif role == "assistant" and tool_calls and isinstance(tool_calls, list):
-                input_items.extend(_reasoning_input_items(msg))
+                input_items.extend(_assistant_replay_items(msg, replay_tool_search))
                 if content:
                     input_items.append(
                         {
@@ -714,6 +915,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                             input_tool_call["name"] = function["name"]
                         if "arguments" in function:
                             input_tool_call["arguments"] = function["arguments"]
+                        if include_namespace and (
+                            namespace := _tool_call_namespace(cast(object, tool_call))  # cast-ok: untyped
+                        ):
+                            input_tool_call["namespace"] = namespace
                         input_items.append(input_tool_call)
                     elif isinstance(custom, dict):
                         input_items.append(
@@ -728,7 +933,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                         raise ValueError(f"tool call not supported: {tool_call}")
             elif content is not None:
                 if role == "assistant":
-                    input_items.extend(_reasoning_input_items(msg))
+                    input_items.extend(_assistant_replay_items(msg, replay_tool_search))
                 input_items.append(
                     {
                         "type": "message",
@@ -739,7 +944,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                     }
                 )
             elif role == "assistant":
-                input_items.extend(_reasoning_input_items(msg))
+                input_items.extend(_assistant_replay_items(msg, replay_tool_search))
 
         return input_items, instructions
 
@@ -747,6 +952,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         self,
         optional_params: dict,
         responses_api_request: "ResponsesAPIOptionalRequestParams",
+        drop_params: bool = False,
     ) -> None:
         """Map optional_params into responses_api_request (mutates in place)."""
         for key, value in optional_params.items():
@@ -756,7 +962,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                 responses_api_request["max_output_tokens"] = value
             elif key == "tools" and value is not None:
                 responses_api_request["tools"] = self._convert_tools_to_responses_format(
-                    cast(list[dict[str, object]], value)
+                    cast(list[dict[str, object]], value), drop_params=drop_params
                 )
             elif key == "response_format":
                 text_format = self._transform_response_format_to_text_format(value)
@@ -845,6 +1051,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         litellm_logging_obj: "LiteLLMLoggingObj",
         client: object | None = None,
     ) -> dict:
+        optional_params = self._extract_extra_body_params(optional_params)
         provider: Final = _served_provider(litellm_params)
         base_model: Final = litellm_params.get("base_model")
         supports_prompt_cache_breakpoint: Final = supports_openai_prompt_cache_breakpoint(model, provider) or (
@@ -856,6 +1063,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
             messages,
             drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params),
             keep_prompt_cache_breakpoints=supports_prompt_cache_breakpoint,
+            replay_tool_search=_requests_tool_search(cast(object, optional_params)),  # cast-ok: upcast
+            include_namespace=True,
         )
         # OpenAI's Responses API rejects an empty input. For a system-only
         # request, carry the system message as a system-role input item instead
@@ -880,8 +1089,6 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         )
         instructions: Final = None if is_system_only_request else converted_instructions
 
-        optional_params = self._extract_extra_body_params(optional_params)
-
         # Build responses API request using the reverse transformation logic
         responses_api_request: Final = ResponsesAPIOptionalRequestParams()
 
@@ -889,7 +1096,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         if instructions:
             responses_api_request["instructions"] = instructions
 
-        self._map_optional_params_to_responses_api_request(optional_params, responses_api_request)
+        self._map_optional_params_to_responses_api_request(
+            optional_params,
+            responses_api_request,
+            drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params),
+        )
 
         stream: Final = optional_params.get("stream") or litellm_params.get("stream", False)
         verbose_logger.debug("Chat provider: Stream parameter: %s", stream)
@@ -1051,7 +1262,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                     # Tool calls accumulate into the single trailing tool_calls choice
                     # like the typed branches above; a choice per call would hide every
                     # call after choices[0] from chat clients
-                    accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index))
+                    accumulated_tool_calls.append(
+                        tool_call_dict_from_output_item(raw_item, tool_call_index, include_namespace=True)
+                    )
                     tool_call_index += 1
                 elif handle_raw_dict_callback is not None:
                     choice, index = handle_raw_dict_callback(item=raw_item, index=index)
@@ -1060,22 +1273,31 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
             else:
                 pass  # don't fail request if item in list is not supported
 
+        tool_search_items: Final = _tool_search_items_from_output_items(output_items)
         if not choices and not accumulated_tool_calls:
-            return choices
+            if not tool_search_items:
+                return choices
+            return [
+                LiteLLMResponsesTransformationHandler._build_empty_choice(
+                    output_items, "stop", tool_search_items=tool_search_items
+                )
+            ]
         return [
             _fold_choices_into_one(
                 choices=choices,
                 tool_calls=accumulated_tool_calls,
                 trailing_reasoning_content=reasoning_content,
                 trailing_reasoning_item=pending_reasoning_item,
+                tool_search_items=tool_search_items,
                 structured_output=structured_output,
             )
         ]
 
     @staticmethod
-    def _build_empty_incomplete_choice(
+    def _build_empty_choice(
         output_items: Sequence[object],
-        finish_reason: Literal["length", "content_filter"],
+        finish_reason: Literal["stop", "length", "content_filter"],
+        tool_search_items: Sequence[Mapping[str, object]] = (),
     ) -> "Choices":
         from litellm.types.utils import Choices, Message
 
@@ -1087,9 +1309,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
             if summary_block.get("text")
         )
         message: Final = Message(
-            content="",
+            content=None if tool_search_items else "",
             reasoning_content=reasoning_content if reasoning_content else None,
             reasoning_items=_as_chat_reasoning_items(reasoning_items),
+            provider_specific_fields=_tool_search_provider_fields(tool_search_items),
         )
         return Choices(message=message, finish_reason=finish_reason, index=0)
 
@@ -1216,7 +1439,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                 raw_response.incomplete_details.reason if raw_response.incomplete_details is not None else None
             )
             if len(choices) == 0:
-                choices.append(self._build_empty_incomplete_choice(output_items, incomplete_finish_reason))
+                choices.append(self._build_empty_choice(output_items, incomplete_finish_reason))
             else:
                 for choice in choices:
                     choice.finish_reason = incomplete_finish_reason
@@ -1232,7 +1455,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         )
 
         model_response.id = _upstream_response_id(raw_response.id) or raw_response.id
-        for key, value in _provider_metadata(raw_response.model_extra).items():
+        for key, value in responses_provider_metadata(raw_response.model_extra).items():
             setattr(model_response, key, value)
 
         # Preserve hidden params from the ResponsesAPIResponse, especially the headers
@@ -1417,14 +1640,21 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
             return result
 
     def _convert_tools_to_responses_format(
-        self, tools: list[dict[str, object]]
+        self, tools: list[dict[str, object]], drop_params: bool = False
     ) -> list["ALL_RESPONSES_API_TOOL_PARAMS"]:
         """Convert chat completion tools to responses API tools format"""
+        tool_search: Final[Literal["hosted", "client", "absent"]] = (
+            "client"
+            if any(tool.get("type") == "tool_search" and tool.get("execution") == "client" for tool in tools)
+            else "hosted"
+            if any(tool.get("type") == "tool_search" or is_anthropic_tool_search_chat_tool(tool) for tool in tools)
+            else "absent"
+        )
         responses_tools: Final[list[ALL_RESPONSES_API_TOOL_PARAMS]] = []
         for tool in tools:
             # convert function tool from chat completion to responses API format
-            if tool.get("type") == "function" and isinstance(tool.get("function"), dict):
-                function_tool = cast(ChatCompletionToolParamFunctionChunk, tool.get("function"))
+            if tool.get("type") == "function" and isinstance(function_payload := tool.get("function"), dict):
+                function_tool = cast(ChatCompletionToolParamFunctionChunk, function_payload)  # cast-ok: client tool
                 responses_tools.append(
                     FunctionToolParam(
                         name=function_tool["name"],
@@ -1432,6 +1662,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                         strict=function_tool.get("strict"),
                         type="function",
                         description=function_tool.get("description"),
+                        **_defer_loading_field(tool, function_tool, tool_search=tool_search, drop_params=drop_params),
                     )
                 )
             elif tool.get("type") == "custom" and isinstance(custom_payload := tool.get("custom"), dict):
@@ -1446,7 +1677,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                     flat_custom["format"] = convert_custom_tool_format_to_responses_shape(custom_payload["format"])
                 responses_tools.append(flat_custom)
             else:
-                responses_tools.append(tool)
+                responses_tools.append(
+                    _without_dropped_defer_loading(tool, tool_search=tool_search, drop_params=drop_params)
+                )
 
         return cast(list["ALL_RESPONSES_API_TOOL_PARAMS"], responses_tools)
 
@@ -1743,7 +1976,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
             # New output item added
             output_item = parsed_chunk.get("item", {})
             if output_item.get("type") in ("function_call", "custom_tool_call"):
-                converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
+                converted: Final = tool_call_dict_from_output_item(
+                    output_item, parsed_chunk.get("output_index", 0), include_namespace=True
+                )
                 provider_specific_fields: Final = converted.get("provider_specific_fields")
 
                 function_chunk: Final = ChatCompletionToolCallFunctionChunk(
@@ -1756,14 +1991,13 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                 tool_call_index = OpenAiResponsesToChatCompletionStreamIterator._sequential_tool_call_index(
                     tool_call_index_map, parsed_chunk.get("output_index", 0)
                 )
-                tool_call_chunk: Final = ChatCompletionToolCallChunk(
+                tool_call_chunk: Final = _ChatToolCallDict(
                     id=converted["id"],
                     index=tool_call_index,
                     type="function",
                     function=function_chunk,
+                    **({"provider_specific_fields": provider_specific_fields} if provider_specific_fields else {}),
                 )
-                if provider_specific_fields:
-                    tool_call_chunk.provider_specific_fields = provider_specific_fields
 
                 return ModelResponseStream(
                     choices=[
@@ -1819,7 +2053,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                                 delta=Delta(
                                     tool_calls=(
                                         tool_call_dict_from_output_item(
-                                            output_item, parsed_chunk.get("output_index", 0)
+                                            output_item, parsed_chunk.get("output_index", 0), include_namespace=True
                                         ),
                                     )
                                 ),
@@ -1905,7 +2139,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                 from litellm.responses.utils import ResponseAPILoggingUtils
 
                 usage = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(response_data.get("usage"))
-            provider_metadata: Final = _provider_metadata(response_data)
+            provider_metadata: Final = responses_provider_metadata(response_data)
             served_service_tier: Final = response_data.get("service_tier")
             return ModelResponseStream(
                 choices=[
@@ -1914,6 +2148,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
                         delta=Delta(
                             content="",
                             reasoning_items=terminal_reasoning_items_typed,
+                            provider_specific_fields=_tool_search_provider_fields(
+                                _tool_search_items_from_output_items(cast(object, output_items))  # cast-ok: untyped
+                            ),
                         ),
                         finish_reason=finish_reason,
                     )

litellm/constants.py

@@ -17,7 +17,7 @@ OAUTH_TOKEN_EXCHANGE_CACHE_SAFETY_MARGIN_SECONDS: Final = 60
 MICROSOFT_365_COPILOT_DEFAULT_TOKEN_EXCHANGE_PROFILE: Final = "jwt_bearer_obo"
 MICROSOFT_365_COPILOT_DEFAULT_TOKEN_EXCHANGE_SCOPE: Final = "https://graph.microsoft.com/.default"
 # Replies depend on the caller's delegated identity, which response-cache keys do not include.
-RESPONSE_CACHE_EXCLUDED_PROVIDERS: Final = frozenset({"microsoft_365_copilot"})
+RESPONSE_CACHE_EXCLUDED_PROVIDERS: Final = frozenset({"microsoft_365_copilot", "databricks_agent"})
 MICROSOFT_365_COPILOT_DEFAULT_TIME_ZONE: Final = "UTC"
 GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth"
 GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type"
@@ -57,6 +57,7 @@ RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset(
         "model_group_retry_policy",
         "model_group_alias",
         "enable_weighted_failover",
+        "enable_mid_stream_fallback_continuation",
         "enable_tag_filtering",
         "tag_routing_prefix",
         "optional_pre_call_checks",
@@ -106,6 +107,7 @@ DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY: Final = get_env_int("DEFAULT_S3_MAX_ADAPTIV
 S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS: Final = get_env_int_in_range("S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS", 30, 1, 3600)
 # https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
 MAX_S3_OBJECT_KEY_BYTES: Final = 1024
+MAX_S3_OBJECT_KEY_FILENAME_BYTES: Final = 255
 S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64
 S3_PREFIX_DIGEST_CHARS: Final = 16
 # s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against
@@ -145,6 +147,7 @@ HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0
 # Streams exceeding this duration are terminated with a Timeout error.
 # None (default) = no limit.  Set env var to a number of seconds to enable globally.
 LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = get_env_float_or_none("LITELLM_MAX_STREAMING_DURATION_SECONDS")
+MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = 1.0
 
 # Maximum number of base64 characters to keep in logging payloads.
 # Data URIs exceeding this are replaced with a size placeholder.
@@ -379,6 +382,7 @@ REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES: Final = get_env_int_or_none("REALTIME
 REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = get_env_float(
     "REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS", 20.0
 )
+BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS: Final = get_env_float("BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS", 5.0)
 
 # RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
 WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
@@ -397,6 +401,7 @@ MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY: Final = "litellm.model_group_alias_resolve
 MCP_PEEKED_BODY_SCOPE_KEY: Final = "litellm_mcp_peeked_body"
 MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS: Final = 5.0
 REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
+REALTIME_METERED_USAGE_KEY: Final = "realtime_metered_usage"
 OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS: Final = 30
 OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS: Final = 5
 OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES: Final = frozenset(
@@ -478,6 +483,7 @@ MINIMUM_PROMPT_CACHE_TOKEN_COUNT: Final = (
     else DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT
 )
 PROMPT_CACHE_LOOKBACK_POSITIONS: Final = 20
+PROMPT_CACHE_PIN_CACHE_MAX_SIZE: Final = 10_000
 DEFAULT_TRIM_RATIO: Final = get_env_float(
     "DEFAULT_TRIM_RATIO", 0.75
 )  # default ratio of tokens to trim from the end of a prompt
@@ -843,6 +849,8 @@ PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO: Final = frozenset(
     }
 )
 
+A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS: Final = frozenset({"databricks_agent"})
+
 LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS: Final = [
     "openai",
     "azure",
@@ -963,7 +971,7 @@ DEFAULT_CHAT_COMPLETION_PARAM_VALUES: Final = {
     "context_management": None,
 }
 
-openai_compatible_endpoints: Final[list] = [
+openai_compatible_endpoints: Final[list[str]] = [
     "api.perplexity.ai",
     "api.endpoints.anyscale.com/v1",
     "api.deepinfra.com/v1/openai",
@@ -1129,6 +1137,7 @@ _openai_like_providers: Final[list[str]] = [
     "databricks",
     "lemonade",
     "watsonx",
+    "inferenceapis",
 ]  # private helper. similar to openai but require some custom auth / endpoint handling, so can't use the openai sdk
 # well supported replicate llms
 replicate_models: Final[set] = set(
@@ -1619,8 +1628,10 @@ CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
 ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
 ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens"
 INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
+MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response"
 SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
 SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"
+CALLER_SESSION_ID_METADATA_KEY: Final = "litellm_caller_session_id"
 LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
 LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (
     "Truncation is a DB storage safeguard. "
@@ -1652,6 +1663,7 @@ DEFAULT_SOFT_BUDGET: Final = get_env_float(
 RATE_LIMIT_ERROR_MESSAGE_FOR_VIRTUAL_KEY: Final = "LiteLLM Virtual Key user_api_key_hash"
 # Prefix of the 401 raised when a submitted virtual key is not shaped like one.
 INVALID_VIRTUAL_KEY_ERROR_MESSAGE: Final = "LiteLLM Virtual Key expected"
+INVALID_API_KEY_ERROR_MESSAGE: Final = "Authentication Error, Invalid API key provided."
 # Attribute stamped on that 401 at its raise site so log routing recognises it by
 # provenance. Message text is caller-influenceable on other 401s, so it must not
 # be used to classify.
@@ -1684,6 +1696,7 @@ ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS: Final = {
 # e.g., 'x-pass-anthropic-beta: value' becomes 'anthropic-beta: value'
 # Works for all LLM pass-through endpoints (Vertex AI, Anthropic, Bedrock, etc.)
 PASS_THROUGH_HEADER_PREFIX: Final = "x-pass-"
+TYPESAFE_DANGER_ZONE_HEADER_PREFIX: Final = "typesafe-danger-zone-"
 INTERNAL_KWARG_PREFIX: Final = "_litellm_"
 CONTROL_OPTIONS_KEY: Final = f"{INTERNAL_KWARG_PREFIX}control"
 
@@ -1695,6 +1708,7 @@ AZURE_SPEECH_FAST_TRANSCRIPTION_PATH: Final = "/speechtotext/transcriptions:tran
 AZURE_SPEECH_STT_DOMAIN: Final = "stt.speech.microsoft.com"
 AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN: Final = "api.cognitive.microsoft.com"
 AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER: Final = "Ocp-Apim-Subscription-Key"
+AWS_ERROR_TYPE_HEADER: Final = "x-amzn-errortype"
 AZURE_SPEECH_SHORT_AUDIO_MODEL: Final = "short-audio"
 AZURE_SPEECH_BATCH_MODEL: Final = "batch-transcription"
 AZURE_SPEECH_FAST_TRANSCRIPTION_MODEL: Final = "fast-transcription"
@@ -1845,6 +1859,7 @@ SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = get_env_int("SPEND_COUNTER_RESEED_L
 PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, get_env_int("PROXY_DB_LOOKUP_MAX_CONCURRENCY", 25))
 PROXY_DB_LOOKUP_DEADLINE_SECONDS: Final = max(0.1, get_env_float("PROXY_DB_LOOKUP_DEADLINE_SECONDS", 10.0))
 PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS: Final = max(0.0, get_env_float("PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", 30.0))
+TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS: Final = 30
 DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = get_env_int("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)  # 1 minute
 PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)
 RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, get_env_int("RESET_BUDGET_JOB_BATCH_SIZE", 500))
@@ -1856,7 +1871,13 @@ RESET_BUDGET_JOB_NAME: Final = "reset_budget_job"
 RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900
 PROXY_BATCH_POLLING_INTERVAL: Final = get_env_int("PROXY_BATCH_POLLING_INTERVAL", 3600)
 MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, get_env_int("MAX_OBJECTS_PER_POLL_CYCLE", 50))
+# raw request and result bytes per batch line-item object, so a large batch is built and stored one part at a time
+BATCH_LINE_ITEMS_MAX_PART_BYTES: Final = get_env_int_in_range(
+    "BATCH_LINE_ITEMS_MAX_PART_BYTES", 16 * 1024 * 1024, 1, 1024 * 1024 * 1024
+)
 MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, get_env_int("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))
+PENDING_RESPONSE_STATUSES: Final = ("queued", "in_progress")
+TERMINAL_RESPONSE_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "incomplete"})
 BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS: Final = 10.0
 BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS: Final = 60
 STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, get_env_int("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000))
@@ -1878,6 +1899,7 @@ _background_interaction_cost_polling_env: Final = os.getenv(
     "BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true"
 ).lower()
 BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_polling_env == "true"
+BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS: Final = 300.0
 PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)
 PROXY_BATCH_WRITE_AT: Final = get_env_int("PROXY_BATCH_WRITE_AT", 10)  # in seconds, increased from 10
 PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30)

litellm/containers/README.md

@@ -120,6 +120,7 @@ class AzureContainerConfig(BaseContainerConfig):
         self,
         headers: dict,
         api_key: Optional[str] = None,
+        litellm_params: Optional[GenericLiteLLMParams] = None,
     ) -> dict:
         """Azure uses api-key header instead of Bearer token."""
         import litellm

litellm/containers/endpoint_factory.py

@@ -36,6 +36,7 @@ RESPONSE_TYPES: Final[dict[str, type]] = {
     "DeleteContainerFileResponse": DeleteContainerFileResponse,
 }
 _ENDPOINTS_DOCUMENT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
+_ENDPOINT_RESULT: Final = TypeAdapter(object)
 
 
 def _load_endpoints_config() -> dict:
@@ -173,7 +174,7 @@ def create_async_endpoint_function(
             init_response: Final = await loop.run_in_executor(None, func_with_context)
 
             if asyncio.iscoroutine(init_response):
-                response = await init_response
+                response = _ENDPOINT_RESULT.validate_python(await init_response)
             else:
                 response = init_response
 

litellm/containers/utils.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from typing import Any, Final, TypeVar
 
 from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
@@ -99,7 +100,7 @@ class ContainerRequestUtils:
         response_obj: T,
         custom_llm_provider: str | None,
         litellm_metadata: dict[str, Any] | None = None,
-        extra_body: dict[str, Any] | None = None,
+        extra_body: Mapping[str, object] | None = None,
     ) -> T:
         """
         Encode container_id in response object with provider/model metadata for routing.

litellm/cost_calculator.py

@@ -10,7 +10,7 @@ from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, cast
 
 from httpx import Response
-from pydantic import BaseModel, ConfigDict, ValidationError
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, TypedDict, assert_never
 
 import litellm
@@ -20,6 +20,7 @@ from litellm.constants import (
     DEFAULT_MAX_LRU_CACHE_SIZE,
     DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND,
     OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES,
+    REALTIME_METERED_USAGE_KEY,
 )
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.cost_discount import resolve_cost_discount
@@ -85,6 +86,9 @@ from litellm.llms.openai.cost_calculation import (
 from litellm.llms.openai.cost_calculation import (
     video_output_cost_per_second,
 )
+from litellm.llms.openai.transcriptions.gpt_transformation import (
+    OpenAIGPTAudioTranscriptionConfig,
+)
 from litellm.llms.perplexity.cost_calculator import (
     cost_per_token as perplexity_cost_per_token,
 )
@@ -117,6 +121,7 @@ from litellm.types.llms.openai import (
     ResponseAPIUsage,
     ResponsesAPIResponse,
 )
+from litellm.types.llms.xai import RealtimeMeteredUsage
 from litellm.types.rerank import RerankBilledUnits, RerankResponse
 from litellm.types.utils import (
     CallTypesLiteral,
@@ -323,6 +328,15 @@ def _transcription_usage_has_token_details(
     return (prompt_tokens_val > 0) or (completion_tokens_val > 0)
 
 
+def _bills_provider_transcription_seconds(model: str | None, custom_llm_provider: str | None) -> bool:
+    if model is None or custom_llm_provider != LlmProviders.OPENAI.value:
+        return False
+    transcription_config: Final = ProviderConfigManager.get_provider_audio_transcription_config(
+        model=model.removeprefix(f"{LlmProviders.OPENAI.value}/"), provider=LlmProviders.OPENAI
+    )
+    return isinstance(transcription_config, OpenAIGPTAudioTranscriptionConfig)
+
+
 OCRPricingField = Literal["ocr_cost_per_page", "ocr_cost_per_credit", "annotation_cost_per_page"]
 
 
@@ -554,12 +568,11 @@ def cost_per_token(
 
     assert custom_llm_provider is not None  # caller-supplied or get_llm_provider
 
-    model_without_prefix = model
     model_parts: Final = model.split("/", 1)
-    if len(model_parts) > 1:
-        model_without_prefix = model_parts[1]
-    else:
-        model_without_prefix = model
+    model_is_provider_relative: Final = (
+        model_with_provider == f"{custom_llm_provider}/{model}" and model_with_provider in model_cost_ref
+    )
+    model_without_prefix: Final = model if len(model_parts) == 1 or model_is_provider_relative else model_parts[1]
     """
     Code block that formats model to lookup in litellm.model_cost
     Option1. model = "bedrock/ap-northeast-1/anthropic.claude-instant-v1". This is the most accurate since it is region based. Should always be option 1
@@ -828,7 +841,14 @@ def _get_hidden_str_for_cost_calc(hidden_params: object, key: str) -> str | None
 
 
 _NON_TOKEN_RATE_FIELDS: Final = frozenset(
-    {"cost_per_second", "input_cost_per_second", "output_cost_per_second", "input_cost_per_query", "tiered_pricing"}
+    {
+        "cost_per_second",
+        "input_cost_per_second",
+        "output_cost_per_second",
+        "input_cost_per_query",
+        "input_cost_per_text_message",
+        "tiered_pricing",
+    }
 )
 
 
@@ -1394,7 +1414,7 @@ def completion_cost(
     total_time: float | None = 0.0,  # used for replicate, sagemaker
     call_type: CallTypesLiteral | None = None,
     ### REGION ###
-    custom_llm_provider=None,
+    custom_llm_provider: str | None = None,
     region_name=None,  # used for bedrock pricing
     ### IMAGE GEN ###
     size: str | None = None,
@@ -1516,6 +1536,9 @@ def completion_cost(
 
         explicit_pricing: Final = custom_pricing is True or base_model is not None
         custom_priced_request_model: Final = model if custom_pricing is True else None
+        configured_cost_per_query: Final = (
+            _configured_search_cost_per_query(litellm_logging_obj) if custom_pricing else None
+        )
         base_model_provider: Final = (
             _base_model_pricing_provider(
                 base_model,
@@ -1537,6 +1560,7 @@ def completion_cost(
             region_name=region_name,
         )
 
+        requested_model: Final = model
         potential_model_names: Final = [
             selected_model,
             _get_response_model(completion_response),
@@ -1696,15 +1720,17 @@ def completion_cost(
                                 video_generation_cost,
                             )
 
-                            return (
-                                video_generation_cost(
-                                    model=model,
-                                    duration_seconds=duration_seconds,
-                                    custom_llm_provider=custom_llm_provider,
-                                    model_info=_video_model_info,
-                                    video_resolution=video_resolution,
-                                )
-                                * video_count
+                            return video_generation_cost(
+                                model=model,
+                                duration_seconds=duration_seconds,
+                                custom_llm_provider=custom_llm_provider,
+                                model_info=_video_model_info,
+                                video_resolution=video_resolution,
+                            ) * video_count + _video_input_image_cost(
+                                requested_model or model,
+                                custom_llm_provider,
+                                _video_model_info,
+                                getattr(completion_response, "usage", None),
                             )
                     # Fallback to default video cost calculation if no duration available
                     return default_video_cost_calculator(
@@ -1725,6 +1751,12 @@ def completion_cost(
                     audio_transcription_file_duration = _hidden.get(
                         "audio_transcription_duration",
                         getattr(completion_response, "duration", 0.0),
+                    ) or (
+                        completion_response.usage.seconds
+                        if isinstance(completion_response, TranscriptionResponse)
+                        and isinstance(completion_response.usage, TranscriptionUsageDurationObject)
+                        and _bills_provider_transcription_seconds(requested_model, custom_llm_provider)
+                        else 0.0
                     )
                 elif call_type in _RERANK_CALL_TYPES:
                     if completion_response is not None and isinstance(completion_response, RerankResponse):
@@ -1762,11 +1794,15 @@ def completion_cost(
                     (
                         prompt_cost,
                         completion_cost_result,
-                    ) = search_provider_cost_per_query(
-                        model=search_model,
-                        custom_llm_provider=custom_llm_provider,
-                        number_of_queries=number_of_queries,
-                        optional_params=optional_params,
+                    ) = (
+                        (number_of_queries * configured_cost_per_query, 0.0)
+                        if configured_cost_per_query is not None
+                        else search_provider_cost_per_query(
+                            model=search_model,
+                            custom_llm_provider=custom_llm_provider,
+                            number_of_queries=number_of_queries,
+                            optional_params=optional_params,
+                        )
                     )
 
                     # Return the total cost (prompt_cost + completion_cost, but for search it's just prompt_cost)
@@ -2152,6 +2188,38 @@ def response_cost_calculator(
         raise e
 
 
+def _video_input_image_cost(
+    model: str,
+    custom_llm_provider: str | None,
+    model_info: ModelInfo | None,
+    usage: object,
+) -> float:
+    match usage:
+        case {"billable_input_image_count": int(billable_input_images)} if billable_input_images > 0:
+            deployment_rate: Final = model_info.get("input_cost_per_image") if model_info is not None else None
+            image_rate: Final = (
+                deployment_rate
+                if deployment_rate is not None
+                else litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider).get(
+                    "input_cost_per_image"
+                )
+            )
+            return billable_input_images * image_rate if isinstance(image_rate, (int, float)) else 0.0
+        case _:
+            return 0.0
+
+
+_LOGGED_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object])
+
+
+def _configured_search_cost_per_query(litellm_logging_obj: LitellmLoggingObject | None) -> float | None:
+    logged_params: Final = getattr(litellm_logging_obj, "litellm_params", None)
+    if logged_params is None:
+        return None
+    cost_per_query: Final = _LOGGED_LITELLM_PARAMS.validate_python(logged_params).get("input_cost_per_query")
+    return float(cost_per_query) if isinstance(cost_per_query, (int, float)) else None
+
+
 def _deployment_model_info(
     litellm_logging_obj: LitellmLoggingObject | None,
     custom_pricing: bool | None,
@@ -2958,6 +3026,7 @@ _RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "re
 
 
 class _ResponsesWsEventResponse(LiteLLMBaseModel):
+    id: str | None = None
     usage: Mapping[str, object] | None = None
     service_tier: str | None = None
 
@@ -3009,6 +3078,23 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
         )
         return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects(list(collected_usage_objects))
 
+    @staticmethod
+    def response_id_from_responses_ws_results(results: Sequence[Mapping[str, object]]) -> str | None:
+        response_ids: Final = tuple(
+            response.id for _, response in _billable_responses_ws_events(results) if response.id
+        )
+        return response_ids[-1] if response_ids else None
+
+    @staticmethod
+    def create_logging_object(results: Sequence[Mapping[str, object]]) -> LiteLLMRealtimeStreamLoggingObject:
+        tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(results)
+        return LiteLLMRealtimeStreamLoggingObject(
+            id=ResponsesWebSocketTokenUsageProcessor.response_id_from_responses_ws_results(results),
+            usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(results),
+            results=cast(OpenAIRealtimeStreamList, list(results)),  # cast-ok: raw WS frames, unvalidated field
+            service_tier=next(iter(tier_partition)) if len(tier_partition) == 1 else None,
+        )
+
 
 _TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"
 
@@ -3101,31 +3187,43 @@ def _live_session_wall_clock_seconds(litellm_logging_obj: LitellmLoggingObject |
     return 0.0
 
 
-def _live_input_cost_per_second(model_name: str, custom_llm_provider: str) -> float | None:
+def _live_model_rate(
+    model_name: str,
+    custom_llm_provider: str,
+    rate_field: Literal["input_cost_per_second", "input_cost_per_text_message"],
+) -> float | None:
     model_info: Final = _lookup_model_info_or_none(model_name, custom_llm_provider)
     if model_info is None:
         return None
-    rate: Final = model_info.get("input_cost_per_second")
+    rate: Final = model_info.get(rate_field)
     if isinstance(rate, (int, float)) and not isinstance(rate, bool) and math.isfinite(rate):
         return float(rate)
     return None
 
 
-def _first_live_input_cost_per_second(
+def _first_live_model_rate(
     potential_model_names: Sequence[str | None],
     custom_llm_provider: str,
+    rate_field: Literal["input_cost_per_second", "input_cost_per_text_message"],
 ) -> float:
     return next(
         (
             rate
             for model_name in potential_model_names
             if model_name is not None
-            and (rate := _live_input_cost_per_second(model_name, custom_llm_provider)) is not None
+            and (rate := _live_model_rate(model_name, custom_llm_provider, rate_field)) is not None
         ),
         0.0,
     )
 
 
+def _xai_metered_usage(litellm_logging_obj: LitellmLoggingObject | None) -> RealtimeMeteredUsage | None:
+    if litellm_logging_obj is None:
+        return None
+    usage: Final = litellm_logging_obj.model_call_details.get(REALTIME_METERED_USAGE_KEY)
+    return usage if isinstance(usage, RealtimeMeteredUsage) else None
+
+
 def _live_response_token_costs(
     response: _LiveResponsePayload,
     custom_llm_provider: str,
@@ -3285,9 +3383,28 @@ def handle_realtime_stream_cost_calculation(
     live_model_names: Final = (
         (custom_pricing_model, base_pricing_model, received_model, litellm_model_name) if live_event_seen else ()
     )
-    live_voice_cost: Final = voice_seconds * _first_live_input_cost_per_second(
+    live_voice_cost: Final = voice_seconds * _first_live_model_rate(
         live_model_names,
         custom_llm_provider,
+        "input_cost_per_second",
+    )
+    xai_metered: Final = (
+        _xai_metered_usage(litellm_logging_obj) if custom_llm_provider == LlmProviders.XAI.value else None
+    )
+    xai_audio_seconds: Final = (
+        xai_metered.audio_seconds_sent + xai_metered.audio_seconds_received
+        if xai_metered is not None and xai_metered.push_to_talk
+        else _live_session_wall_clock_seconds(litellm_logging_obj)
+        if custom_llm_provider == LlmProviders.XAI.value
+        else 0.0
+    )
+    audio_duration_cost: Final = xai_audio_seconds * _first_live_model_rate(
+        potential_model_names,
+        custom_llm_provider,
+        "input_cost_per_second",
+    )
+    text_message_cost: Final = (xai_metered.text_messages if xai_metered is not None else 0) * _first_live_model_rate(
+        potential_model_names, custom_llm_provider, "input_cost_per_text_message"
     )
     live_response_costs: Final = tuple(
         _live_response_token_costs(
@@ -3301,10 +3418,17 @@ def handle_realtime_stream_cost_calculation(
     live_response_output_cost: Final = sum(costs[1] for costs in live_response_costs)
     input_cost: Final = input_cost_per_token + live_response_input_cost
     output_cost: Final = output_cost_per_token + live_response_output_cost
-    total_cost: Final = input_cost + output_cost + transcription_cost + live_voice_cost
+    total_cost: Final = (
+        input_cost + output_cost + transcription_cost + live_voice_cost + audio_duration_cost + text_message_cost
+    )
     additional_costs: Final = tuple(
         (cost_name, cost)
-        for cost_name, cost in (("transcription_cost", transcription_cost), ("live_voice_cost", live_voice_cost))
+        for cost_name, cost in (
+            ("transcription_cost", transcription_cost),
+            ("live_voice_cost", live_voice_cost),
+            ("audio_duration_cost", audio_duration_cost),
+            ("text_message_cost", text_message_cost),
+        )
         if cost > 0
     )
 

litellm/decisions/call.py

@@ -2,14 +2,15 @@ from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field, replace
 from functools import partial
 from types import MappingProxyType
-from typing import Final, TypeAlias
+from typing import Final, Literal, TypeAlias
 
 import httpx
-from pydantic import ConfigDict, TypeAdapter, ValidationError
+from pydantic import BaseModel, ConfigDict, StrictBool, TypeAdapter, ValidationError
 from typing_extensions import assert_never
 
 import litellm
 from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER, normalize_drop_params
+from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.prompt_templates.image_handling import (
     async_convert_urls_to_base64,
@@ -31,10 +32,17 @@ from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
 from litellm.llms.openai.decisions.transformation import ir_to_openai_response, openai_request_to_ir
 from litellm.types.decisions import (
     DecisionQuestion,
+    DecisionsAsset,
+    DecisionsAssetReference,
+    DecisionsAssets,
     DecisionsAudio,
     DecisionsBase64Image,
     DecisionsImage,
+    DecisionsIRAnswer,
+    DecisionsIRBoxQuestion,
     DecisionsIRMessages,
+    DecisionsIRQuestion,
+    DecisionsIRRefusal,
     DecisionsIRRequest,
     DecisionsIRResponse,
     DecisionsIRState,
@@ -48,12 +56,13 @@ from litellm.types.decisions import (
     OpenAIDecisionResponse,
     UnsupportedDecisionsRequest,
 )
-from litellm.types.utils import LlmProviders
+from litellm.types.utils import LlmProviders, ProviderSpecificHeader
 from litellm.utils import ProviderConfigManager
 
 SystemOneQuestions: TypeAlias = Mapping[str, DecisionQuestion | Mapping[str, object]]
 OpenAIDecisionQuestions: TypeAlias = Sequence[OpenAIDecisionQuestion | Mapping[str, object]]
 DecisionsImages: TypeAlias = Sequence[str | DecisionsBase64Image | Mapping[str, object]]
+DecisionsAssetEntries: TypeAlias = Mapping[str, DecisionsAsset | Mapping[str, object]]
 DecisionsAudioClips: TypeAlias = Sequence[DecisionsAudio | Mapping[str, object]]
 DecisionsVideoClips: TypeAlias = Sequence[DecisionsVideo | Mapping[str, object]]
 DecisionsRequestFormat: TypeAlias = DecisionsRequestBody | OpenAIDecisionRequestBody
@@ -66,12 +75,22 @@ _SAFETY_IDENTIFIER_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(
 _IMAGES_ADAPTER: Final[TypeAdapter[tuple[DecisionsImage, ...] | None]] = TypeAdapter(
     tuple[DecisionsImage, ...] | None, config=ConfigDict(title="images")
 )
+_PROVIDER_SPECIFIC_HEADER_ADAPTER: Final[
+    TypeAdapter[ProviderSpecificHeader | tuple[ProviderSpecificHeader, ...] | None]
+] = TypeAdapter(ProviderSpecificHeader | tuple[ProviderSpecificHeader, ...] | None)
 _HANDLER: Final = BaseLLMHTTPHandler()
 _BASE64_MARKER: Final = ";base64,"
 
 
 _IMAGE_ADAPTER: Final[TypeAdapter[DecisionsImage]] = TypeAdapter(DecisionsImage, config=ConfigDict(title="images"))
 _IMAGE_ENTRIES_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...])
+_ASSETS_ADAPTER: Final[TypeAdapter[DecisionsAssets | None]] = TypeAdapter(
+    DecisionsAssets | None, config=ConfigDict(title="assets")
+)
+_ASSET_ENTRY_ADAPTER: Final[TypeAdapter[tuple[DecisionsAssetReference, DecisionsAsset]]] = TypeAdapter(
+    tuple[DecisionsAssetReference, DecisionsAsset]
+)
+_ASSET_ENTRIES_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
 
 
 def _valid_image(image: object) -> DecisionsImage | None:
@@ -96,6 +115,29 @@ def _no_image(_image: DecisionsImage) -> bool:
     return False
 
 
+def _valid_asset(entry: tuple[str, object]) -> tuple[str, DecisionsAsset] | None:
+    try:
+        return _ASSET_ENTRY_ADAPTER.validate_python(entry)
+    except ValidationError:
+        return None
+
+
+def _asset_entries(assets: object) -> Mapping[str, object]:
+    try:
+        return _ASSET_ENTRIES_ADAPTER.validate_python(assets)
+    except ValidationError:
+        return MappingProxyType({})
+
+
+def _request_assets(assets: object, drops_params: bool) -> Mapping[str, DecisionsAsset] | None:
+    if not drops_params:
+        return _ASSETS_ADAPTER.validate_python(assets)
+    valid: Final = MappingProxyType(
+        dict(entry for entry in map(_valid_asset, _asset_entries(assets).items()) if entry is not None)
+    )
+    return valid or None
+
+
 def _request_images(images: object, drops_params: bool) -> tuple[DecisionsImage, ...]:
     if not drops_params:
         return _IMAGES_ADAPTER.validate_python(images) or ()
@@ -107,6 +149,7 @@ class SystemOneRequestFields:
     state: DecisionsJSON | None
     questions: SystemOneQuestions | None
     images: object = None
+    assets: object = None
     audio: DecisionsAudioClips | None = None
     videos: DecisionsVideoClips | None = None
 
@@ -116,6 +159,7 @@ class SystemOneRequestFields:
                 "state": self.state,
                 "questions": self.questions,
                 "images": _request_images(self.images, drops_params) or None,
+                "assets": _request_assets(self.assets, drops_params),
                 "audio": self.audio,
                 "videos": self.videos,
             }
@@ -158,6 +202,7 @@ class DecisionsCall:
     headers: Mapping[str, str]
     timeout: float | httpx.Timeout | None
     litellm_params: Mapping[str, object] = field(repr=False)
+    caller_questions: tuple[DecisionsIRQuestion, ...]
 
 
 def supported_decisions_providers() -> tuple[str, ...]:
@@ -229,8 +274,9 @@ class _UnsupportedParams:
 
 
 @dataclass(frozen=True, slots=True)
-class _NoInputLeftAfterDrop:
+class _NothingLeftAfterDrop:
     names: tuple[str, ...]
+    missing: Literal["input to decide on", "questions to answer"]
 
 
 @dataclass(frozen=True, slots=True)
@@ -239,12 +285,21 @@ class _MalformedImages:
 
 
 def _unsupported_params(
-    ir_request: DecisionsIRRequest, provider_config: BaseDecisionsConfig, model: str, custom_llm_provider: str
+    ir_request: DecisionsIRRequest,
+    provider_config: BaseDecisionsConfig,
+    model: str,
+    custom_llm_provider: str,
+    headers: Mapping[str, str],
+    supports_box_questions: bool,
 ) -> tuple[str, ...]:
     refuses_safety_identifier: Final = (
         ir_request.safety_identifier is not None and not provider_config.supports_safety_identifier
     )
-    refused_image_param: Final = None if provider_config.supports_images else image_param(ir_request.input)
+    refused_image_param: Final = (
+        None if provider_config.supports_images_in(ir_request.input, headers) else image_param(ir_request.input)
+    )
+    refuses_assets: Final = bool(ir_request.assets) and not provider_config.supports_assets(headers)
+    refuses_box_questions: Final = _has_box_question(ir_request.questions) and not supports_box_questions
     refuses_audio: Final = bool(ir_request.audio) and not provider_config.supports_audio_input(
         model, custom_llm_provider
     )
@@ -254,18 +309,41 @@ def _unsupported_params(
     return (
         *(("safety_identifier",) if refuses_safety_identifier else ()),
         *(() if refused_image_param is None else (refused_image_param,)),
+        *(("assets",) if refuses_assets else ()),
+        *(("box questions",) if refuses_box_questions else ()),
         *(("audio",) if refuses_audio else ()),
         *(("videos",) if refuses_videos else ()),
     )
 
 
+def _has_box_question(questions: Sequence[DecisionsIRQuestion]) -> bool:
+    return any(isinstance(question, DecisionsIRBoxQuestion) for question in questions)
+
+
+def _is_not_box_question(question: DecisionsIRQuestion) -> bool:
+    return not isinstance(question, DecisionsIRBoxQuestion)
+
+
 def _without_unsupported_params(
-    ir_request: DecisionsIRRequest, provider_config: BaseDecisionsConfig
+    ir_request: DecisionsIRRequest,
+    provider_config: BaseDecisionsConfig,
+    headers: Mapping[str, str],
+    supports_box_questions: bool,
 ) -> DecisionsIRRequest:
     return replace(
         ir_request,
-        input=ir_request.input if provider_config.supports_images else with_images_where(ir_request.input, _no_image),
+        input=(
+            ir_request.input
+            if provider_config.supports_images_in(ir_request.input, headers)
+            else with_images_where(ir_request.input, _no_image)
+        ),
         safety_identifier=ir_request.safety_identifier if provider_config.supports_safety_identifier else None,
+        assets=ir_request.assets if provider_config.supports_assets(headers) else MappingProxyType({}),
+        questions=(
+            ir_request.questions
+            if supports_box_questions
+            else tuple(filter(_is_not_box_question, ir_request.questions))
+        ),
     )
 
 
@@ -285,29 +363,58 @@ def _without_malformed_images(
     return well_formed
 
 
+class _DeploymentBoxQuestions(BaseModel):
+    model_config = ConfigDict(frozen=True, from_attributes=True)
+
+    supports_box_questions: StrictBool | None = None
+
+
+def _deployment_supports_box_questions(kwargs: Mapping[str, object]) -> bool | None:
+    try:
+        return _DeploymentBoxQuestions.model_validate(kwargs.get("model_info") or {}).supports_box_questions
+    except ValidationError:
+        return None
+
+
 def _provider_ir_request(
-    request: DecisionsRequestFormat,
+    ir_request: DecisionsIRRequest,
     *,
-    safety_identifier: str | None,
     provider_config: BaseDecisionsConfig,
     model: str,
     custom_llm_provider: str,
+    headers: Mapping[str, str],
     kwargs: Mapping[str, object],
-) -> DecisionsIRRequest | _UnsupportedParams | _NoInputLeftAfterDrop | _MalformedImages:
-    ir_request: Final = _without_malformed_images(_ir_request(request, safety_identifier), _drops_params(kwargs))
-    if isinstance(ir_request, _MalformedImages):
-        return ir_request
-    unsupported: Final = _unsupported_params(ir_request, provider_config, model, custom_llm_provider)
+) -> DecisionsIRRequest | _UnsupportedParams | _NothingLeftAfterDrop:
+    supports_box_questions: Final = provider_config.supports_box_questions(
+        model, custom_llm_provider, headers, _deployment_supports_box_questions(kwargs)
+    )
+    unsupported: Final = _unsupported_params(
+        ir_request, provider_config, model, custom_llm_provider, headers, supports_box_questions
+    )
     if not unsupported:
         return ir_request
     if not _drops_params(kwargs):
         return _UnsupportedParams(names=unsupported)
-    supported: Final = _without_unsupported_params(ir_request, provider_config)
+    supported: Final = _without_unsupported_params(ir_request, provider_config, headers, supports_box_questions)
     if _leaves_no_input(supported.input):
-        return _NoInputLeftAfterDrop(names=unsupported)
+        return _NothingLeftAfterDrop(names=unsupported, missing="input to decide on")
+    if not supported.questions:
+        return _NothingLeftAfterDrop(names=unsupported, missing="questions to answer")
     return supported
 
 
+def _outbound_headers(
+    extra_headers: Mapping[str, str] | None, custom_llm_provider: str, kwargs: Mapping[str, object]
+) -> Mapping[str, str]:
+    forwarded: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers(
+        provider_specific_header=_PROVIDER_SPECIFIC_HEADER_ADAPTER.validate_python(
+            kwargs.get("provider_specific_header")
+        ),
+        custom_llm_provider=custom_llm_provider,
+    )
+    return MappingProxyType({**(extra_headers or {}), **forwarded})
+
+
 def prepare_call(
     *,
     model: str,
@@ -338,6 +445,7 @@ def prepare_call(
     except ValueError as error:
         raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error
     try:
+        headers: Final = _outbound_headers(extra_headers, provider, kwargs)
         request_safety_identifier: Final = _request_safety_identifier(safety_identifier, provider_config, kwargs)
         request: Final = _fields_without_unsupported_media(
             fields, provider_config, canonical_model, provider, kwargs
@@ -350,7 +458,7 @@ def prepare_call(
         ) from error
 
     try:
-        resolved_api_base, resolved_api_key = provider_config.resolve_connection(
+        resolved_api_base, static_api_key = provider_config.resolve_connection(
             dynamic_api_base or api_base, dynamic_api_key or api_key
         )
     except ValueError as error:
@@ -361,6 +469,7 @@ def prepare_call(
             model=model,
             llm_provider=provider,
         )
+    resolved_api_key: Final = provider_config.resolve_credential(static_api_key, resolved_api_base, kwargs)
     if resolved_api_key is None and provider_config.api_key_required:
         raise litellm.AuthenticationError(
             message=f"Missing API key for Decisions provider '{provider}'",
@@ -368,12 +477,24 @@ def prepare_call(
             llm_provider=provider,
         )
 
+    caller_ir_request: Final = _without_malformed_images(
+        _ir_request(request, request_safety_identifier), _drops_params(kwargs)
+    )
+    if isinstance(caller_ir_request, _MalformedImages):
+        raise litellm.BadRequestError(
+            message=(
+                f"Invalid Decisions request: {caller_ir_request.count} input_image part(s) have an image_url that is "
+                "neither a data:image/...;base64, URL nor an http(s) URL"
+            ),
+            model=model,
+            llm_provider=provider,
+        )
     ir_request: Final = _provider_ir_request(
-        request,
-        safety_identifier=request_safety_identifier,
+        caller_ir_request,
         provider_config=provider_config,
         model=canonical_model,
         custom_llm_provider=provider,
+        headers=headers,
         kwargs=kwargs,
     )
     if isinstance(ir_request, _UnsupportedParams):
@@ -386,20 +507,11 @@ def prepare_call(
             model=model,
             llm_provider=provider,
         )
-    if isinstance(ir_request, _MalformedImages):
-        raise litellm.BadRequestError(
-            message=(
-                f"Invalid Decisions request: {ir_request.count} input_image part(s) have an image_url that is "
-                "neither a data:image/...;base64, URL nor an http(s) URL"
-            ),
-            model=model,
-            llm_provider=provider,
-        )
-    if isinstance(ir_request, _NoInputLeftAfterDrop):
+    if isinstance(ir_request, _NothingLeftAfterDrop):
         raise litellm.BadRequestError(
             message=(
                 f"{provider} does not support parameters: {list(ir_request.names)}, for model={model}, "
-                "and dropping them leaves no input to decide on"
+                f"and dropping them leaves no {ir_request.missing}"
             ),
             model=model,
             llm_provider=provider,
@@ -425,9 +537,10 @@ def prepare_call(
         api_base=resolved_api_base,
         api_key=resolved_api_key,
         logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
-        headers=extra_headers or {},
+        headers=headers,
         timeout=timeout,
         litellm_params=kwargs,
+        caller_questions=caller_ir_request.questions,
     )
 
 
@@ -561,14 +674,24 @@ def _format_response(response: DecisionsIRResponse, call: DecisionsCall) -> Deci
     return formatted
 
 
+def _caller_answers(call: DecisionsCall, answers: Sequence[DecisionsIRAnswer]) -> tuple[DecisionsIRAnswer, ...]:
+    sent_answers: Final = iter(answers)
+    return tuple(
+        next(sent_answers) if question in call.ir_request.questions else DecisionsIRRefusal()
+        for question in call.caller_questions
+    )
+
+
 def _formatted_response(
     response: DecisionsIRResponse, call: DecisionsCall
 ) -> DecisionsResponse | OpenAIDecisionResponse:
+    caller_request: Final = replace(call.ir_request, questions=call.caller_questions)
+    caller_response: Final = replace(response, answers=_caller_answers(call, response.answers))
     match call.request:
         case DecisionsRequestBody():
-            return ir_to_systemone_response(response, call.ir_request)
+            return ir_to_systemone_response(caller_response, caller_request)
         case OpenAIDecisionRequestBody():
-            return ir_to_openai_response(response, call.ir_request, call.requested_model)
+            return ir_to_openai_response(caller_response, caller_request, call.requested_model)
         case _:
             assert_never(call.request)
 

litellm/decisions/main.py

@@ -1,5 +1,7 @@
+import asyncio
 import warnings
 from collections.abc import Mapping
+from functools import partial
 from typing import Final
 
 import httpx
@@ -27,7 +29,7 @@ _SYSTEMONE_QUESTIONS_ADAPTER: Final[TypeAdapter[SystemOneQuestions | None]] = Ty
 _OPENAI_QUESTIONS_ADAPTER: Final[TypeAdapter[OpenAIDecisionQuestions | None]] = TypeAdapter(
     OpenAIDecisionQuestions | None
 )
-_SYSTEMONE_FIELDS: Final = frozenset({"state", "images"})
+_SYSTEMONE_FIELDS: Final = frozenset({"state", "images", "assets"})
 
 
 def _request_fields(
@@ -55,6 +57,7 @@ def _request_fields(
         state=_STATE_ADAPTER.validate_python(kwargs["state"]),
         questions=_SYSTEMONE_QUESTIONS_ADAPTER.validate_python(questions),
         images=kwargs.get("images"),
+        assets=kwargs.get("assets"),
         audio=audio,
         videos=videos,
     )
@@ -103,19 +106,22 @@ async def adecisions(
     videos: DecisionsVideoClips | None = None,
     **kwargs: object,
 ) -> OpenAIDecisionResponse | DecisionsResponse:
-    call: Final = _prepare_decisions_call(
-        model=model,
-        decision_input=input,
-        questions=questions,
-        safety_identifier=safety_identifier,
-        audio=audio,
-        videos=videos,
-        api_key=api_key,
-        api_base=api_base,
-        timeout=timeout,
-        custom_llm_provider=custom_llm_provider,
-        extra_headers=extra_headers,
-        kwargs=kwargs,
+    call: Final = await asyncio.to_thread(
+        partial(
+            _prepare_decisions_call,
+            model=model,
+            decision_input=input,
+            questions=questions,
+            safety_identifier=safety_identifier,
+            audio=audio,
+            videos=videos,
+            api_key=api_key,
+            api_base=api_base,
+            timeout=timeout,
+            custom_llm_provider=custom_llm_provider,
+            extra_headers=extra_headers,
+            kwargs=kwargs,
+        )
     )
     return await asend(call)
 

litellm/exceptions.py

@@ -10,7 +10,7 @@
 ## LiteLLM versions of the OpenAI Exception Types
 
 import enum
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from typing import Any, Final
 
 import httpx
@@ -313,6 +313,10 @@ class VectorStoreSearchError(BadRequestError):
         )
 
 
+class MalformedRequestFieldError(BadRequestError):
+    pass
+
+
 class UnprocessableEntityError(openai.UnprocessableEntityError):
     def __init__(
         self,
@@ -468,7 +472,7 @@ class RateLimitError(openai.RateLimitError):
         category: str | RateLimitErrorCategory = (RateLimitErrorCategory.VENDOR_RATE_LIMIT),
         rate_limit_type: str | RateLimitType | None = None,
         headers: dict[str, str] | None = None,
-        detail: Any = None,
+        detail: object = None,
         body: object | None = None,
     ):
         self.status_code = 429
@@ -1217,6 +1221,7 @@ class MidStreamFallbackError(ServiceUnavailableError):
         num_retries: int | None = None,
         generated_content: str = "",
         is_pre_first_chunk: bool = False,
+        emitted_disqualifying_content: bool = False,
     ):
         original_status: Final = getattr(original_exception, "status_code", None)
         self.status_code = int(original_status) if original_status is not None else 503
@@ -1229,6 +1234,7 @@ class MidStreamFallbackError(ServiceUnavailableError):
         self.num_retries = num_retries
         self.generated_content = generated_content
         self.is_pre_first_chunk = is_pre_first_chunk
+        self.emitted_disqualifying_content = emitted_disqualifying_content
 
         # Create a response if one wasn't provided
         if response is None:
@@ -1267,6 +1273,10 @@ class MidStreamFallbackError(ServiceUnavailableError):
         self.message = _saved_message
         self.args = (_saved_message,)
 
+        original_headers: Final = getattr(original_exception, "litellm_response_headers", None)
+        if original_headers is not None:
+            self.litellm_response_headers = original_headers
+
     def __str__(self):
         _message = self.message
         if self.num_retries:
@@ -1300,8 +1310,8 @@ class ModifyResponseException(Exception):
         model: str,
         request_data: dict[str, Any],
         guardrail_name: str | None = None,
-        detection_info: dict[str, Any] | None = None,
-        original_response: Any | None = None,
+        detection_info: Mapping[str, object] | None = None,
+        original_response: object = None,
     ):
         self.message = message
         self.model = model
@@ -1334,7 +1344,7 @@ class SensitiveDataRouteException(Exception):
         route_to_model: str,
         session_id: str,
         guardrail_name: str | None = None,
-        detection_info: dict[str, Any] | None = None,
+        detection_info: Mapping[str, object] | None = None,
         message: str | None = None,
         sticky_session_routing: bool = True,
     ):

litellm/experimental_mcp_client/client.py

@@ -5,6 +5,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
 import asyncio
 import base64
 import hashlib
+import inspect
 import json
 import os
 import time
@@ -37,6 +38,9 @@ _TransportContext: TypeAlias = AbstractAsyncContextManager[_TransportStreams]
 _SessionCallbackName: TypeAlias = Literal["sampling_callback", "elicitation_callback", "logging_callback"]
 
 
+SESSION_CLEANUP_TIMEOUT_MESSAGE: Final = "MCP session cleanup timed out"
+
+
 class _SessionCallbackMap(TypedDict):
     sampling_callback: ReadOnly[SamplingFnT | None]
     elicitation_callback: ReadOnly[ElicitationFnT | None]
@@ -185,6 +189,23 @@ def _strip_header_whitespace(headers: dict[str, str]) -> dict[str, str]:
     }
 
 
+async def _receive_session_message(
+    message: ServerNotification | Exception,
+    *,
+    stream_error: asyncio.Future[Exception],
+    on_notification: Callable[[ServerNotification], Awaitable[None]] | None,
+) -> None:
+    if not isinstance(message, Exception):
+        if on_notification is not None:
+            await on_notification(message)
+        return
+    if not isinstance(message, (ValueError, httpx2.HTTPError, OSError)):
+        return
+    if not stream_error.done():
+        stream_error.set_result(message)
+    raise RuntimeError("MCP response stream failed")
+
+
 def _first_non_cancelled_cause(exc: BaseException, cleanup_errors: tuple[Exception, ...] = ()) -> BaseException | None:
     queue: Final[list[BaseException]] = [exc]
     while queue:
@@ -533,6 +554,17 @@ class MCPClient:
     async def discovery_auth_fingerprint(self) -> str:
         return self._hash_discovery_auth(await self.prepare_request_auth(), self._aws_identity())
 
+    async def subscription_auth_fingerprint(self) -> str:
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import (
+            ClientCredentialsBearerAuth,
+        )
+
+        request: Final = await self.prepare_request_auth()
+        auth: Final = self._resolved_auth
+        if isinstance(auth, ClientCredentialsBearerAuth) and auth.subscription_identity is not None:
+            request.headers[auth.header_name] = "managed-grant:" + auth.subscription_identity
+        return self._hash_discovery_auth(request, self._aws_identity())
+
     def _aws_identity(self) -> tuple[str, str, str | None, str | None] | None:
         return self._aws_auth.identity if isinstance(self._aws_auth, MCPSigV4Auth) else None
 
@@ -696,6 +728,7 @@ class MCPClient:
         *,
         on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None,
         on_cleanup: Callable[[], None] | None = None,
+        on_notification: Callable[[ServerNotification], Awaitable[None]] | None = None,
     ) -> TSessionResult:
         """
         Execute an operation within a transport and session context.
@@ -716,16 +749,6 @@ class MCPClient:
                     stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future()
                     _notify_stream_error(on_stream_error, stream_error)
 
-                    async def receive_message(
-                        message: ServerNotification | Exception,
-                    ) -> None:
-                        if not isinstance(message, (ValueError, httpx2.HTTPError, OSError)):
-                            return
-                        if not stream_error.done():
-                            stream_error.set_result(message)
-                        # The SDK closes pending requests when its message handler raises.
-                        raise RuntimeError("MCP response stream failed")
-
                     session_kwargs: Final = (
                         self._session_callbacks_override
                         if self._session_callbacks_override is not None
@@ -737,7 +760,9 @@ class MCPClient:
                         read_stream,
                         write_stream,
                         read_timeout_seconds=self.timeout,
-                        message_handler=receive_message,
+                        message_handler=partial(
+                            _receive_session_message, stream_error=stream_error, on_notification=on_notification
+                        ),
                         sampling_callback=session_kwargs["sampling_callback"],
                         elicitation_callback=session_kwargs["elicitation_callback"],
                         logging_callback=session_kwargs["logging_callback"],
@@ -793,7 +818,7 @@ class MCPClient:
             raise (
                 in_flight_error
                 if in_flight_error is not None
-                else asyncio.CancelledError("MCP session cleanup timed out")
+                else asyncio.CancelledError(SESSION_CLEANUP_TIMEOUT_MESSAGE)
             )
         return result
 
@@ -804,6 +829,7 @@ class MCPClient:
         quiet_on_error: bool = False,
         on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None,
         on_cleanup: Callable[[], None] | None = None,
+        on_notification: Callable[[ServerNotification], Awaitable[None]] | None = None,
     ) -> TSessionResult:
         """Open a session, run the provided coroutine, and clean up.
 
@@ -822,8 +848,9 @@ class MCPClient:
                     http_client=http_client,
                     on_stream_error=on_stream_error,
                     on_cleanup=on_cleanup,
+                    on_notification=on_notification,
                 )
-                if on_stream_error is not None or on_cleanup is not None
+                if on_stream_error is not None or on_cleanup is not None or on_notification is not None
                 else self._execute_session_operation(transport_ctx, operation, http_client=http_client)
             )
             result: Final = await execution
@@ -1094,7 +1121,7 @@ class MCPClient:
     async def call_tool(
         self,
         call_tool_request_params: MCPCallToolRequestParams,
-        host_progress_callback: Callable | None = None,
+        host_progress_callback: Callable[..., Awaitable[None]] | None = None,
         raise_on_error: bool = False,
         allow_input_required: bool = False,
         persistent_session: "PersistentMCPSession | None" = None,
@@ -1126,7 +1153,12 @@ class MCPClient:
             # Forward to Host if callback provided
             if host_progress_callback:
                 try:
-                    await host_progress_callback(progress, total)
+                    try:
+                        inspect.signature(host_progress_callback).bind(progress, total, message)
+                    except (TypeError, ValueError):
+                        await host_progress_callback(progress, total)
+                    else:
+                        await host_progress_callback(progress, total, message)
                 except Exception as e:
                     verbose_logger.warning("Failed to forward to Host: %s", e)
 
@@ -1579,6 +1611,20 @@ def _deliver_persistent_outcome(
     )
 
 
+async def _await_persistent_operation(operation: Awaitable[object], result: asyncio.Future[object]) -> object:
+    running: Final = asyncio.ensure_future(operation)
+
+    def cancel_operation(_: asyncio.Future[object]) -> None:
+        running.cancel()
+
+    result.add_done_callback(cancel_operation)
+    try:
+        (outcome,) = await asyncio.gather(running, return_exceptions=True)
+        return outcome
+    finally:
+        result.remove_done_callback(cancel_operation)
+
+
 class UpstreamSessionClosedError(RuntimeError):
     def __init__(self) -> None:
         super().__init__("upstream MCP session closed")
@@ -1707,12 +1753,14 @@ class PersistentMCPSession:
                     if callbacks is not None:
                         self._current_callbacks = callbacks
                     async with self._admission() if self._admission is not None else nullcontext():
-                        if not future.done():
-                            (outcome,) = await asyncio.gather(operation(session), return_exceptions=True)
-                            if _deliver_persistent_outcome(future, outcome, self._stream_error):
-                                self._active = None
-                                self._ended = True
-                                return
+                        if not future.done() and _deliver_persistent_outcome(
+                            future,
+                            await _await_persistent_operation(operation(session), future),
+                            self._stream_error,
+                        ):
+                            self._active = None
+                            self._ended = True
+                            return
                     self._active = None
                 if self._retiring and self._queue.empty():
                     return

litellm/files/main.py

@@ -11,7 +11,7 @@ import time
 import uuid as uuid_module
 from collections.abc import Coroutine
 from functools import partial
-from typing import Any, Final, Literal, cast
+from typing import Final, Literal, cast
 
 import httpx
 from openai import AsyncOpenAI, OpenAI
@@ -163,7 +163,7 @@ def create_file(
     extra_headers: dict[str, str] | None = None,
     extra_body: dict[str, str] | None = None,
     **kwargs,
-) -> OpenAIFileObject | Coroutine[Any, Any, OpenAIFileObject]:
+) -> OpenAIFileObject | Coroutine[object, None, OpenAIFileObject]:
     """
     Files are used to upload documents that can be used with features like Assistants, Fine-tuning, and Batch API.
 
@@ -259,6 +259,7 @@ def create_file(
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
                 create_file_data=_create_file_request,
+                litellm_params=litellm_params_dict,
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -366,7 +367,15 @@ def file_retrieve(
 
         _is_async: Final = kwargs.pop("is_async", False) is True
 
-        if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
+        if custom_llm_provider == LlmProviders.XAI.value and is_xai_batch_results_id(file_id):
+            response = xai_batch_results_instance.batch_results_file(  # rebind-ok: shares the result slot with the provider branches below
+                _is_async=_is_async,
+                batch_id=file_id,
+                api_base=optional_params.api_base,
+                api_key=optional_params.api_key,
+                timeout=timeout,
+            )
+        elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
             openai_creds: Final = get_openai_credentials(
                 api_base=optional_params.api_base,
                 api_key=optional_params.api_key,
@@ -380,6 +389,7 @@ def file_retrieve(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -561,6 +571,7 @@ def file_delete(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
+                litellm_params=litellm_params_dict,
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -767,6 +778,7 @@ def file_list(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -859,8 +871,8 @@ def file_content(
 ) -> (
     HttpxBinaryResponseContent
     | FileContentStreamingResult
-    | Coroutine[Any, Any, HttpxBinaryResponseContent]
-    | Coroutine[Any, Any, FileContentStreamingResult]
+    | Coroutine[object, None, HttpxBinaryResponseContent]
+    | Coroutine[object, None, FileContentStreamingResult]
 ):
     """
     Returns the contents of the specified file.
@@ -964,6 +976,7 @@ def file_content(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 organization=openai_creds.organization,
+                litellm_params=litellm_params_dict,
             )
         elif custom_llm_provider == "azure":
             azure_creds: Final = get_azure_credentials(
@@ -1085,6 +1098,7 @@ def file_content_streaming(
             organization=openai_creds.organization,
             chunk_size=chunk_size,
             client=client if isinstance(client, (OpenAI, AsyncOpenAI)) else None,
+            litellm_params=optional_params.model_dump(exclude_none=True),
         )
     elif custom_llm_provider == LlmProviders.VERTEX_AI.value:
         if not _is_async:

litellm/fine_tuning/main.py

@@ -239,6 +239,7 @@ def create_fine_tuning_job(
                 client=kwargs.get(
                     "client", None
                 ),  # note, when we add this to `GenericLiteLLMParams` it impacts a lot of other tests + linting
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         # Azure OpenAI
         elif custom_llm_provider == "azure":
@@ -436,6 +437,7 @@ def cancel_fine_tuning_job(
                 max_retries=optional_params.max_retries,
                 _is_async=_is_async,
                 client=kwargs.get("client", None),
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         # Azure OpenAI
         elif custom_llm_provider == "azure":
@@ -591,6 +593,7 @@ def list_fine_tuning_jobs(
                 max_retries=optional_params.max_retries,
                 _is_async=_is_async,
                 client=kwargs.get("client", None),
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         # Azure OpenAI
         elif custom_llm_provider == "azure":
@@ -730,6 +733,7 @@ def retrieve_fine_tuning_job(
                 max_retries=optional_params.max_retries,
                 _is_async=_is_async,
                 client=kwargs.get("client", None),
+                litellm_params=optional_params.model_dump(exclude_none=True),
             )
         # Azure OpenAI
         elif custom_llm_provider == "azure":

litellm/google_genai/adapters/handler.py

@@ -1,7 +1,10 @@
 from collections.abc import AsyncIterator, Coroutine, Mapping
 from typing import Final, cast
 
+from openai.types.chat import ChatCompletionStreamOptionsParam
+
 import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.types.google_genai.adapters import GenerateContentCompletionKwargs
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import ModelResponse
@@ -12,6 +15,23 @@ from .transformation import GoogleGenAIAdapter
 GOOGLE_GENAI_ADAPTER: Final = GoogleGenAIAdapter()
 
 
+def _mark_failure_logged_by_inner_completion(adapter_kwargs: Mapping[str, object]) -> None:
+    outer_logging_obj: Final = adapter_kwargs.get("litellm_logging_obj")
+    if isinstance(outer_logging_obj, LiteLLMLoggingObj):
+        outer_logging_obj.mark_logging_complete(event_type="sync_failure")
+        outer_logging_obj.mark_logging_complete(event_type="async_failure")
+
+
+def _adapter_stream_options(requested: object) -> ChatCompletionStreamOptionsParam | None:
+    match requested:
+        case {"include_usage": True}:
+            return {"include_usage": True}
+        case Mapping():
+            return None
+        case _:
+            return {"include_usage": True}
+
+
 class GenerateContentToCompletionHandler:
     """Handler for transforming generate_content calls to completion format when provider config is None"""
 
@@ -47,9 +67,14 @@ class GenerateContentToCompletionHandler:
                 completion_kwargs["extra_headers"] = extra_kwargs["extra_headers"]
             if "proxy_server_request" in extra_kwargs:
                 completion_kwargs["proxy_server_request"] = extra_kwargs["proxy_server_request"]
+            if "litellm_metadata" in extra_kwargs:
+                completion_kwargs["litellm_metadata"] = extra_kwargs["litellm_metadata"]
 
         if stream:
             completion_kwargs["stream"] = stream
+            stream_options: Final = _adapter_stream_options((extra_kwargs or {}).get("stream_options"))
+            if stream_options is not None:
+                completion_kwargs["stream_options"] = stream_options
 
         return GenerateContentCompletionKwargs(**completion_kwargs)
 
@@ -75,33 +100,33 @@ class GenerateContentToCompletionHandler:
 
         try:
             completion_response: Final = await litellm.acompletion(**completion_kwargs)
+        except Exception:
+            _mark_failure_logged_by_inner_completion(kwargs)
+            raise
 
-            if stream:
-                # Check if completion_response is actually a stream or a ModelResponse
-                # This can happen in error cases or when stream is not properly supported
-                if not hasattr(completion_response, "__aiter__"):
-                    # If it's not a stream, treat it as a regular response
-                    generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
-                        cast(ModelResponse, completion_response)
-                    )
-                    return generate_content_response
-                else:
-                    # Transform streaming completion response to generate_content format
-                    transformed_stream: Final = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
-                        completion_response
-                    )
-                    if transformed_stream is not None:
-                        return transformed_stream
-                    raise ValueError("Failed to transform streaming response")
-            else:
-                # Transform completion response back to generate_content format
+        if stream:
+            # Check if completion_response is actually a stream or a ModelResponse
+            # This can happen in error cases or when stream is not properly supported
+            if not hasattr(completion_response, "__aiter__"):
+                # If it's not a stream, treat it as a regular response
                 generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
                     cast(ModelResponse, completion_response)
                 )
                 return generate_content_response
-
-        except Exception as e:
-            raise ValueError(f"Error calling litellm.acompletion for generate_content: {e}")
+            else:
+                # Transform streaming completion response to generate_content format
+                transformed_stream: Final = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
+                    completion_response
+                )
+                if transformed_stream is not None:
+                    return transformed_stream
+                raise ValueError("Failed to transform streaming response")
+        else:
+            # Transform completion response back to generate_content format
+            generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
+                cast(ModelResponse, completion_response)
+            )
+            return generate_content_response
 
     @staticmethod
     def generate_content_handler(
@@ -136,30 +161,30 @@ class GenerateContentToCompletionHandler:
 
         try:
             completion_response: Final = litellm.completion(**completion_kwargs)
+        except Exception:
+            _mark_failure_logged_by_inner_completion(kwargs)
+            raise
 
-            if stream:
-                # Check if completion_response is actually a stream or a ModelResponse
-                # This can happen in error cases or when stream is not properly supported
-                if not hasattr(completion_response, "__iter__"):
-                    # If it's not a stream, treat it as a regular response
-                    generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
-                        cast(ModelResponse, completion_response)
-                    )
-                    return generate_content_response
-                else:
-                    # Transform streaming completion response to generate_content format
-                    transformed_stream: Final = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
-                        completion_response
-                    )
-                    if transformed_stream is not None:
-                        return transformed_stream
-                    raise ValueError("Failed to transform streaming response")
-            else:
-                # Transform completion response back to generate_content format
+        if stream:
+            # Check if completion_response is actually a stream or a ModelResponse
+            # This can happen in error cases or when stream is not properly supported
+            if not hasattr(completion_response, "__iter__"):
+                # If it's not a stream, treat it as a regular response
                 generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
                     cast(ModelResponse, completion_response)
                 )
                 return generate_content_response
-
-        except Exception as e:
-            raise ValueError(f"Error calling litellm.completion for generate_content: {e}")
+            else:
+                # Transform streaming completion response to generate_content format
+                transformed_stream: Final = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
+                    completion_response
+                )
+                if transformed_stream is not None:
+                    return transformed_stream
+                raise ValueError("Failed to transform streaming response")
+        else:
+            # Transform completion response back to generate_content format
+            generate_content_response = GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
+                cast(ModelResponse, completion_response)
+            )
+            return generate_content_response

litellm/google_genai/adapters/transformation.py

@@ -158,8 +158,17 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
         self.sent_first_chunk = False
         self.accumulated_tool_calls = dict[int, _ToolCallAccumulator]()
         self._returned_response = False
+        self.finish_chunk_awaiting_usage: Mapping[str, object] | None = None
         super().__init__(completion_stream)
 
+    def hold_finish_chunk_for_usage(self, finish_chunk: Mapping[str, object]) -> None:
+        self.finish_chunk_awaiting_usage = finish_chunk
+
+    def take_finish_chunk_awaiting_usage(self) -> Mapping[str, object] | None:
+        finish_chunk: Final = self.finish_chunk_awaiting_usage
+        self.finish_chunk_awaiting_usage = None
+        return finish_chunk
+
     def __next__(self):
         try:
             if not hasattr(self.completion_stream, "__iter__"):
@@ -176,6 +185,9 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
                 if transformed_chunk:
                     return transformed_chunk
 
+            finish_chunk: Final = self.take_finish_chunk_awaiting_usage()
+            if finish_chunk:
+                return finish_chunk
             raise StopIteration
         except StopIteration:
             raise
@@ -198,6 +210,10 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
                 if transformed_chunk:
                     return transformed_chunk
 
+            finish_chunk: Final = self.take_finish_chunk_awaiting_usage()
+            if finish_chunk:
+                return finish_chunk
+
             # After the stream is exhausted, check for any remaining accumulated tool calls
             if self.accumulated_tool_calls:
                 try:
@@ -286,6 +302,10 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
                 else:
                     yield str(chunk).encode()
 
+        finish_chunk: Final = self.take_finish_chunk_awaiting_usage()
+        if finish_chunk:
+            yield f"data: {json.dumps(finish_chunk)}\n\n".encode()
+
 
 class GoogleGenAIAdapter:
     """Adapter for transforming Google GenAI generate_content requests to/from litellm.completion format"""
@@ -657,8 +677,7 @@ class GoogleGenAIAdapter:
         # Extract the main response content from streaming chunk
         choice: Final = response.choices[0] if response.choices else None
         if not choice:
-            # Return empty chunk if no choices
-            return None
+            return self._usage_only_streaming_chunk(response, wrapper)
 
         # Handle streaming choice
         if isinstance(choice, StreamingChoices):
@@ -675,7 +694,7 @@ class GoogleGenAIAdapter:
 
         # Only create response chunk if we have parts or it's the final chunk
         if not parts and not finish_reason:
-            return None
+            return self._usage_only_streaming_chunk(response, wrapper)
 
         # Create Google GenAI streaming format response
         streaming_chunk: Final[_JsonDict] = {
@@ -689,20 +708,34 @@ class GoogleGenAIAdapter:
             ]
         }
 
-        # Add usage metadata only in the final chunk (when finish_reason is present)
-        if finish_reason:
-            usage_metadata: Final = (
-                self._map_usage(getattr(response, "usage", None))
-                if hasattr(response, "usage") and getattr(response, "usage", None)
-                else {
-                    "promptTokenCount": 0,
-                    "candidatesTokenCount": 0,
-                    "totalTokenCount": 0,
-                }
-            )
-            streaming_chunk["usageMetadata"] = usage_metadata
+        if not finish_reason:
+            return streaming_chunk
 
-        return streaming_chunk
+        finish_usage: Final[object] = getattr(response, "usage", None)
+        if finish_usage:
+            return {**streaming_chunk, "usageMetadata": self._map_usage(finish_usage)}
+
+        wrapper.hold_finish_chunk_for_usage(
+            {
+                **streaming_chunk,
+                "usageMetadata": {"promptTokenCount": 0, "candidatesTokenCount": 0, "totalTokenCount": 0},
+            }
+        )
+        return None
+
+    def _usage_only_streaming_chunk(
+        self, response: ModelResponse | ModelResponseStream, wrapper: GoogleGenAIStreamWrapper
+    ) -> Mapping[str, object] | None:
+        usage: Final[object] = getattr(response, "usage", None)
+        if not usage:
+            return None
+        finish_chunk: Final = wrapper.take_finish_chunk_awaiting_usage()
+        if finish_chunk:
+            return {**finish_chunk, "usageMetadata": self._map_usage(usage)}
+        return {
+            "candidates": [{"content": {"parts": [{"text": ""}], "role": "model"}, "index": 0}],
+            "usageMetadata": self._map_usage(usage),
+        }
 
     def _transform_openai_message_to_google_genai_parts(
         self,

litellm/google_genai/main.py

@@ -1,11 +1,10 @@
 import asyncio
 import contextvars
-from collections.abc import Iterator
 from functools import partial
 from typing import TYPE_CHECKING, Any, ClassVar, Final
 
 import httpx
-from pydantic import ConfigDict
+from pydantic import ConfigDict, TypeAdapter
 
 import litellm
 from litellm.constants import request_timeout
@@ -39,6 +38,7 @@ else:
 ####### ENVIRONMENT VARIABLES ###################
 # Initialize any necessary instances or variables here
 base_llm_http_handler = BaseLLMHTTPHandler()
+_GENERATE_CONTENT_RESULT: Final = TypeAdapter(object)
 #################################################
 
 
@@ -145,6 +145,7 @@ class GenerateContentHelper:
             ProviderConfigManager.get_provider_google_genai_generate_content_config(
                 model=model,
                 provider=litellm.LlmProviders(custom_llm_provider),
+                litellm_params=litellm_params.model_dump(include={"base_model", "api_base"}),
             )
         )
 
@@ -247,7 +248,7 @@ async def agenerate_content(
     # LiteLLM specific params,
     custom_llm_provider: str | None = None,
     **kwargs,
-) -> Any:
+) -> object:
     """
     Async: Generate content using Google GenAI
     """
@@ -285,7 +286,7 @@ async def agenerate_content(
         init_response: Final = await loop.run_in_executor(None, func_with_context)
 
         if asyncio.iscoroutine(init_response):
-            response = await init_response
+            response = _GENERATE_CONTENT_RESULT.validate_python(await init_response)
         else:
             response = init_response
 
@@ -315,7 +316,7 @@ def generate_content(
     # LiteLLM specific params,
     custom_llm_provider: str | None = None,
     **kwargs,
-) -> Any:
+) -> object:
     """
     Generate content using Google GenAI
     """
@@ -361,22 +362,24 @@ def generate_content(
             )
 
         # Call the standard handler
-        response: Final = base_llm_http_handler.generate_content_handler(
-            model=setup_result.model,
-            contents=contents,
-            tools=tools,
-            generate_content_provider_config=setup_result.generate_content_provider_config,
-            generate_content_config_dict=setup_result.generate_content_config_dict,
-            custom_llm_provider=setup_result.custom_llm_provider,
-            litellm_params=setup_result.litellm_params,
-            logging_obj=setup_result.litellm_logging_obj,
-            extra_headers=extra_headers,
-            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"),
-            litellm_metadata=kwargs.get("litellm_metadata", {}),
-            system_instruction=system_instruction,
+        response: Final = _GENERATE_CONTENT_RESULT.validate_python(
+            base_llm_http_handler.generate_content_handler(
+                model=setup_result.model,
+                contents=contents,
+                tools=tools,
+                generate_content_provider_config=setup_result.generate_content_provider_config,
+                generate_content_config_dict=setup_result.generate_content_config_dict,
+                custom_llm_provider=setup_result.custom_llm_provider,
+                litellm_params=setup_result.litellm_params,
+                logging_obj=setup_result.litellm_logging_obj,
+                extra_headers=extra_headers,
+                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"),
+                litellm_metadata=kwargs.get("litellm_metadata", {}),
+                system_instruction=system_instruction,
+            )
         )
 
         return response
@@ -405,7 +408,7 @@ async def agenerate_content_stream(
     # LiteLLM specific params,
     custom_llm_provider: str | None = None,
     **kwargs,
-) -> Any:
+) -> object:
     """
     Async: Generate content using Google GenAI with streaming response
     """
@@ -456,23 +459,25 @@ async def agenerate_content_stream(
 
         # Call the handler with async enabled and streaming
         # Return the coroutine directly for the router to handle
-        return await base_llm_http_handler.generate_content_handler(
-            model=setup_result.model,
-            contents=contents,
-            generate_content_provider_config=setup_result.generate_content_provider_config,
-            generate_content_config_dict=setup_result.generate_content_config_dict,
-            tools=tools,
-            custom_llm_provider=setup_result.custom_llm_provider,
-            litellm_params=setup_result.litellm_params,
-            logging_obj=setup_result.litellm_logging_obj,
-            extra_headers=extra_headers,
-            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"),
-            stream=True,
-            litellm_metadata=kwargs.get("litellm_metadata", {}),
-            system_instruction=system_instruction,
+        return _GENERATE_CONTENT_RESULT.validate_python(
+            await base_llm_http_handler.generate_content_handler(
+                model=setup_result.model,
+                contents=contents,
+                generate_content_provider_config=setup_result.generate_content_provider_config,
+                generate_content_config_dict=setup_result.generate_content_config_dict,
+                tools=tools,
+                custom_llm_provider=setup_result.custom_llm_provider,
+                litellm_params=setup_result.litellm_params,
+                logging_obj=setup_result.litellm_logging_obj,
+                extra_headers=extra_headers,
+                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"),
+                stream=True,
+                litellm_metadata=kwargs.get("litellm_metadata", {}),
+                system_instruction=system_instruction,
+            )
         )
 
     except Exception as e:
@@ -500,7 +505,7 @@ def generate_content_stream(
     # LiteLLM specific params,
     custom_llm_provider: str | None = None,
     **kwargs,
-) -> Iterator[Any]:
+) -> object:
     """
     Generate content using Google GenAI with streaming response
     """
@@ -545,23 +550,25 @@ def generate_content_stream(
             )
 
         # Call the handler with streaming enabled (sync version)
-        return base_llm_http_handler.generate_content_handler(
-            model=setup_result.model,
-            contents=contents,
-            generate_content_provider_config=setup_result.generate_content_provider_config,
-            generate_content_config_dict=setup_result.generate_content_config_dict,
-            tools=tools,
-            custom_llm_provider=setup_result.custom_llm_provider,
-            litellm_params=setup_result.litellm_params,
-            logging_obj=setup_result.litellm_logging_obj,
-            extra_headers=extra_headers,
-            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,
+        return _GENERATE_CONTENT_RESULT.validate_python(
+            base_llm_http_handler.generate_content_handler(
+                model=setup_result.model,
+                contents=contents,
+                generate_content_provider_config=setup_result.generate_content_provider_config,
+                generate_content_config_dict=setup_result.generate_content_config_dict,
+                tools=tools,
+                custom_llm_provider=setup_result.custom_llm_provider,
+                litellm_params=setup_result.litellm_params,
+                logging_obj=setup_result.litellm_logging_obj,
+                extra_headers=extra_headers,
+                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:

litellm/google_genai/streaming_iterator.py

@@ -67,7 +67,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
         request_body: dict,
         model: str,
         custom_llm_provider: str,
-        hidden_params: dict[str, Any] | None = None,
+        hidden_params: dict[str, object] | None = None,
     ):
         self.litellm_logging_obj = litellm_logging_obj
         self.request_body = request_body
@@ -118,7 +118,7 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent
         litellm_metadata: dict,
         custom_llm_provider: str,
         request_body: dict | None = None,
-        hidden_params: dict[str, Any] | None = None,
+        hidden_params: dict[str, object] | None = None,
     ):
         super().__init__(
             litellm_logging_obj=logging_obj,
@@ -169,7 +169,7 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
         litellm_metadata: dict,
         custom_llm_provider: str,
         request_body: dict | None = None,
-        hidden_params: dict[str, Any] | None = None,
+        hidden_params: dict[str, object] | None = None,
     ):
         super().__init__(
             litellm_logging_obj=logging_obj,

litellm/harness/handlers/deepagents_handler.py

@@ -13,8 +13,8 @@ import asyncio
 import importlib
 from collections.abc import AsyncIterator, Callable, Mapping
 from dataclasses import dataclass
-from types import MappingProxyType, ModuleType
-from typing import TYPE_CHECKING, Any
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Protocol
 
 from litellm.harness.context import SessionContext
 from litellm.harness.errors import HarnessError, HarnessInstallFailed
@@ -51,10 +51,17 @@ if TYPE_CHECKING:
     from langgraph.types import Command
 
     from litellm.llms.base_llm.harness.transformation import BaseHarnessConfig
+    from litellm.llms.deepagents.harness.sandbox_backend import SandboxBackend, ToolFilterMiddleware, UsageCallback
 
 _MODEL_NODE = "model"
 
 
+class _SandboxBackendModule(Protocol):
+    SandboxBackend: type[SandboxBackend]
+    UsageCallback: type[UsageCallback]
+    ToolFilterMiddleware: type[ToolFilterMiddleware]
+
+
 @dataclass(frozen=True)
 class DeepAgentsDeps:
     """The optional-dependency entrypoints this handler uses."""
@@ -65,7 +72,7 @@ class DeepAgentsDeps:
     command_cls: Callable[..., Command]
     subagent_defaults: Mapping[str, object]
     convert_to_openai_messages: Any
-    backend: ModuleType
+    backend: _SandboxBackendModule
 
 
 def load_deps() -> DeepAgentsDeps:
@@ -77,7 +84,7 @@ def load_deps() -> DeepAgentsDeps:
         memory = importlib.import_module("langgraph.checkpoint.memory")
         lg_types = importlib.import_module("langgraph.types")
         messages = importlib.import_module("langchain_core.messages")
-        backend = importlib.import_module("litellm.llms.deepagents.harness.sandbox_backend")
+        from litellm.llms.deepagents.harness import sandbox_backend as backend
     except ImportError as e:
         raise HarnessInstallFailed(f"{INSTALL_HINT} ({e})") from e
     return DeepAgentsDeps(

litellm/images/main.py

@@ -5,6 +5,8 @@ from collections.abc import Coroutine
 from functools import partial
 from typing import TYPE_CHECKING, Final, Literal, Optional, cast, overload
 
+from pydantic import TypeAdapter, ValidationError
+
 if TYPE_CHECKING:
     from litellm.images.utils import ImageEditRequestUtils
 
@@ -17,6 +19,9 @@ from litellm import client
 from litellm.constants import DEFAULT_IMAGE_ENDPOINT_MODEL
 from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT
 from litellm.exceptions import LiteLLMUnknownProvider
+from litellm.litellm_core_utils.get_llm_provider_logic import (
+    invalid_provider_routing_hints_message,
+)
 from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.llm_request_utils import flatten_form_field_values
@@ -34,6 +39,7 @@ from openai.types.audio.transcription_create_params import FileTypes
 # BFL handlers
 from litellm.llms.black_forest_labs.image_edit.handler import bfl_image_edit
 from litellm.llms.black_forest_labs.image_generation.handler import bfl_image_generation
+from litellm.llms.openai.workload_identity import resolve_openai_workload_identity_config
 from litellm.main import (
     azure_chat_completions,
     base_llm_aiohttp_handler,
@@ -52,6 +58,7 @@ from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import (
     LITELLM_IMAGE_VARIATION_PROVIDERS,
     LlmProviders,
+    ProviderRoutingHints,
 )
 from litellm.utils import (
     ImageResponse,
@@ -60,6 +67,8 @@ from litellm.utils import (
     get_optional_params_image_gen,
 )
 
+_PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints)
+
 # Cache for ImageEditRequestUtils to avoid repeated __getattr__ calls
 _ImageEditRequestUtils_cache: Optional["ImageEditRequestUtils"] = None
 
@@ -92,6 +101,12 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse:
     model: Final = args[0] if len(args) > 0 else kwargs["model"]
     ### PASS ARGS TO Image Generation ###
     kwargs["aimg_generation"] = True
+    try:
+        provider_hints: Final = _PROVIDER_ROUTING_HINTS.validate_python(kwargs)
+    except ValidationError as error:
+        raise litellm.BadRequestError(
+            message=invalid_provider_routing_hints_message(error), model=model, llm_provider=""
+        ) from error
     custom_llm_provider = None
     try:
         # Use a partial function to pass your keyword arguments
@@ -101,7 +116,11 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse:
         ctx: Final = contextvars.copy_context()
         func_with_context: Final = partial(ctx.run, func)
 
-        _, custom_llm_provider, _, _ = get_llm_provider(model=model, api_base=kwargs.get("api_base", None))
+        _, custom_llm_provider, _, _ = get_llm_provider(
+            model=model,
+            custom_llm_provider=provider_hints.get("custom_llm_provider"),
+            api_base=provider_hints.get("api_base"),
+        )
 
         # Await normally
         init_response: Final = await loop.run_in_executor(None, func_with_context)
@@ -487,6 +506,7 @@ def image_generation(
                 aimg_generation=aimg_generation,
                 client=client,
                 headers=headers,
+                litellm_params=litellm_params_dict,
             )
         elif custom_llm_provider == "bedrock":
             if model is None:
@@ -663,10 +683,13 @@ def image_variation(
     api_base = provider_config.get_api_base(litellm_params.get("api_base", None))
 
     if image_variation_provider == LITELLM_IMAGE_VARIATION_PROVIDERS.OPENAI:
-        if api_key is None:
-            raise ValueError("API key is required for OpenAI image variations")
         if api_base is None:
             raise ValueError("API base is required for OpenAI image variations")
+        if api_key is None and (
+            resolve_openai_workload_identity_config(api_key=api_key, api_base=api_base, litellm_params=litellm_params)
+            is None
+        ):
+            raise ValueError("API key is required for OpenAI image variations")
 
         response = openai_image_variations.image_variations(
             model_response=model_response,

litellm/integrations/anthropic_cache_control_hook.py

@@ -10,9 +10,10 @@ Supported for both `v1/chat/completions` (via the prompt-management hook) and
 """
 
 import copy
+import itertools
 import os
 import re
-from collections.abc import Callable, Iterable, Mapping, Sequence
+from collections.abc import Callable, Container, Iterable, Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, cast
 from urllib.parse import urlparse
 
@@ -26,6 +27,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
     cache_control_as_prompt_cache_breakpoint,
     with_prompt_cache_breakpoint,
 )
+from litellm.litellm_core_utils.prompt_templates.factory import (
+    tool_call_ids_answered_by_client,
+    tool_call_ids_answered_in,
+)
 from litellm.llms.anthropic.common_utils import (
     is_claude_code_one_shot_subagent_request,
     supports_anthropic_cache_control,
@@ -38,6 +43,7 @@ from litellm.types.integrations.anthropic_cache_control_hook import (
     CacheControlMessageInjectionPoint,
 )
 from litellm.types.llms.anthropic import (
+    ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX,
     ANTHROPIC_TOOL_SEARCH_TOOL_TYPES,
     AllAnthropicToolsValues,
     AnthropicSystemMessageContent,
@@ -74,6 +80,7 @@ OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE")
 _OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
 _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
 _CONTENT_PART_ADAPTER: Final = TypeAdapter(dict[str, object])
+_ATTRIBUTE_VALUE: Final = TypeAdapter(object)
 
 AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues
 
@@ -215,7 +222,7 @@ def _carries_cache_breakpoint(block: object) -> bool:
 
 def _attribute_or_key(value: object, key: str) -> object | None:
     if hasattr(value, key):
-        return getattr(value, key)
+        return _ATTRIBUTE_VALUE.validate_python(getattr(value, key))
     if isinstance(value, Mapping):
         return value.get(key)
     return None
@@ -227,12 +234,83 @@ def _as_object_list(value: object | None) -> list[object] | None:
     return _validated_object_list(value)
 
 
-def _tool_call_carries_cache_breakpoint(tool_call: object, message: object) -> bool:
+def _tool_call_carries_cache_breakpoint(tool_call: object, message: object, answered_tool_ids: Container[str]) -> bool:
     if _attribute_or_key(tool_call, "cache_control") is None:
         return False
 
     return not tool_call_is_rebuilt_as_server_tool_use(
-        _attribute_or_key(tool_call, "id"), _attribute_or_key(message, "provider_specific_fields")
+        _attribute_or_key(tool_call, "id"), _attribute_or_key(message, "provider_specific_fields"), answered_tool_ids
+    )
+
+
+def _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(
+    messages: Sequence[object], model: str | None = None, custom_llm_provider: str | None = None
+) -> frozenset[str]:
+    if not _marks_a_server_tool_call(messages):
+        return frozenset()
+    if model is not None and not _builds_anthropic_messages(model, custom_llm_provider):
+        return tool_call_ids_answered_in(messages)
+    return tool_call_ids_answered_by_client(messages, _keeps_mid_conversation_system(model, custom_llm_provider))
+
+
+def _marks_a_server_tool_call(messages: Sequence[object]) -> bool:
+    tool_calls: Final = itertools.chain.from_iterable(
+        _as_object_list(_attribute_or_key(message, "tool_calls")) or () for message in messages
+    )
+    return any(_is_marked_server_tool_call(tool_call) for tool_call in tool_calls)
+
+
+def _builds_anthropic_messages(model: str, custom_llm_provider: str | None) -> bool:
+    from litellm.exceptions import BadRequestError
+    from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, get_llm_provider
+    from litellm.llms.anthropic.chat.transformation import AnthropicConfig
+    from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
+    from litellm.types.utils import LlmProviders
+    from litellm.utils import ProviderConfigManager
+
+    if declared_authenticating_provider(model, custom_llm_provider) is not None:
+        return False
+    try:
+        served_model, provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider)
+        chat_config: Final = ProviderConfigManager.get_provider_chat_config(
+            model=served_model, provider=LlmProviders(provider)
+        )
+    except (BadRequestError, ValueError):
+        return True
+    return isinstance(chat_config, AnthropicConfig) and not isinstance(chat_config, OpenAIGPTConfig)
+
+
+def _custom_llm_provider(non_default_params: Mapping[str, object], point: Mapping[str, object]) -> str | None:
+    provider: Final = non_default_params.get("custom_llm_provider") or point.get(CUSTOM_LLM_PROVIDER_STAMP)
+    return provider if isinstance(provider, str) else None
+
+
+def _keeps_mid_conversation_system(model: str | None, custom_llm_provider: str | None) -> bool:
+    if model is None:
+        return True
+    from litellm.utils import supports_mid_conversation_system
+
+    if supports_mid_conversation_system(model=model, custom_llm_provider=custom_llm_provider):
+        return True
+    return _model_info_flags_mid_conversation_system(model, custom_llm_provider)
+
+
+def _model_info_flags_mid_conversation_system(model: str, custom_llm_provider: str | None) -> bool:
+    from litellm.utils import get_model_info
+
+    try:
+        model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
+    except Exception:
+        return False
+    return model_info.get("supports_mid_conversation_system") is True
+
+
+def _is_marked_server_tool_call(tool_call: object) -> bool:
+    tool_call_id: Final = _attribute_or_key(tool_call, "id")
+    return (
+        isinstance(tool_call_id, str)
+        and tool_call_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX)
+        and _attribute_or_key(tool_call, "cache_control") is not None
     )
 
 
@@ -261,6 +339,7 @@ def _may_carry_prompt_cache_breakpoint(message: Mapping[str, object]) -> bool:
 CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_points"
 
 EXTERNAL_BREAKPOINTS_STAMP: Final = "_litellm_external_breakpoints"
+CUSTOM_LLM_PROVIDER_STAMP: Final = "_litellm_custom_llm_provider"
 
 
 class AnthropicCacheControlHook(CustomPromptManagement):
@@ -334,13 +413,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         )
         remaining_points: Final = tuple(point for point in injection_points if point.get("location") != "message")
 
+        custom_llm_provider: Final = _custom_llm_provider(non_default_params, injection_points[0])
         stamped_dialect: Final = injection_points[0].get("_litellm_openai_dialect")
         openai_dialect: Final = (
             stamped_dialect
             if isinstance(stamped_dialect, bool)
             else AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(
                 model,
-                non_default_params.get("custom_llm_provider"),
+                custom_llm_provider,
                 non_default_params.get("api_base") or non_default_params.get("base_url"),
                 non_default_params.get("prompt_cache_options"),
             )
@@ -363,16 +443,25 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         reserved_blocks: Final = AnthropicCacheControlHook._blocks_reserved_outside_messages(
             remaining_points, external_breakpoints, openai_dialect
         )
-        breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages)
+        answered_tool_ids: Final = _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(
+            processed_messages, model, custom_llm_provider
+        )
+        breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(
+            processed_messages, answered_tool_ids=answered_tool_ids
+        )
         processed_messages = self._apply_message_injections(
             points=applied_message_points,
             messages=processed_messages,
             max_blocks=MAX_CACHE_CONTROL_BLOCKS - reserved_blocks,
             openai_dialect=openai_dialect,
+            answered_tool_ids=answered_tool_ids,
         )
         if (
             openai_dialect
-            and AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages) > breakpoints_before
+            and AnthropicCacheControlHook.count_request_cache_breakpoints(
+                processed_messages, answered_tool_ids=answered_tool_ids
+            )
+            > breakpoints_before
             and non_default_params.get("prompt_cache_options") is None
         ):
             non_default_params["prompt_cache_options"] = PromptCacheOptions(mode="implicit")
@@ -385,7 +474,10 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         carried_points: Final[Sequence[CacheControlInjectionPoint]] = (
             *AnthropicCacheControlHook._points_with_a_slot_left(
                 remaining_points,
-                AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages) + external_breakpoints,
+                AnthropicCacheControlHook.count_request_cache_breakpoints(
+                    processed_messages, answered_tool_ids=answered_tool_ids
+                )
+                + external_breakpoints,
                 openai_dialect,
             ),
             *carried_message_points,
@@ -420,11 +512,21 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         return _served_provider_for_model(model)
 
     @staticmethod
-    def count_request_cache_breakpoints(messages: Iterable[object], system: object = None) -> int:
+    def count_request_cache_breakpoints(
+        messages: Iterable[object], system: object = None, answered_tool_ids: frozenset[str] | None = None
+    ) -> int:
         system_blocks: Final = (
             sum(1 for block in system if _carries_cache_breakpoint(block)) if isinstance(system, list) else 0
         )
-        return system_blocks + sum(AnthropicCacheControlHook._count_cache_control_blocks(msg) for msg in messages)
+        message_list: Final = tuple(messages)
+        client_answers: Final = (
+            _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(message_list)
+            if answered_tool_ids is None
+            else answered_tool_ids
+        )
+        return system_blocks + sum(
+            AnthropicCacheControlHook._count_cache_control_blocks(msg, client_answers) for msg in message_list
+        )
 
     @staticmethod
     def count_external_cache_breakpoints(
@@ -498,6 +600,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         messages: list[AllMessageValues],
         max_blocks: int,
         openai_dialect: bool = False,
+        answered_tool_ids: frozenset[str] | None = None,
     ) -> list[AllMessageValues]:
         """Apply message-level cache control injection points in order.
 
@@ -508,7 +611,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         ``max_blocks`` is reached. Injection points are honored in config order,
         so earlier points win when slots are scarce.
         """
-        used_blocks = AnthropicCacheControlHook.count_request_cache_breakpoints(messages)
+        client_answers: Final = (
+            _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(messages)
+            if answered_tool_ids is None
+            else answered_tool_ids
+        )
+        used_blocks = AnthropicCacheControlHook.count_request_cache_breakpoints(
+            messages, answered_tool_ids=client_answers
+        )
 
         limit_reached = False
         for point in points:
@@ -525,14 +635,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
                     limit_reached = True
                     break
 
-                if AnthropicCacheControlHook._message_has_cache_control(messages[target_index]):
+                if AnthropicCacheControlHook._message_has_cache_control(messages[target_index], client_answers):
                     # Client already marked this message; don't overwrite it.
                     continue
 
                 messages[target_index] = AnthropicCacheControlHook._safe_insert_cache_control_in_message(
                     messages[target_index], control, openai_dialect
                 )
-                if AnthropicCacheControlHook._message_has_cache_control(messages[target_index]):
+                if AnthropicCacheControlHook._message_has_cache_control(messages[target_index], client_answers):
                     used_blocks += 1
 
             if limit_reached:
@@ -586,22 +696,26 @@ class AnthropicCacheControlHook(CustomPromptManagement):
         return []
 
     @staticmethod
-    def _count_cache_control_blocks(message: object) -> int:
+    def _count_cache_control_blocks(message: object, answered_tool_ids: Container[str] = frozenset()) -> int:
         message_count: Final = 1 if _carries_cache_breakpoint(message) else 0
         content: Final = _as_object_list(_attribute_or_key(message, "content"))
         content_count: Final = sum(1 for block in content if _carries_cache_breakpoint(block)) if content else 0
         tool_calls: Final = _as_object_list(_attribute_or_key(message, "tool_calls"))
         tool_call_count: Final = (
-            sum(1 for tool_call in tool_calls if _tool_call_carries_cache_breakpoint(tool_call, message))
+            sum(
+                1
+                for tool_call in tool_calls
+                if _tool_call_carries_cache_breakpoint(tool_call, message, answered_tool_ids)
+            )
             if tool_calls
             else 0
         )
         return message_count + content_count + tool_call_count
 
     @staticmethod
-    def _message_has_cache_control(message: AllMessageValues) -> bool:
+    def _message_has_cache_control(message: AllMessageValues, answered_tool_ids: Container[str]) -> bool:
         """Return True if the message already carries any cache_control."""
-        return AnthropicCacheControlHook._count_cache_control_blocks(message) > 0
+        return AnthropicCacheControlHook._count_cache_control_blocks(message, answered_tool_ids) > 0
 
     @staticmethod
     def _safe_insert_cache_control_in_message(
@@ -973,7 +1087,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
 
     @staticmethod
     def maybe_seed_default_injection_points(
-        non_default_params: dict[str, Any],
+        non_default_params: dict[str, object],
         messages: list[AllMessageValues],
         model: str,
         custom_llm_provider: str | None,
@@ -997,7 +1111,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
             tools_keeping_marks: Final = tuple(
                 tool for tool in tools or () if not _chat_transform_drops_tool_cache_control(tool)
             )
-            non_default_params["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_for_prompt_hook(
+            stamped: Final = AnthropicCacheControlHook._stamped_for_prompt_hook(
                 configured,
                 AnthropicCacheControlHook.count_external_cache_breakpoints(
                     tools_keeping_marks, non_default_params.get("cache_control"), non_default_params
@@ -1007,6 +1121,11 @@ class AnthropicCacheControlHook(CustomPromptManagement):
                 api_base,
                 non_default_params.get("prompt_cache_options"),
             )
+            non_default_params["cache_control_injection_points"] = (
+                AnthropicCacheControlHook._stamped(stamped, CUSTOM_LLM_PROVIDER_STAMP, custom_llm_provider)
+                if custom_llm_provider is not None and _marks_a_server_tool_call(messages)
+                else stamped
+            )
             return
         points: Final = AnthropicCacheControlHook.get_default_injection_points(
             messages=messages,
@@ -1099,7 +1218,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
     def maybe_inject_cache_control(
         messages: list[dict],
         system: str | list | None,
-        kwargs: dict[str, Any],
+        kwargs: dict[str, object],
         model: str | None = None,
         custom_llm_provider: str | None = None,
         tools: list[dict] | None = None,

litellm/integrations/arize/arize_phoenix.py

@@ -1,6 +1,7 @@
 import os
 import threading
 from collections import OrderedDict
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, Final
 
 from litellm._logging import verbose_logger
@@ -22,7 +23,7 @@ if TYPE_CHECKING:
 
     Protocol = _Protocol
     OpenTelemetryConfig = _OpenTelemetryConfig
-    Span = _Span | Any
+    Span = _Span
     OpenTelemetry = _OpenTelemetry
     LITELLM_TRACER_NAME: str
 else:
@@ -100,7 +101,7 @@ class ArizePhoenixLogger(OpenTelemetry):
         if getattr(self, "_use_injected_tracer_provider", False):
             return
 
-        shared_processor: Final = getattr(self, "_shared_span_processor", None)
+        shared_processor: Final[SpanProcessor | None] = getattr(self, "_shared_span_processor", None)
         if shared_processor is not None:
             try:
                 shared_processor.force_flush()
@@ -111,7 +112,8 @@ class ArizePhoenixLogger(OpenTelemetry):
                 )
 
         with getattr(self, "_project_providers_lock", threading.Lock()):
-            providers: Final = list(getattr(self, "_project_providers", {}).values())
+            project_providers: Final[Mapping[str, TracerProvider]] = getattr(self, "_project_providers", {})
+            providers: Final = list(project_providers.values())
 
         for provider in providers:
             try:
@@ -193,7 +195,7 @@ class ArizePhoenixLogger(OpenTelemetry):
         primary ``otel`` callback which handles proxy-level parent spans.
         """
 
-    def set_attributes(self, span: Span, kwargs, response_obj: Any | None):
+    def set_attributes(self, span: Span, kwargs, response_obj: object):
         ArizePhoenixLogger.set_arize_phoenix_attributes(span, kwargs, response_obj)
 
     @staticmethod