Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_internal_copy_36281

.github/workflows/test-rust.yml

@@ -9,6 +9,7 @@ on:
       - "rust-toolchain.toml"
       - ".github/scripts/smoke_test_native_wheel.py"
       - ".github/scripts/verify_linux_native_wheel.py"
+      - "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
       - ".github/workflows/test-rust.yml"
   pull_request:
     branches:
@@ -23,6 +24,7 @@ on:
       - "rust-toolchain.toml"
       - ".github/scripts/smoke_test_native_wheel.py"
       - ".github/scripts/verify_linux_native_wheel.py"
+      - "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
       - ".github/workflows/test-rust.yml"
 
 permissions:
@@ -121,3 +123,6 @@ jobs:
         env:
           RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
         run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
+
+      - name: Test native route wheel
+        run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl

.github/workflows/test-unit.yml

@@ -96,6 +96,7 @@ jobs:
           - shard: misc
             artifact-name: misc
             test-path: >-
+              tests/sdk_function_trace
               tests/test_litellm/batches
               tests/test_litellm/secret_managers
               tests/test_litellm/a2a_protocol

basedpyright-code-budget.json

@@ -3,7 +3,7 @@
     "limit": 14074
   },
   "reportArgumentType": {
-    "limit": 2216
+    "limit": 2215
   },
   "reportAssignmentType": {
     "limit": 319
@@ -24,7 +24,7 @@
     "limit": 19
   },
   "reportExplicitAny": {
-    "limit": 4125
+    "limit": 4124
   },
   "reportFunctionMemberAccess": {
     "limit": 7
@@ -57,7 +57,7 @@
     "limit": 5601
   },
   "reportMissingTypeArgument": {
-    "limit": 15306
+    "limit": 15290
   },
   "reportMissingTypeStubs": {
     "limit": 40
@@ -99,31 +99,31 @@
     "limit": 0
   },
   "reportUnknownArgumentType": {
-    "limit": 44364
+    "limit": 44362
   },
   "reportUnknownLambdaType": {
     "limit": 109
   },
   "reportUnknownMemberType": {
-    "limit": 38350
+    "limit": 38332
   },
   "reportUnknownParameterType": {
     "limit": 19625
   },
   "reportUnknownVariableType": {
-    "limit": 29877
+    "limit": 29861
   },
   "reportUnnecessaryCast": {
     "limit": 111
   },
   "reportUnnecessaryComparison": {
-    "limit": 692
+    "limit": 687
   },
   "reportUnnecessaryContains": {
-    "limit": 5
+    "limit": 4
   },
   "reportUnnecessaryIsInstance": {
-    "limit": 826
+    "limit": 823
   },
   "reportUntypedBaseClass": {
     "limit": 0

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -43,6 +43,8 @@ def _get_prisma_env() -> dict:
 
 _MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
 
+_MIGRATION_DEADLOCK_MARKER = "deadlock detected"
+
 _SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
 _SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(
     r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE
@@ -265,6 +267,50 @@ class ProxyExtrasDBManager:
             env=prisma_env,
         )
 
+    @staticmethod
+    def _roll_back_migration_best_effort(migration_name: str) -> None:
+        """Mark a migration rolled back, tolerating a concurrent resolver
+        having already done it."""
+        try:
+            ProxyExtrasDBManager._roll_back_migration(migration_name)
+        except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
+            pass
+
+    @staticmethod
+    def _failed_migration_logs(migration_name: str) -> Optional[str]:
+        """Return failed migration logs, or None if the ledger is unavailable."""
+        database_url = os.getenv("DATABASE_URL")
+        if not database_url:
+            return None
+
+        try:
+            import psycopg
+        except ImportError:
+            return None
+
+        cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
+        ledger_table = psycopg.sql.SQL("{}.{}").format(
+            psycopg.sql.Identifier(
+                ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
+            ),
+            psycopg.sql.Identifier("_prisma_migrations"),
+        )
+        try:
+            with psycopg.connect(
+                cleaned_url, connect_timeout=10, autocommit=True
+            ) as conn:
+                row = conn.execute(
+                    psycopg.sql.SQL(
+                        "SELECT logs FROM {} "
+                        "WHERE migration_name = %s AND finished_at IS NULL "
+                        "AND rolled_back_at IS NULL"
+                    ).format(ledger_table),
+                    (migration_name,),
+                ).fetchone()
+        except (psycopg.OperationalError, psycopg.DatabaseError):
+            return None
+        return (row[0] or "") if row else ""
+
     @staticmethod
     def _resolve_specific_migration(migration_name: str):
         """Mark a specific migration as applied"""
@@ -661,7 +707,8 @@ class ProxyExtrasDBManager:
         v2 migration resolver (opt-in via --use_v2_migration_resolver).
 
         Runs `prisma migrate deploy` and handles standard recovery paths
-        (P3005 baseline, P3009/P3018 idempotent errors). Critically, it does
+        (P3005 baseline, P3009/P3018 idempotent errors, deadlocks against a
+        concurrent migrate deploy). Critically, it does
         NOT call `_resolve_all_migrations` — the diff-and-force recovery that
         caused schema thrashing when two LiteLLM versions contended for the
         same DB during rolling deploys.
@@ -772,6 +819,20 @@ class ProxyExtrasDBManager:
                                     f"Detail: {resolve_err}"
                                 ) from resolve_err
                             continue
+                        if migration_match:
+                            migration_name = migration_match.group(1)
+                            ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name)
+                            if ledger_logs is not None and (
+                                ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs
+                            ):
+                                logger.info(
+                                    "Migration %s failed in a concurrent migrate deploy "
+                                    "deadlock race, rolling its ledger row back and retrying",
+                                    migration_name,
+                                )
+                                ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name)
+                                time.sleep(random.randrange(5, 15))
+                                continue
                         raise RuntimeError(
                             "Database migration failed and cannot be auto-recovered. "
                             f"Manual intervention required.\n\nPrisma error:\n{stderr}"
@@ -817,11 +878,42 @@ class ProxyExtrasDBManager:
                                 ) from resolve_err
                             continue
 
+                        if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr:
+                            logger.info(
+                                "Migration %s deadlocked against a concurrent "
+                                "migrate deploy, rolling its ledger row back "
+                                "and retrying",
+                                migration_match.group(1),
+                            )
+                            ProxyExtrasDBManager._roll_back_migration_best_effort(
+                                migration_match.group(1)
+                            )
+                            time.sleep(random.randrange(5, 15))
+                            continue
+
                         raise RuntimeError(
                             "Database migration failed and cannot be auto-recovered. "
                             f"Manual intervention required.\n\nPrisma error:\n{stderr}"
                         ) from e
 
+                    if _MIGRATION_DEADLOCK_MARKER in stderr:
+                        logger.info(
+                            "prisma migrate deploy attempt %s deadlocked against "
+                            "a concurrent migrate deploy, retrying",
+                            attempt + 1,
+                        )
+                        time.sleep(random.randrange(5, 15))
+                        continue
+
+                    if "P1002" in stderr and "advisory lock" in stderr:
+                        logger.info(
+                            "prisma migrate deploy attempt %s timed out waiting for "
+                            "the advisory lock a concurrent migrate deploy holds, retrying",
+                            attempt + 1,
+                        )
+                        time.sleep(random.randrange(5, 15))
+                        continue
+
                     raise RuntimeError(
                         "Database migration failed and cannot be auto-recovered. "
                         f"Manual intervention required.\n\nPrisma error:\n{stderr}"
@@ -829,9 +921,9 @@ class ProxyExtrasDBManager:
 
             raise RuntimeError(
                 "Database migration failed after 4 attempts (retry loop "
-                "exhausted by timeouts or repeated idempotent-recovery "
-                "continues). Check database connectivity, load, and "
-                "_prisma_migrations ledger state, and raise "
+                "exhausted by timeouts, deadlock retries, or repeated "
+                "idempotent-recovery continues). Check database connectivity, "
+                "load, and _prisma_migrations ledger state, and raise "
                 f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
             )
         finally:

litellm-proxy-extras/tests/test_setup_database_fail_fast.py

@@ -240,3 +240,223 @@ def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
     ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
     assert ok is True
     assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery"
+
+
+_DEADLOCK_P3018_STDERR = (
+    "Error: P3018\n"
+    "Migration name: 20260415120000_health_check_latest_per_model_index\n"
+    "Database error code: 40P01\n"
+    "deadlock detected"
+)
+
+
+def _stub_v2_env(monkeypatch, tmp_path):
+    monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
+    monkeypatch.setattr(
+        ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
+    )
+    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
+    (tmp_path / "schema.prisma").write_text("// stub")
+    monkeypatch.setattr("time.sleep", lambda _: None)
+
+
+def _succeed_after(failures: int, stderr: str):
+    calls = {"n": 0}
+
+    class _OkResult:
+        stdout = "Applied migration.\n"
+        stderr = ""
+
+    def _run(*args, **kwargs):
+        if "deploy" not in args[0]:
+            return _OkResult()
+        calls["n"] += 1
+        if calls["n"] <= failures:
+            raise subprocess.CalledProcessError(
+                returncode=1, cmd=args[0], stderr=stderr, output=""
+            )
+        return _OkResult()
+
+    return _run
+
+
+def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path):
+    """v2: losing the migrate deploy deadlock race against a concurrent
+    instance rolls the ledger row back and retries instead of dying."""
+    _stub_v2_env(monkeypatch, tmp_path)
+
+    rolled_back = []
+    monkeypatch.setattr(
+        ProxyExtrasDBManager,
+        "_roll_back_migration",
+        lambda name: rolled_back.append(name),
+    )
+    monkeypatch.setattr(
+        ProxyExtrasDBManager,
+        "_resolve_specific_migration",
+        lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
+    )
+    monkeypatch.setattr("subprocess.run", _succeed_after(1, _DEADLOCK_P3018_STDERR))
+
+    ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
+    assert ok is True
+    assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]
+
+
+def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path):
+    """v2: a deadlock on every attempt still fails after the retry budget."""
+    _stub_v2_env(monkeypatch, tmp_path)
+    monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda name: None)
+
+    with patch(
+        "subprocess.run",
+        side_effect=_fake_migrate_deploy_failure(1, _DEADLOCK_P3018_STDERR),
+    ):
+        with pytest.raises(RuntimeError, match="after 4 attempts"):
+            ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
+
+
+def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_path):
+    """v2: the surviving instance sees the victim's failed ledger row as P3009.
+    When that row's logs show a deadlock, roll it back and retry."""
+    _stub_v2_env(monkeypatch, tmp_path)
+
+    stderr = (
+        "Error: P3009\n"
+        "migrate found failed migrations in the target database\n"
+        "The `20260415120000_health_check_latest_per_model_index` migration "
+        "started at 2026-09-01 18:46:13 UTC failed"
+    )
+    monkeypatch.setattr(
+        ProxyExtrasDBManager,
+        "_failed_migration_logs",
+        lambda name: "ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock",
+    )
+    rolled_back = []
+    monkeypatch.setattr(
+        ProxyExtrasDBManager,
+        "_roll_back_migration",
+        lambda name: rolled_back.append(name),
+    )
+    monkeypatch.setattr(
+        ProxyExtrasDBManager,
+        "_resolve_specific_migration",
+        lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
+    )
+    monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr))
+
+    ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
+    assert ok is True
+    assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]
+
+
+def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path):
+    """v2: empty failed ledger logs mean a concurrent deploy moved it on."""
+    _stub_v2_env(monkeypatch, tmp_path)
+
+    stderr = (
+        "Error: P3009\n"
+        "migrate found failed migrations in the target database\n"
+        "The `20260415120000_health_check_latest_per_model_index` migration "
+        "started at 2026-09-01 18:46:13 UTC failed"
+    )
+    monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "")
+    rolled_back = []
+    monkeypatch.setattr(
+        ProxyExtrasDBManager,
+        "_roll_back_migration",
+        lambda name: rolled_back.append(name),
+    )
+    monkeypatch.setattr(
+        ProxyExtrasDBManager,
+        "_resolve_specific_migration",
+        lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
+    )
+    monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr))
+
+    ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
+    assert ok is True
+    assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]
+
+
+def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path):
+    """v2: an unreadable ledger cannot establish that P3009 was a deadlock."""
+    _stub_v2_env(monkeypatch, tmp_path)
+
+    stderr = (
+        "Error: P3009\n"
+        "migrate found failed migrations in the target database\n"
+        "The `20260415120000_health_check_latest_per_model_index` migration "
+        "started at 2026-09-01 18:46:13 UTC failed"
+    )
+    monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: None)
+    monkeypatch.setattr(
+        ProxyExtrasDBManager,
+        "_roll_back_migration",
+        lambda name: pytest.fail("an unreadable ledger must not trigger a retry"),
+    )
+    monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr))
+
+    with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
+        ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
+
+
+def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path):
+    """v2: a failed ledger row whose logs show a real SQL error stays fatal."""
+    _stub_v2_env(monkeypatch, tmp_path)
+
+    stderr = (
+        "Error: P3009\n"
+        "migrate found failed migrations in the target database\n"
+        "The `20260101000000_genuinely_broken` migration started at "
+        "2026-09-01 18:46:13 UTC failed"
+    )
+    monkeypatch.setattr(
+        ProxyExtrasDBManager,
+        "_failed_migration_logs",
+        lambda name: 'ERROR: syntax error at or near "BRKN"',
+    )
+
+    with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
+        with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
+            ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
+
+
+def test_v2_bare_deadlock_stderr_retries(monkeypatch, tmp_path):
+    """v2: a deadlock reported without a Prisma error code (the advisory-lock
+    waiter as victim) is retried, not fatal."""
+    _stub_v2_env(monkeypatch, tmp_path)
+    monkeypatch.setattr(
+        "subprocess.run", _succeed_after(1, "Database error: deadlock detected")
+    )
+
+    ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
+    assert ok is True
+
+
+_P1002_ADVISORY_LOCK_STDERR = (
+    "Error: P1002\n\n"
+    "The database server at `127.0.0.1`:`45743` was reached but timed out.\n\n"
+    "Context: Timed out trying to acquire a postgres advisory lock "
+    "(SELECT pg_advisory_lock(72707369)). Elapsed: 10000ms."
+)
+
+
+def test_v2_advisory_lock_timeout_retries(monkeypatch, tmp_path):
+    """v2: the advisory-lock waiter that times out while a peer's retry holds
+    the lock retries instead of dying."""
+    _stub_v2_env(monkeypatch, tmp_path)
+    monkeypatch.setattr("subprocess.run", _succeed_after(2, _P1002_ADVISORY_LOCK_STDERR))
+
+    ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
+    assert ok is True
+
+
+def test_v2_p1002_without_advisory_lock_context_still_raises(monkeypatch, tmp_path):
+    """v2: a plain P1002 (database unreachable) stays fatal."""
+    _stub_v2_env(monkeypatch, tmp_path)
+    stderr = "Error: P1002\n\nThe database server at `db`:`5432` was reached but timed out."
+    monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr))
+
+    with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
+        ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)

litellm-rust/Cargo.lock

@@ -1392,6 +1392,12 @@ dependencies = [
  "wasm-bindgen",
 ]
 
+[[package]]
+name = "lazy_static"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
+
 [[package]]
 name = "libc"
 version = "0.2.186"
@@ -1416,6 +1422,7 @@ dependencies = [
  "tokio",
  "tokio-tungstenite",
  "tower",
+ "tracing",
 ]
 
 [[package]]
@@ -1435,6 +1442,7 @@ dependencies = [
  "sha2 0.10.9",
  "thiserror 2.0.19",
  "tokio",
+ "tracing",
 ]
 
 [[package]]
@@ -1442,13 +1450,18 @@ name = "litellm-python-bridge"
 version = "0.1.0"
 dependencies = [
  "criterion",
+ "futures-util",
  "litellm-ai-gateway",
  "litellm-core",
  "litellm-python-interop",
  "pyo3",
  "pyo3-async-runtimes",
+ "serde",
  "serde_json",
  "tokio",
+ "tokio-tungstenite",
+ "tracing",
+ "tracing-subscriber",
 ]
 
 [[package]]
@@ -1669,9 +1682,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3"
-version = "0.29.0"
+version = "0.29.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c"
+checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
 dependencies = [
  "libc",
  "once_cell",
@@ -1697,18 +1710,18 @@ dependencies = [
 
 [[package]]
 name = "pyo3-build-config"
-version = "0.29.0"
+version = "0.29.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078"
+checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9"
 dependencies = [
  "target-lexicon",
 ]
 
 [[package]]
 name = "pyo3-ffi"
-version = "0.29.0"
+version = "0.29.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b"
+checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327"
 dependencies = [
  "libc",
  "pyo3-build-config",
@@ -1716,9 +1729,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-macros"
-version = "0.29.0"
+version = "0.29.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771"
+checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c"
 dependencies = [
  "proc-macro2",
  "pyo3-macros-backend",
@@ -1728,9 +1741,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-macros-backend"
-version = "0.29.0"
+version = "0.29.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362"
+checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952"
 dependencies = [
  "heck",
  "proc-macro2",
@@ -2276,6 +2289,15 @@ dependencies = [
  "digest 0.11.3",
 ]
 
+[[package]]
+name = "sharded-slab"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6"
+dependencies = [
+ "lazy_static",
+]
+
 [[package]]
 name = "shlex"
 version = "2.0.1"
@@ -2414,6 +2436,15 @@ dependencies = [
  "syn 3.0.0",
 ]
 
+[[package]]
+name = "thread_local"
+version = "1.1.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070"
+dependencies = [
+ "cfg-if",
+]
+
 [[package]]
 name = "time"
 version = "0.3.53"
@@ -2662,6 +2693,17 @@ dependencies = [
  "once_cell",
 ]
 
+[[package]]
+name = "tracing-subscriber"
+version = "0.3.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319"
+dependencies = [
+ "sharded-slab",
+ "thread_local",
+ "tracing-core",
+]
+
 [[package]]
 name = "try-lock"
 version = "0.2.5"

litellm-rust/Cargo.toml

@@ -14,11 +14,13 @@ license = "MIT"
 repository = "https://github.com/BerriAI/litellm"
 
 [workspace.dependencies]
+tracing = "0.1"
+tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
 litellm-core = { path = "crates/core" }
 litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
 litellm-python-interop = { path = "crates/python-interop" }
 axum = "0.7"
-pyo3 = "0.29.0"
+pyo3 = "0.29.2"
 pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
 pythonize = "0.29.0"
 rand = "0.8"

litellm-rust/crates/ai-gateway/Cargo.toml

@@ -14,6 +14,7 @@ path = "src/main.rs"
 required-features = ["server"]
 
 [dependencies]
+tracing.workspace = true
 litellm-core = { workspace = true, features = ["bedrock-auth"] }
 # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
 # Python proxy callbacks API.

litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs

@@ -32,6 +32,7 @@ pub(super) fn truncate_error_body(body: &str) -> String {
     format!("{truncated}... (truncated)")
 }
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub(super) fn ocr_provider_config(
     provider: &str,
     model: &str,
@@ -73,12 +74,6 @@ pub(super) fn string_headers(
         .collect()
 }
 
-pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool {
-    headers
-        .iter()
-        .any(|(key, _)| key.eq_ignore_ascii_case(name))
-}
-
 fn document_url_field(document: &Value) -> Result<Option<(&str, &str)>, Error> {
     let Some(object) = document.as_object() else {
         return Ok(None);

litellm-rust/crates/ai-gateway/src/ocr/handler.rs

@@ -1,12 +1,19 @@
 use litellm_core::error::Error;
+use litellm_core::http_utils::http_request;
 use litellm_core::ocr::transformation::OcrResponseHandling;
 use serde_json::Value;
 
 use super::common_utils::{poll_document_intelligence, truncate_error_body};
-use super::types::ProviderOcrRequest;
+use super::hooks::OcrLifecycleHooks;
+use super::types::PreparedOcrRequest;
 use crate::client::http_client;
 
-pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Result<Value, Error> {
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+pub(crate) async fn execute_ocr_provider_call(
+    request: PreparedOcrRequest,
+    hooks: &OcrLifecycleHooks,
+) -> Result<Value, Error> {
+    let request = hooks.prepare_provider_request(request).await?;
     let mut request_builder = http_client().post(&request.url).json(&request.body);
     for (key, value) in &request.upstream_headers {
         request_builder = request_builder.header(key, value);
@@ -15,8 +22,7 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Re
         request_builder = request_builder.timeout(duration);
     }
 
-    let response = request_builder
-        .send()
+    let response = http_request(request_builder)
         .await
         .map_err(|err| Error::Network(err.to_string()))?;
 

litellm-rust/crates/ai-gateway/src/ocr/hooks.rs

@@ -1,13 +1,10 @@
 use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
 use litellm_core::error::Error;
-use litellm_core::ocr::transformation::OcrAuthStrategy;
 use serde_json::{Map, Value, json};
 use std::future::Future;
 use std::pin::Pin;
 
-use super::common_utils::{
-    convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers,
-};
+use super::common_utils::{convert_document_url_to_data_uri, string_headers};
 use super::types::{PreparedOcrRequest, ProviderOcrRequest};
 use crate::integrations::custom_guardrail::{
     CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
@@ -62,6 +59,10 @@ impl OcrLifecycleHooks {
             .await
             .map_err(guardrail_error_to_core_error)?;
         let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?;
+        let optional_params = match &request.config {
+            Ok(config) => config.map_ocr_params(&optional_params),
+            Err(_) => optional_params,
+        };
         Ok(PreparedOcrRequest {
             document,
             optional_params,
@@ -69,25 +70,23 @@ impl OcrLifecycleHooks {
         })
     }
 
-    async fn prepare_provider_request(
+    pub(crate) async fn prepare_provider_request(
         &self,
         request: PreparedOcrRequest,
     ) -> Result<ProviderOcrRequest, Error> {
-        let config = ocr_provider_config(&request.custom_llm_provider, &request.model)
-            .ok_or_else(|| Error::InvalidProvider(request.custom_llm_provider.clone()))?;
+        let config = request.config?;
         let env_lookup = |key: &str| std::env::var(key).ok();
-        let headers = string_headers(request.extra_headers)?;
-        let auth_strategy = config.auth_strategy();
-        let api_key = (!has_header(&headers, auth_strategy.header_name()))
-            .then(|| config.resolve_api_key(request.api_key.as_deref(), &env_lookup))
-            .transpose()?;
+        let upstream_headers = config.validate_environment(
+            string_headers(request.extra_headers)?,
+            request.api_key.as_deref(),
+            &env_lookup,
+        )?;
         let url = config.complete_url(
             request.api_base.as_deref(),
             &request.model,
             &request.optional_params,
             &env_lookup,
         )?;
-        let filtered_params = config.map_ocr_params(&request.optional_params);
         let model = request.model.clone();
         let custom_llm_provider = request.custom_llm_provider.clone();
         let document = if config.requires_data_uri_document() {
@@ -96,9 +95,8 @@ impl OcrLifecycleHooks {
             request.document
         };
         let body = config
-            .transform_ocr_request(&request.model, document, filtered_params)?
+            .transform_ocr_request(&request.model, document, request.optional_params)?
             .data;
-        let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref());
         let body = self
             .run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
             .await?;
@@ -167,9 +165,9 @@ impl OcrLifecycleHooks {
     }
 }
 
-impl CallLifecycleHooks<PreparedOcrRequest, ProviderOcrRequest, Value> for OcrLifecycleHooks {
+impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLifecycleHooks {
     type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
-    type DuringCallFuture<'a> = OcrFuture<'a, ProviderOcrRequest>;
+    type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
     type SuccessFuture<'a> = OcrLogFuture<'a>;
     type FailureFuture<'a> = OcrLogFuture<'a>;
 
@@ -186,7 +184,7 @@ impl CallLifecycleHooks<PreparedOcrRequest, ProviderOcrRequest, Value> for OcrLi
         _context: &'a CallLifecycleContext,
         request: PreparedOcrRequest,
     ) -> Self::DuringCallFuture<'a> {
-        Box::pin(async move { self.prepare_provider_request(request).await })
+        Box::pin(async move { Ok(request) })
     }
 
     fn async_log_success_event<'a>(
@@ -247,21 +245,6 @@ impl CallLifecycleHooks<PreparedOcrRequest, ProviderOcrRequest, Value> for OcrLi
     }
 }
 
-fn upstream_headers(
-    headers: &[(String, String)],
-    auth_strategy: OcrAuthStrategy,
-    api_key: Option<&str>,
-) -> Vec<(String, String)> {
-    api_key
-        .map(|api_key| match auth_strategy {
-            OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")),
-            OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()),
-        })
-        .into_iter()
-        .chain(headers.iter().cloned())
-        .collect()
-}
-
 fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
     GuardrailContext {
         call_type: CallType::Ocr,

litellm-rust/crates/ai-gateway/src/ocr/mod.rs

@@ -13,10 +13,13 @@ pub use types::OcrRequest;
 use handler::execute_ocr_provider_call;
 use prepare::{PreparedOcrCall, prepare_ocr_call};
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub async fn ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
     let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);
     CallLifecycle::default()
-        .run_request(request, &hooks, execute_ocr_provider_call)
+        .run_request(request, &hooks, |request| {
+            execute_ocr_provider_call(request, &hooks)
+        })
         .await
 }
 

litellm-rust/crates/ai-gateway/src/ocr/prepare.rs

@@ -3,6 +3,7 @@ use std::time::{SystemTime, UNIX_EPOCH};
 
 use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
 
+use super::common_utils::ocr_provider_config;
 use super::hooks::OcrLifecycleHooks;
 use super::types::{OcrRequest, PreparedOcrRequest};
 use crate::integrations::custom_guardrail::CustomGuardrailRunner;
@@ -13,6 +14,7 @@ pub(crate) struct PreparedOcrCall {
     pub(crate) hooks: OcrLifecycleHooks,
 }
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
     let call_id = request
         .litellm_call_id
@@ -25,9 +27,25 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
         });
     let model = provider_info.model.to_string();
     let custom_llm_provider = provider_info.custom_llm_provider.to_string();
+    let config = ocr_provider_config(&custom_llm_provider, &model)
+        .ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone()));
+    let optional_params = match &config {
+        Ok(config) => {
+            let supported = config.supported_ocr_params();
+            config.map_ocr_params(
+                &request
+                    .optional_params
+                    .into_iter()
+                    .filter(|(name, _)| supported.contains(&name.as_str()))
+                    .collect(),
+            )
+        }
+        Err(_) => request.optional_params,
+    };
 
     PreparedOcrCall {
         request: PreparedOcrRequest {
+            config,
             model,
             custom_llm_provider,
             litellm_call_id: call_id,
@@ -35,7 +53,7 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
             api_key: request.api_key.map(str::to_string),
             api_base: request.api_base.map(str::to_string),
             extra_headers: request.extra_headers,
-            optional_params: request.optional_params,
+            optional_params,
             timeout: request.timeout,
         },
         hooks: OcrLifecycleHooks::new(

litellm-rust/crates/ai-gateway/src/ocr/tests.rs

@@ -2,12 +2,13 @@ use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
 use litellm_core::error::Error;
+use litellm_core::http_utils::has_header;
 use litellm_core::ocr::transformation::OcrResponseHandling;
 use serde_json::{Map, Value, json};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::{TcpListener, TcpStream};
 
-use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body};
+use super::common_utils::{ocr_provider_config, string_headers, truncate_error_body};
 use super::{OcrRequest, ocr};
 use crate::integrations::custom_guardrail::{
     CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook,

litellm-rust/crates/ai-gateway/src/ocr/types.rs

@@ -25,6 +25,7 @@ pub struct OcrRequest<'a> {
 }
 
 pub(crate) struct PreparedOcrRequest {
+    pub(crate) config: Result<&'static dyn OcrProviderConfig, litellm_core::Error>,
     pub(crate) model: String,
     pub(crate) custom_llm_provider: String,
     pub(crate) litellm_call_id: String,

litellm-rust/crates/core/Cargo.toml

@@ -11,6 +11,7 @@ reqwest.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 thiserror.workspace = true
+tracing.workspace = true
 sha2.workspace = true
 aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
 aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }

litellm-rust/crates/core/src/audio_transcription/handler.rs

@@ -1,11 +1,12 @@
 use serde_json::Value;
 
 use crate::error::Error;
-use crate::http_utils::truncate_error_body;
+use crate::http_utils::{http_request, truncate_error_body};
 
 use super::client::http_client;
 use super::types::ProviderAudioTranscriptionRequest;
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub async fn execute_audio_transcription_provider_call(
     request: ProviderAudioTranscriptionRequest,
 ) -> Result<Value, Error> {
@@ -19,8 +20,7 @@ pub async fn execute_audio_transcription_provider_call(
     if let Some(duration) = request.timeout {
         request_builder = request_builder.timeout(duration);
     }
-    let response = request_builder
-        .send()
+    let response = http_request(request_builder)
         .await
         .map_err(|error| Error::Network(error.to_string()))?;
     let status = response.status();

litellm-rust/crates/core/src/audio_transcription/mod.rs

@@ -11,6 +11,7 @@ pub use handler::execute_audio_transcription_provider_call;
 pub use prepare::prepare_audio_transcription_provider_call;
 pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
     execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?)
         .await

litellm-rust/crates/core/src/audio_transcription/prepare.rs

@@ -7,6 +7,7 @@ use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}
 use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
 use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
     #[cfg(feature = "bedrock-auth")]
     if provider == "bedrock" {
@@ -16,6 +17,7 @@ fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProv
     None
 }
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub fn prepare_audio_transcription_provider_call(
     request: AudioTranscriptionRequest<'_>,
 ) -> Result<ProviderAudioTranscriptionRequest, Error> {

litellm-rust/crates/core/src/audio_transcription/transformation.rs

@@ -15,6 +15,7 @@ pub enum AudioTranscriptionAuth {
 pub trait AudioTranscriptionProviderConfig: Sync {
     fn supported_transcription_params(&self) -> &'static [&'static str];
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn map_transcription_params(&self, params: &Map<String, Value>) -> Map<String, Value> {
         params
             .iter()

litellm-rust/crates/core/src/chat_completions/common_utils.rs

@@ -7,6 +7,7 @@ use super::transformation::ChatCompletionsProviderConfig;
 
 const HEADER_CONTEXT: &str = "chat completions";
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub(super) fn chat_completions_provider_config(
     provider: &str,
 ) -> Option<&'static dyn ChatCompletionsProviderConfig> {

litellm-rust/crates/core/src/chat_completions/handler.rs

@@ -1,17 +1,21 @@
 use serde_json::Value;
 
 use crate::error::Error;
-use crate::http_utils::truncate_error_body;
+use crate::http_utils::{http_request, truncate_error_body};
 
 use super::client::http_client;
+use super::prepare::prepare_provider_request;
 use super::transformation::ChatCompletionsAuth;
 use super::types::{
     ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData,
+    ResolvedChatCompletionsRequest,
 };
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub(super) async fn execute_chat_completions_provider_call(
-    request: ProviderChatCompletionsRequest,
+    request: ResolvedChatCompletionsRequest<'_>,
 ) -> Result<ChatCompletionsResponse, Error> {
+    let request = prepare_provider_request(request)?;
     let body = serde_json::to_vec(&request.body).map_err(|err| {
         Error::InvalidRequest(format!(
             "failed to serialize chat completions request: {err}"
@@ -27,7 +31,7 @@ pub(super) async fn execute_chat_completions_provider_call(
         request_builder = request_builder.timeout(duration);
     }
 
-    let response = request_builder.send().await.map_err(|err| {
+    let response = http_request(request_builder).await.map_err(|err| {
         // Failing to establish the connection means the request never went out,
         // so the host can still serve it. Everything else here, a timeout
         // above all, may have reached the provider and been answered.

litellm-rust/crates/core/src/chat_completions/mod.rs

@@ -19,13 +19,14 @@ pub mod types;
 use serde_json::{Map, Value};
 
 use handler::execute_chat_completions_provider_call;
-use prepare::{parse_messages, prepare_chat_completions_call, resolve_provider_config};
+use prepare::{parse_messages, resolve_provider_config, resolve_request};
 use types::{ChatCompletionsRequest, ChatCompletionsResponse};
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub async fn chat_completions(
     request: ChatCompletionsRequest<'_>,
 ) -> Result<ChatCompletionsResponse, Error> {
-    execute_chat_completions_provider_call(prepare_chat_completions_call(request)?).await
+    execute_chat_completions_provider_call(resolve_request(request)?).await
 }
 
 /// Whether the core would accept this request, without resolving credentials or

litellm-rust/crates/core/src/chat_completions/prepare.rs

@@ -6,7 +6,10 @@ use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}
 
 use super::common_utils::{chat_completions_provider_config, string_headers};
 use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
-use super::types::{ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest};
+use super::types::{
+    ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest,
+    ResolvedChatCompletionsRequest,
+};
 
 pub(super) fn resolve_provider_config<'a>(
     model: &'a str,
@@ -34,12 +37,10 @@ pub(super) fn parse_messages(messages: Value) -> Result<Vec<ChatMessage>, Error>
         .map_err(|err| Error::InvalidRequest(format!("invalid chat completions messages: {err}")))
 }
 
-pub(super) fn prepare_chat_completions_call(
+pub(super) fn resolve_request(
     request: ChatCompletionsRequest<'_>,
-) -> Result<ProviderChatCompletionsRequest, Error> {
+) -> Result<ResolvedChatCompletionsRequest<'_>, Error> {
     let (model, config) = resolve_provider_config(request.model, request.custom_llm_provider)?;
-    let env_lookup = |key: &str| std::env::var(key).ok();
-
     let messages = parse_messages(request.messages)?;
     if messages.is_empty() {
         return Err(Error::InvalidRequest(
@@ -49,11 +50,29 @@ pub(super) fn prepare_chat_completions_call(
     if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) {
         return Err(Error::Unsupported(reason.0));
     }
+    Ok(ResolvedChatCompletionsRequest {
+        model,
+        config,
+        messages,
+        optional_params: request.optional_params,
+        api_key: request.api_key,
+        api_base: request.api_base,
+        extra_headers: request.extra_headers,
+        timeout: request.timeout,
+    })
+}
 
-    let mut headers = string_headers(request.extra_headers)?;
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+fn validate_environment(
+    request: &ResolvedChatCompletionsRequest<'_>,
+    model: &str,
+    config: &dyn ChatCompletionsProviderConfig,
+) -> Result<(Vec<(String, String)>, ChatCompletionsAuth), Error> {
+    let env_lookup = |key: &str| std::env::var(key).ok();
+    let mut headers = string_headers(request.extra_headers.clone())?;
     let auth = config.auth(
         request.api_key,
-        &model,
+        model,
         &request.optional_params,
         &env_lookup,
     )?;
@@ -94,7 +113,16 @@ pub(super) fn prepare_chat_completions_call(
             headers.push(((*name).to_string(), (*value).to_string()));
         }
     }
+    Ok((headers, auth))
+}
 
+pub(super) fn prepare_provider_request(
+    request: ResolvedChatCompletionsRequest<'_>,
+) -> Result<ProviderChatCompletionsRequest, Error> {
+    let (headers, auth) = validate_environment(&request, &request.model, request.config)?;
+    let model = request.model;
+    let config = request.config;
+    let env_lookup = |key: &str| std::env::var(key).ok();
     let url = config.complete_url(
         request.api_base,
         &model,
@@ -102,7 +130,7 @@ pub(super) fn prepare_chat_completions_call(
         &env_lookup,
     )?;
     let transformed =
-        config.transform_request(&model, messages, request.optional_params.clone())?;
+        config.transform_request(&model, request.messages, request.optional_params.clone())?;
 
     Ok(ProviderChatCompletionsRequest {
         model,

litellm-rust/crates/core/src/chat_completions/tests.rs

@@ -2,9 +2,15 @@ use serde_json::{Map, Value, json};
 
 use crate::error::Error;
 
-use super::prepare::prepare_chat_completions_call;
+use super::prepare::{prepare_provider_request, resolve_request};
 use super::transformation::ChatCompletionsAuth;
-use super::types::ChatCompletionsRequest;
+use super::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest};
+
+fn prepare_chat_completions_call(
+    request: ChatCompletionsRequest<'_>,
+) -> Result<ProviderChatCompletionsRequest, Error> {
+    prepare_provider_request(resolve_request(request)?)
+}
 
 fn request<'a>(
     model: &'a str,

litellm-rust/crates/core/src/chat_completions/transformation.rs

@@ -62,9 +62,8 @@ pub trait ChatCompletionsProviderConfig: Sync {
         false
     }
 
-    /// Provider parameter names (post-mapping) the Rust path knows how to place
-    /// in the upstream body. Anything outside this set declines the request.
-    fn supported_params(&self) -> &'static [&'static str];
+    /// Supported OpenAI parameter names paired with their provider names.
+    fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)];
 
     /// Parameters consumed as call configuration (credentials, endpoints)
     /// rather than placed in the body. Accepted, never serialized.
@@ -78,7 +77,7 @@ pub trait ChatCompletionsProviderConfig: Sync {
         optional_params: &Map<String, Value>,
     ) -> Option<Unsupported> {
         unsupported_param(
-            self.supported_params(),
+            self.supported_openai_params(),
             self.config_params(),
             optional_params,
         )
@@ -100,7 +99,7 @@ pub trait ChatCompletionsProviderConfig: Sync {
 }
 
 pub fn unsupported_param(
-    supported: &'static [&'static str],
+    supported: &'static [(&'static str, &'static str)],
     config: &'static [&'static str],
     optional_params: &Map<String, Value>,
 ) -> Option<Unsupported> {
@@ -115,7 +114,9 @@ pub fn unsupported_param(
         .keys()
         .any(|key| {
             key != STREAM_PARAM
-                && !supported.contains(&key.as_str())
+                && !supported
+                    .iter()
+                    .any(|(_, provider_name)| *provider_name == key)
                 && !config.contains(&key.as_str())
         })
         .then_some(Unsupported("unrecognized request parameter"))

litellm-rust/crates/core/src/chat_completions/types.rs

@@ -22,6 +22,17 @@ pub struct ChatCompletionsRequest<'a> {
     pub timeout: Option<Duration>,
 }
 
+pub(super) struct ResolvedChatCompletionsRequest<'a> {
+    pub(super) model: String,
+    pub(super) config: &'static dyn ChatCompletionsProviderConfig,
+    pub(super) messages: Vec<ChatMessage>,
+    pub(super) optional_params: Map<String, Value>,
+    pub(super) api_key: Option<&'a str>,
+    pub(super) api_base: Option<&'a str>,
+    pub(super) extra_headers: Option<Map<String, Value>>,
+    pub(super) timeout: Option<Duration>,
+}
+
 pub(super) struct ProviderChatCompletionsRequest {
     pub(super) model: String,
     pub(super) config: &'static dyn ChatCompletionsProviderConfig,

litellm-rust/crates/core/src/http_utils.rs

@@ -5,6 +5,13 @@ use serde_json::{Map, Value};
 use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS;
 use crate::error::{Error, json_type_name};
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+pub async fn http_request(
+    request: reqwest::RequestBuilder,
+) -> Result<reqwest::Response, reqwest::Error> {
+    request.send().await
+}
+
 /// Bound an upstream error body before it crosses a host boundary, so provider
 /// bodies stay data-minimized.
 pub fn truncate_error_body(body: &str) -> String {

litellm-rust/crates/core/src/messages/common_utils.rs

@@ -10,6 +10,7 @@ pub(super) use crate::http_utils::{has_bearer_auth, has_header, truncate_error_b
 
 const HEADER_CONTEXT: &str = "messages";
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub(super) fn messages_provider_config(
     provider: &str,
 ) -> Option<&'static dyn AnthropicMessagesProviderConfig> {

litellm-rust/crates/core/src/messages/handler.rs

@@ -1,13 +1,17 @@
 use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
 use crate::error::Error;
+use crate::http_utils::http_request;
 
 use super::client::http_client;
 use super::common_utils::truncate_error_body;
-use super::types::{AnthropicMessagesResponse, ProviderMessagesRequest};
+use super::prepare::prepare_provider_request;
+use super::types::{AnthropicMessagesResponse, MessagesRequest};
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub(super) async fn execute_messages_provider_call(
-    request: ProviderMessagesRequest,
+    request: MessagesRequest<'_>,
 ) -> Result<AnthropicMessagesResponse, Error> {
+    let request = prepare_provider_request(request)?;
     let mut request_builder = http_client().post(&request.url).json(&request.body);
     for (key, value) in &request.upstream_headers {
         request_builder = request_builder.header(key, value);
@@ -16,8 +20,7 @@ pub(super) async fn execute_messages_provider_call(
         request_builder = request_builder.timeout(duration);
     }
 
-    let response = request_builder
-        .send()
+    let response = http_request(request_builder)
         .await
         .map_err(|err| Error::Network(err.to_string()))?;
 
@@ -40,8 +43,9 @@ pub(super) async fn execute_messages_provider_call(
 }
 
 pub(super) async fn execute_messages_provider_stream(
-    request: ProviderMessagesRequest,
+    request: MessagesRequest<'_>,
 ) -> Result<reqwest::Response, Error> {
+    let request = prepare_provider_request(request)?;
     if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
         return Err(Error::InvalidRequest(
             "streaming messages is not supported for this provider".to_string(),
@@ -56,8 +60,7 @@ pub(super) async fn execute_messages_provider_stream(
         request_builder = request_builder.timeout(duration);
     }
 
-    let response = request_builder
-        .send()
+    let response = http_request(request_builder)
         .await
         .map_err(|err| Error::Network(err.to_string()))?;
     let status = response.status();

litellm-rust/crates/core/src/messages/mod.rs

@@ -16,15 +16,15 @@ pub mod transformation;
 pub mod types;
 
 use handler::{execute_messages_provider_call, execute_messages_provider_stream};
-use prepare::prepare_messages_call;
 use types::{AnthropicMessagesResponse, MessagesRequest};
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub async fn messages(request: MessagesRequest<'_>) -> Result<AnthropicMessagesResponse, Error> {
-    execute_messages_provider_call(prepare_messages_call(request)?).await
+    execute_messages_provider_call(request).await
 }
 
 pub async fn messages_stream(request: MessagesRequest<'_>) -> Result<reqwest::Response, Error> {
-    execute_messages_provider_stream(prepare_messages_call(request)?).await
+    execute_messages_provider_stream(request).await
 }
 
 #[cfg(test)]

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

@@ -2,10 +2,11 @@ use crate::error::Error;
 use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
 
 use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers};
-use super::transformation::MessagesAuthStrategy;
+use super::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
 use super::types::{MessagesRequest, ProviderMessagesRequest};
+use serde_json::{Map, Value};
 
-pub(super) fn prepare_messages_call(
+pub(super) fn prepare_provider_request(
     request: MessagesRequest<'_>,
 ) -> Result<ProviderMessagesRequest, Error> {
     let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
@@ -29,13 +30,46 @@ pub(super) fn prepare_messages_call(
         .ok_or_else(|| Error::InvalidProvider(provider.to_string()))?;
     let env_lookup = |key: &str| std::env::var(key).ok();
 
-    let mut headers = string_headers(request.extra_headers)?;
+    let headers =
+        validate_environment(config, request.extra_headers, request.api_key, &env_lookup)?;
+
+    let typed_request = serde_json::from_value(request.body).map_err(|err| {
+        Error::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
+    })?;
+    let transformed = config.transform_request(typed_request)?;
+    let body = serde_json::to_value(transformed).map_err(|err| {
+        Error::InvalidRequest(format!(
+            "failed to serialize Anthropic messages request: {err}"
+        ))
+    })?;
+
+    let url = config.complete_url(request.api_base, &model, &env_lookup)?;
+
+    Ok(ProviderMessagesRequest {
+        provider: provider.to_string(),
+        model,
+        config,
+        url,
+        body,
+        upstream_headers: headers,
+        timeout: request.timeout,
+    })
+}
+
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+fn validate_environment(
+    config: &dyn AnthropicMessagesProviderConfig,
+    extra_headers: Option<Map<String, Value>>,
+    api_key: Option<&str>,
+    env_lookup: &dyn Fn(&str) -> Option<String>,
+) -> Result<Vec<(String, String)>, Error> {
+    let mut headers = string_headers(extra_headers)?;
 
     let auth_strategy = config.auth_strategy();
     let already_authorized = has_header(&headers, auth_strategy.header_name())
         || (config.accepts_bearer_auth() && has_bearer_auth(&headers));
     if !already_authorized {
-        let api_key = config.resolve_api_key(request.api_key, &env_lookup)?;
+        let api_key = config.resolve_api_key(api_key, env_lookup)?;
         let auth_header = match auth_strategy {
             MessagesAuthStrategy::Bearer => {
                 ("authorization".to_string(), format!("Bearer {api_key}"))
@@ -51,24 +85,5 @@ pub(super) fn prepare_messages_call(
         }
     }
 
-    let url = config.complete_url(request.api_base, &model, &env_lookup)?;
-    let typed_request = serde_json::from_value(request.body).map_err(|err| {
-        Error::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
-    })?;
-    let transformed = config.transform_request(typed_request)?;
-    let body = serde_json::to_value(transformed).map_err(|err| {
-        Error::InvalidRequest(format!(
-            "failed to serialize Anthropic messages request: {err}"
-        ))
-    })?;
-
-    Ok(ProviderMessagesRequest {
-        provider: provider.to_string(),
-        model,
-        config,
-        url,
-        body,
-        upstream_headers: headers,
-        timeout: request.timeout,
-    })
+    Ok(headers)
 }

litellm-rust/crates/core/src/messages/transformation.rs

@@ -45,6 +45,7 @@ pub trait AnthropicMessagesProviderConfig: Sync {
         ]
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn transform_request(
         &self,
         request: AnthropicMessagesRequest,
@@ -52,6 +53,7 @@ pub trait AnthropicMessagesProviderConfig: Sync {
         Ok(request)
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn transform_response(
         &self,
         _model: &str,

litellm-rust/crates/core/src/ocr/transformation.rs

@@ -27,6 +27,7 @@ pub enum OcrResponseHandling {
 pub trait OcrProviderConfig: Sync {
     fn supported_ocr_params(&self) -> &'static [&'static str];
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
         let mut mapped_params = Map::new();
         for (param, value) in non_default_params {
@@ -64,6 +65,25 @@ pub trait OcrProviderConfig: Sync {
         env_lookup: &dyn Fn(&str) -> Option<String>,
     ) -> Result<String, Error>;
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+    fn validate_environment(
+        &self,
+        headers: Vec<(String, String)>,
+        api_key: Option<&str>,
+        env_lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> Result<Vec<(String, String)>, Error> {
+        let strategy = self.auth_strategy();
+        if crate::http_utils::has_header(&headers, strategy.header_name()) {
+            return Ok(headers);
+        }
+        let api_key = self.resolve_api_key(api_key, env_lookup)?;
+        let auth_header = match strategy {
+            OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")),
+            OcrAuthStrategy::Header(name) => (name.to_string(), api_key),
+        };
+        Ok(std::iter::once(auth_header).chain(headers).collect())
+    }
+
     fn auth_strategy(&self) -> OcrAuthStrategy {
         OcrAuthStrategy::Bearer
     }

litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs

@@ -27,7 +27,12 @@ use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage
 /// per-model gate inside `transform_request`, the function this route replaces.
 /// Forwarding it would send `top_k` to a model that removed sampling params and
 /// take a 400 after the call, where Python drops it and succeeds.
-const SUPPORTED_PARAMS: &[&str] = &["max_tokens", "temperature", "top_p", "stop_sequences"];
+const SUPPORTED_PARAMS: &[(&str, &str)] = &[
+    ("max_tokens", "max_tokens"),
+    ("temperature", "temperature"),
+    ("top_p", "top_p"),
+    ("stop", "stop_sequences"),
+];
 
 pub struct AnthropicChatCompletionsConfig;
 
@@ -112,7 +117,8 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
         })
     }
 
-    fn supported_params(&self) -> &'static [&'static str] {
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+    fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)] {
         SUPPORTED_PARAMS
     }
 
@@ -121,7 +127,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
         messages: &[ChatMessage],
         optional_params: &Map<String, Value>,
     ) -> Option<Unsupported> {
-        unsupported_param(SUPPORTED_PARAMS, &[], optional_params)
+        unsupported_param(self.supported_openai_params(), &[], optional_params)
             .or_else(|| messages.iter().find_map(unsupported_message))
             // Anthropic rejects a request whose first turn is not a user turn.
             // Python only repairs that under `litellm.modify_params`, which the
@@ -132,6 +138,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
             })
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn transform_request(
         &self,
         model: &str,
@@ -143,6 +150,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
         })
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn transform_response(
         &self,
         _model: &str,

litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs

@@ -47,6 +47,7 @@ pub fn complete_anthropic_url(
 }
 
 impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn complete_url(
         &self,
         api_base: Option<&str>,

litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs

@@ -142,6 +142,7 @@ fn fold_system_role_messages(request: AnthropicMessagesRequest) -> AnthropicMess
 }
 
 impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig {
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn complete_url(
         &self,
         api_base: Option<&str>,

litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs

@@ -46,10 +46,12 @@ fn optional_string<'a>(params: &'a Map<String, Value>, key: &str) -> Option<&'a
 }
 
 impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn supported_transcription_params(&self) -> &'static [&'static str] {
         SUPPORTED_PARAMS
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn transform_transcription_request(
         &self,
         _model: &str,
@@ -83,6 +85,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
         })
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn transform_transcription_response(
         &self,
         _model: &str,

litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs

@@ -23,11 +23,12 @@ use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT
 /// `additionalModelRequestFields` for Anthropic base models and to
 /// `inferenceConfig` otherwise, and that branch reads the model catalog the
 /// core cannot see.
-const SUPPORTED_PARAMS: &[&str] = &["maxTokens", "temperature", "topP", "stopSequences"];
-
-/// Params that belong in `inferenceConfig`, in the order Python's
-/// `AmazonConverseConfig` declares them, so bodies compare cleanly.
-const INFERENCE_CONFIG_PARAMS: &[&str] = SUPPORTED_PARAMS;
+const SUPPORTED_PARAMS: &[(&str, &str)] = &[
+    ("max_tokens", "maxTokens"),
+    ("temperature", "temperature"),
+    ("top_p", "topP"),
+    ("stop", "stopSequences"),
+];
 
 const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = "aws_bedrock_runtime_endpoint";
 
@@ -66,7 +67,7 @@ fn converse_body(conversation: &Conversation, params: &Map<String, Value>) -> Va
         })
         .collect();
 
-    let inference_config = Map::from_iter(INFERENCE_CONFIG_PARAMS.iter().filter_map(|name| {
+    let inference_config = Map::from_iter(SUPPORTED_PARAMS.iter().filter_map(|(_, name)| {
         params
             .get(*name)
             .map(|value| ((*name).to_string(), value.clone()))
@@ -162,7 +163,8 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
         &[("Content-Type", "application/json")]
     }
 
-    fn supported_params(&self) -> &'static [&'static str] {
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+    fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)] {
         SUPPORTED_PARAMS
     }
 
@@ -175,32 +177,36 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
         messages: &[ChatMessage],
         optional_params: &Map<String, Value>,
     ) -> Option<Unsupported> {
-        unsupported_param(SUPPORTED_PARAMS, CONFIG_PARAMS, optional_params)
-            .or_else(|| messages.iter().find_map(unsupported_message))
-            // Python's Converse translation drops blank text blocks instead of
-            // substituting the placeholder the shared conversation builder
-            // applies, so decline blank text rather than diverge.
-            .or_else(|| {
-                messages
-                    .iter()
-                    .any(has_blank_text)
-                    .then_some(Unsupported("blank message text"))
-            })
-            // Converse has no assistant prefill: Python inserts a continue turn
-            // when a conversation opens or closes on an assistant message, and
-            // only under `litellm.modify_params`, which the core cannot see.
-            // Declining both ends also keeps the shared builder's final
-            // assistant right-strip (an Anthropic rule) unreachable here.
-            .or_else(|| {
-                let conversation = build_conversation(messages);
-                let ends_on_assistant = conversation
-                    .turns
-                    .last()
-                    .is_some_and(|turn| turn.role == TurnRole::Assistant);
-                (!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported(
-                    "conversation does not run user turn to user turn",
-                ))
-            })
+        unsupported_param(
+            self.supported_openai_params(),
+            CONFIG_PARAMS,
+            optional_params,
+        )
+        .or_else(|| messages.iter().find_map(unsupported_message))
+        // Python's Converse translation drops blank text blocks instead of
+        // substituting the placeholder the shared conversation builder
+        // applies, so decline blank text rather than diverge.
+        .or_else(|| {
+            messages
+                .iter()
+                .any(has_blank_text)
+                .then_some(Unsupported("blank message text"))
+        })
+        // Converse has no assistant prefill: Python inserts a continue turn
+        // when a conversation opens or closes on an assistant message, and
+        // only under `litellm.modify_params`, which the core cannot see.
+        // Declining both ends also keeps the shared builder's final
+        // assistant right-strip (an Anthropic rule) unreachable here.
+        .or_else(|| {
+            let conversation = build_conversation(messages);
+            let ends_on_assistant = conversation
+                .turns
+                .last()
+                .is_some_and(|turn| turn.role == TurnRole::Assistant);
+            (!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported(
+                "conversation does not run user turn to user turn",
+            ))
+        })
     }
 
     fn transform_request(

litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs

@@ -70,10 +70,12 @@ pub struct MistralOcrConfig;
 pub const MISTRAL_OCR_CONFIG: MistralOcrConfig = MistralOcrConfig;
 
 impl OcrProviderConfig for MistralOcrConfig {
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn supported_ocr_params(&self) -> &'static [&'static str] {
         SUPPORTED_OCR_PARAMS
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn transform_ocr_request(
         &self,
         model: &str,
@@ -100,6 +102,7 @@ impl OcrProviderConfig for MistralOcrConfig {
         })
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn transform_ocr_response(
         &self,
         model: &str,
@@ -134,6 +137,7 @@ impl OcrProviderConfig for MistralOcrConfig {
         })
     }
 
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
     fn complete_url(
         &self,
         api_base: Option<&str>,
@@ -153,6 +157,7 @@ impl OcrProviderConfig for MistralOcrConfig {
     }
 }
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub fn supported_ocr_params() -> &'static [&'static str] {
     MISTRAL_OCR_CONFIG.supported_ocr_params()
 }
@@ -161,6 +166,7 @@ pub fn map_ocr_params(non_default_params: &Map<String, Value>) -> Map<String, Va
     MISTRAL_OCR_CONFIG.map_ocr_params(non_default_params)
 }
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub fn transform_ocr_request(
     model: &str,
     document: Value,
@@ -169,6 +175,7 @@ pub fn transform_ocr_request(
     MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
 }
 
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub fn transform_ocr_response(model: &str, response_json: Value) -> Result<OcrResponseData, Error> {
     MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
 }

litellm-rust/crates/python-bridge/Cargo.toml

@@ -16,16 +16,21 @@ extension-module = ["pyo3/extension-module"]
 panic-test = []
 
 [dependencies]
+futures-util.workspace = true
+tracing.workspace = true
+tracing-subscriber.workspace = true
 litellm-core = { workspace = true, features = ["bedrock-auth"] }
 litellm-ai-gateway = { workspace = true, default-features = false }
 litellm-python-interop.workspace = true
 pyo3.workspace = true
 pyo3-async-runtimes.workspace = true
+serde.workspace = true
 serde_json.workspace = true
 tokio.workspace = true
 
 [dev-dependencies]
 criterion = "0.8.2"
+tokio-tungstenite.workspace = true
 
 [[bench]]
 name = "serialization"

litellm-rust/crates/python-bridge/src/constants.rs

@@ -0,0 +1 @@
+pub(crate) const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";

litellm-rust/crates/python-bridge/src/diagnostics.rs

@@ -0,0 +1,23 @@
+use litellm_python_interop::release_count;
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+
+#[pyfunction]
+fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
+    let stats = PyDict::new(py);
+    stats.set_item("releases", release_count())?;
+    Ok(stats.into_any().unbind())
+}
+
+#[cfg(feature = "panic-test")]
+#[pyfunction]
+fn _panic_for_test() {
+    panic!("intentional PyO3 panic smoke test");
+}
+
+pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    module.add_function(wrap_pyfunction!(gil_stats, module)?)?;
+    #[cfg(feature = "panic-test")]
+    module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?;
+    Ok(())
+}

litellm-rust/crates/python-bridge/src/errors.rs

@@ -0,0 +1,61 @@
+use litellm_core::error::Error;
+use pyo3::exceptions::{PyRuntimeError, PyValueError};
+use pyo3::prelude::*;
+
+pyo3::create_exception!(
+    _native,
+    RustBridgeDeclined,
+    pyo3::exceptions::PyException,
+    "The route declined before calling the provider, so the host may retry on its own path."
+);
+
+pyo3::create_exception!(
+    _native,
+    RustUpstreamError,
+    pyo3::exceptions::PyException,
+    "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response."
+);
+
+pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr {
+    match err {
+        Error::Auth(message) => PyValueError::new_err(message),
+        Error::InvalidProvider(_)
+        | Error::InvalidRequest(_)
+        | Error::InvalidType { .. }
+        | Error::MissingField(_) => PyValueError::new_err(err.to_string()),
+        other => PyRuntimeError::new_err(other.to_string()),
+    }
+}
+
+/// Map a core error for a route whose host keeps a Python implementation.
+///
+/// The distinction the host needs is whether the provider was already called.
+/// Everything raised before the request goes out is safe for the host to retry
+/// on its own path; anything after it is not, because the provider has already
+/// done the work and billed for it.
+pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr {
+    match err {
+        Error::Unsupported(_)
+        | Error::Auth(_)
+        | Error::InvalidProvider(_)
+        | Error::InvalidRequest(_)
+        | Error::InvalidType { .. }
+        | Error::MissingField(_)
+        | Error::Routing(_)
+        // Nothing reached the provider, so serving it on Python cannot double
+        // bill and is the only way the caller gets an answer at all.
+        | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()),
+        Error::Http { status, body } => {
+            RustUpstreamError::new_err((status, format!("{status}: {body}")))
+        }
+        Error::Network(message) | Error::InvalidResponse(message) => {
+            RustUpstreamError::new_err((0u16, message))
+        }
+    }
+}
+
+pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    let py = module.py();
+    module.add("RustBridgeDeclined", py.get_type::<RustBridgeDeclined>())?;
+    module.add("RustUpstreamError", py.get_type::<RustUpstreamError>())
+}

litellm-rust/crates/python-bridge/src/execution.rs

@@ -0,0 +1,423 @@
+use std::future::Future;
+use std::panic::AssertUnwindSafe;
+use std::time::Duration;
+
+use futures_util::FutureExt;
+use litellm_core::error::Error;
+use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil};
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
+use serde::Serialize;
+use tokio::runtime::{Handle, Runtime};
+use tokio::time::{self, MissedTickBehavior};
+
+pub(crate) fn run_sync<T, F>(
+    py: Python<'_>,
+    future: F,
+    map_error: fn(Error) -> PyErr,
+) -> PyResult<Py<PyAny>>
+where
+    T: Serialize + Send + 'static,
+    F: Future<Output = Result<T, Error>> + Send + 'static,
+{
+    run_sync_on(
+        py,
+        pyo3_async_runtimes::tokio::get_runtime(),
+        future,
+        map_error,
+    )
+}
+
+fn run_sync_on<T, F>(
+    py: Python<'_>,
+    runtime: &Runtime,
+    future: F,
+    map_error: fn(Error) -> PyErr,
+) -> PyResult<Py<PyAny>>
+where
+    T: Serialize + Send + 'static,
+    F: Future<Output = Result<T, Error>> + Send + 'static,
+{
+    if Handle::try_current().is_ok() {
+        return Err(PyRuntimeError::new_err(
+            "synchronous native routes cannot run from a Tokio context; use the async route",
+        ));
+    }
+
+    let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;
+    let result = map_core_result(result, map_error)?;
+    Pythonized(result).into_pyobject(py).map(Bound::unbind)
+}
+
+pub(crate) fn run_async<T, F>(
+    py: Python<'_>,
+    future: F,
+    map_error: fn(Error) -> PyErr,
+) -> PyResult<Bound<'_, PyAny>>
+where
+    T: Serialize + Send + 'static,
+    F: Future<Output = Result<T, Error>> + Send + 'static,
+{
+    pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        let result = catch_future_panic(future).await?;
+        let result = map_core_result(result, map_error)?;
+        Ok(Pythonized(result))
+    })
+}
+
+fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
+    match result {
+        Ok(value) => Ok(value),
+        Err(error) => Err(
+            std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error)))
+                .map_err(panic_to_pyerr)?,
+        ),
+    }
+}
+
+async fn catch_future_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
+where
+    F: Future<Output = Result<T, Error>>,
+{
+    AssertUnwindSafe(future)
+        .catch_unwind()
+        .await
+        .map_err(panic_to_pyerr)
+}
+
+async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
+where
+    F: Future<Output = Result<T, Error>>,
+{
+    let future = catch_future_panic(future);
+    tokio::pin!(future);
+
+    let signal_interval = Duration::from_millis(50);
+    let mut signal_checks =
+        time::interval_at(time::Instant::now() + signal_interval, signal_interval);
+    signal_checks.set_missed_tick_behavior(MissedTickBehavior::Delay);
+    loop {
+        tokio::select! {
+            result = &mut future => return result,
+            _ = signal_checks.tick() => Python::attach(|py| py.check_signals())?,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::ffi::CString;
+    use std::future::poll_fn;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::sync::{Arc, mpsc};
+    use std::task::Poll;
+    use std::thread;
+    use std::time::Instant;
+
+    use pyo3::panic::PanicException;
+    use pyo3::types::{PyDict, PyModule};
+    use serde::Serializer;
+    use tokio::runtime::Builder;
+
+    use super::*;
+
+    fn runtime_error(error: Error) -> PyErr {
+        PyRuntimeError::new_err(error.to_string())
+    }
+
+    fn panicking_error_mapper(_error: Error) -> PyErr {
+        panic!("error mapper panicked")
+    }
+
+    struct PanickingOutput;
+
+    static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0);
+
+    impl Serialize for PanickingOutput {
+        fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
+        where
+            S: Serializer,
+        {
+            panic!("serializer panicked")
+        }
+    }
+
+    #[pyfunction]
+    fn async_serialization_panic(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
+        run_async(py, async { Ok(PanickingOutput) }, runtime_error)
+    }
+
+    #[pyfunction]
+    fn async_runtime_probe(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
+        run_async(
+            py,
+            async {
+                ASYNC_PROBE_COMPLETED.fetch_add(1, Ordering::SeqCst);
+                Ok(true)
+            },
+            runtime_error,
+        )
+    }
+
+    #[pyfunction]
+    fn runtime_worker_count() -> usize {
+        pyo3_async_runtimes::tokio::get_runtime()
+            .metrics()
+            .num_workers()
+    }
+
+    #[pyfunction]
+    fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
+        let completion_deadline = Instant::now() + Duration::from_secs(2);
+        while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
+            if Instant::now() >= completion_deadline {
+                return false;
+            }
+            thread::sleep(Duration::from_millis(1));
+        }
+
+        let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
+        pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
+            let _ = heartbeat_tx.send(());
+        });
+        heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
+    }
+
+    fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
+        result
+            .expect("route should complete")
+            .bind(py)
+            .extract()
+            .expect("result should convert")
+    }
+
+    #[test]
+    fn sync_runner_polls_future_on_the_caller_thread() {
+        Python::initialize();
+        Python::attach(|py| {
+            let caller_thread = std::thread::current().id();
+            let result = run_sync(
+                py,
+                async move { Ok(std::thread::current().id() == caller_thread) },
+                runtime_error,
+            );
+
+            assert!(extract_bool(py, result));
+        });
+    }
+
+    #[test]
+    fn sync_runner_releases_gil_while_waiting() {
+        Python::initialize();
+        Python::attach(|py| {
+            let result = run_sync(
+                py,
+                async {
+                    let gil_acquired = tokio::time::timeout(
+                        Duration::from_secs(2),
+                        tokio::task::spawn_blocking(|| Python::attach(|_| true)),
+                    )
+                    .await;
+                    Ok(matches!(gil_acquired, Ok(Ok(true))))
+                },
+                runtime_error,
+            );
+
+            assert!(extract_bool(py, result));
+        });
+    }
+
+    #[test]
+    fn sync_runner_rejects_calls_from_a_tokio_context() {
+        Python::initialize();
+        let runtime = Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .expect("runtime should build");
+
+        let error = runtime.block_on(async {
+            Python::attach(|py| {
+                run_sync::<bool, _>(py, async { Ok(true) }, runtime_error)
+                    .expect_err("sync route should reject a nested Tokio runtime")
+            })
+        });
+
+        assert_eq!(
+            error.to_string(),
+            "RuntimeError: synchronous native routes cannot run from a Tokio context; use the async route"
+        );
+    }
+
+    #[test]
+    fn sync_runner_can_drive_a_current_thread_runtime() {
+        Python::initialize();
+        let runtime = Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .expect("runtime should build");
+        Python::attach(|py| {
+            let result = run_sync_on(
+                py,
+                &runtime,
+                async {
+                    tokio::task::yield_now().await;
+                    Ok(true)
+                },
+                runtime_error,
+            );
+            assert!(extract_bool(py, result));
+        });
+    }
+
+    #[test]
+    fn sync_runner_maps_a_panicked_future() {
+        Python::initialize();
+        Python::attach(|py| {
+            let error = run_sync::<bool, _>(
+                py,
+                poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
+                runtime_error,
+            )
+            .expect_err("panicked route should become a Python exception");
+
+            assert!(error.is_instance_of::<PanicException>(py));
+            assert_eq!(error.to_string(), "PanicException: route future panicked");
+        });
+    }
+
+    #[test]
+    fn sync_runner_maps_a_panicked_error_mapper() {
+        Python::initialize();
+        Python::attach(|py| {
+            let error = run_sync::<bool, _>(
+                py,
+                async { Err(Error::InvalidRequest("invalid".to_string())) },
+                panicking_error_mapper,
+            )
+            .expect_err("panicked mapper should become a Python exception");
+
+            assert!(error.is_instance_of::<PanicException>(py));
+            assert_eq!(error.to_string(), "PanicException: error mapper panicked");
+        });
+    }
+
+    #[test]
+    fn sync_runner_surfaces_serializer_panics() {
+        Python::initialize();
+        Python::attach(|py| {
+            let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error)
+                .expect_err("serializer panic should become a Python exception");
+
+            assert!(error.is_instance_of::<PanicException>(py));
+            assert_eq!(error.to_string(), "PanicException: serializer panicked");
+        });
+    }
+
+    #[test]
+    fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() {
+        Python::initialize();
+        let barrier = Arc::new(tokio::sync::Barrier::new(2));
+        let callers: Vec<_> = (0..2)
+            .map(|_| {
+                let barrier = Arc::clone(&barrier);
+                thread::spawn(move || {
+                    Python::attach(|py| {
+                        extract_bool(
+                            py,
+                            run_sync(
+                                py,
+                                async move {
+                                    Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait())
+                                        .await
+                                        .is_ok())
+                                },
+                                runtime_error,
+                            ),
+                        )
+                    })
+                })
+            })
+            .collect();
+        let results: Vec<_> = callers
+            .into_iter()
+            .map(|caller| caller.join().expect("caller should not panic"))
+            .collect();
+
+        assert_eq!(results, vec![true, true]);
+    }
+
+    #[test]
+    fn async_runner_surfaces_serializer_panics() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = PyModule::new(py, "runtime").expect("module should be created");
+            module
+                .add_function(
+                    wrap_pyfunction!(async_serialization_panic, &module)
+                        .expect("function should wrap"),
+                )
+                .expect("function should register");
+            let locals = PyDict::new(py);
+            locals
+                .set_item("runtime", &module)
+                .expect("module should enter Python locals");
+            let code = CString::new(
+                r#"
+import asyncio
+
+async def exercise():
+    try:
+        await runtime.async_serialization_panic()
+    except BaseException as error:
+        assert type(error).__name__ == "PanicException"
+        assert str(error) == "serializer panicked"
+    else:
+        raise AssertionError("serializer panic was not raised")
+
+asyncio.run(exercise())
+"#,
+            )
+            .expect("Python source should not contain null bytes");
+            py.run(&code, Some(&locals), Some(&locals))
+                .expect("serializer panic should reach the Python awaiter");
+        });
+    }
+
+    #[test]
+    fn async_result_delivery_does_not_stall_tokio_workers() {
+        Python::initialize();
+        ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst);
+        Python::attach(|py| {
+            let module = PyModule::new(py, "runtime").expect("module should be created");
+            for function in [
+                wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"),
+                wrap_pyfunction!(runtime_worker_count, &module).expect("function should wrap"),
+                wrap_pyfunction!(runtime_is_responsive, &module).expect("function should wrap"),
+            ] {
+                module
+                    .add_function(function)
+                    .expect("function should register");
+            }
+            let locals = PyDict::new(py);
+            locals
+                .set_item("runtime", &module)
+                .expect("module should enter Python locals");
+            let code = CString::new(
+                r#"
+import asyncio
+
+async def exercise():
+    worker_count = runtime.runtime_worker_count()
+    awaitables = [runtime.async_runtime_probe() for _ in range(worker_count)]
+    assert runtime.runtime_is_responsive(worker_count)
+    assert await asyncio.gather(*awaitables) == [True] * worker_count
+
+asyncio.run(exercise())
+"#,
+            )
+            .expect("Python source should not contain null bytes");
+            py.run(&code, Some(&locals), Some(&locals))
+                .expect("result delivery should leave Tokio workers responsive");
+        });
+    }
+}

litellm-rust/crates/python-bridge/src/function_trace.rs

@@ -0,0 +1,216 @@
+use std::future::Future;
+use std::sync::{Arc, Mutex};
+
+use serde::Serialize;
+use tracing::instrument::WithSubscriber;
+use tracing::span::{Attributes, Id};
+use tracing::{Dispatch, Level, Subscriber};
+use tracing_subscriber::filter::{LevelFilter, filter_fn};
+use tracing_subscriber::layer::Context;
+use tracing_subscriber::prelude::*;
+use tracing_subscriber::registry::LookupSpan;
+use tracing_subscriber::{Layer, Registry};
+
+use crate::constants::FUNCTION_TRACE_TARGET;
+
+#[derive(Serialize)]
+#[serde(untagged)]
+pub(crate) enum TraceResponse<T> {
+    Plain(T),
+    Traced {
+        response: T,
+        trace: Vec<FunctionTraceEvent>,
+    },
+}
+
+pub(crate) async fn trace_call<T, E>(
+    future: impl Future<Output = Result<T, E>>,
+    enabled: bool,
+) -> Result<TraceResponse<T>, E> {
+    if !enabled {
+        return future.await.map(TraceResponse::Plain);
+    }
+    let trace = FunctionTrace::default();
+    let response = future.with_subscriber(trace.dispatcher()).await?;
+    Ok(TraceResponse::Traced {
+        response,
+        trace: trace.events(),
+    })
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize)]
+pub struct FunctionTraceEvent {
+    pub function: &'static str,
+    pub depth: usize,
+}
+
+#[derive(Clone, Default)]
+pub struct FunctionTrace {
+    events: Arc<Mutex<Vec<FunctionTraceEvent>>>,
+}
+
+impl FunctionTrace {
+    pub fn dispatcher(&self) -> Dispatch {
+        let filter = filter_fn(|metadata| {
+            metadata.is_span()
+                && metadata.target() == FUNCTION_TRACE_TARGET
+                && *metadata.level() == Level::TRACE
+        })
+        .with_max_level_hint(LevelFilter::TRACE);
+        Dispatch::new(
+            Registry::default().with(
+                FunctionTraceLayer {
+                    trace: self.clone(),
+                }
+                .with_filter(filter),
+            ),
+        )
+    }
+
+    pub fn events(&self) -> Vec<FunctionTraceEvent> {
+        self.events
+            .lock()
+            .unwrap_or_else(|error| error.into_inner())
+            .clone()
+    }
+}
+
+struct FunctionTraceLayer {
+    trace: FunctionTrace,
+}
+
+impl<S> Layer<S> for FunctionTraceLayer
+where
+    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
+{
+    fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) {
+        let depth = context
+            .span(id)
+            .map(|span| span.scope().skip(1).count())
+            .unwrap_or_default();
+        self.trace
+            .events
+            .lock()
+            .unwrap_or_else(|error| error.into_inner())
+            .push(FunctionTraceEvent {
+                function: attributes.metadata().name(),
+                depth,
+            });
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+    async fn outer() {
+        tokio::task::yield_now().await;
+        inner().await;
+    }
+
+    #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+    async fn inner() {
+        tokio::task::yield_now().await;
+    }
+
+    #[tokio::test]
+    async fn concurrent_futures_keep_separate_traces_across_yields() {
+        use tracing::instrument::WithSubscriber;
+
+        let first = FunctionTrace::default();
+        let second = FunctionTrace::default();
+        let outside = FunctionTrace::default();
+
+        async {
+            tokio::join!(
+                outer().with_subscriber(first.dispatcher()),
+                inner().with_subscriber(second.dispatcher()),
+            );
+            inner().await;
+        }
+        .with_subscriber(outside.dispatcher())
+        .await;
+
+        assert_eq!(
+            first.events(),
+            vec![
+                FunctionTraceEvent {
+                    function: "outer",
+                    depth: 0
+                },
+                FunctionTraceEvent {
+                    function: "inner",
+                    depth: 1
+                },
+            ],
+        );
+        assert_eq!(
+            second.events(),
+            vec![FunctionTraceEvent {
+                function: "inner",
+                depth: 0
+            }],
+        );
+        assert_eq!(
+            outside.events(),
+            vec![FunctionTraceEvent {
+                function: "inner",
+                depth: 0
+            }],
+        );
+    }
+
+    #[test]
+    fn records_matching_spans_in_creation_order() {
+        let trace = FunctionTrace::default();
+        let dispatch = trace.dispatcher();
+
+        tracing::dispatcher::with_default(&dispatch, || {
+            let _ignored = tracing::trace_span!(target: "other", "ignored");
+            let _first = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
+            let _wrong_level = tracing::debug_span!(target: FUNCTION_TRACE_TARGET, "wrong_level");
+            let _second = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
+        });
+
+        assert_eq!(
+            trace.events(),
+            vec![
+                FunctionTraceEvent {
+                    function: "same_name",
+                    depth: 0,
+                },
+                FunctionTraceEvent {
+                    function: "same_name",
+                    depth: 0,
+                },
+            ]
+        );
+    }
+
+    #[test]
+    fn records_matching_span_nesting_depth() {
+        let trace = FunctionTrace::default();
+        let dispatch = trace.dispatcher();
+
+        tracing::dispatcher::with_default(&dispatch, || {
+            let outer = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "outer");
+            let _outer_guard = outer.enter();
+            let _inner = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "inner");
+        });
+
+        assert_eq!(
+            trace.events(),
+            vec![
+                FunctionTraceEvent {
+                    function: "outer",
+                    depth: 0,
+                },
+                FunctionTraceEvent {
+                    function: "inner",
+                    depth: 1,
+                },
+            ]
+        );
+    }
+}

litellm-rust/crates/python-bridge/src/lib.rs

@@ -1,142 +1,18 @@
-use std::collections::HashMap;
-use std::time::Duration;
+mod constants;
+mod diagnostics;
+mod errors;
+mod execution;
+pub mod function_trace;
+mod marshal;
+mod routes;
 
-use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
 use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
-use litellm_core::audio_transcription::{
-    AudioTranscriptionRequest, audio_transcription as run_audio_transcription,
-};
-use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse};
-use litellm_core::chat_completions::{
-    chat_completions as run_chat_completions, chat_completions_decline_reason,
-};
-use litellm_core::error::Error;
-use litellm_core::messages::messages as run_messages;
-use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest};
-use litellm_python_interop::{from_py, release_count, release_gil, to_py};
-use pyo3::exceptions::{PyRuntimeError, PyValueError};
 use pyo3::prelude::*;
-use pyo3::types::{PyAny, PyDict};
-use serde_json::{Map, Value};
+use pyo3::types::PyAny;
+use serde_json::Value;
 
-pyo3::create_exception!(
-    _native,
-    RustBridgeDeclined,
-    pyo3::exceptions::PyException,
-    "The route declined before calling the provider, so the host may retry on its own path."
-);
-
-pyo3::create_exception!(
-    _native,
-    RustUpstreamError,
-    pyo3::exceptions::PyException,
-    "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response."
-);
-
-type MarshaledOcrInputs = (
-    Value,
-    Option<Map<String, Value>>,
-    Map<String, Value>,
-    Option<Duration>,
-);
-
-fn messages_response_to_py(
-    py: Python<'_>,
-    response: AnthropicMessagesResponse,
-) -> PyResult<Py<PyAny>> {
-    to_py(py, &response)
-}
-
-fn chat_completions_response_to_py(
-    py: Python<'_>,
-    response: ChatCompletionsResponse,
-) -> PyResult<Py<PyAny>> {
-    to_py(py, &response)
-}
-
-fn core_error_to_pyerr(err: Error) -> PyErr {
-    match err {
-        Error::Auth(message) => PyValueError::new_err(message),
-        Error::InvalidProvider(_)
-        | Error::InvalidRequest(_)
-        | Error::InvalidType { .. }
-        | Error::MissingField(_) => PyValueError::new_err(err.to_string()),
-        other => PyRuntimeError::new_err(other.to_string()),
-    }
-}
-
-/// Map a core error for a route whose host keeps a Python implementation.
-///
-/// The distinction the host needs is whether the provider was already called.
-/// Everything raised before the request goes out is safe for the host to retry
-/// on its own path; anything after it is not, because the provider has already
-/// done the work and billed for it.
-fn chat_completions_error_to_pyerr(err: Error) -> PyErr {
-    match err {
-        Error::Unsupported(_)
-        | Error::Auth(_)
-        | Error::InvalidProvider(_)
-        | Error::InvalidRequest(_)
-        | Error::InvalidType { .. }
-        | Error::MissingField(_)
-        | Error::Routing(_)
-        // Nothing reached the provider, so serving it on Python cannot double
-        // bill and is the only way the caller gets an answer at all.
-        | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()),
-        Error::Http { status, body } => {
-            RustUpstreamError::new_err((status, format!("{status}: {body}")))
-        }
-        Error::Network(message) | Error::InvalidResponse(message) => {
-            RustUpstreamError::new_err((0u16, message))
-        }
-    }
-}
-
-fn optional_object_to_map(
-    py: Python<'_>,
-    name: &'static str,
-    value: Option<Py<PyAny>>,
-) -> PyResult<Map<String, Value>> {
-    match value {
-        Some(value) => match from_py(value.bind(py))? {
-            Value::Object(map) => Ok(map),
-            _ => Err(PyValueError::new_err(format!("{name} must be a dict"))),
-        },
-        None => Ok(Map::new()),
-    }
-}
-
-fn optional_timeout(timeout_seconds: Option<f64>) -> Option<Duration> {
-    timeout_seconds.and_then(|secs| {
-        if secs.is_finite() && secs > 0.0 {
-            Some(Duration::from_secs_f64(secs))
-        } else {
-            None
-        }
-    })
-}
-
-fn marshal_headers(
-    py: Python<'_>,
-    headers: Option<Py<PyAny>>,
-) -> PyResult<HashMap<String, String>> {
-    let value = match headers {
-        Some(headers) => from_py(headers.bind(py))?,
-        None => Value::Object(Map::new()),
-    };
-    let Value::Object(headers) = value else {
-        return Err(PyValueError::new_err("headers must be a dict"));
-    };
-    headers
-        .into_iter()
-        .map(|(name, value)| {
-            value
-                .as_str()
-                .map(|value| (name, value.to_string()))
-                .ok_or_else(|| PyValueError::new_err("header values must be strings"))
-        })
-        .collect()
-}
+use crate::errors::core_error_to_pyerr;
+use crate::marshal::{marshal_headers, optional_timeout};
 
 #[pyclass]
 struct ResponsesWebSocketConnection {
@@ -151,16 +27,16 @@ impl ResponsesWebSocketConnection {
         _cls: &Bound<'py, pyo3::types::PyType>,
         py: Python<'py>,
         url: String,
-        headers: Option<Py<PyAny>>,
+        #[pyo3(from_py_with = litellm_python_interop::from_py)] headers: Option<Value>,
         timeout_seconds: Option<f64>,
     ) -> PyResult<Bound<'py, PyAny>> {
-        let headers = marshal_headers(py, headers)?;
+        let headers = marshal_headers(headers)?;
         let timeout = optional_timeout(timeout_seconds);
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
             let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout)
                 .await
                 .map_err(core_error_to_pyerr)?;
-            Python::attach(|py| Py::new(py, ResponsesWebSocketConnection { inner }))
+            Ok(ResponsesWebSocketConnection { inner })
         })
     }
 
@@ -186,445 +62,126 @@ impl ResponsesWebSocketConnection {
     }
 }
 
-fn marshal_inputs(
-    py: Python<'_>,
-    document: Py<PyAny>,
-    extra_headers: Option<Py<PyAny>>,
-    optional_params: Option<Py<PyAny>>,
-    timeout_seconds: Option<f64>,
-) -> PyResult<MarshaledOcrInputs> {
-    let document = from_py(document.bind(py))?;
-    let extra_headers = match extra_headers {
-        Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
-        None => None,
-    };
-    let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
-    let timeout = optional_timeout(timeout_seconds);
+#[pymodule(gil_used = false)]
+mod _native {
+    use pyo3::prelude::*;
 
-    Ok((document, extra_headers, optional_params, timeout))
-}
-
-#[pyfunction]
-#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
-#[allow(clippy::too_many_arguments)]
-fn ocr(
-    py: Python<'_>,
-    model: String,
-    document: Py<PyAny>,
-    api_key: Option<String>,
-    api_base: Option<String>,
-    custom_llm_provider: Option<String>,
-    extra_headers: Option<Py<PyAny>>,
-    optional_params: Option<Py<PyAny>>,
-    timeout_seconds: Option<f64>,
-) -> PyResult<Py<PyAny>> {
-    let (document, extra_headers, optional_params, timeout) = marshal_inputs(
-        py,
-        document,
-        extra_headers,
-        optional_params,
-        timeout_seconds,
-    )?;
-
-    let result = release_gil(py, || {
-        pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest {
-            model: &model,
-            document,
-            api_key: api_key.as_deref(),
-            api_base: api_base.as_deref(),
-            custom_llm_provider: custom_llm_provider.as_deref(),
-            extra_headers,
-            optional_params,
-            timeout,
-            callbacks: Vec::new(),
-            guardrails: Vec::new(),
-            request_metadata: Default::default(),
-            litellm_call_id: None,
-        }))
-    });
-
-    match result {
-        Ok(value) => to_py(py, &value),
-        Err(err) => Err(core_error_to_pyerr(err)),
+    #[pymodule_init]
+    fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
+        super::errors::register(module)?;
+        super::routes::register(module)?;
+        module.add_class::<super::ResponsesWebSocketConnection>()?;
+        super::diagnostics::register(module)
     }
 }
 
-#[pyfunction]
-#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
-#[allow(clippy::too_many_arguments)]
-fn aocr(
-    py: Python<'_>,
-    model: String,
-    document: Py<PyAny>,
-    api_key: Option<String>,
-    api_base: Option<String>,
-    custom_llm_provider: Option<String>,
-    extra_headers: Option<Py<PyAny>>,
-    optional_params: Option<Py<PyAny>>,
-    timeout_seconds: Option<f64>,
-) -> PyResult<Bound<'_, PyAny>> {
-    let (document, extra_headers, optional_params, timeout) = marshal_inputs(
-        py,
-        document,
-        extra_headers,
-        optional_params,
-        timeout_seconds,
-    )?;
+#[cfg(test)]
+mod tests {
+    use std::ffi::CString;
+    use std::time::Duration;
 
-    pyo3_async_runtimes::tokio::future_into_py(py, async move {
-        let value = run_ocr(OcrRequest {
-            model: &model,
-            document,
-            api_key: api_key.as_deref(),
-            api_base: api_base.as_deref(),
-            custom_llm_provider: custom_llm_provider.as_deref(),
-            extra_headers,
-            optional_params,
-            timeout,
-            callbacks: Vec::new(),
-            guardrails: Vec::new(),
-            request_metadata: Default::default(),
-            litellm_call_id: None,
-        })
-        .await
-        .map_err(core_error_to_pyerr)?;
+    use futures_util::{SinkExt, StreamExt};
+    use pyo3::types::PyDict;
+    use tokio::net::TcpListener;
+    use tokio_tungstenite::{accept_async, tungstenite::Message};
 
-        Python::attach(|py| to_py(py, &value))
-    })
-}
+    use super::*;
 
-#[pyfunction]
-#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
-#[allow(clippy::too_many_arguments)]
-fn transcription(
-    py: Python<'_>,
-    model: String,
-    audio: Py<PyAny>,
-    api_key: Option<String>,
-    api_base: Option<String>,
-    custom_llm_provider: Option<String>,
-    extra_headers: Option<Py<PyAny>>,
-    optional_params: Option<Py<PyAny>>,
-    timeout_seconds: Option<f64>,
-) -> PyResult<Py<PyAny>> {
-    let audio = from_py(audio.bind(py))?;
-    let extra_headers = match extra_headers {
-        Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
-        None => None,
-    };
-    let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
-    let timeout = optional_timeout(timeout_seconds);
-    let result = release_gil(py, || {
-        pyo3_async_runtimes::tokio::get_runtime().block_on(run_audio_transcription(
-            AudioTranscriptionRequest {
-                model: &model,
-                audio,
-                api_key: api_key.as_deref(),
-                api_base: api_base.as_deref(),
-                custom_llm_provider: custom_llm_provider.as_deref(),
-                extra_headers,
-                optional_params,
-                timeout,
-            },
-        ))
-    });
-    match result {
-        Ok(value) => to_py(py, &value),
-        Err(err) => Err(core_error_to_pyerr(err)),
+    #[test]
+    fn module_registration_preserves_the_public_surface() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = pyo3::wrap_pymodule!(_native)(py).into_bound(py);
+
+            let expected = [
+                "RustBridgeDeclined",
+                "RustUpstreamError",
+                "ocr",
+                "aocr",
+                "transcription",
+                "atranscription",
+                "messages",
+                "amessages",
+                "chat_completions_decline",
+                "chat_completions",
+                "achat_completions",
+                "ResponsesWebSocketConnection",
+                "gil_stats",
+            ];
+
+            let public_names: Vec<String> = module
+                .dict()
+                .keys()
+                .extract::<Vec<String>>()
+                .expect("module names should be strings")
+                .into_iter()
+                .filter(|name| !name.starts_with("__"))
+                .collect();
+            assert_eq!(public_names, expected);
+        });
+    }
+
+    #[test]
+    fn responses_websocket_connection_round_trips_through_python() {
+        Python::initialize();
+        let runtime = pyo3_async_runtimes::tokio::get_runtime();
+        let listener = runtime
+            .block_on(TcpListener::bind("127.0.0.1:0"))
+            .expect("listener should bind");
+        let address = listener
+            .local_addr()
+            .expect("listener should have an address");
+        let server = runtime.spawn(async move {
+            let (stream, _) = listener.accept().await.expect("server should accept");
+            let mut socket = accept_async(stream)
+                .await
+                .expect("handshake should succeed");
+
+            let message = socket
+                .next()
+                .await
+                .expect("client should send a frame")
+                .expect("client frame should be valid");
+            assert_eq!(message, Message::Text("from-python".into()));
+            socket
+                .send(Message::Text("from-server".into()))
+                .await
+                .expect("server should reply");
+            assert!(matches!(socket.next().await, Some(Ok(Message::Close(_)))));
+        });
+
+        Python::attach(|py| {
+            let module = pyo3::wrap_pymodule!(_native)(py).into_bound(py);
+            let locals = PyDict::new(py);
+            locals
+                .set_item("native", &module)
+                .expect("module should enter Python locals");
+            locals
+                .set_item("url", format!("ws://{address}"))
+                .expect("URL should enter Python locals");
+            let code = CString::new(
+                r#"
+import asyncio
+
+async def exercise():
+    connection = await native.ResponsesWebSocketConnection.connect(url)
+    assert type(connection) is native.ResponsesWebSocketConnection
+    await connection.send_text("from-python")
+    assert await connection.recv_text() == "from-server"
+    await connection.close()
+    assert await connection.recv_text() is None
+
+asyncio.run(asyncio.wait_for(exercise(), timeout=5))
+"#,
+            )
+            .expect("Python source should not contain null bytes");
+            py.run(&code, Some(&locals), Some(&locals))
+                .expect("Python WebSocket methods should round trip");
+        });
+
+        runtime
+            .block_on(async { tokio::time::timeout(Duration::from_secs(5), server).await })
+            .expect("server should finish")
+            .expect("server task should not panic");
     }
 }
-
-#[pyfunction]
-#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
-#[allow(clippy::too_many_arguments)]
-fn atranscription(
-    py: Python<'_>,
-    model: String,
-    audio: Py<PyAny>,
-    api_key: Option<String>,
-    api_base: Option<String>,
-    custom_llm_provider: Option<String>,
-    extra_headers: Option<Py<PyAny>>,
-    optional_params: Option<Py<PyAny>>,
-    timeout_seconds: Option<f64>,
-) -> PyResult<Bound<'_, PyAny>> {
-    let audio = from_py(audio.bind(py))?;
-    let extra_headers = match extra_headers {
-        Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
-        None => None,
-    };
-    let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
-    let timeout = optional_timeout(timeout_seconds);
-    pyo3_async_runtimes::tokio::future_into_py(py, async move {
-        let value = run_audio_transcription(AudioTranscriptionRequest {
-            model: &model,
-            audio,
-            api_key: api_key.as_deref(),
-            api_base: api_base.as_deref(),
-            custom_llm_provider: custom_llm_provider.as_deref(),
-            extra_headers,
-            optional_params,
-            timeout,
-        })
-        .await
-        .map_err(core_error_to_pyerr)?;
-        Python::attach(|py| to_py(py, &value))
-    })
-}
-
-type MarshaledMessagesInputs = (Value, Option<Map<String, Value>>, Option<Duration>);
-
-fn marshal_messages_inputs(
-    py: Python<'_>,
-    body: Py<PyAny>,
-    extra_headers: Option<Py<PyAny>>,
-    timeout_seconds: Option<f64>,
-) -> PyResult<MarshaledMessagesInputs> {
-    let body: Value = from_py(body.bind(py))?;
-    if !body.is_object() {
-        return Err(PyValueError::new_err("body must be a dict"));
-    }
-    let extra_headers = match extra_headers {
-        Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
-        None => None,
-    };
-    Ok((body, extra_headers, optional_timeout(timeout_seconds)))
-}
-
-#[pyfunction]
-#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
-#[allow(clippy::too_many_arguments)]
-fn messages(
-    py: Python<'_>,
-    model: String,
-    body: Py<PyAny>,
-    api_key: Option<String>,
-    api_base: Option<String>,
-    custom_llm_provider: Option<String>,
-    extra_headers: Option<Py<PyAny>>,
-    timeout_seconds: Option<f64>,
-) -> PyResult<Py<PyAny>> {
-    let (body, extra_headers, timeout) =
-        marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?;
-
-    let result = release_gil(py, || {
-        pyo3_async_runtimes::tokio::get_runtime().block_on(run_messages(MessagesRequest {
-            model: &model,
-            body,
-            api_key: api_key.as_deref(),
-            api_base: api_base.as_deref(),
-            custom_llm_provider: custom_llm_provider.as_deref(),
-            extra_headers,
-            timeout,
-        }))
-    });
-
-    match result {
-        Ok(response) => messages_response_to_py(py, response),
-        Err(err) => Err(core_error_to_pyerr(err)),
-    }
-}
-
-#[pyfunction]
-#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
-#[allow(clippy::too_many_arguments)]
-fn amessages(
-    py: Python<'_>,
-    model: String,
-    body: Py<PyAny>,
-    api_key: Option<String>,
-    api_base: Option<String>,
-    custom_llm_provider: Option<String>,
-    extra_headers: Option<Py<PyAny>>,
-    timeout_seconds: Option<f64>,
-) -> PyResult<Bound<'_, PyAny>> {
-    let (body, extra_headers, timeout) =
-        marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?;
-
-    pyo3_async_runtimes::tokio::future_into_py(py, async move {
-        let response = run_messages(MessagesRequest {
-            model: &model,
-            body,
-            api_key: api_key.as_deref(),
-            api_base: api_base.as_deref(),
-            custom_llm_provider: custom_llm_provider.as_deref(),
-            extra_headers,
-            timeout,
-        })
-        .await
-        .map_err(core_error_to_pyerr)?;
-
-        Python::attach(|py| messages_response_to_py(py, response))
-    })
-}
-
-type MarshaledChatCompletionsInputs = (
-    Value,
-    Map<String, Value>,
-    Option<Map<String, Value>>,
-    Option<Duration>,
-);
-
-fn marshal_chat_completions_inputs(
-    py: Python<'_>,
-    messages: Py<PyAny>,
-    optional_params: Option<Py<PyAny>>,
-    extra_headers: Option<Py<PyAny>>,
-    timeout_seconds: Option<f64>,
-) -> PyResult<MarshaledChatCompletionsInputs> {
-    let messages: Value = from_py(messages.bind(py))?;
-    if !messages.is_array() {
-        return Err(PyValueError::new_err("messages must be a list"));
-    }
-    let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
-    let extra_headers = match extra_headers {
-        Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
-        None => None,
-    };
-    Ok((
-        messages,
-        optional_params,
-        extra_headers,
-        optional_timeout(timeout_seconds),
-    ))
-}
-
-/// The decline reason for this request, or `None` when the Rust path accepts
-/// it. Resolves no credentials and performs no I/O, so a host can ask before
-/// committing to either path.
-#[pyfunction]
-#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))]
-fn chat_completions_decline(
-    py: Python<'_>,
-    model: String,
-    messages: Py<PyAny>,
-    optional_params: Option<Py<PyAny>>,
-    custom_llm_provider: Option<String>,
-) -> PyResult<Option<String>> {
-    let messages = from_py(messages.bind(py))?;
-    let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
-    Ok(chat_completions_decline_reason(
-        &model,
-        custom_llm_provider.as_deref(),
-        messages,
-        &optional_params,
-    )
-    .map(str::to_string))
-}
-
-#[pyfunction]
-#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
-#[allow(clippy::too_many_arguments)]
-fn chat_completions(
-    py: Python<'_>,
-    model: String,
-    messages: Py<PyAny>,
-    optional_params: Option<Py<PyAny>>,
-    api_key: Option<String>,
-    api_base: Option<String>,
-    custom_llm_provider: Option<String>,
-    extra_headers: Option<Py<PyAny>>,
-    timeout_seconds: Option<f64>,
-) -> PyResult<Py<PyAny>> {
-    let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs(
-        py,
-        messages,
-        optional_params,
-        extra_headers,
-        timeout_seconds,
-    )?;
-
-    let result = release_gil(py, || {
-        pyo3_async_runtimes::tokio::get_runtime().block_on(run_chat_completions(
-            ChatCompletionsRequest {
-                model: &model,
-                messages,
-                optional_params,
-                api_key: api_key.as_deref(),
-                api_base: api_base.as_deref(),
-                custom_llm_provider: custom_llm_provider.as_deref(),
-                extra_headers,
-                timeout,
-            },
-        ))
-    });
-
-    match result {
-        Ok(response) => chat_completions_response_to_py(py, response),
-        Err(err) => Err(chat_completions_error_to_pyerr(err)),
-    }
-}
-
-#[pyfunction]
-#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
-#[allow(clippy::too_many_arguments)]
-fn achat_completions(
-    py: Python<'_>,
-    model: String,
-    messages: Py<PyAny>,
-    optional_params: Option<Py<PyAny>>,
-    api_key: Option<String>,
-    api_base: Option<String>,
-    custom_llm_provider: Option<String>,
-    extra_headers: Option<Py<PyAny>>,
-    timeout_seconds: Option<f64>,
-) -> PyResult<Bound<'_, PyAny>> {
-    let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs(
-        py,
-        messages,
-        optional_params,
-        extra_headers,
-        timeout_seconds,
-    )?;
-
-    pyo3_async_runtimes::tokio::future_into_py(py, async move {
-        let response = run_chat_completions(ChatCompletionsRequest {
-            model: &model,
-            messages,
-            optional_params,
-            api_key: api_key.as_deref(),
-            api_base: api_base.as_deref(),
-            custom_llm_provider: custom_llm_provider.as_deref(),
-            extra_headers,
-            timeout,
-        })
-        .await
-        .map_err(chat_completions_error_to_pyerr)?;
-
-        Python::attach(|py| chat_completions_response_to_py(py, response))
-    })
-}
-
-#[pyfunction]
-fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
-    let stats = PyDict::new(py);
-    stats.set_item("releases", release_count())?;
-    Ok(stats.into_any().unbind())
-}
-
-#[cfg(feature = "panic-test")]
-#[pyfunction]
-fn _panic_for_test() {
-    panic!("intentional PyO3 panic smoke test");
-}
-
-#[pymodule]
-fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> {
-    let py = module.py();
-    module.add_function(wrap_pyfunction!(ocr, module)?)?;
-    module.add_function(wrap_pyfunction!(aocr, module)?)?;
-    module.add_function(wrap_pyfunction!(transcription, module)?)?;
-    module.add_function(wrap_pyfunction!(atranscription, module)?)?;
-    module.add_function(wrap_pyfunction!(messages, module)?)?;
-    module.add_function(wrap_pyfunction!(amessages, module)?)?;
-    module.add("RustBridgeDeclined", py.get_type::<RustBridgeDeclined>())?;
-    module.add("RustUpstreamError", py.get_type::<RustUpstreamError>())?;
-    module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?;
-    module.add_function(wrap_pyfunction!(chat_completions, module)?)?;
-    module.add_function(wrap_pyfunction!(achat_completions, module)?)?;
-    module.add_class::<ResponsesWebSocketConnection>()?;
-    module.add_function(wrap_pyfunction!(gil_stats, module)?)?;
-    #[cfg(feature = "panic-test")]
-    module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?;
-    Ok(())
-}

litellm-rust/crates/python-bridge/src/marshal.rs

@@ -0,0 +1,104 @@
+use std::collections::HashMap;
+use std::time::Duration;
+
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use serde_json::{Map, Value};
+
+pub(crate) struct RouteOptions {
+    pub(crate) model: String,
+    pub(crate) api_key: Option<String>,
+    pub(crate) api_base: Option<String>,
+    pub(crate) custom_llm_provider: Option<String>,
+    pub(crate) extra_headers: Option<Map<String, Value>>,
+    pub(crate) timeout: Option<Duration>,
+}
+
+pub(crate) struct RouteOptionsInputs {
+    pub(crate) model: String,
+    pub(crate) api_key: Option<String>,
+    pub(crate) api_base: Option<String>,
+    pub(crate) custom_llm_provider: Option<String>,
+    pub(crate) extra_headers: Option<Value>,
+    pub(crate) timeout_seconds: Option<f64>,
+}
+
+impl RouteOptions {
+    pub(crate) fn from_python(inputs: RouteOptionsInputs) -> PyResult<Self> {
+        Ok(Self {
+            model: inputs.model,
+            api_key: inputs.api_key,
+            api_base: inputs.api_base,
+            custom_llm_provider: inputs.custom_llm_provider,
+            extra_headers: optional_object("extra_headers", inputs.extra_headers)?,
+            timeout: optional_timeout(inputs.timeout_seconds),
+        })
+    }
+}
+
+pub(crate) fn required_value(
+    name: &'static str,
+    value: Value,
+    expected: fn(&Value) -> bool,
+    expected_name: &'static str,
+) -> PyResult<Value> {
+    if expected(&value) {
+        return Ok(value);
+    }
+    Err(PyValueError::new_err(format!(
+        "{name} must be a {expected_name}"
+    )))
+}
+
+pub(crate) fn object_or_empty(
+    name: &'static str,
+    value: Option<Value>,
+) -> PyResult<Map<String, Value>> {
+    match value {
+        Some(value) => object(name, value),
+        None => Ok(Map::new()),
+    }
+}
+
+fn optional_object(
+    name: &'static str,
+    value: Option<Value>,
+) -> PyResult<Option<Map<String, Value>>> {
+    value.map(|value| object(name, value)).transpose()
+}
+
+fn object(name: &'static str, value: Value) -> PyResult<Map<String, Value>> {
+    match value {
+        Value::Object(map) => Ok(map),
+        _ => Err(PyValueError::new_err(format!("{name} must be a dict"))),
+    }
+}
+
+pub(crate) fn optional_timeout(timeout_seconds: Option<f64>) -> Option<Duration> {
+    timeout_seconds.and_then(|secs| {
+        if secs.is_finite() && secs > 0.0 {
+            Some(Duration::from_secs_f64(secs))
+        } else {
+            None
+        }
+    })
+}
+
+pub(crate) fn marshal_headers(headers: Option<Value>) -> PyResult<HashMap<String, String>> {
+    let value = match headers {
+        Some(headers) => headers,
+        None => Value::Object(Map::new()),
+    };
+    let Value::Object(headers) = value else {
+        return Err(PyValueError::new_err("headers must be a dict"));
+    };
+    headers
+        .into_iter()
+        .map(|(name, value)| {
+            value
+                .as_str()
+                .map(|value| (name, value.to_string()))
+                .ok_or_else(|| PyValueError::new_err("header values must be strings"))
+        })
+        .collect()
+}

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

@@ -0,0 +1,71 @@
+use litellm_core::Error;
+use std::future::Future;
+
+use litellm_core::audio_transcription::{
+    AudioTranscriptionRequest, audio_transcription as run_audio_transcription,
+};
+use pyo3::prelude::*;
+use serde_json::Value;
+
+use crate::errors::core_error_to_pyerr;
+use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
+
+fn prepare_transcription(
+    inputs: AudioTranscriptionInputs,
+) -> PyResult<impl Future<Output = Result<Value, Error>> + Send + 'static> {
+    let audio = inputs.audio;
+    let options = RouteOptions::from_python(RouteOptionsInputs {
+        model: inputs.model,
+        api_key: inputs.api_key,
+        api_base: inputs.api_base,
+        custom_llm_provider: inputs.custom_llm_provider,
+        extra_headers: inputs.extra_headers,
+        timeout_seconds: inputs.timeout_seconds,
+    })?;
+    let optional_params = object_or_empty("optional_params", inputs.optional_params)?;
+
+    Ok(async move {
+        let RouteOptions {
+            model,
+            api_key,
+            api_base,
+            custom_llm_provider,
+            extra_headers,
+            timeout,
+        } = options;
+        run_audio_transcription(AudioTranscriptionRequest {
+            model: &model,
+            audio,
+            api_key: api_key.as_deref(),
+            api_base: api_base.as_deref(),
+            custom_llm_provider: custom_llm_provider.as_deref(),
+            extra_headers,
+            optional_params,
+            timeout,
+        })
+        .await
+    })
+}
+
+bridge_route! {
+    sync = transcription,
+    asynchronous = atranscription,
+    inputs = AudioTranscriptionInputs,
+    required = {
+        model: String,
+        #[pyo3(from_py_with = litellm_python_interop::from_py)]
+        audio: Value,
+    },
+    optional = {
+        api_key: Option<String>,
+        api_base: Option<String>,
+        custom_llm_provider: Option<String>,
+        #[pyo3(from_py_with = litellm_python_interop::from_py)]
+        extra_headers: Option<Value>,
+        #[pyo3(from_py_with = litellm_python_interop::from_py)]
+        optional_params: Option<Value>,
+        timeout_seconds: Option<f64>,
+    },
+    prepare = prepare_transcription,
+    errors = core_error_to_pyerr,
+}

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

@@ -0,0 +1,91 @@
+use litellm_core::Error;
+use std::future::Future;
+
+use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse};
+use litellm_core::chat_completions::{
+    chat_completions as run_chat_completions, chat_completions_decline_reason,
+};
+use pyo3::prelude::*;
+use serde_json::Value;
+
+use crate::errors::chat_completions_error_to_pyerr;
+use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_value};
+
+fn prepare_chat_completions(
+    inputs: ChatCompletionsInputs,
+) -> PyResult<impl Future<Output = Result<ChatCompletionsResponse, Error>> + Send + 'static> {
+    let messages = required_value("messages", inputs.messages, Value::is_array, "list")?;
+    let optional_params = object_or_empty("optional_params", inputs.optional_params)?;
+    let options = RouteOptions::from_python(RouteOptionsInputs {
+        model: inputs.model,
+        api_key: inputs.api_key,
+        api_base: inputs.api_base,
+        custom_llm_provider: inputs.custom_llm_provider,
+        extra_headers: inputs.extra_headers,
+        timeout_seconds: inputs.timeout_seconds,
+    })?;
+
+    Ok(async move {
+        let RouteOptions {
+            model,
+            api_key,
+            api_base,
+            custom_llm_provider,
+            extra_headers,
+            timeout,
+        } = options;
+        run_chat_completions(ChatCompletionsRequest {
+            model: &model,
+            messages,
+            optional_params,
+            api_key: api_key.as_deref(),
+            api_base: api_base.as_deref(),
+            custom_llm_provider: custom_llm_provider.as_deref(),
+            extra_headers,
+            timeout,
+        })
+        .await
+    })
+}
+
+#[pyfunction]
+#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))]
+fn chat_completions_decline(
+    model: String,
+    #[pyo3(from_py_with = litellm_python_interop::from_py)] messages: Value,
+    #[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option<Value>,
+    custom_llm_provider: Option<String>,
+) -> PyResult<Option<String>> {
+    let optional_params = object_or_empty("optional_params", optional_params)?;
+    Ok(chat_completions_decline_reason(
+        &model,
+        custom_llm_provider.as_deref(),
+        messages,
+        &optional_params,
+    )
+    .map(str::to_string))
+}
+
+bridge_route! {
+    sync = chat_completions,
+    asynchronous = achat_completions,
+    inputs = ChatCompletionsInputs,
+    required = {
+        model: String,
+        #[pyo3(from_py_with = litellm_python_interop::from_py)]
+        messages: Value,
+    },
+    optional = {
+        #[pyo3(from_py_with = litellm_python_interop::from_py)]
+        optional_params: Option<Value>,
+        api_key: Option<String>,
+        api_base: Option<String>,
+        custom_llm_provider: Option<String>,
+        #[pyo3(from_py_with = litellm_python_interop::from_py)]
+        extra_headers: Option<Value>,
+        timeout_seconds: Option<f64>,
+    },
+    prepare = prepare_chat_completions,
+    errors = chat_completions_error_to_pyerr,
+    extra = [chat_completions_decline],
+}

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

@@ -0,0 +1,429 @@
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
+use pyo3::types::PyCFunction;
+
+macro_rules! bridge_route {
+    (
+        sync = $sync_name:ident,
+        asynchronous = $async_name:ident,
+        inputs = $inputs:ident,
+        required = { $($(#[$required_attr:meta])* $required_name:ident: $required_type:ty),+ $(,)? },
+        optional = { $($(#[$optional_attr:meta])* $optional_name:ident: $optional_type:ty),* $(,)? },
+        prepare = $prepare:path,
+        errors = $map_error:path
+        $(, extra = [$($extra:ident),* $(,)?])?
+        $(,)?
+    ) => {
+        struct $inputs {
+            $($required_name: $required_type,)*
+            $($optional_name: $optional_type),*
+        }
+
+        #[pyfunction]
+        #[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))]
+        #[allow(clippy::too_many_arguments)]
+        fn $sync_name(
+            py: pyo3::Python<'_>,
+            $($(#[$required_attr])* $required_name: $required_type,)*
+            $($(#[$optional_attr])* $optional_name: $optional_type,)*
+            trace: bool,
+        ) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {
+            let future = $prepare($inputs {
+                $($required_name,)*
+                $($optional_name),*
+            })?;
+            $crate::execution::run_sync(
+                py,
+                $crate::function_trace::trace_call(future, trace),
+                $map_error,
+            )
+        }
+
+        #[pyfunction]
+        #[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))]
+        #[allow(clippy::too_many_arguments)]
+        fn $async_name(
+            py: pyo3::Python<'_>,
+            $($(#[$required_attr])* $required_name: $required_type,)*
+            $($(#[$optional_attr])* $optional_name: $optional_type,)*
+            trace: bool,
+        ) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
+            let future = $prepare($inputs {
+                $($required_name,)*
+                $($optional_name),*
+            })?;
+            $crate::execution::run_async(
+                py,
+                $crate::function_trace::trace_call(future, trace),
+                $map_error,
+            )
+        }
+
+        pub(super) fn register(
+            module: &pyo3::Bound<'_, pyo3::types::PyModule>,
+        ) -> pyo3::PyResult<()> {
+            $($($crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($extra, module)?)?;)*)?
+            $crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($sync_name, module)?)?;
+            $crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($async_name, module)?)?;
+            Ok(())
+        }
+    };
+}
+
+pub(super) fn add_function(
+    module: &Bound<'_, PyModule>,
+    function: Bound<'_, PyCFunction>,
+) -> PyResult<()> {
+    let name: String = function.getattr("__name__")?.extract()?;
+    if module.hasattr(&name)? {
+        return Err(PyRuntimeError::new_err(format!(
+            "duplicate native route: {name}"
+        )));
+    }
+    module.add_function(function)
+}
+
+#[cfg(test)]
+mod tests {
+    use std::ffi::CString;
+    use std::sync::atomic::{AtomicBool, Ordering};
+
+    use litellm_core::error::Error;
+    use pyo3::exceptions::PyLookupError;
+    use pyo3::types::{PyDict, PyList};
+
+    use super::*;
+
+    mod synthetic {
+        use std::future::{Future, pending};
+
+        use super::*;
+
+        static FUTURE_DROPPED: AtomicBool = AtomicBool::new(false);
+
+        struct DropGuard;
+
+        impl Drop for DropGuard {
+            fn drop(&mut self) {
+                FUTURE_DROPPED.store(true, Ordering::SeqCst);
+            }
+        }
+
+        #[pyfunction]
+        fn future_dropped() -> bool {
+            FUTURE_DROPPED.load(Ordering::SeqCst)
+        }
+
+        bridge_route! {
+            sync = echo,
+            asynchronous = aecho,
+            inputs = EchoInputs,
+            required = { value: String },
+            optional = {},
+            prepare = prepare_echo,
+            errors = map_error,
+            extra = [future_dropped],
+        }
+
+        fn prepare_echo(
+            inputs: EchoInputs,
+        ) -> PyResult<impl Future<Output = Result<String, Error>> + Send + 'static> {
+            FUTURE_DROPPED.store(false, Ordering::SeqCst);
+            let drop_guard = (inputs.value == "pending").then_some(DropGuard);
+            Ok(async move {
+                let _drop_guard = drop_guard;
+                tokio::task::yield_now().await;
+                match inputs.value.as_str() {
+                    "error" => Err(Error::InvalidRequest("synthetic error".to_string())),
+                    "map_panic" => Err(Error::InvalidRequest("panic in mapper".to_string())),
+                    "panic" => panic!("synthetic panic"),
+                    "pending" => {
+                        pending::<()>().await;
+                        unreachable!()
+                    }
+                    _ => Ok(inputs.value),
+                }
+            })
+        }
+
+        fn map_error(error: Error) -> PyErr {
+            if matches!(&error, Error::InvalidRequest(message) if message == "panic in mapper") {
+                panic!("synthetic mapper panic")
+            }
+            PyLookupError::new_err(error.to_string())
+        }
+    }
+
+    #[test]
+    fn sync_and_async_route_signatures_match_the_python_contract() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = PyModule::new(py, "routes").expect("module should be created");
+            crate::routes::register(&module).expect("routes should register");
+            let routes = [
+                (
+                    "ocr",
+                    "aocr",
+                    "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)",
+                ),
+                (
+                    "transcription",
+                    "atranscription",
+                    "(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)",
+                ),
+                (
+                    "messages",
+                    "amessages",
+                    "(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)",
+                ),
+                (
+                    "chat_completions",
+                    "achat_completions",
+                    "(model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)",
+                ),
+            ];
+
+            for (sync_name, async_name, expected) in routes {
+                let sync_signature: String = module
+                    .getattr(sync_name)
+                    .and_then(|function| function.getattr("__text_signature__"))
+                    .and_then(|signature| signature.extract())
+                    .expect("sync signature should be available");
+                let async_signature: String = module
+                    .getattr(async_name)
+                    .and_then(|function| function.getattr("__text_signature__"))
+                    .and_then(|signature| signature.extract())
+                    .expect("async signature should be available");
+
+                assert_eq!(sync_signature, expected);
+                assert_eq!(async_signature, expected);
+            }
+        });
+    }
+
+    #[test]
+    fn sync_and_async_routes_apply_the_same_input_validation() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = PyModule::new(py, "routes").expect("module should be created");
+            crate::routes::register(&module).expect("routes should register");
+
+            let invalid_messages = PyDict::new(py);
+            let sync_chat_error = module
+                .getattr("chat_completions")
+                .and_then(|function| function.call1(("model", &invalid_messages)))
+                .expect_err("sync chat should reject a non-list messages value");
+            let async_chat_error = module
+                .getattr("achat_completions")
+                .and_then(|function| function.call1(("model", &invalid_messages)))
+                .expect_err("async chat should reject a non-list messages value");
+
+            assert_eq!(
+                sync_chat_error.to_string(),
+                "ValueError: messages must be a list"
+            );
+            assert_eq!(async_chat_error.to_string(), sync_chat_error.to_string());
+
+            let invalid_body = PyList::empty(py);
+            let sync_messages_error = module
+                .getattr("messages")
+                .and_then(|function| function.call1(("model", &invalid_body)))
+                .expect_err("sync Messages should reject a non-dict body");
+            let async_messages_error = module
+                .getattr("amessages")
+                .and_then(|function| function.call1(("model", &invalid_body)))
+                .expect_err("async Messages should reject a non-dict body");
+
+            assert_eq!(
+                sync_messages_error.to_string(),
+                "ValueError: body must be a dict"
+            );
+            assert_eq!(
+                async_messages_error.to_string(),
+                sync_messages_error.to_string()
+            );
+
+            let invalid_headers = PyList::empty(py);
+            let kwargs = PyDict::new(py);
+            kwargs
+                .set_item("extra_headers", &invalid_headers)
+                .expect("kwargs should accept extra_headers");
+            let document = PyDict::new(py);
+
+            for (sync_name, async_name) in [("ocr", "aocr"), ("transcription", "atranscription")] {
+                let sync_error = module
+                    .getattr(sync_name)
+                    .and_then(|function| function.call(("model", &document), Some(&kwargs)))
+                    .expect_err("sync route should reject non-dict extra_headers");
+                let async_error = module
+                    .getattr(async_name)
+                    .and_then(|function| function.call(("model", &document), Some(&kwargs)))
+                    .expect_err("async route should reject non-dict extra_headers");
+
+                assert_eq!(
+                    sync_error.to_string(),
+                    "ValueError: extra_headers must be a dict"
+                );
+                assert_eq!(async_error.to_string(), sync_error.to_string());
+            }
+        });
+    }
+
+    #[test]
+    fn route_input_validation_preserves_left_to_right_order() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = PyModule::new(py, "routes").expect("module should be created");
+            crate::routes::register(&module).expect("routes should register");
+            let invalid = PyList::empty(py);
+
+            let chat_kwargs = PyDict::new(py);
+            chat_kwargs
+                .set_item("optional_params", &invalid)
+                .expect("kwargs should accept optional_params");
+            chat_kwargs
+                .set_item("extra_headers", &invalid)
+                .expect("kwargs should accept extra_headers");
+            let invalid_messages = PyDict::new(py);
+            let error = module
+                .getattr("chat_completions")
+                .and_then(|function| {
+                    function.call(("model", &invalid_messages), Some(&chat_kwargs))
+                })
+                .expect_err("messages should be validated first");
+            assert_eq!(error.to_string(), "ValueError: messages must be a list");
+
+            let valid_messages = PyList::empty(py);
+            let error = module
+                .getattr("chat_completions")
+                .and_then(|function| function.call(("model", &valid_messages), Some(&chat_kwargs)))
+                .expect_err("optional_params should be validated before headers");
+            assert_eq!(
+                error.to_string(),
+                "ValueError: optional_params must be a dict"
+            );
+
+            let headers_kwargs = PyDict::new(py);
+            headers_kwargs
+                .set_item("extra_headers", &invalid)
+                .expect("kwargs should accept extra_headers");
+            let invalid_body = PyList::empty(py);
+            let error = module
+                .getattr("messages")
+                .and_then(|function| function.call(("model", &invalid_body), Some(&headers_kwargs)))
+                .expect_err("body should be validated before headers");
+            assert_eq!(error.to_string(), "ValueError: body must be a dict");
+
+            let invalid_payload =
+                PyModule::new(py, "invalid_payload").expect("invalid payload should be created");
+            for name in ["ocr", "transcription"] {
+                let error = module
+                    .getattr(name)
+                    .and_then(|function| {
+                        function.call(("model", &invalid_payload), Some(&headers_kwargs))
+                    })
+                    .expect_err("payload should be validated before headers");
+                assert!(!error.to_string().contains("extra_headers"));
+            }
+        });
+    }
+
+    #[test]
+    fn generated_routes_execute_sync_and_async_contracts() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = PyModule::new(py, "synthetic").expect("module should be created");
+            synthetic::register(&module).expect("routes should register");
+
+            let sync_value: String = module
+                .getattr("echo")
+                .and_then(|function| function.call1(("sync",)))
+                .and_then(|value| value.extract())
+                .expect("sync route should return its value");
+            assert_eq!(sync_value, "sync");
+
+            let sync_error = module
+                .getattr("echo")
+                .and_then(|function| function.call1(("error",)))
+                .expect_err("sync route should map its error");
+            assert!(sync_error.is_instance_of::<PyLookupError>(py));
+            assert_eq!(
+                sync_error.to_string(),
+                "LookupError: invalid request: synthetic error"
+            );
+
+            let locals = PyDict::new(py);
+            locals
+                .set_item("routes", &module)
+                .expect("module should enter Python locals");
+            let code = CString::new(
+                r#"
+import asyncio
+
+async def exercise():
+    assert await routes.aecho("async") == "async"
+
+    try:
+        await routes.aecho("error")
+    except LookupError as error:
+        assert str(error) == "invalid request: synthetic error"
+    else:
+        raise AssertionError("mapped error was not raised")
+
+    try:
+        await routes.aecho("panic")
+    except BaseException as error:
+        assert type(error).__name__ == "PanicException"
+        assert str(error) == "synthetic panic"
+    else:
+        raise AssertionError("panic was not raised")
+
+    try:
+        await routes.aecho("map_panic")
+    except BaseException as error:
+        assert type(error).__name__ == "PanicException"
+        assert str(error) == "synthetic mapper panic"
+    else:
+        raise AssertionError("mapper panic was not raised")
+
+    task = asyncio.ensure_future(routes.aecho("pending"))
+    await asyncio.sleep(0)
+    task.cancel()
+    try:
+        await task
+    except asyncio.CancelledError:
+        pass
+    else:
+        raise AssertionError("cancelled route completed")
+
+    for _ in range(100):
+        if routes.future_dropped():
+            break
+        await asyncio.sleep(0.001)
+    assert routes.future_dropped()
+
+asyncio.run(exercise())
+"#,
+            )
+            .expect("Python source should not contain null bytes");
+            py.run(&code, Some(&locals), Some(&locals))
+                .expect("async route contract should hold");
+        });
+    }
+
+    #[test]
+    fn route_registration_rejects_duplicate_python_names() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = PyModule::new(py, "synthetic").expect("module should be created");
+            synthetic::register(&module).expect("first registration should succeed");
+            let error = synthetic::register(&module)
+                .expect_err("duplicate registration should be rejected");
+
+            assert_eq!(
+                error.to_string(),
+                "RuntimeError: duplicate native route: future_dropped"
+            );
+        });
+    }
+}

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

@@ -0,0 +1,65 @@
+use litellm_core::Error;
+use litellm_core::messages::messages as run_messages;
+use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest};
+use pyo3::prelude::*;
+use serde_json::Value;
+use std::future::Future;
+
+use crate::errors::core_error_to_pyerr;
+use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value};
+
+fn prepare_messages(
+    inputs: MessagesInputs,
+) -> PyResult<impl Future<Output = Result<AnthropicMessagesResponse, Error>> + Send + 'static> {
+    let body = required_value("body", inputs.body, Value::is_object, "dict")?;
+    let options = RouteOptions::from_python(RouteOptionsInputs {
+        model: inputs.model,
+        api_key: inputs.api_key,
+        api_base: inputs.api_base,
+        custom_llm_provider: inputs.custom_llm_provider,
+        extra_headers: inputs.extra_headers,
+        timeout_seconds: inputs.timeout_seconds,
+    })?;
+
+    Ok(async move {
+        let RouteOptions {
+            model,
+            api_key,
+            api_base,
+            custom_llm_provider,
+            extra_headers,
+            timeout,
+        } = options;
+        run_messages(MessagesRequest {
+            model: &model,
+            body,
+            api_key: api_key.as_deref(),
+            api_base: api_base.as_deref(),
+            custom_llm_provider: custom_llm_provider.as_deref(),
+            extra_headers,
+            timeout,
+        })
+        .await
+    })
+}
+
+bridge_route! {
+    sync = messages,
+    asynchronous = amessages,
+    inputs = MessagesInputs,
+    required = {
+        model: String,
+        #[pyo3(from_py_with = litellm_python_interop::from_py)]
+        body: Value,
+    },
+    optional = {
+        api_key: Option<String>,
+        api_base: Option<String>,
+        custom_llm_provider: Option<String>,
+        #[pyo3(from_py_with = litellm_python_interop::from_py)]
+        extra_headers: Option<Value>,
+        timeout_seconds: Option<f64>,
+    },
+    prepare = prepare_messages,
+    errors = core_error_to_pyerr,
+}

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

@@ -0,0 +1,16 @@
+use pyo3::prelude::*;
+
+#[macro_use]
+mod definition;
+
+mod audio_transcription;
+mod chat_completions;
+mod messages;
+mod ocr;
+
+pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    ocr::register(module)?;
+    audio_transcription::register(module)?;
+    messages::register(module)?;
+    chat_completions::register(module)
+}

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

@@ -0,0 +1,73 @@
+use litellm_core::Error;
+use std::future::Future;
+
+use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
+use pyo3::prelude::*;
+use serde_json::Value;
+
+use crate::errors::core_error_to_pyerr;
+use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty};
+
+fn prepare_ocr(
+    inputs: OcrInputs,
+) -> PyResult<impl Future<Output = Result<Value, Error>> + Send + 'static> {
+    let document = inputs.document;
+    let options = RouteOptions::from_python(RouteOptionsInputs {
+        model: inputs.model,
+        api_key: inputs.api_key,
+        api_base: inputs.api_base,
+        custom_llm_provider: inputs.custom_llm_provider,
+        extra_headers: inputs.extra_headers,
+        timeout_seconds: inputs.timeout_seconds,
+    })?;
+    let optional_params = object_or_empty("optional_params", inputs.optional_params)?;
+
+    Ok(async move {
+        let RouteOptions {
+            model,
+            api_key,
+            api_base,
+            custom_llm_provider,
+            extra_headers,
+            timeout,
+        } = options;
+        run_ocr(OcrRequest {
+            model: &model,
+            document,
+            api_key: api_key.as_deref(),
+            api_base: api_base.as_deref(),
+            custom_llm_provider: custom_llm_provider.as_deref(),
+            extra_headers,
+            optional_params,
+            timeout,
+            callbacks: Vec::new(),
+            guardrails: Vec::new(),
+            request_metadata: Default::default(),
+            litellm_call_id: None,
+        })
+        .await
+    })
+}
+
+bridge_route! {
+    sync = ocr,
+    asynchronous = aocr,
+    inputs = OcrInputs,
+    required = {
+        model: String,
+        #[pyo3(from_py_with = litellm_python_interop::from_py)]
+        document: Value,
+    },
+    optional = {
+        api_key: Option<String>,
+        api_base: Option<String>,
+        custom_llm_provider: Option<String>,
+        #[pyo3(from_py_with = litellm_python_interop::from_py)]
+        extra_headers: Option<Value>,
+        #[pyo3(from_py_with = litellm_python_interop::from_py)]
+        optional_params: Option<Value>,
+        timeout_seconds: Option<f64>,
+    },
+    prepare = prepare_ocr,
+    errors = core_error_to_pyerr,
+}

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

@@ -0,0 +1,423 @@
+use std::future::Future;
+use std::panic::AssertUnwindSafe;
+use std::time::Duration;
+
+use futures_util::FutureExt;
+use litellm_core::error::Error;
+use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil};
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
+use serde::Serialize;
+use tokio::runtime::{Handle, Runtime};
+use tokio::time::{self, MissedTickBehavior};
+
+pub(super) fn run_sync<T, F>(
+    py: Python<'_>,
+    future: F,
+    map_error: fn(Error) -> PyErr,
+) -> PyResult<Py<PyAny>>
+where
+    T: Serialize + Send + 'static,
+    F: Future<Output = Result<T, Error>> + Send + 'static,
+{
+    run_sync_on(
+        py,
+        pyo3_async_runtimes::tokio::get_runtime(),
+        future,
+        map_error,
+    )
+}
+
+fn run_sync_on<T, F>(
+    py: Python<'_>,
+    runtime: &Runtime,
+    future: F,
+    map_error: fn(Error) -> PyErr,
+) -> PyResult<Py<PyAny>>
+where
+    T: Serialize + Send + 'static,
+    F: Future<Output = Result<T, Error>> + Send + 'static,
+{
+    if Handle::try_current().is_ok() {
+        return Err(PyRuntimeError::new_err(
+            "synchronous native routes cannot run from a Tokio context; use the async route",
+        ));
+    }
+
+    let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;
+    let result = map_core_result(result, map_error)?;
+    Pythonized(result).into_pyobject(py).map(Bound::unbind)
+}
+
+pub(super) fn run_async<T, F>(
+    py: Python<'_>,
+    future: F,
+    map_error: fn(Error) -> PyErr,
+) -> PyResult<Bound<'_, PyAny>>
+where
+    T: Serialize + Send + 'static,
+    F: Future<Output = Result<T, Error>> + Send + 'static,
+{
+    pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        let result = catch_route_panic(future).await?;
+        let result = map_core_result(result, map_error)?;
+        Ok(Pythonized(result))
+    })
+}
+
+fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
+    match result {
+        Ok(value) => Ok(value),
+        Err(error) => Err(
+            std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error)))
+                .map_err(panic_to_pyerr)?,
+        ),
+    }
+}
+
+async fn catch_route_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
+where
+    F: Future<Output = Result<T, Error>>,
+{
+    AssertUnwindSafe(future)
+        .catch_unwind()
+        .await
+        .map_err(panic_to_pyerr)
+}
+
+async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
+where
+    F: Future<Output = Result<T, Error>>,
+{
+    let future = catch_route_panic(future);
+    tokio::pin!(future);
+
+    let signal_interval = Duration::from_millis(50);
+    let mut signal_checks =
+        time::interval_at(time::Instant::now() + signal_interval, signal_interval);
+    signal_checks.set_missed_tick_behavior(MissedTickBehavior::Delay);
+    loop {
+        tokio::select! {
+            result = &mut future => return result,
+            _ = signal_checks.tick() => Python::attach(|py| py.check_signals())?,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::ffi::CString;
+    use std::future::poll_fn;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::sync::{Arc, mpsc};
+    use std::task::Poll;
+    use std::thread;
+    use std::time::Instant;
+
+    use pyo3::panic::PanicException;
+    use pyo3::types::{PyDict, PyModule};
+    use serde::Serializer;
+    use tokio::runtime::Builder;
+
+    use super::*;
+
+    fn runtime_error(error: Error) -> PyErr {
+        PyRuntimeError::new_err(error.to_string())
+    }
+
+    fn panicking_error_mapper(_error: Error) -> PyErr {
+        panic!("error mapper panicked")
+    }
+
+    struct PanickingOutput;
+
+    static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0);
+
+    impl Serialize for PanickingOutput {
+        fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
+        where
+            S: Serializer,
+        {
+            panic!("serializer panicked")
+        }
+    }
+
+    #[pyfunction]
+    fn async_serialization_panic(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
+        run_async(py, async { Ok(PanickingOutput) }, runtime_error)
+    }
+
+    #[pyfunction]
+    fn async_runtime_probe(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
+        run_async(
+            py,
+            async {
+                ASYNC_PROBE_COMPLETED.fetch_add(1, Ordering::SeqCst);
+                Ok(true)
+            },
+            runtime_error,
+        )
+    }
+
+    #[pyfunction]
+    fn runtime_worker_count() -> usize {
+        pyo3_async_runtimes::tokio::get_runtime()
+            .metrics()
+            .num_workers()
+    }
+
+    #[pyfunction]
+    fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
+        let completion_deadline = Instant::now() + Duration::from_secs(2);
+        while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
+            if Instant::now() >= completion_deadline {
+                return false;
+            }
+            thread::sleep(Duration::from_millis(1));
+        }
+
+        let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
+        pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
+            let _ = heartbeat_tx.send(());
+        });
+        heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
+    }
+
+    fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
+        result
+            .expect("route should complete")
+            .bind(py)
+            .extract()
+            .expect("result should convert")
+    }
+
+    #[test]
+    fn sync_runner_polls_future_on_the_caller_thread() {
+        Python::initialize();
+        Python::attach(|py| {
+            let caller_thread = std::thread::current().id();
+            let result = run_sync(
+                py,
+                async move { Ok(std::thread::current().id() == caller_thread) },
+                runtime_error,
+            );
+
+            assert!(extract_bool(py, result));
+        });
+    }
+
+    #[test]
+    fn sync_runner_releases_gil_while_waiting() {
+        Python::initialize();
+        Python::attach(|py| {
+            let result = run_sync(
+                py,
+                async {
+                    let gil_acquired = tokio::time::timeout(
+                        Duration::from_secs(2),
+                        tokio::task::spawn_blocking(|| Python::attach(|_| true)),
+                    )
+                    .await;
+                    Ok(matches!(gil_acquired, Ok(Ok(true))))
+                },
+                runtime_error,
+            );
+
+            assert!(extract_bool(py, result));
+        });
+    }
+
+    #[test]
+    fn sync_runner_rejects_calls_from_a_tokio_context() {
+        Python::initialize();
+        let runtime = Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .expect("runtime should build");
+
+        let error = runtime.block_on(async {
+            Python::attach(|py| {
+                run_sync::<bool, _>(py, async { Ok(true) }, runtime_error)
+                    .expect_err("sync route should reject a nested Tokio runtime")
+            })
+        });
+
+        assert_eq!(
+            error.to_string(),
+            "RuntimeError: synchronous native routes cannot run from a Tokio context; use the async route"
+        );
+    }
+
+    #[test]
+    fn sync_runner_can_drive_a_current_thread_runtime() {
+        Python::initialize();
+        let runtime = Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .expect("runtime should build");
+        Python::attach(|py| {
+            let result = run_sync_on(
+                py,
+                &runtime,
+                async {
+                    tokio::task::yield_now().await;
+                    Ok(true)
+                },
+                runtime_error,
+            );
+            assert!(extract_bool(py, result));
+        });
+    }
+
+    #[test]
+    fn sync_runner_maps_a_panicked_future() {
+        Python::initialize();
+        Python::attach(|py| {
+            let error = run_sync::<bool, _>(
+                py,
+                poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
+                runtime_error,
+            )
+            .expect_err("panicked route should become a Python exception");
+
+            assert!(error.is_instance_of::<PanicException>(py));
+            assert_eq!(error.to_string(), "PanicException: route future panicked");
+        });
+    }
+
+    #[test]
+    fn sync_runner_maps_a_panicked_error_mapper() {
+        Python::initialize();
+        Python::attach(|py| {
+            let error = run_sync::<bool, _>(
+                py,
+                async { Err(Error::InvalidRequest("invalid".to_string())) },
+                panicking_error_mapper,
+            )
+            .expect_err("panicked mapper should become a Python exception");
+
+            assert!(error.is_instance_of::<PanicException>(py));
+            assert_eq!(error.to_string(), "PanicException: error mapper panicked");
+        });
+    }
+
+    #[test]
+    fn sync_runner_surfaces_serializer_panics() {
+        Python::initialize();
+        Python::attach(|py| {
+            let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error)
+                .expect_err("serializer panic should become a Python exception");
+
+            assert!(error.is_instance_of::<PanicException>(py));
+            assert_eq!(error.to_string(), "PanicException: serializer panicked");
+        });
+    }
+
+    #[test]
+    fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() {
+        Python::initialize();
+        let barrier = Arc::new(tokio::sync::Barrier::new(2));
+        let callers: Vec<_> = (0..2)
+            .map(|_| {
+                let barrier = Arc::clone(&barrier);
+                thread::spawn(move || {
+                    Python::attach(|py| {
+                        extract_bool(
+                            py,
+                            run_sync(
+                                py,
+                                async move {
+                                    Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait())
+                                        .await
+                                        .is_ok())
+                                },
+                                runtime_error,
+                            ),
+                        )
+                    })
+                })
+            })
+            .collect();
+        let results: Vec<_> = callers
+            .into_iter()
+            .map(|caller| caller.join().expect("caller should not panic"))
+            .collect();
+
+        assert_eq!(results, vec![true, true]);
+    }
+
+    #[test]
+    fn async_runner_surfaces_serializer_panics() {
+        Python::initialize();
+        Python::attach(|py| {
+            let module = PyModule::new(py, "runtime").expect("module should be created");
+            module
+                .add_function(
+                    wrap_pyfunction!(async_serialization_panic, &module)
+                        .expect("function should wrap"),
+                )
+                .expect("function should register");
+            let locals = PyDict::new(py);
+            locals
+                .set_item("runtime", &module)
+                .expect("module should enter Python locals");
+            let code = CString::new(
+                r#"
+import asyncio
+
+async def exercise():
+    try:
+        await runtime.async_serialization_panic()
+    except BaseException as error:
+        assert type(error).__name__ == "PanicException"
+        assert str(error) == "serializer panicked"
+    else:
+        raise AssertionError("serializer panic was not raised")
+
+asyncio.run(exercise())
+"#,
+            )
+            .expect("Python source should not contain null bytes");
+            py.run(&code, Some(&locals), Some(&locals))
+                .expect("serializer panic should reach the Python awaiter");
+        });
+    }
+
+    #[test]
+    fn async_result_delivery_does_not_stall_tokio_workers() {
+        Python::initialize();
+        ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst);
+        Python::attach(|py| {
+            let module = PyModule::new(py, "runtime").expect("module should be created");
+            for function in [
+                wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"),
+                wrap_pyfunction!(runtime_worker_count, &module).expect("function should wrap"),
+                wrap_pyfunction!(runtime_is_responsive, &module).expect("function should wrap"),
+            ] {
+                module
+                    .add_function(function)
+                    .expect("function should register");
+            }
+            let locals = PyDict::new(py);
+            locals
+                .set_item("runtime", &module)
+                .expect("module should enter Python locals");
+            let code = CString::new(
+                r#"
+import asyncio
+
+async def exercise():
+    worker_count = runtime.runtime_worker_count()
+    awaitables = [runtime.async_runtime_probe() for _ in range(worker_count)]
+    assert runtime.runtime_is_responsive(worker_count)
+    assert await asyncio.gather(*awaitables) == [True] * worker_count
+
+asyncio.run(exercise())
+"#,
+            )
+            .expect("Python source should not contain null bytes");
+            py.run(&code, Some(&locals), Some(&locals))
+                .expect("result delivery should leave Tokio workers responsive");
+        });
+    }
+}

litellm-rust/crates/python-interop/src/lib.rs

@@ -2,4 +2,4 @@ mod gil;
 mod marshal;
 
 pub use gil::{release_count, release_gil};
-pub use marshal::{from_py, to_py};
+pub use marshal::{Pythonized, from_py, panic_to_pyerr, to_py};

litellm-rust/crates/python-interop/src/marshal.rs

@@ -1,4 +1,8 @@
+use std::any::Any;
+use std::panic::{AssertUnwindSafe, catch_unwind};
+
 use pyo3::exceptions::PyValueError;
+use pyo3::panic::PanicException;
 use pyo3::prelude::*;
 use serde::Serialize;
 use serde::de::DeserializeOwned;
@@ -18,3 +22,71 @@ where
         .map(Bound::unbind)
         .map_err(|error| PyValueError::new_err(error.to_string()))
 }
+
+pub struct Pythonized<T>(pub T);
+
+impl<'py, T> IntoPyObject<'py> for Pythonized<T>
+where
+    T: Serialize,
+{
+    type Target = PyAny;
+    type Output = Bound<'py, PyAny>;
+    type Error = PyErr;
+
+    fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
+        catch_unwind(AssertUnwindSafe(|| pythonize::pythonize(py, &self.0)))
+            .map_err(panic_to_pyerr)?
+            .map_err(|error| PyValueError::new_err(error.to_string()))
+    }
+}
+
+pub fn panic_to_pyerr(payload: Box<dyn Any + Send>) -> PyErr {
+    let message = payload
+        .downcast_ref::<String>()
+        .map(String::as_str)
+        .or_else(|| payload.downcast_ref::<&str>().copied())
+        .unwrap_or("panic from Rust code");
+    PanicException::new_err(message.to_string())
+}
+
+#[cfg(test)]
+mod tests {
+    use serde::Serializer;
+
+    use super::*;
+
+    struct PanickingSerializer;
+
+    impl Serialize for PanickingSerializer {
+        fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
+        where
+            S: Serializer,
+        {
+            panic!("serializer panicked")
+        }
+    }
+
+    #[test]
+    fn pythonized_converts_on_the_attached_thread() {
+        Python::initialize();
+        Python::attach(|py| {
+            let value: Vec<i32> = Pythonized(vec![1, 2, 3])
+                .into_pyobject(py)
+                .and_then(|value| value.extract())
+                .expect("value should convert");
+            assert_eq!(value, vec![1, 2, 3]);
+        });
+    }
+
+    #[test]
+    fn pythonized_maps_serializer_panics_to_a_base_exception() {
+        Python::initialize();
+        Python::attach(|py| {
+            let error = Pythonized(PanickingSerializer)
+                .into_pyobject(py)
+                .expect_err("serializer panic should become a Python exception");
+            assert!(error.is_instance_of::<PanicException>(py));
+            assert_eq!(error.to_string(), "PanicException: serializer panicked");
+        });
+    }
+}

litellm/__init__.py

@@ -1417,7 +1417,7 @@ from .skills.main import (
 )
 from .containers.main import *
 from .ocr.main import *
-from .rust_bridge.ocr import use_litellm_rust
+from .rust_bridge import use_litellm_rust
 from .rag.main import *
 from .sandbox.main import *
 from .search.main import *

litellm/integrations/SlackAlerting/slack_alerting.py

@@ -1955,11 +1955,10 @@ Model Info:
         if not thresholds_enabled and not anomalies_enabled:
             return
 
-        if prisma_client is None:
-            from litellm.proxy.proxy_server import prisma_client as global_prisma_client
+        from litellm.proxy.proxy_server import prisma_client as global_prisma_client
 
-            prisma_client = global_prisma_client  # rebind-ok: fall back to the proxy's global client
-        if prisma_client is None:
+        client: Final = prisma_client if prisma_client is not None else global_prisma_client
+        if client is None:
             return
 
         from litellm.integrations.SlackAlerting.user_spend_alerts import (
@@ -1970,7 +1969,7 @@ Model Info:
         try:
             today: Final = datetime.datetime.now(datetime.timezone.utc).date()
             rows: Final = await fetch_user_spend_rows(
-                prisma_client=prisma_client,
+                prisma_client=client,
                 today=today,
                 baseline_days=self.alerting_args.spend_anomaly_baseline_days,
             )

litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py

@@ -5,6 +5,7 @@ This hook is called before making an LLM request when a vector store is configur
 It searches the vector store for relevant context and appends it to the messages.
 """
 
+from collections.abc import Awaitable, Callable
 from typing import TYPE_CHECKING, Any, Final, cast
 
 import litellm
@@ -80,10 +81,17 @@ class VectorStorePreCallHook(CustomLogger):
 
             # Get prisma_client for database fallback
             prisma_client = None
+            llm_router = None
             try:
-                from litellm.proxy.proxy_server import prisma_client as _prisma_client
+                from litellm.proxy.proxy_server import (
+                    llm_router as _llm_router,
+                )
+                from litellm.proxy.proxy_server import (
+                    prisma_client as _prisma_client,
+                )
 
                 prisma_client = _prisma_client
+                llm_router = _llm_router
             except ImportError:
                 pass
 
@@ -114,12 +122,26 @@ class VectorStorePreCallHook(CustomLogger):
                 vector_store_id = vector_store_to_run.get("vector_store_id", "")
                 custom_llm_provider = vector_store_to_run.get("custom_llm_provider")
                 litellm_params_for_vector_store = vector_store_to_run.get("litellm_params", {}) or {}
-                # Call litellm.vector_stores.search() with the required parameters
-                search_response = await litellm.vector_stores.asearch(
+                request_litellm_params = litellm_logging_obj.model_call_details.get("litellm_params", {})
+                request_metadata = (
+                    request_litellm_params.get("metadata", {}) if isinstance(request_litellm_params, dict) else {}
+                )
+                if llm_router is not None:
+                    search_function = cast(  # cast-ok: normalize router search callable
+                        Callable[..., Awaitable[VectorStoreSearchResponse]],
+                        llm_router.avector_store_search,
+                    )
+                else:
+                    search_function = cast(  # cast-ok: normalize SDK search callable
+                        Callable[..., Awaitable[VectorStoreSearchResponse]],
+                        litellm.vector_stores.asearch,
+                    )
+                search_response = await search_function(
                     **{
                         "vector_store_id": vector_store_id,
                         "query": query,
                         "custom_llm_provider": custom_llm_provider,
+                        "metadata": request_metadata,
                         **litellm_params_for_vector_store,
                     },
                 )

litellm/integrations/websearch_interception/handler.py

@@ -419,7 +419,6 @@ class WebSearchInterceptionLogger(CustomLogger):
         if call_type in (CallTypes.responses, CallTypes.aresponses):
             return self._convert_responses_tools(kwargs=kwargs, tools=tools)
 
-        # Check if any tool is a web search tool (native or already LiteLLM standard)
         has_websearch: Final = any(is_web_search_tool(t) for t in tools)
 
         if not has_websearch:

litellm/llms/anthropic/chat/transformation.py

@@ -1565,6 +1565,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
                         optional_params.pop("thinking", None)
                 else:
                     optional_params["thinking"] = value
+                    AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
+                        model=model, optional_params=optional_params, custom_llm_provider=self._resolved_provider
+                    )
             elif param == "reasoning_effort":
                 # Accept both string ("low") and dict ({"effort": "low",
                 # "summary": "concise"}). The Responses->Chat parser keeps the

litellm/llms/anthropic/common_utils.py

@@ -13,7 +13,12 @@ import httpx
 from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 
 import litellm
-from litellm.constants import DEFAULT_MODEL_CREATED_AT_TIME
+from litellm.constants import (
+    DEFAULT_MODEL_CREATED_AT_TIME,
+    DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
+    DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
+    DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
+)
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     get_file_ids_from_messages,
 )
@@ -534,6 +539,51 @@ class AnthropicModelInfo(BaseLLMModelInfo):
         )
         optional_params.pop("thinking", None)
 
+    @staticmethod
+    def translate_legacy_thinking_for_adaptive_model(
+        model: str,
+        optional_params: MutableMapping[str, object],  # mutable-ok: in-place out-param like the sibling helpers
+        custom_llm_provider: str,
+    ) -> None:
+        """Translate legacy ``thinking.type=enabled`` to adaptive for the
+        adaptive-thinking models that reject it (4.7+ and the 5 families).
+        Models flagged ``supports_legacy_thinking`` (the 4.6 family) accept the
+        legacy shape natively, so it is forwarded verbatim and the caller's
+        ``budget_tokens`` cap keeps applying. Caller-provided
+        ``output_config.effort`` is never overridden.
+        """
+        if not AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider):
+            return
+        if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider):
+            return
+        thinking: Final = optional_params.get("thinking")
+        if not isinstance(thinking, dict) or thinking.get("type") != "enabled":
+            return
+
+        effort: Final = AnthropicModelInfo._legacy_budget_to_effort(
+            model=model,
+            budget_tokens=int(thinking.get("budget_tokens") or 0),
+            custom_llm_provider=custom_llm_provider,
+        )
+        existing_output_config: Final = optional_params.get("output_config")
+        optional_params["thinking"] = {"type": "adaptive"}
+        optional_params["output_config"] = {
+            "effort": effort,
+            **(existing_output_config if isinstance(existing_output_config, dict) else MappingProxyType({})),
+        }
+
+    @staticmethod
+    def _legacy_budget_to_effort(model: str, budget_tokens: int, custom_llm_provider: str) -> str:
+        if budget_tokens >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and (
+            AnthropicModelInfo._supports_model_capability(model, "supports_xhigh_reasoning_effort", custom_llm_provider)
+        ):
+            return "xhigh"
+        if budget_tokens >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET:
+            return "high"
+        if budget_tokens >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET:
+            return "medium"
+        return "low"
+
     def is_effort_used(
         self,
         optional_params: dict | None,
@@ -1361,6 +1411,97 @@ def flatten_unencrypted_web_search_results_in_anthropic_messages(  # mutable-ok:
     return [_flatten_web_search_results_in_message(m) for m in messages]  # mutable-ok: JSON wire format
 
 
+def _normalized_cache_control(cache_control: object) -> dict[str, str] | None:  # mutable-ok: JSON wire format
+    if not isinstance(cache_control, Mapping):
+        return None
+    cache_type: Final = cache_control.get("type")
+    return {"type": cache_type if isinstance(cache_type, str) else "ephemeral"}  # mutable-ok: JSON wire format
+
+
+def _with_portable_cache_control(block: Mapping[str, object]) -> dict[str, object]:  # mutable-ok: JSON wire format
+    if "cache_control" not in block:
+        return dict(block)  # mutable-ok: JSON wire format
+    normalized: Final = _normalized_cache_control(block["cache_control"])
+    rest: Final = {key: value for key, value in block.items() if key != "cache_control"}  # mutable-ok: JSON wire format
+    return rest if normalized is None else {**rest, "cache_control": normalized}  # mutable-ok: JSON wire format
+
+
+def _with_portable_cache_control_in_blocks(blocks: object) -> object:
+    if isinstance(blocks, str) or not isinstance(blocks, Sequence):
+        return blocks
+    return [  # mutable-ok: JSON wire format
+        _with_portable_cache_control(block) if isinstance(block, Mapping) else block for block in blocks
+    ]
+
+
+def _with_portable_cache_control_in_content_block(block: object) -> object:
+    if not isinstance(block, Mapping):
+        return block
+    portable: Final = _with_portable_cache_control(block)
+    if portable.get("type") != "tool_result" or "content" not in portable:
+        return portable
+    return {  # mutable-ok: JSON wire format
+        **portable,
+        "content": _with_portable_cache_control_in_blocks(portable["content"]),
+    }
+
+
+def _with_portable_cache_control_in_message(message: object) -> object:
+    if not isinstance(message, Mapping) or "content" not in message:
+        return message
+    content: Final = message["content"]
+    if isinstance(content, str) or not isinstance(content, Sequence):
+        return message
+    return {  # mutable-ok: JSON wire format
+        **message,
+        "content": [  # mutable-ok: JSON wire format
+            _with_portable_cache_control_in_content_block(block) for block in content
+        ],
+    }
+
+
+def _with_portable_cache_control_in_messages(messages: object) -> object:
+    if isinstance(messages, str) or not isinstance(messages, Sequence):
+        return messages
+    return [  # mutable-ok: JSON wire format
+        _with_portable_cache_control_in_message(message) for message in messages
+    ]
+
+
+def _with_portable_cache_control_in_scoped_value(key: str, value: object) -> object:
+    match key:
+        case "system" | "tools":
+            return _with_portable_cache_control_in_blocks(value)
+        case "messages":
+            return _with_portable_cache_control_in_messages(value)
+        case _:
+            return value
+
+
+def normalize_cache_control_in_anthropic_payload(
+    payload: Mapping[str, object],
+) -> dict[str, object]:  # mutable-ok: JSON wire format
+    """
+    Return a copy of an Anthropic /v1/messages payload with every
+    ``cache_control`` entry reduced to ``{"type": <its type, or "ephemeral">}``
+    at the places the Messages API defines it: the request itself, system
+    blocks, tools, message content blocks, and ``tool_result`` content blocks.
+    Application data such as ``tool_use.input`` and tool ``input_schema`` is
+    never touched, even when it happens to contain a ``cache_control`` key.
+
+    Anthropic itself accepts prompt-caching extensions such as ``ttl``, but
+    strict non-Anthropic implementations of the Messages API validate the field
+    literally and reject the whole request (``cache_control.ttl: 1h is not
+    supported``, ``cache_control.type is required``), which 400s clients like
+    Claude Code that send cache hints. Non-dict ``cache_control`` values are
+    dropped entirely. The caller's payload is never mutated.
+    """
+    portable: Final = _with_portable_cache_control(payload)
+    return {  # mutable-ok: JSON wire format
+        key: _with_portable_cache_control_in_scoped_value(key, value) for key, value in portable.items()
+    }
+
+
 def process_anthropic_headers(headers: httpx.Headers | dict) -> dict:
     openai_headers: Final = {}
     if "anthropic-ratelimit-requests-limit" in headers:

litellm/llms/anthropic/experimental_pass_through/messages/handler.py

@@ -99,6 +99,10 @@ def _deployment_passes_through_anthropic_messages(model_info: object) -> bool:
     return isinstance(supported_endpoints, (list, tuple)) and "/v1/messages" in supported_endpoints
 
 
+def _deployment_supports_cache_control_ttl(model_info: object) -> bool:
+    return isinstance(model_info, dict) and model_info.get("cache_control_ttl") is True
+
+
 ####### ENVIRONMENT VARIABLES ###################
 # Initialize any necessary instances or variables here
 base_llm_http_handler = BaseLLMHTTPHandler()
@@ -568,7 +572,9 @@ def anthropic_messages_handler(
             OpenAILikeAnthropicMessagesConfig,
         )
 
-        anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig()
+        anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig(
+            cache_control_ttl=_deployment_supports_cache_control_ttl(kwargs.get("model_info")),
+        )
     if anthropic_messages_provider_config is None:
         # Route to Responses API for OpenAI / Azure, chat/completions for everything else.
         if _should_route_to_responses_api(custom_llm_provider, original_model, model):

litellm/llms/anthropic/experimental_pass_through/messages/transformation.py

@@ -3,11 +3,6 @@ from typing import Any, Final
 
 import httpx
 
-from litellm.constants import (
-    DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
-    DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
-    DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
-)
 from litellm.exceptions import AuthenticationError
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.litellm_logging import verbose_logger
@@ -400,46 +395,6 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
             existing_output_config.setdefault("effort", mapped_effort)
             optional_params["output_config"] = existing_output_config
 
-    @staticmethod
-    def _translate_legacy_thinking_for_adaptive_model(
-        model: str, optional_params: dict, custom_llm_provider: str
-    ) -> None:
-        """Translate legacy ``thinking.type=enabled`` to adaptive for the
-        adaptive-thinking models that reject it (4.7+ and the 5 families).
-        Models flagged ``supports_legacy_thinking`` (the 4.6 family) accept the
-        legacy shape natively, so it is forwarded verbatim and the caller's
-        ``budget_tokens`` cap keeps applying. Caller-provided
-        ``output_config.effort`` is never overridden.
-        """
-        from litellm.llms.anthropic.chat.transformation import AnthropicConfig
-
-        if not AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider):
-            return
-        if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider):
-            return
-        thinking: Final = optional_params.get("thinking")
-        if not isinstance(thinking, dict) or thinking.get("type") != "enabled":
-            return
-
-        budget: Final = int(thinking.get("budget_tokens") or 0)
-        if budget >= DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET and (
-            AnthropicConfig._supports_effort_level(model, "xhigh", custom_llm_provider)
-        ):
-            effort = "xhigh"
-        elif budget >= DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET:
-            effort = "high"
-        elif budget >= DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET:
-            effort = "medium"
-        else:
-            effort = "low"
-
-        optional_params["thinking"] = {"type": "adaptive"}
-        existing_output_config = optional_params.get("output_config")
-        if not isinstance(existing_output_config, dict):
-            existing_output_config = {}
-        existing_output_config.setdefault("effort", effort)
-        optional_params["output_config"] = existing_output_config
-
     @staticmethod
     def _translate_adaptive_effort_for_non_adaptive_model(
         model: str, optional_params: dict, max_tokens: int | None, custom_llm_provider: str
@@ -606,7 +561,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
             custom_llm_provider=self._resolved_provider,
         )
 
-        self._translate_legacy_thinking_for_adaptive_model(
+        AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
             model=model,
             optional_params=anthropic_messages_optional_request_params,
             custom_llm_provider=self._resolved_provider,

litellm/llms/azure_ai/anthropic/transformation.py

@@ -17,13 +17,14 @@ def _promote_extra_body_to_optional_params(optional_params: dict) -> None:
     ``output_config`` get auto-routed into ``extra_body`` by
     ``add_provider_specific_params_to_optional_params``. For the Azure→Anthropic
     route those keys must reach the request body and be validated, so promote
-    them. ``setdefault`` keeps explicit top-level values authoritative.
+    them. The caller's values overwrite mapped top-level duplicates, matching
+    the native ``anthropic`` provider, where the same passthrough lands on
+    top-level ``optional_params`` after mapping.
     """
     extra_body: Final = optional_params.get("extra_body")
     if not isinstance(extra_body, dict) or not extra_body:
         return
-    for k, v in extra_body.items():
-        optional_params.setdefault(k, v)
+    optional_params.update(extra_body)
     optional_params.pop("extra_body", None)
 
 

litellm/llms/azure_ai/vector_stores/transformation.py

@@ -1,10 +1,15 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final
 
 import httpx
 
-import litellm
 from litellm.llms.azure.common_utils import BaseAzureLLM
-from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
+from litellm.llms.base_llm.vector_store.transformation import (
+    BaseQueryEmbeddingVectorStoreConfig,
+    VectorStoreEmbeddingExecutor,
+)
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.vector_stores import (
     BaseVectorStoreAuthCredentials,
@@ -26,7 +31,7 @@ else:
     LiteLLMLoggingObj = Any
 
 
-class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM):
+class AzureAIVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAzureLLM):
     """
     Configuration for Azure AI Search Vector Store
 
@@ -110,83 +115,73 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM):
     def transform_search_vector_store_request(
         self,
         vector_store_id: str,
-        query: str | list[str],
+        query: str | Sequence[str],
         vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
         api_base: str,
         litellm_logging_obj: LiteLLMLoggingObj,
-        litellm_params: dict,
-        extra_body: dict[str, Any] | None = None,
-        router: "Router | None" = None,
-    ) -> tuple[str, dict[str, Any]]:
-        """
-        Transform search request for Azure AI Search API
+        litellm_params: Mapping[str, object],
+        extra_body: Mapping[str, object] | None = None,
+        router: Router | None = None,
+        embedding_executor: VectorStoreEmbeddingExecutor | None = None,
+    ) -> tuple[str, dict[str, object]]:
+        query_text: Final = self.query_text(query)
+        query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor, router)
+        return self._search_request(
+            vector_store_id,
+            query_text,
+            query_vector,
+            vector_store_search_optional_params,
+            api_base,
+            litellm_logging_obj,
+            litellm_params,
+        )
 
-        Generates embeddings using litellm.embeddings and constructs Azure AI Search request
-        """
-        # Convert query to string if it's a list
-        if isinstance(query, list):
-            query = " ".join(query)
+    async def atransform_search_vector_store_request(
+        self,
+        vector_store_id: str,
+        query: str | Sequence[str],
+        vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
+        api_base: str,
+        litellm_logging_obj: LiteLLMLoggingObj,
+        litellm_params: Mapping[str, object],
+        extra_body: Mapping[str, object] | None = None,
+        router: Router | None = None,
+        embedding_executor: VectorStoreEmbeddingExecutor | None = None,
+    ) -> tuple[str, dict[str, object]]:
+        query_text: Final = self.query_text(query)
+        query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor, router)
+        return self._search_request(
+            vector_store_id,
+            query_text,
+            query_vector,
+            vector_store_search_optional_params,
+            api_base,
+            litellm_logging_obj,
+            litellm_params,
+        )
 
-        # Get embedding model from litellm_params (required)
-        embedding_model: Final = litellm_params.get("litellm_embedding_model")
-        if not embedding_model:
-            raise ValueError(
-                "embedding_model is required in litellm_params for Azure AI Search. "
-                "Example: litellm_params['embedding_model'] = 'azure/text-embedding-3-large'"
-            )
-
-        embedding_config: Final = litellm_params.get("litellm_embedding_config", {})
-        if not embedding_config:
-            raise ValueError(
-                "embedding_config is required in litellm_params for Azure AI Search. "
-                "Example: litellm_params['embedding_config'] = {'api_base': 'https://krris-mh44uf7y-eastus2.cognitiveservices.azure.com/', 'api_key': 'os.environ/AZURE_API_KEY', 'api_version': '2025-09-01'}"
-            )
-
-        # Get vector field name (defaults to contentVector)
+    @staticmethod
+    def _search_request(
+        vector_store_id: str,
+        query_text: str,
+        query_vector: Sequence[float],
+        vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
+        api_base: str,
+        litellm_logging_obj: LiteLLMLoggingObj,
+        litellm_params: Mapping[str, object],
+    ) -> tuple[str, dict[str, object]]:
         vector_field: Final = litellm_params.get("azure_search_vector_field", "contentVector")
-
-        # Get top_k (number of results to return)
         top_k: Final = vector_store_search_optional_params.get("top_k", 10)
-
-        # Generate embedding for the query using litellm.embeddings
-        try:
-            embedding_response: Final = litellm.embedding(
-                model=embedding_model,
-                input=[query],
-                **embedding_config,
-            )
-            query_vector: Final = embedding_response.data[0]["embedding"]
-        except Exception as e:
-            raise Exception(f"Failed to generate embedding for query: {e}")
-
-        # Azure AI Search endpoint for search
-        index_name: Final = vector_store_id  # vector_store_id is the index name
-        url: Final = f"{api_base}/indexes/{index_name}/docs/search?api-version=2024-07-01"
-
-        # Build the request body for Azure AI Search with vector search
-        request_body: Final = {
-            "search": "*",  # Get all documents (filtered by vector similarity)
-            "vectorQueries": [
-                {
-                    "vector": query_vector,
-                    "fields": vector_field,
-                    "kind": "vector",
-                    "k": top_k,  # Number of nearest neighbors to return
-                }
-            ],
-            "select": "id,content",  # Fields to return (customize based on schema)
+        litellm_logging_obj.model_call_details["input"] = query_text
+        litellm_logging_obj.model_call_details["embedding_model"] = litellm_params.get("litellm_embedding_model")
+        litellm_logging_obj.model_call_details["top_k"] = top_k
+        return f"{api_base}/indexes/{vector_store_id}/docs/search?api-version=2024-07-01", {
+            "search": "*",
+            "vectorQueries": [{"vector": query_vector, "fields": vector_field, "kind": "vector", "k": top_k}],
+            "select": "id,content",
             "top": top_k,
         }
 
-        #########################################################
-        # Update logging object with details of the request
-        #########################################################
-        litellm_logging_obj.model_call_details["input"] = query
-        litellm_logging_obj.model_call_details["embedding_model"] = embedding_model
-        litellm_logging_obj.model_call_details["top_k"] = top_k
-
-        return url, request_body
-
     def transform_search_vector_store_response(
         self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj
     ) -> VectorStoreSearchResponse:

litellm/llms/base_llm/vector_store/transformation.py

@@ -1,10 +1,16 @@
+from __future__ import annotations
+
 from abc import abstractmethod
 from collections.abc import Mapping, Sequence
-from typing import TYPE_CHECKING, Any, NoReturn
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, runtime_checkable
 
 import httpx
+from pydantic import TypeAdapter
 
 from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import EmbeddingResponse
 from litellm.types.vector_stores import (
     VECTOR_STORE_OPENAI_PARAMS,
     BaseVectorStoreAuthCredentials,
@@ -28,6 +34,95 @@ else:
     BaseLLMException = Any
 
 
+@runtime_checkable
+class VectorStoreEmbeddingExecutor(Protocol):
+    def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: ...
+
+    async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: ...
+
+
+@dataclass(frozen=True, slots=True)
+class LiteLLMVectorStoreEmbeddingExecutor:
+    def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse:
+        import litellm
+
+        return litellm.embedding(  # pyright: ignore[reportCallIssue, reportUnknownMemberType, reportUnknownVariableType]  # provider kwargs are intentionally dynamic
+            model=model,
+            input=[query],  # mutable-ok: LiteLLM embedding requires a mutable input list
+            **dict(configuration),  # pyright: ignore[reportArgumentType]  # provider-specific embedding config is validated downstream  # mutable-ok: kwargs require a concrete dict
+        )
+
+    async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse:
+        import litellm
+
+        return await litellm.aembedding(  # pyright: ignore[reportUnknownMemberType]  # provider kwargs are intentionally dynamic
+            model=model,
+            input=[query],  # mutable-ok: LiteLLM embedding requires a mutable input list
+            **dict(configuration),  # pyright: ignore[reportArgumentType]  # provider-specific embedding config is validated downstream  # mutable-ok: kwargs require a concrete dict
+        )
+
+
+_REQUEST_METADATA: Final = TypeAdapter(dict[str, object])
+
+
+def vector_store_request_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]:
+    litellm_metadata: Final = kwargs.get("litellm_metadata")
+    if isinstance(litellm_metadata, dict):
+        return _REQUEST_METADATA.validate_python(litellm_metadata)
+    metadata: Final = kwargs.get("metadata")
+    if isinstance(metadata, dict):
+        return _REQUEST_METADATA.validate_python(metadata)
+    return MappingProxyType({})
+
+
+@dataclass(frozen=True, slots=True)
+class RouterVectorStoreEmbeddingExecutor:
+    router: Router
+    metadata: Mapping[str, object]
+
+    def _embedding_kwargs(self, configuration: Mapping[str, object]) -> Mapping[str, object]:
+        configured_metadata: Final = configuration.get("metadata")
+        metadata: Final = {
+            **(configured_metadata if isinstance(configured_metadata, Mapping) else {}),
+            **self.metadata,
+        }
+        return {
+            **{key: value for key, value in configuration.items() if key not in ("input", "metadata", "model")},
+            "metadata": metadata,
+        }
+
+    def _router_serves(self, model: str) -> bool:
+        team_id: Final = self.metadata.get("user_api_key_team_id")
+        resolved: Final = self.router.resolved_litellm_models(model, team_id if isinstance(team_id, str) else None)
+        deployment_models: Final = (
+            deployment.get("litellm_params", {}).get("model") for deployment in self.router.get_model_list() or ()
+        )
+        return bool(resolved) or model in deployment_models
+
+    def _embeds_through_sdk(self, model: str, configuration: Mapping[str, object]) -> bool:
+        return bool(configuration) and not self._router_serves(model)
+
+    def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse:
+        embedding_kwargs: Final = self._embedding_kwargs(configuration)
+        if self._embeds_through_sdk(model, configuration):
+            return LiteLLMVectorStoreEmbeddingExecutor().embed(model, query, embedding_kwargs)
+        return self.router.embedding(  # pyright: ignore[reportUnknownMemberType]  # Router embedding input retains a legacy untyped list
+            model=model,
+            input=[query],  # mutable-ok: Router embedding requires a mutable input list
+            **embedding_kwargs,  # pyright: ignore[reportArgumentType]  # provider kwargs are intentionally dynamic
+        )
+
+    async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse:
+        embedding_kwargs: Final = self._embedding_kwargs(configuration)
+        if self._embeds_through_sdk(model, configuration):
+            return await LiteLLMVectorStoreEmbeddingExecutor().aembed(model, query, embedding_kwargs)
+        return await self.router.aembedding(  # pyright: ignore[reportUnknownMemberType]  # Router embedding input retains a legacy untyped list
+            model=model,
+            input=[query],  # mutable-ok: Router embedding requires a mutable input list
+            **embedding_kwargs,  # pyright: ignore[reportArgumentType]  # provider kwargs are intentionally dynamic
+        )
+
+
 class BaseVectorStoreConfig:
     def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]:
         return []
@@ -58,7 +153,7 @@ class BaseVectorStoreConfig:
         litellm_logging_obj: LiteLLMLoggingObj,
         litellm_params: dict,
         extra_body: dict[str, Any] | None = None,
-        router: "Router | None" = None,
+        router: Router | None = None,
     ) -> tuple[str, dict]:
         pass
 
@@ -71,7 +166,7 @@ class BaseVectorStoreConfig:
         litellm_logging_obj: LiteLLMLoggingObj,
         litellm_params: dict,
         extra_body: dict[str, Any] | None = None,
-        router: "Router | None" = None,
+        router: Router | None = None,
     ) -> tuple[str, dict]:
         """
         Optional async version of transform_search_vector_store_request.
@@ -161,6 +256,116 @@ class BaseVectorStoreConfig:
         return 0.0, 0.0
 
 
+_EMPTY_EMBEDDING_CONFIGURATION: Final[Mapping[str, object]] = MappingProxyType({})
+_QUERY_VECTOR: Final = TypeAdapter(list[float])
+
+
+class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig):
+    @abstractmethod
+    def transform_search_vector_store_request(
+        self,
+        vector_store_id: str,
+        query: str | Sequence[str],
+        vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
+        api_base: str,
+        litellm_logging_obj: LiteLLMLoggingObj,
+        litellm_params: Mapping[str, object],
+        extra_body: Mapping[str, object] | None = None,
+        router: Router | None = None,
+        embedding_executor: VectorStoreEmbeddingExecutor | None = None,
+    ) -> tuple[str, dict[str, object]]:
+        pass
+
+    async def atransform_search_vector_store_request(
+        self,
+        vector_store_id: str,
+        query: str | Sequence[str],
+        vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
+        api_base: str,
+        litellm_logging_obj: LiteLLMLoggingObj,
+        litellm_params: Mapping[str, object],
+        extra_body: Mapping[str, object] | None = None,
+        router: Router | None = None,
+        embedding_executor: VectorStoreEmbeddingExecutor | None = None,
+    ) -> tuple[str, dict[str, object]]:
+        return self.transform_search_vector_store_request(
+            vector_store_id=vector_store_id,
+            query=query,
+            vector_store_search_optional_params=vector_store_search_optional_params,
+            api_base=api_base,
+            litellm_logging_obj=litellm_logging_obj,
+            litellm_params=litellm_params,
+            extra_body=extra_body,
+            router=router,
+            embedding_executor=embedding_executor,
+        )
+
+    @staticmethod
+    def query_text(query: str | Sequence[str]) -> str:
+        return query if isinstance(query, str) else " ".join(query)
+
+    @staticmethod
+    def query_embedding_model(litellm_params: Mapping[str, object]) -> str:
+        embedding_model: Final = litellm_params.get("litellm_embedding_model")
+        if isinstance(embedding_model, str) and embedding_model:
+            return embedding_model
+        raise ValueError(
+            "litellm_embedding_model is required in litellm_params for this vector store. "
+            "Example: litellm_params['litellm_embedding_model'] = 'openai/text-embedding-3-small'"
+        )
+
+    @staticmethod
+    def query_embedding_configuration(litellm_params: Mapping[str, object]) -> Mapping[str, object]:
+        configuration: Final = litellm_params.get("litellm_embedding_config")
+        if isinstance(configuration, Mapping):
+            return {str(key): value for key, value in configuration.items()}  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]  # litellm_params is an untyped dict, keys are re-validated as str here
+        return _EMPTY_EMBEDDING_CONFIGURATION
+
+    @staticmethod
+    def query_embedding_executor(
+        embedding_executor: VectorStoreEmbeddingExecutor | None,
+        router: Router | None,
+        request_metadata: Mapping[str, object] = MappingProxyType({}),
+    ) -> VectorStoreEmbeddingExecutor:
+        if embedding_executor is not None:
+            return embedding_executor
+        if router is not None:
+            return RouterVectorStoreEmbeddingExecutor(router=router, metadata=request_metadata)
+        return LiteLLMVectorStoreEmbeddingExecutor()
+
+    def embed_query(
+        self,
+        query_text: str,
+        litellm_params: Mapping[str, object],
+        embedding_executor: VectorStoreEmbeddingExecutor | None,
+        router: Router | None = None,
+    ) -> Sequence[float]:
+        model: Final = self.query_embedding_model(litellm_params)
+        configuration: Final = self.query_embedding_configuration(litellm_params)
+        executor: Final = self.query_embedding_executor(embedding_executor, router)
+        try:
+            response: Final = executor.embed(model, query_text, configuration)
+        except Exception as e:
+            raise Exception(f"Failed to generate embedding for query: {e}")
+        return _QUERY_VECTOR.validate_python(response.data[0]["embedding"])  # pyright: ignore[reportUnknownMemberType]  # EmbeddingResponse.data is an untyped list, the vector is validated here
+
+    async def aembed_query(
+        self,
+        query_text: str,
+        litellm_params: Mapping[str, object],
+        embedding_executor: VectorStoreEmbeddingExecutor | None,
+        router: Router | None = None,
+    ) -> Sequence[float]:
+        model: Final = self.query_embedding_model(litellm_params)
+        configuration: Final = self.query_embedding_configuration(litellm_params)
+        executor: Final = self.query_embedding_executor(embedding_executor, router)
+        try:
+            response: Final = await executor.aembed(model, query_text, configuration)
+        except Exception as e:
+            raise Exception(f"Failed to generate embedding for query: {e}")
+        return _QUERY_VECTOR.validate_python(response.data[0]["embedding"])  # pyright: ignore[reportUnknownMemberType]  # EmbeddingResponse.data is an untyped list, the vector is validated here
+
+
 class BaseDirectVectorStoreConfig(BaseVectorStoreConfig):
     """
     Base config for vector store providers whose datastore has no HTTP API
@@ -176,6 +381,7 @@ class BaseDirectVectorStoreConfig(BaseVectorStoreConfig):
         vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
         litellm_logging_obj: LiteLLMLoggingObj,
         litellm_params: Mapping[str, object],
+        embedding_executor: VectorStoreEmbeddingExecutor | None = None,
         timeout: float | httpx.Timeout | None = None,
     ) -> VectorStoreSearchResponse:
         pass
@@ -188,6 +394,7 @@ class BaseDirectVectorStoreConfig(BaseVectorStoreConfig):
         vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
         litellm_logging_obj: LiteLLMLoggingObj,
         litellm_params: Mapping[str, object],
+        embedding_executor: VectorStoreEmbeddingExecutor | None = None,
         timeout: float | httpx.Timeout | None = None,
     ) -> VectorStoreSearchResponse:
         pass
@@ -201,7 +408,7 @@ class BaseDirectVectorStoreConfig(BaseVectorStoreConfig):
         litellm_logging_obj: LiteLLMLoggingObj,
         litellm_params: Mapping[str, object],
         extra_body: Mapping[str, object] | None = None,
-        router: "Router | None" = None,
+        router: Router | None = None,
     ) -> NoReturn:
         raise NotImplementedError("Direct vector store providers execute the search themselves; no HTTP request shape")
 

litellm/llms/bedrock/chat/converse_transformation.py

@@ -943,6 +943,9 @@ class AmazonConverseConfig(BaseConfig):
                         litellm.verbose_logger.warning(DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, model)
                 else:
                     optional_params["thinking"] = value
+                    AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
+                        model=model, optional_params=optional_params, custom_llm_provider="bedrock"
+                    )
             elif param == "reasoning_effort" and isinstance(value, str):
                 self._handle_reasoning_effort_parameter(
                     model=model, reasoning_effort=value, optional_params=optional_params
@@ -1334,6 +1337,7 @@ class AmazonConverseConfig(BaseConfig):
             )
 
         additional_request_params.pop("parallel_tool_calls", None)
+        additional_request_params.pop("client_metadata", None)
 
         # Only set the topK value in for models that support it
         additional_request_params.update(self._handle_top_k_value(model, inference_params, drop_params))

litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py

@@ -107,6 +107,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
         # Restore original model name
         model = original_model
 
+        AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
+            model=original_model, optional_params=optional_params, custom_llm_provider="bedrock"
+        )
+
         # The stub model hides the original model from the parent's forced-tool-use backstop
         response_format_tool_choice: Final = optional_params.get("tool_choice")
         if (

litellm/llms/bedrock/common_utils.py

@@ -748,6 +748,15 @@ def strip_bedrock_throughput_suffix(model: str) -> str:
 
 
 MANTLE_MESSAGES_PATH: Final = "/anthropic/v1/messages"
+_MANTLE_OPENAI_BASE_SUFFIXES: Final = ("/openai/v1", "/v1")
+
+
+def _mantle_api_base_from_env() -> str | None:
+    env_base: Final = get_secret_str("BEDROCK_MANTLE_API_BASE")
+    if env_base is None:
+        return None
+    base: Final = env_base.rstrip("/")
+    return next((base[: -len(suffix)] for suffix in _MANTLE_OPENAI_BASE_SUFFIXES if base.endswith(suffix)), base)
 
 
 def build_mantle_messages_url(
@@ -758,12 +767,15 @@ def build_mantle_messages_url(
     """Build the bedrock-mantle Anthropic /messages URL.
 
     Honors an explicit endpoint override (``api_base``, then
-    ``aws_bedrock_runtime_endpoint``) so private VPC / VPCE / GovCloud Mantle
-    endpoints are reachable; otherwise falls back to the public regional host.
+    ``aws_bedrock_runtime_endpoint``, then ``BEDROCK_MANTLE_API_BASE``) so
+    private VPC / VPCE / GovCloud Mantle endpoints are reachable; otherwise
+    falls back to the public regional host.
     The mantle messages path is appended unless the override already carries it,
-    so callers can pass either the host or the full messages URL.
+    so callers can pass either the host or the full messages URL. The env var is
+    shared with the OpenAI-surface ``bedrock_mantle/*`` routes, which need it to
+    carry their ``/v1`` or ``/openai/v1`` base, so that suffix is dropped first.
     """
-    override: Final = api_base or aws_bedrock_runtime_endpoint
+    override: Final = api_base or aws_bedrock_runtime_endpoint or _mantle_api_base_from_env()
     if override:
         base: Final = override.rstrip("/")
         if base.endswith(MANTLE_MESSAGES_PATH):

litellm/llms/custom_httpx/http_handler.py

@@ -9,13 +9,15 @@ import threading
 import time
 from collections.abc import AsyncIterable, Callable, Iterable, Mapping
 from http.cookiejar import CookieJar, DefaultCookiePolicy
-from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, Optional, TypeAlias, TypedDict
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, Optional, TypeAlias, TypedDict, TypeVar
 
 import certifi
 import httpx
 from aiohttp import ClientSession, DummyCookieJar, TCPConnector
 from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport
-from httpx._types import RequestFiles
+from httpx._types import CertTypes, RequestFiles
+from httpx._utils import get_environment_proxies
 
 import litellm
 from litellm._logging import verbose_logger
@@ -66,6 +68,22 @@ _AddrInfo: TypeAlias = tuple[int | socket.AddressFamily, int | socket.SocketKind
 
 _RequestContent: TypeAlias = str | bytes | Iterable[bytes] | AsyncIterable[bytes]
 
+_IPV4_LOCAL_ADDRESS: Final = "0.0.0.0"
+
+_HttpxTransportT = TypeVar("_HttpxTransportT", HTTPTransport, AsyncHTTPTransport)
+
+
+def _environment_proxy_mounts(
+    build_proxy_transport: Callable[[str], _HttpxTransportT],
+) -> Mapping[str, _HttpxTransportT | None]:
+    """httpx skips its own HTTP(S)_PROXY / NO_PROXY mounts whenever an explicit `transport=` is passed."""
+    return MappingProxyType(
+        {
+            pattern: None if proxy_url is None else build_proxy_transport(proxy_url)
+            for pattern, proxy_url in get_environment_proxies().items()
+        }
+    )
+
 
 class _TCPConnectorKwargs(TypedDict, total=False):
     local_addr: tuple[str, int] | None
@@ -607,6 +625,7 @@ class AsyncHTTPHandler:
 
         return httpx.AsyncClient(
             transport=transport,
+            mounts=AsyncHTTPHandler._create_httpx_proxy_mounts(transport, verify=ssl_config, cert=cert),
             event_hooks=event_hooks,
             timeout=timeout,
             verify=ssl_config,
@@ -1191,10 +1210,22 @@ class AsyncHTTPHandler:
         - [Default] If force_ipv4 is False, it will return None
         """
         if litellm.force_ipv4:
-            return AsyncHTTPTransport(local_address="0.0.0.0")
+            return AsyncHTTPTransport(local_address=_IPV4_LOCAL_ADDRESS)
         else:
             return None
 
+    @staticmethod
+    def _create_httpx_proxy_mounts(
+        transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | None,
+        verify: VerifyTypes,
+        cert: CertTypes | None,
+    ) -> Mapping[str, AsyncHTTPTransport | None] | None:
+        if not isinstance(transport, AsyncHTTPTransport):
+            return None
+        return _environment_proxy_mounts(
+            lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert)
+        )
+
 
 class HTTPHandler:
     def __init__(
@@ -1227,6 +1258,7 @@ class HTTPHandler:
         # Create a client with a connection pool
         return httpx.Client(
             transport=self._create_sync_transport(),
+            mounts=self._create_sync_proxy_mounts(verify=ssl_config, cert=cert),
             timeout=self.timeout if self.timeout is not None else _DEFAULT_TIMEOUT,
             verify=ssl_config,
             cert=cert,
@@ -1507,10 +1539,19 @@ class HTTPHandler:
         Some users have seen httpx ConnectionError when using ipv6 - forcing ipv4 resolves the issue for them
         """
         if litellm.force_ipv4:
-            return HTTPTransport(local_address="0.0.0.0")
+            return HTTPTransport(local_address=_IPV4_LOCAL_ADDRESS)
         else:
             return getattr(litellm, "sync_transport", None)
 
+    @staticmethod
+    def _create_sync_proxy_mounts(
+        verify: VerifyTypes,
+        cert: CertTypes | None,
+    ) -> Mapping[str, HTTPTransport | None] | None:
+        if not litellm.force_ipv4:
+            return None
+        return _environment_proxy_mounts(lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert))
+
 
 def get_async_httpx_client(
     llm_provider: LlmProviders | httpxSpecialProvider,

litellm/llms/custom_httpx/llm_http_handler.py

@@ -1,6 +1,5 @@
 import asyncio
 import json
-import os
 import ssl
 from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
 from contextlib import asynccontextmanager
@@ -29,6 +28,7 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import (
     SUBTITLE_RESPONSE_FORMATS,
     synthesize_subtitle_document,
 )
+from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
 from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields
 from litellm.litellm_core_utils.realtime_errors import realtime_error_event, websocket_close_reason
 from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming
@@ -69,7 +69,9 @@ from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig
 from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig
 from litellm.llms.base_llm.vector_store.transformation import (
     BaseDirectVectorStoreConfig,
+    BaseQueryEmbeddingVectorStoreConfig,
     BaseVectorStoreConfig,
+    VectorStoreEmbeddingExecutor,
 )
 from litellm.llms.base_llm.vector_store_files.transformation import (
     BaseVectorStoreFilesConfig,
@@ -159,7 +161,11 @@ def _rust_responses_websocket_enabled(
     custom_llm_provider: str | None,
     litellm_params: GenericLiteLLMParams,
 ) -> bool:
-    return custom_llm_provider == "openai" and litellm_params.get("rust") is True
+    from litellm.rust_bridge.configuration import rust_enabled
+
+    raw_request_override: Final = litellm_params.get("rust")
+    request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
+    return custom_llm_provider == "openai" and rust_enabled(request_override=request_override)
 
 
 from .http_handler import get_shared_realtime_ssl_context
@@ -271,6 +277,16 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool:
     return False
 
 
+def _aws_signing_overrides(optional_params: Mapping[str, Any], litellm_params: Mapping[str, Any]) -> Mapping[str, Any]:
+    return MappingProxyType(
+        {
+            key: litellm_params[key]
+            for key in AWS_CREDENTIAL_KWARGS_KEYS
+            if optional_params.get(key) is None and litellm_params.get(key) is not None
+        }
+    )
+
+
 def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]:
     """Duck-type discover proxy hooks exposing per-frame project ITPM/OTPM
     enforcement, so the Responses WebSocket loop can charge every
@@ -535,7 +551,10 @@ class BaseLLMHTTPHandler:
 
         headers, signed_json_body = provider_config.sign_request(
             headers=headers,
-            optional_params=optional_params,
+            optional_params={
+                **optional_params,
+                **_aws_signing_overrides(optional_params, litellm_params),
+            },
             request_data=data,
             api_base=api_base,
             api_key=api_key,
@@ -2364,10 +2383,6 @@ class BaseLLMHTTPHandler:
             "anthropic_messages",
         )
 
-    @staticmethod
-    def _rust_env_enabled() -> bool:
-        return os.getenv("LITELLM_RUST", "").strip().lower() in {"1", "true", "yes", "on"}
-
     @staticmethod
     async def _maybe_rust_anthropic_messages(
         *,
@@ -2383,7 +2398,11 @@ class BaseLLMHTTPHandler:
     ) -> AnthropicMessagesResponse | None:
         if custom_llm_provider not in ("azure_ai", "anthropic"):
             return None
-        if litellm_params.get("rust") is not True and not BaseLLMHTTPHandler._rust_env_enabled():
+        from litellm.rust_bridge.configuration import rust_enabled
+
+        raw_request_override: Final = litellm_params.get("rust")
+        request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
+        if not rust_enabled(request_override=request_override):
             return None
         if has_agentic_hook:
             return None
@@ -9684,6 +9703,7 @@ class BaseLLMHTTPHandler:
         custom_llm_provider: str,
         litellm_params: GenericLiteLLMParams,
         logging_obj: LiteLLMLoggingObj,
+        embedding_executor: VectorStoreEmbeddingExecutor | None = None,
         extra_headers: dict[str, object] | None = None,
         extra_body: dict[str, object] | None = None,
         timeout: float | httpx.Timeout | None = None,
@@ -9704,6 +9724,7 @@ class BaseLLMHTTPHandler:
                 vector_store_search_optional_params=vector_store_search_optional_params,
                 litellm_logging_obj=logging_obj,
                 litellm_params=dict(litellm_params),  # mutable-ok: snapshot GenericLiteLLMParams into the Mapping shape
+                embedding_executor=embedding_executor,
                 timeout=timeout,
             )
 
@@ -9727,8 +9748,7 @@ class BaseLLMHTTPHandler:
             litellm_params=dict(litellm_params),
         )
 
-        # Check if provider has async transform method
-        if hasattr(vector_store_provider_config, "atransform_search_vector_store_request"):
+        if isinstance(vector_store_provider_config, BaseQueryEmbeddingVectorStoreConfig):
             (
                 url,
                 request_body,
@@ -9741,12 +9761,13 @@ class BaseLLMHTTPHandler:
                 litellm_params=dict(litellm_params),
                 extra_body=extra_body,
                 router=router,
+                embedding_executor=embedding_executor,
             )
         else:
             (
                 url,
                 request_body,
-            ) = vector_store_provider_config.transform_search_vector_store_request(
+            ) = await vector_store_provider_config.atransform_search_vector_store_request(
                 vector_store_id=vector_store_id,
                 query=query,
                 vector_store_search_optional_params=vector_store_search_optional_params,
@@ -9801,6 +9822,7 @@ class BaseLLMHTTPHandler:
         custom_llm_provider: str,
         litellm_params: GenericLiteLLMParams,
         logging_obj: LiteLLMLoggingObj,
+        embedding_executor: VectorStoreEmbeddingExecutor | None = None,
         extra_headers: dict[str, object] | None = None,
         extra_body: dict[str, object] | None = None,
         timeout: float | httpx.Timeout | None = None,
@@ -9817,6 +9839,7 @@ class BaseLLMHTTPHandler:
                 litellm_params=litellm_params,
                 logging_obj=logging_obj,
                 custom_llm_provider=custom_llm_provider,
+                embedding_executor=embedding_executor,
                 extra_headers=extra_headers,
                 extra_body=extra_body,
                 timeout=timeout,
@@ -9837,6 +9860,7 @@ class BaseLLMHTTPHandler:
                 vector_store_search_optional_params=vector_store_search_optional_params,
                 litellm_logging_obj=logging_obj,
                 litellm_params=dict(litellm_params),  # mutable-ok: snapshot GenericLiteLLMParams into the Mapping shape
+                embedding_executor=embedding_executor,
                 timeout=timeout,
             )
 
@@ -9857,19 +9881,35 @@ class BaseLLMHTTPHandler:
             litellm_params=dict(litellm_params),
         )
 
-        (
-            url,
-            request_body,
-        ) = vector_store_provider_config.transform_search_vector_store_request(
-            vector_store_id=vector_store_id,
-            query=query,
-            vector_store_search_optional_params=vector_store_search_optional_params,
-            api_base=api_base,
-            litellm_logging_obj=logging_obj,
-            litellm_params=dict(litellm_params),
-            extra_body=extra_body,
-            router=router,
-        )
+        if isinstance(vector_store_provider_config, BaseQueryEmbeddingVectorStoreConfig):
+            (
+                url,
+                request_body,
+            ) = vector_store_provider_config.transform_search_vector_store_request(
+                vector_store_id=vector_store_id,
+                query=query,
+                vector_store_search_optional_params=vector_store_search_optional_params,
+                api_base=api_base,
+                litellm_logging_obj=logging_obj,
+                litellm_params=dict(litellm_params),
+                extra_body=extra_body,
+                router=router,
+                embedding_executor=embedding_executor,
+            )
+        else:
+            (
+                url,
+                request_body,
+            ) = vector_store_provider_config.transform_search_vector_store_request(
+                vector_store_id=vector_store_id,
+                query=query,
+                vector_store_search_optional_params=vector_store_search_optional_params,
+                api_base=api_base,
+                litellm_logging_obj=logging_obj,
+                litellm_params=dict(litellm_params),
+                extra_body=extra_body,
+                router=router,
+            )
 
         all_optional_params: Final[dict[str, object]] = dict(litellm_params)
         all_optional_params.update(vector_store_search_optional_params or {})

litellm/llms/databricks/chat/transformation.py

@@ -330,6 +330,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
     ) -> dict:
         is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params)
         mapped_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params)
+        if "claude" in model:
+            AnthropicConfig.translate_legacy_thinking_for_adaptive_model(
+                model=model, optional_params=mapped_params, custom_llm_provider="databricks"
+            )
         if "tools" in mapped_params:
             mapped_params["tools"] = self._map_openai_to_dbrx_tool(model=model, tools=mapped_params["tools"])
         if "max_completion_tokens" in non_default_params and replace_max_completion_tokens_with_max_tokens:

litellm/llms/gigachat/authenticator.py

@@ -8,6 +8,7 @@ Based on official GigaChat SDK authentication flow.
 import time
 import uuid
 from collections.abc import Mapping
+from types import MappingProxyType
 from typing import Final
 
 import httpx
@@ -32,8 +33,8 @@ GIGACHAT_SCOPE: Final = "GIGACHAT_API_PERS"
 # Token expiry buffer in milliseconds (refresh token 60s before expiry)
 TOKEN_EXPIRY_BUFFER_MS: Final = 60000
 
-# Cache for access tokens
 _token_cache: Final = InMemoryCache()
+_NO_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
 
 
 class GigaChatAuthError(BaseLLMException):
@@ -80,10 +81,9 @@ def get_access_token(
     Raises:
         GigaChatAuthError: If authentication fails
     """
-    if not litellm_params:
-        litellm_params = {}  # mutable-ok: empty dict default; rebind-ok: provide default
+    params: Final = litellm_params or _NO_LITELLM_PARAMS
 
-    access_token: Final = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN")
+    access_token: Final = params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN")
     if access_token:
         return access_token
 
@@ -94,24 +94,20 @@ def get_access_token(
             message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.",
         )
 
-    effective_scope: Final = scope or litellm_params.get("gigachat_scope") or _get_scope()
-    effective_auth_url: Final = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url()
+    effective_scope: Final = scope or params.get("gigachat_scope") or _get_scope()
+    effective_auth_url: Final = auth_url or params.get("gigachat_auth_url") or _get_auth_url()
 
-    # Check cache
     cache_key: Final = f"gigachat_token:{effective_credentials[:16]}"
     cached: Final = _token_cache.get_cache(cache_key)
     if cached:
         _token, _expires_at = cached
-        # Check if token is still valid (with buffer)
         if time.time() * 1000 < _expires_at - TOKEN_EXPIRY_BUFFER_MS:
             verbose_logger.debug("Using cached GigaChat access token")
             return _token
 
-    # Request new token
     new_token, new_expires_at = _request_token_sync(effective_credentials, effective_scope, effective_auth_url)  # pyright: ignore[reportArgumentType]  # credential keys may be broader than str
 
     if new_expires_at:
-        # Cache token
         ttl_seconds: Final = max(0, (new_expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000)
         if ttl_seconds > 0:
             _token_cache.set_cache(cache_key, (new_token, new_expires_at), ttl=ttl_seconds)
@@ -126,10 +122,9 @@ async def get_access_token_async(
     litellm_params: Mapping[str, object] | None = None,
 ) -> str:
     """Async version of get_access_token."""
-    if not litellm_params:
-        litellm_params = {}  # mutable-ok: empty dict default; rebind-ok: provide default
+    params: Final = litellm_params or _NO_LITELLM_PARAMS
 
-    access_token: Final = litellm_params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN")
+    access_token: Final = params.get("gigachat_access_token") or get_secret_str("GIGACHAT_ACCESS_TOKEN")
     if access_token:
         return access_token
 
@@ -140,10 +135,9 @@ async def get_access_token_async(
             message="GigaChat credentials not provided. Set GIGACHAT_CREDENTIALS or GIGACHAT_API_KEY environment variable.",
         )
 
-    effective_scope: Final = scope or litellm_params.get("gigachat_scope") or _get_scope()
-    effective_auth_url: Final = auth_url or litellm_params.get("gigachat_auth_url") or _get_auth_url()
+    effective_scope: Final = scope or params.get("gigachat_scope") or _get_scope()
+    effective_auth_url: Final = auth_url or params.get("gigachat_auth_url") or _get_auth_url()
 
-    # Check cache
     cache_key: Final = f"gigachat_token:{effective_credentials[:16]}"
     cached: Final = _token_cache.get_cache(cache_key)
     if cached:
@@ -152,11 +146,9 @@ async def get_access_token_async(
             verbose_logger.debug("Using cached GigaChat access token")
             return _token
 
-    # Request new token
     new_token, new_expires_at = await _request_token_async(effective_credentials, effective_scope, effective_auth_url)  # pyright: ignore[reportArgumentType]  # credential keys may be broader than str
 
     if new_expires_at:
-        # Cache token
         ttl_seconds: Final = max(0, (new_expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000)
         if ttl_seconds > 0:
             _token_cache.set_cache(cache_key, (new_token, new_expires_at), ttl=ttl_seconds)

litellm/llms/gigachat/chat/streaming.py

@@ -52,7 +52,6 @@ class GigaChatModelResponseIterator:
         tool_use: ChatCompletionToolCallChunk | None = None  # rebind-ok: conditionally assigned on function_call
         finish_reason: str | None = chunk_finish_reason
 
-        # Handle function_call in stream
         raw_function_call: Final = delta.get("function_call")
         if chunk_finish_reason == "function_call" and isinstance(raw_function_call, Mapping) and raw_function_call:
             func_call: Final[Mapping[str, object]] = raw_function_call

litellm/llms/gigachat/chat/transformation.py

@@ -10,6 +10,7 @@ import json
 import time
 import uuid
 from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final
 
 import httpx
@@ -34,6 +35,9 @@ else:
     LiteLLMLoggingObj = Any
 
 
+_EMPTY_FUNCTION: Final[Mapping[str, object]] = MappingProxyType({})
+
+
 def is_valid_json(value: str) -> bool:
     """Checks whether the value passed is a valid serialized JSON string"""
     try:
@@ -111,11 +115,9 @@ class GigaChatConfig(BaseConfig):
         """
         Set up headers with OAuth token.
         """
-        # Get access token
         credentials: Final = api_key or get_secret_str("GIGACHAT_CREDENTIALS") or get_secret_str("GIGACHAT_API_KEY")
         access_token: Final = get_access_token(credentials=credentials, litellm_params=litellm_params)
 
-        # Store credentials for image uploads
         self._current_credentials = credentials
         self._current_api_base = api_base
 
@@ -208,18 +210,18 @@ class GigaChatConfig(BaseConfig):
 
     def _convert_tools_to_functions(self, tools: Sequence) -> Sequence[dict]:
         """Convert OpenAI tools format to GigaChat functions format."""
-        functions: Final[list[dict]] = []  # mutable-ok: accumulator for building functions list
-        for tool in tools:
-            if isinstance(tool, dict) and tool.get("type") == "function":
-                func = tool.get("function", {})
-                functions.append(
-                    {
-                        "name": func.get("name", ""),
-                        "description": func.get("description", ""),
-                        "parameters": func.get("parameters", {}),
-                    }
-                )
-        return functions
+        return [
+            {
+                "name": function.get("name", ""),
+                "description": function.get("description", ""),
+                "parameters": function.get("parameters", {}),
+            }
+            for function in (
+                tool.get("function", _EMPTY_FUNCTION)
+                for tool in tools
+                if isinstance(tool, dict) and tool.get("type") == "function"
+            )
+        ]
 
     def _map_tool_choice(self, tool_choice: str | Mapping[str, object]) -> str | Mapping[str, object] | None:
         """
@@ -299,7 +301,6 @@ class GigaChatConfig(BaseConfig):
                 if part.get("type") == "text":
                     texts.append(part.get("text", ""))
                 elif part.get("type") == "image_url":
-                    # Extract image URL and upload to GigaChat
                     image_url: object = part.get("image_url", {})
                     upload_url: str
                     if isinstance(image_url, str):
@@ -322,16 +323,13 @@ class GigaChatConfig(BaseConfig):
         headers: Mapping[str, object],
     ) -> dict:  # mutable-ok: request payload sent to httpx
         """Transform OpenAI request to GigaChat format."""
-        # Transform messages
         giga_messages: Final = self._transform_messages(messages)
 
-        # Build request
         request_data: Final[dict[str, object]] = {
             "model": model.replace("gigachat/", ""),
             "messages": giga_messages,
         }
 
-        # Add optional params
         for key in [
             "temperature",
             "top_p",
@@ -343,7 +341,6 @@ class GigaChatConfig(BaseConfig):
             if key in optional_params:
                 request_data[key] = optional_params[key]
 
-        # Add functions if present
         if "functions" in optional_params:
             request_data["functions"] = optional_params["functions"]
         if "function_call" in optional_params:
@@ -358,10 +355,8 @@ class GigaChatConfig(BaseConfig):
         for i, msg in enumerate(messages):
             message = dict(msg)
 
-            # Remove unsupported fields
             message.pop("name", None)
 
-            # Transform roles
             role = message.get("role", "user")
             if role == "developer":
                 message["role"] = "system"
@@ -374,18 +369,15 @@ class GigaChatConfig(BaseConfig):
                 if not isinstance(content, str) or not is_valid_json(content):
                     message["content"] = json.dumps(content, ensure_ascii=False)
 
-            # Handle None content
             if message.get("content") is None:
                 message["content"] = ""
 
-            # Handle list content (multimodal) - extract text and images
             content = message.get("content")
             if isinstance(content, list):
                 message["content"], attachments = self._transform_list_content(content)
                 if attachments:
                     message["attachments"] = attachments
 
-            # Transform tool_calls to function_call
             tool_calls = message.get("tool_calls")
             if tool_calls and isinstance(tool_calls, list) and len(tool_calls) > 0:
                 tool_call = tool_calls[0]
@@ -436,13 +428,11 @@ class GigaChatConfig(BaseConfig):
             message_data = choice.get("message", {})
             finish_reason = choice.get("finish_reason", "stop")
 
-            # Transform function_call to tool_calls or content
             if finish_reason == "function_call" and message_data.get("function_call"):
                 func_call = message_data["function_call"]
                 args = func_call.get("arguments", {})
 
                 if is_structured_output:
-                    # Convert to content for structured output
                     if isinstance(args, dict):
                         content = json.dumps(args, ensure_ascii=False)
                     else:
@@ -452,7 +442,6 @@ class GigaChatConfig(BaseConfig):
                     message_data.pop("functions_state_id", None)
                     finish_reason = "stop"
                 else:
-                    # Convert to tool_calls format
                     if isinstance(args, dict):
                         args = json.dumps(args, ensure_ascii=False)
                     message_data["tool_calls"] = [
@@ -468,7 +457,6 @@ class GigaChatConfig(BaseConfig):
                     message_data.pop("function_call", None)
                     finish_reason = "tool_calls"
 
-            # Clean up GigaChat-specific fields
             message_data.pop("functions_state_id", None)
 
             choices.append(

litellm/llms/gigachat/embedding/transformation.py

@@ -112,18 +112,10 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig):
             "input": ["text1", "text2", ...]
         }
         """
-        # Normalize input to list
-        if isinstance(input, str):
-            input_list: list = [input]  # rebind-ok: locally scoped conversion
-        else:
-            input_list = input
-
-        # Remove gigachat/ prefix from model if present
-        model = model.removeprefix("gigachat/")  # rebind-ok: parameter reassignment for normalization
-
+        normalized_input: Final = [input] if isinstance(input, str) else input  # mutable-ok: preserve list API
         return {
-            "model": model,
-            "input": input_list,
+            "model": model.removeprefix("gigachat/"),
+            "input": normalized_input,
         }
 
     def transform_embedding_response(

litellm/llms/gigachat/passthrough/transformation.py

@@ -60,7 +60,6 @@ class GigaChatPassthroughConfig(BasePassthroughConfig):
         """
         Set up headers with OAuth token.
         """
-        # Get access token
         access_token: Final = get_access_token(credentials=api_key, litellm_params=litellm_params)
 
         headers["Authorization"] = f"Bearer {access_token}"  # rebind-ok: mutating for OAuth setup
@@ -82,7 +81,6 @@ class GigaChatPassthroughConfig(BasePassthroughConfig):
         from litellm.types.utils import LlmProviders, ModelResponse
         from litellm.utils import ProviderConfigManager
 
-        # cost tracking only for completions and embeddings
         if "completions" in endpoint:
             provider_chat_config: Final = ProviderConfigManager.get_provider_chat_config(
                 provider=LlmProviders(custom_llm_provider),

litellm/llms/gigachat/utils.py

@@ -4,7 +4,6 @@ from typing import Final
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.utils import PromptTokensDetailsWrapper, Usage
 
-# GigaChat API endpoint
 GIGACHAT_BASE_URL: Final = "https://gigachat.devices.sberbank.ru/api/v1"
 
 

litellm/llms/hosted_vllm/rerank/transformation.py

@@ -3,16 +3,20 @@ Transformation logic for Hosted VLLM rerank
 """
 
 from collections.abc import Mapping
+from types import MappingProxyType
 from typing import Any, Final
 
 import httpx
+from pydantic import ValidationError
 
 from litellm._uuid import uuid
+from litellm.exceptions import UnsupportedParamsError
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.rerank import (
+    HostedVLLMRerankTruncationParams,
     OptionalRerankParams,
     RerankBilledUnits,
     RerankRequest,
@@ -34,6 +38,13 @@ class HostedVLLMRerankError(BaseLLMException):
         super().__init__(status_code=status_code, message=message, headers=headers)
 
 
+def validated_truncation_params(non_default_params: Mapping[str, object] | None) -> HostedVLLMRerankTruncationParams:
+    try:
+        return HostedVLLMRerankTruncationParams.model_validate(non_default_params or MappingProxyType({}))
+    except ValidationError as error:
+        raise UnsupportedParamsError(status_code=400, message=f"hosted_vllm rerank: {error}") from error
+
+
 class HostedVLLMRerankConfig(BaseRerankConfig):
     def __init__(self) -> None:
         pass
@@ -62,7 +73,11 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
             "top_n",
             "rank_fields",
             "return_documents",
+            "max_tokens_per_doc",
             "instruction",
+            "truncate_prompt_tokens",
+            "truncation_side",
+            "max_tokens_per_query",
         ]
 
     def map_cohere_rerank_params(
@@ -100,7 +115,15 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
         if instruction is not None:
             mapped_params["instruction"] = instruction
 
-        return dict(mapped_params)
+        truncation: Final = validated_truncation_params(non_default_params)
+        forwarded: Final[OptionalRerankParams] = {
+            **mapped_params,
+            "max_tokens_per_doc": max_tokens_per_doc,
+            "truncate_prompt_tokens": truncation.truncate_prompt_tokens,
+            "truncation_side": truncation.truncation_side,
+            "max_tokens_per_query": truncation.max_tokens_per_query,
+        }
+        return dict(forwarded)
 
     def validate_environment(
         self,
@@ -138,6 +161,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
         if "documents" not in optional_rerank_params:
             raise ValueError("documents is required for Hosted VLLM rerank")
 
+        truncation: Final = HostedVLLMRerankTruncationParams.model_validate(optional_rerank_params)
         rerank_request: Final = RerankRequest(
             model=model,
             query=optional_rerank_params["query"],
@@ -146,6 +170,10 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
             rank_fields=optional_rerank_params.get("rank_fields", None),
             return_documents=optional_rerank_params.get("return_documents", None),
             instruction=optional_rerank_params.get("instruction", None),
+            max_tokens_per_doc=truncation.max_tokens_per_doc,
+            truncate_prompt_tokens=truncation.truncate_prompt_tokens,
+            truncation_side=truncation.truncation_side,
+            max_tokens_per_query=truncation.max_tokens_per_query,
         )
         return rerank_request.model_dump(exclude_none=True)
 

litellm/llms/milvus/vector_stores/transformation.py

@@ -1,9 +1,14 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final
 
 import httpx
 
-import litellm
-from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
+from litellm.llms.base_llm.vector_store.transformation import (
+    BaseQueryEmbeddingVectorStoreConfig,
+    VectorStoreEmbeddingExecutor,
+)
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.vector_stores import (
@@ -37,7 +42,7 @@ MILVUS_OPTIONAL_PARAMS: Final = {
 }
 
 
-class MilvusVectorStoreConfig(BaseVectorStoreConfig):
+class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig):
     """
     Configuration for Milvus Vector Store
 
@@ -118,78 +123,79 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig):
     def transform_search_vector_store_request(
         self,
         vector_store_id: str,
-        query: str | list[str],
+        query: str | Sequence[str],
         vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
         api_base: str,
         litellm_logging_obj: LiteLLMLoggingObj,
-        litellm_params: dict,
-        extra_body: dict[str, Any] | None = None,
-        router: "Router | None" = None,
-    ) -> tuple[str, dict[str, Any]]:
-        """
-        Transform search request for Azure AI Search API
+        litellm_params: Mapping[str, object],
+        extra_body: Mapping[str, object] | None = None,
+        router: Router | None = None,
+        embedding_executor: VectorStoreEmbeddingExecutor | None = None,
+    ) -> tuple[str, dict[str, object]]:
+        query_text: Final = self.query_text(query)
+        query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor, router)
+        return self._search_request(
+            vector_store_id,
+            query_text,
+            query_vector,
+            vector_store_search_optional_params,
+            api_base,
+            litellm_logging_obj,
+            litellm_params,
+        )
 
-        Generates embeddings using litellm.embeddings and constructs Azure AI Search request
-        """
-        # Convert query to string if it's a list
-        if isinstance(query, list):
-            query = " ".join(query)
+    async def atransform_search_vector_store_request(
+        self,
+        vector_store_id: str,
+        query: str | Sequence[str],
+        vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
+        api_base: str,
+        litellm_logging_obj: LiteLLMLoggingObj,
+        litellm_params: Mapping[str, object],
+        extra_body: Mapping[str, object] | None = None,
+        router: Router | None = None,
+        embedding_executor: VectorStoreEmbeddingExecutor | None = None,
+    ) -> tuple[str, dict[str, object]]:
+        query_text: Final = self.query_text(query)
+        query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor, router)
+        return self._search_request(
+            vector_store_id,
+            query_text,
+            query_vector,
+            vector_store_search_optional_params,
+            api_base,
+            litellm_logging_obj,
+            litellm_params,
+        )
 
-        # Get embedding model from litellm_params (required)
-        embedding_model: Final = litellm_params.get("litellm_embedding_model")
-        if not embedding_model:
-            raise ValueError(
-                "embedding_model is required in litellm_params for Milvus. You can call any litellm embedding model."
-                "Example: litellm_params['embedding_model'] = 'azure/text-embedding-3-large'"
+    @staticmethod
+    def _search_request(
+        vector_store_id: str,
+        query_text: str,
+        query_vector: Sequence[float],
+        vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
+        api_base: str,
+        litellm_logging_obj: LiteLLMLoggingObj,
+        litellm_params: Mapping[str, object],
+    ) -> tuple[str, dict[str, object]]:
+        scope: Final = {
+            key: value
+            for key, value in (
+                ("dbName", litellm_params.get("milvus_db_name")),
+                ("partitionNames", litellm_params.get("milvus_partition_names")),
             )
-
-        embedding_config: Final = litellm_params.get("litellm_embedding_config", {})
-        if not embedding_config:
-            raise ValueError(
-                "embedding_config is required in litellm_params for Milvus. You can call any litellm embedding model."
-                "Example: litellm_params['embedding_config'] = {'api_base': 'https://krris-mh44uf7y-eastus2.cognitiveservices.azure.com/', 'api_key': 'os.environ/AZURE_API_KEY', 'api_version': '2025-09-01'}"
-            )
-
-        # Get top_k (number of results to return)
-        # Generate embedding for the query using litellm.embeddings
-        try:
-            embedding_response: Final = litellm.embedding(
-                model=embedding_model,
-                input=[query],
-                **embedding_config,
-            )
-            query_vector: Final = embedding_response.data[0]["embedding"]
-        except Exception as e:
-            raise Exception(f"Failed to generate embedding for query: {e}")
-
-        # Azure AI Search endpoint for search
-        index_name: Final = vector_store_id  # vector_store_id is the index name
-        url: Final = f"{api_base}/v2/vectordb/entities/search"
-
-        # Build the request body for Azure AI Search with vector search
-        request_body: Final[dict[str, Any]] = {
-            "collectionName": index_name,
+            if value
+        }
+        litellm_logging_obj.model_call_details["input"] = query_text
+        litellm_logging_obj.model_call_details["embedding_model"] = litellm_params.get("litellm_embedding_model")
+        return f"{api_base}/v2/vectordb/entities/search", {
+            "collectionName": vector_store_id,
             "data": [query_vector],
             "annsField": "book_intro_vector",
             **vector_store_search_optional_params,
+            **scope,
         }
 
-        db_name: Final = litellm_params.get("milvus_db_name")
-        if db_name:
-            request_body["dbName"] = db_name
-
-        partition_names: Final = litellm_params.get("milvus_partition_names")
-        if partition_names:
-            request_body["partitionNames"] = partition_names
-
-        #########################################################
-        # Update logging object with details of the request
-        #########################################################
-        litellm_logging_obj.model_call_details["input"] = query
-        litellm_logging_obj.model_call_details["embedding_model"] = embedding_model
-
-        return url, request_body
-
     def transform_search_vector_store_response(
         self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj
     ) -> VectorStoreSearchResponse:

litellm/llms/ollama/chat/transformation.py

@@ -423,6 +423,7 @@ class OllamaChatConfig(BaseConfig):
 class OllamaChatCompletionResponseIterator(BaseModelResponseIterator):
     started_reasoning_content: bool = False
     finished_reasoning_content: bool = False
+    seen_tool_calls: bool = False
 
     def _is_function_call_complete(self, function_args: str | dict) -> bool:
         if isinstance(function_args, dict):
@@ -468,6 +469,7 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator):
             # process tool calls - if complete function arg - add id to tool call
             tool_calls: Final = chunk["message"].get("tool_calls")
             if tool_calls is not None:
+                self.seen_tool_calls = True
                 for tool_call in tool_calls:
                     function_args = tool_call.get("function").get("arguments")
                     if function_args is not None and len(function_args) > 0:
@@ -508,9 +510,10 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator):
 
             if chunk["done"] is True:
                 finish_reason = chunk.get("done_reason") or "stop"
-                # Override finish_reason when tool_calls are present
+                # Override finish_reason when tool_calls appeared in any chunk
                 # Fixes: https://github.com/BerriAI/litellm/issues/18922
-                if tool_calls is not None:
+                # Fixes: https://github.com/BerriAI/litellm/issues/34692
+                if self.seen_tool_calls:
                     finish_reason = "tool_calls"
                 choices = [
                     StreamingChoices(

litellm/llms/openai/common_utils.py

@@ -305,14 +305,16 @@ class BaseOpenAILLM:
 
         # Get unified SSL configuration
         ssl_config: Final = get_ssl_configuration()
+        transport: Final = AsyncHTTPHandler._create_async_transport(
+            ssl_context=(ssl_config if isinstance(ssl_config, ssl.SSLContext) else None),
+            ssl_verify=ssl_config if isinstance(ssl_config, bool) else None,
+            shared_session=shared_session,
+        )
 
         return httpx.AsyncClient(
             verify=ssl_config,
-            transport=AsyncHTTPHandler._create_async_transport(
-                ssl_context=(ssl_config if isinstance(ssl_config, ssl.SSLContext) else None),
-                ssl_verify=ssl_config if isinstance(ssl_config, bool) else None,
-                shared_session=shared_session,
-            ),
+            transport=transport,
+            mounts=AsyncHTTPHandler._create_httpx_proxy_mounts(transport, verify=ssl_config, cert=None),
             follow_redirects=True,
         )
 

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -28,6 +28,7 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
     - text: str
 """
 
+import copy
 import time
 import uuid
 from collections.abc import Mapping, Sequence
@@ -36,7 +37,6 @@ from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Union, cast
 
 from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
-from openai.types.responses.tool_param import FunctionToolParam
 from pydantic import BaseModel, TypeAdapter
 from typing_extensions import ReadOnly, TypedDict
 
@@ -49,6 +49,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     blocked_responses_stream_usage,
     stream_item_field,
 )
+from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
 from litellm.responses.litellm_completion_transformation.transformation import (
     LiteLLMCompletionResponsesConfig,
 )
@@ -62,7 +63,6 @@ from litellm.types.llms.openai import (
     ContentPartDonePartOutputText,
     ErrorEvent,
     ErrorEventError,
-    OpenAIMcpServerTool,
     OutputItemAddedEvent,
     OutputItemDoneEvent,
     OutputTextDeltaEvent,
@@ -157,23 +157,31 @@ class OpenAIResponsesHandler(BaseTranslation):
         Handles both string input and list of message objects.
         """
         input_data: Final[str | ResponseInputParam | None] = data.get("input")
-        tools_to_check: Final[list[ChatCompletionToolParam]] = []
         if input_data is None:
             return data
 
         structured_messages: Final = self.get_structured_messages(data)
+        raw_tools: Final = data.get("tools")
+        original_tools: Final[tuple[Mapping[str, object], ...]] = (
+            tuple(raw_tools) if isinstance(raw_tools, list) else ()
+        )
+        flattened_tool_groups: Final = tuple(
+            form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools)
+        )
+        flattened_tools: Final = tuple(
+            cast(ChatCompletionToolParam, tool)  # cast-ok: mcp tools ride along in the guardrail's tool list
+            for group in flattened_tool_groups
+            for tool in group
+        )
+        tools_to_check: Final[list[ChatCompletionToolParam]] = list(  # mutable-ok: guardrail inputs want a list
+            copy.deepcopy(flattened_tools)
+        )
 
         # Handle simple string input
         if isinstance(input_data, str):
             inputs = GenericGuardrailAPIInputs(texts=[input_data])
-            original_tools: list[dict[str, object]] = []
-
-            # Extract and transform tools if present
-            if "tools" in data and data["tools"]:
-                original_tools = list(data["tools"])
-                self._extract_and_transform_tools(data["tools"], tools_to_check)
-                if tools_to_check:
-                    inputs["tools"] = tools_to_check
+            if tools_to_check:
+                inputs["tools"] = tools_to_check
             if structured_messages:
                 inputs["structured_messages"] = structured_messages
             # Include model information if available
@@ -189,7 +197,9 @@ class OpenAIResponsesHandler(BaseTranslation):
             )
             guardrailed_texts = guardrailed_inputs.get("texts", [])
             data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data
-            self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools"))
+            self._apply_guardrailed_tools_to_data(
+                data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools")
+            )
             verbose_proxy_logger.debug("OpenAI Responses API: Processed string input")
             return data
 
@@ -200,7 +210,6 @@ class OpenAIResponsesHandler(BaseTranslation):
         texts_to_check: Final[list[str]] = []
         images_to_check: Final[list[str]] = []
         task_mappings: Final[list[tuple[int, int | None]]] = []
-        original_tools_list: Final[list[dict[str, object]]] = list(data.get("tools") or [])
 
         # Step 1: Extract all text content, images, and tools
         for msg_idx, message in enumerate(input_data):
@@ -212,10 +221,6 @@ class OpenAIResponsesHandler(BaseTranslation):
                 task_mappings=task_mappings,
             )
 
-        # Extract and transform tools if present
-        if "tools" in data and data["tools"]:
-            self._extract_and_transform_tools(data["tools"], tools_to_check)
-
         # Step 2: Apply guardrail to all texts in batch
         if texts_to_check:
             inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
@@ -238,9 +243,7 @@ class OpenAIResponsesHandler(BaseTranslation):
 
             guardrailed_texts = guardrailed_inputs.get("texts", [])
             self._apply_guardrailed_tools_to_data(
-                data,
-                original_tools_list,
-                guardrailed_inputs.get("tools"),
+                data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools")
             )
 
             # Step 3: Map guardrail responses back to original input structure
@@ -267,73 +270,18 @@ class OpenAIResponsesHandler(BaseTranslation):
                 names.append(str(tool["server_label"]))
         return names
 
-    def _extract_and_transform_tools(
-        self,
-        tools: list[FunctionToolParam | OpenAIMcpServerTool],
-        tools_to_check: list[ChatCompletionToolParam],
-    ) -> None:
-        """
-        Extract and transform tools from Responses API format to Chat Completion format.
-
-        Uses the LiteLLM transformation function to convert Responses API tools
-        to Chat Completion tools that can be passed to guardrails.
-        """
-        if tools is not None and isinstance(tools, list):
-            # Transform Responses API tools to Chat Completion tools
-            (
-                transformed_tools,
-                _,
-            ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools)
-            tools_to_check.extend(cast(list[ChatCompletionToolParam], transformed_tools))
-
-    def _remap_tools_to_responses_api_format(self, guardrailed_tools: list[Any]) -> list[dict[str, object]]:
-        """
-        Remap guardrail-returned tools (Chat Completion format) back to
-        Responses API request tool format.
-        """
-        return LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools(
-            guardrailed_tools
-        )
-
-    def _merge_tools_after_guardrail(
-        self,
-        original_tools: list[dict[str, object]],
-        remapped: list[dict[str, object]],
-    ) -> list[dict[str, object]]:
-        """
-        Merge remapped guardrailed tools with original tools that were not sent
-        to the guardrail (e.g. web_search, web_search_preview), preserving order.
-        Tools a guardrail appended (``remapped`` longer than ``original_tools``)
-        have no original slot and are kept so an injected tool is not dropped.
-        """
-        if not original_tools:
-            return remapped
-        result: Final[list[dict[str, object]]] = []
-        j = 0
-        for tool in original_tools:
-            if isinstance(tool, dict) and tool.get("type") in (
-                "web_search",
-                "web_search_preview",
-            ):
-                result.append(tool)
-            else:
-                if j < len(remapped):
-                    result.append(remapped[j])
-                    j += 1
-        # Keep guardrail-appended tools that matched no original slot above.
-        result.extend(remapped[j:])
-        return result
-
     def _apply_guardrailed_tools_to_data(
         self,
         data: dict,
-        original_tools: list[dict[str, object]],
-        guardrailed_tools: list[ChatCompletionToolParam] | None,
+        original_tools: Sequence[Mapping[str, object]],
+        flattened_tool_groups: Sequence[Sequence[Mapping[str, object]]],
+        guardrailed_tools: Sequence[ChatCompletionToolParam] | None,
     ) -> None:
-        """Remap guardrailed tools to Responses API format and merge with original, then set data['tools']."""
-        if guardrailed_tools is not None:
-            remapped: Final = self._remap_tools_to_responses_api_format(guardrailed_tools)
-            data["tools"] = self._merge_tools_after_guardrail(original_tools, remapped)
+        if guardrailed_tools is None:
+            return
+        data["tools"] = list(  # mutable-ok: downstream wants a list  # rebind-ok: in-place request rewrite
+            merge_guardrailed_tools(original_tools, flattened_tool_groups, guardrailed_tools)
+        )
 
     def _extract_input_text_and_images(
         self,

litellm/llms/openai/responses/guardrail_translation/tool_merge.py

@@ -0,0 +1,182 @@
+from collections.abc import Iterable, Mapping, Sequence
+from itertools import accumulate, chain, groupby
+from types import MappingProxyType
+from typing import Final, TypeAlias
+
+from pydantic import BaseModel, TypeAdapter, ValidationError
+
+from litellm._logging import verbose_logger
+from litellm.responses.litellm_completion_transformation.transformation import (
+    NAMESPACE_DESCRIPTION_SEPARATOR,
+    LiteLLMCompletionResponsesConfig,
+)
+
+Tool: TypeAlias = Mapping[str, object]
+IndexedKey: TypeAlias = tuple[str, int]
+
+_TOOL_ADAPTER: Final = TypeAdapter(dict[str, object])
+_CHAT_TOOL_TOP_LEVEL_KEYS: Final = frozenset({"type", "function"})
+
+
+def _as_tool(value: object) -> Tool | None:
+    candidate: Final = value.model_dump(exclude_unset=True) if isinstance(value, BaseModel) else value
+    try:
+        return _TOOL_ADAPTER.validate_python(candidate)
+    except ValidationError:
+        return None
+
+
+def _validated_tools(values: Iterable[object]) -> tuple[Tool, ...]:
+    validated: Final = tuple(map(_as_tool, values))
+    dropped: Final = sum(tool is None for tool in validated)
+    if dropped:
+        verbose_logger.warning("Dropping %d guardrail-returned tools that are not objects", dropped)
+    return tuple(tool for tool in validated if tool is not None)
+
+
+def _is_function(tool: Tool) -> bool:
+    return tool.get("type") == "function"
+
+
+def _chat_tool_key(tool: Tool) -> str:
+    tool_type: Final = str(tool.get("type") or "")
+    function: Final = _as_tool(tool.get("function"))
+    if function is not None:
+        return f"{tool_type}:{function.get('name') or ''}"
+    return f"{tool_type}:{tool.get('server_label') or tool.get('name') or ''}"
+
+
+def _indexed_keys(tools: Sequence[Tool]) -> tuple[IndexedKey, ...]:
+    keys: Final = tuple(_chat_tool_key(tool) for tool in tools)
+    positions_by_key: Final = groupby(sorted(range(len(keys)), key=keys.__getitem__), key=keys.__getitem__)
+    ordinal_by_position: Final = MappingProxyType(
+        {position: ordinal for _, positions in positions_by_key for ordinal, position in enumerate(positions)}
+    )
+    return tuple((key, ordinal_by_position[position]) for position, key in enumerate(keys))
+
+
+def _namespace_members(namespace: Tool) -> tuple[Tool, ...]:
+    members: Final = namespace.get("tools")
+    if not isinstance(members, Sequence) or isinstance(members, (str, bytes)):
+        return ()
+    return tuple(member for member in map(_as_tool, members) if member is not None)
+
+
+def _function_fields(tool: Tool) -> Tool:
+    function: Final = _as_tool(tool.get("function"))
+    return function if function is not None else MappingProxyType({})
+
+
+def _without_namespace_prefix(key: str, value: object, prefix: str) -> object:
+    if key != "description" or not isinstance(value, str) or not value.startswith(prefix):
+        return value
+    return value[len(prefix) :]
+
+
+def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_description: str) -> Tool:
+    flattened_function: Final = _function_fields(flattened)
+    prefix: Final = f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}" if namespace_description else ""
+    changed_function: Final = MappingProxyType(
+        {
+            key: _without_namespace_prefix(key, value, prefix)
+            for key, value in _function_fields(guardrailed).items()
+            if flattened_function.get(key) != value
+        }
+    )
+    changed_extras: Final = MappingProxyType(
+        {
+            key: value
+            for key, value in guardrailed.items()
+            if key not in _CHAT_TOOL_TOP_LEVEL_KEYS and flattened.get(key) != value
+        }
+    )
+    return {**member, **changed_extras, **changed_function}  # mutable-ok: json.dumps rejects MappingProxyType
+
+
+def _rebuilt_function_members(
+    function_members: Sequence[Tool],
+    flattened_group: Sequence[Tool],
+    group_keys: Sequence[IndexedKey],
+    guardrailed_by_key: Mapping[IndexedKey, Tool],
+    namespace_description: str,
+) -> tuple[Tool | None, ...]:
+    return tuple(
+        None
+        if key not in guardrailed_by_key
+        else member
+        if guardrailed_by_key[key] == flattened
+        else _rebuilt_member(member, flattened, guardrailed_by_key[key], namespace_description)
+        for member, flattened, key in zip(function_members, flattened_group, group_keys)
+    )
+
+
+def _rebuilt_namespace(
+    original: Tool,
+    members: Sequence[Tool],
+    flattened_group: Sequence[Tool],
+    group_keys: Sequence[IndexedKey],
+    guardrailed_by_key: Mapping[IndexedKey, Tool],
+) -> tuple[Tool, ...]:
+    namespace_description: Final = str(original.get("description") or "")
+    rebuilt_functions: Final = iter(
+        _rebuilt_function_members(
+            tuple(member for member in members if _is_function(member)),
+            flattened_group,
+            group_keys,
+            guardrailed_by_key,
+            namespace_description,
+        )
+    )
+    rebuilt_members: Final = tuple(
+        rebuilt
+        for rebuilt in (next(rebuilt_functions) if _is_function(member) else member for member in members)
+        if rebuilt is not None
+    )
+    if not rebuilt_members:
+        return ()
+    return ({**original, "tools": list(rebuilt_members)},)  # mutable-ok: json.dumps needs a plain dict and list
+
+
+def _merged_original(
+    original: Tool,
+    flattened_group: Sequence[Tool],
+    group_keys: Sequence[IndexedKey],
+    guardrailed_by_key: Mapping[IndexedKey, Tool],
+) -> tuple[Tool, ...]:
+    if not group_keys:
+        return (original,)
+    guardrailed_group: Final = tuple(guardrailed_by_key[key] for key in group_keys if key in guardrailed_by_key)
+    if guardrailed_group == tuple(flattened_group):
+        return (original,)
+    members: Final = _namespace_members(original) if original.get("type") == "namespace" else ()
+    if members and sum(map(_is_function, members)) == len(flattened_group):
+        return _rebuilt_namespace(original, members, flattened_group, group_keys, guardrailed_by_key)
+    if not guardrailed_group:
+        return ()
+    return tuple(
+        LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools(guardrailed_group)
+    )
+
+
+def merge_guardrailed_tools(
+    original_tools: Sequence[Tool],
+    flattened_groups: Sequence[Sequence[Tool]],
+    guardrailed_tools: Iterable[object],
+) -> tuple[Tool, ...]:
+    guardrailed: Final = _validated_tools(guardrailed_tools)
+    flattened_keys: Final = _indexed_keys(tuple(chain.from_iterable(flattened_groups)))
+    guardrailed_keys: Final = _indexed_keys(guardrailed)
+    guardrailed_by_key: Final = MappingProxyType(dict(zip(guardrailed_keys, guardrailed)))
+    group_ends: Final = tuple(accumulate(len(group) for group in flattened_groups))
+    group_key_slices: Final = tuple(
+        flattened_keys[end - len(group) : end] for group, end in zip(flattened_groups, group_ends)
+    )
+    merged_originals: Final = chain.from_iterable(
+        _merged_original(original, group, group_keys, guardrailed_by_key)
+        for original, group, group_keys in zip(original_tools, flattened_groups, group_key_slices)
+    )
+    owned_keys: Final = frozenset(flattened_keys)
+    appended: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools(
+        tuple(tool for key, tool in zip(guardrailed_keys, guardrailed) if key not in owned_keys)
+    )
+    return tuple(chain(merged_originals, appended))

litellm/llms/openai_like/README.md

@@ -54,7 +54,10 @@ That's it! The provider will be automatically loaded and available.
     "constraints": {
       "temperature_max": 1.0,
       "temperature_min": 0.0,
-      "temperature_min_with_n_gt_1": 0.3
+      "temperature_min_with_n_gt_1": 0.3,
+      // /v1/messages providers only: keep Anthropic cache_control extensions
+      // such as ttl instead of stripping them down to {"type": ...}
+      "cache_control_ttl": true
     },
     
     // Optional: Special handling flags

litellm/llms/openai_like/messages/transformation.py

@@ -1,11 +1,13 @@
 from typing import Any, Final
 
 import litellm
+from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload
 from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
     AnthropicMessagesConfig,
 )
 from litellm.llms.openai_like.json_loader import SimpleProviderConfig
 from litellm.secret_managers.main import get_secret_str
+from litellm.types.router import GenericLiteLLMParams
 
 DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01"
 
@@ -19,10 +21,17 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig):
     ``"/v1/messages"``. The inbound Anthropic payload (system, cache_control,
     thinking, tools, ...) is forwarded essentially unchanged to
     ``{api_base}/v1/messages``, so Anthropic-only features that the
-    Anthropic->OpenAI translation would otherwise drop are preserved. Response
-    parsing and streaming are inherited from the native Anthropic config.
+    Anthropic->OpenAI translation would otherwise drop are preserved. The one
+    exception is ``cache_control``, whose Anthropic-only extensions (``ttl``)
+    are stripped unless the deployment opts in with
+    ``model_info.cache_control_ttl: true``. Response parsing and streaming are
+    inherited from the native Anthropic config.
     """
 
+    def __init__(self, cache_control_ttl: bool = False) -> None:
+        super().__init__()
+        self._cache_control_ttl: Final = cache_control_ttl
+
     def validate_anthropic_messages_environment(
         self,
         headers: dict[str, str],
@@ -53,6 +62,35 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig):
     def should_filter_anthropic_beta_headers(self) -> bool:
         return False
 
+    def supports_cache_control_ttl(self) -> bool:
+        return self._cache_control_ttl
+
+    def transform_anthropic_messages_request(
+        self,
+        model: str,
+        messages: list[dict],  # mutable-ok: matches dict-typed base signature
+        anthropic_messages_optional_request_params: dict,  # mutable-ok: matches dict-typed base signature
+        litellm_params: GenericLiteLLMParams,
+        headers: dict,  # mutable-ok: matches dict-typed base signature
+    ) -> dict:  # mutable-ok: matches dict-typed base signature
+        """
+        Anthropic ignores prompt-caching hints it cannot honor, but strict
+        non-Anthropic implementations of the Messages API 400 the whole request
+        on Anthropic-only ``cache_control`` extensions (``cache_control.ttl: 1h
+        is not supported``), so unless the provider declares ttl support the
+        hints are reduced to their portable ``{"type": ...}`` core.
+        """
+        request: Final = super().transform_anthropic_messages_request(
+            model=model,
+            messages=messages,
+            anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
+            litellm_params=litellm_params,
+            headers=headers,
+        )
+        if self.supports_cache_control_ttl():
+            return request
+        return normalize_cache_control_in_anthropic_payload(request)
+
     def get_complete_url(
         self,
         api_base: str | None,
@@ -81,7 +119,7 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig):
     """
 
     def __init__(self, provider: SimpleProviderConfig):
-        super().__init__()
+        super().__init__(cache_control_ttl=bool(provider.constraints.get("cache_control_ttl")))
         self._provider = provider
 
     @property

litellm/llms/valkey/vector_stores/transformation.py

@@ -15,7 +15,10 @@ import httpx
 from pydantic import BaseModel, ConfigDict
 
 import litellm
-from litellm.llms.base_llm.vector_store.transformation import BaseDirectVectorStoreConfig
+from litellm.llms.base_llm.vector_store.transformation import (
+    BaseDirectVectorStoreConfig,
+    VectorStoreEmbeddingExecutor,
+)
 from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector
 from litellm.types.utils import EmbeddingResponse
 from litellm.types.vector_stores import (
@@ -213,6 +216,7 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig):
         vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
         litellm_logging_obj: "LiteLLMLoggingObj",
         litellm_params: Mapping[str, object],
+        embedding_executor: VectorStoreEmbeddingExecutor | None = None,
         timeout: float | httpx.Timeout | None = None,
     ) -> VectorStoreSearchResponse:
         params: Final = _ValkeySearchParams.model_validate(litellm_params)
@@ -222,10 +226,18 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig):
             embedding_field=params.embedding_field,
             text_field=params.text_field,
         )
-        embedding_response: Final = self.embedding_fn(
-            model=params.require_embedding_model(),
-            input=[query_text],  # mutable-ok: litellm.embedding's input contract is a list
-            **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG),
+        embedding_response: Final = (
+            embedding_executor.embed(
+                params.require_embedding_model(),
+                query_text,
+                params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG,
+            )
+            if embedding_executor is not None
+            else self.embedding_fn(
+                model=params.require_embedding_model(),
+                input=[query_text],  # mutable-ok: the injected embedding callable requires list input
+                **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG),
+            )
         )
         vec_params: Final = {"vec": pack_vector(embedding_response.data[0]["embedding"])}  # mutable-ok: redis-py API
 
@@ -252,6 +264,7 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig):
         vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
         litellm_logging_obj: "LiteLLMLoggingObj",
         litellm_params: Mapping[str, object],
+        embedding_executor: VectorStoreEmbeddingExecutor | None = None,
         timeout: float | httpx.Timeout | None = None,
     ) -> VectorStoreSearchResponse:
         params: Final = _ValkeySearchParams.model_validate(litellm_params)
@@ -261,10 +274,18 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig):
             embedding_field=params.embedding_field,
             text_field=params.text_field,
         )
-        embedding_response: Final = await self.aembedding_fn(
-            model=params.require_embedding_model(),
-            input=[query_text],  # mutable-ok: litellm.embedding's input contract is a list
-            **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG),
+        embedding_response: Final = (
+            await embedding_executor.aembed(
+                params.require_embedding_model(),
+                query_text,
+                params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG,
+            )
+            if embedding_executor is not None
+            else await self.aembedding_fn(
+                model=params.require_embedding_model(),
+                input=[query_text],  # mutable-ok: the injected embedding callable requires list input
+                **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG),
+            )
         )
         vec_params: Final = {"vec": pack_vector(embedding_response.data[0]["embedding"])}  # mutable-ok: redis-py API
 

litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py

@@ -177,6 +177,10 @@ class VertexAIAnthropicConfig(AnthropicConfig):
         # Restore original model name for any other processing
         model = original_model
 
+        AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
+            model=original_model, optional_params=optional_params, custom_llm_provider="vertex_ai"
+        )
+
         return optional_params
 
     def transform_response(

litellm/ocr/main.py

@@ -194,6 +194,12 @@ def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool:
     return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS
 
 
+def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool:
+    raw_request_override: Final = prepared_request.litellm_params.get("rust")
+    request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
+    return rust_ocr_bridge.rust_ocr_enabled(request_override=request_override)
+
+
 def _rust_bridge_optional_params(
     prepared_request: _PreparedOCRRequest,
     resolve_secret: Callable[[str], str | None],
@@ -422,7 +428,7 @@ async def aocr(
         custom_llm_provider = prepared.custom_llm_provider
         completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
 
-        if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled():
+        if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared):
             from litellm.secret_managers.main import get_secret_str
 
             rust_response: Final = await _run_rust_aocr(
@@ -694,7 +700,7 @@ def ocr(
         custom_llm_provider = prepared.custom_llm_provider
         completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
 
-        if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled():
+        if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared):
             from litellm.secret_managers.main import get_secret_str
 
             rust_response: Final = _run_rust_ocr(

litellm/passthrough/main.py

@@ -113,10 +113,8 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]):
                 )
             )
 
-            # Compliant: Save a strong reference to prevent GC
             self._background_tasks.add(task)
 
-            # Remove the task from the set when it finishes to avoid memory leaks
             task.add_done_callback(self._background_tasks.discard)
         except Exception as e:  # noqa: BLE001 # Safe catch-all for verbose logging
             verbose_logger.exception(
@@ -578,7 +576,6 @@ def llm_passthrough_route(
             else:
                 return response
     except Exception as e:
-        # provider_config is guaranteed non-None here due to the earlier guard
         assert provider_config is not None
         raise base_llm_http_handler._handle_error(
             e=e,

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -74,7 +74,13 @@ _MCP_GUARDRAIL_REJECTIONS: Final = (
 )
 
 
-def _connection_error_message(exc: BaseException) -> str:
+def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str:
+    if isinstance(exc, TimeoutError):
+        return (
+            f"Failed to connect to MCP server: no response from {url or 'the server'} "
+            f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL "
+            "from its network (DNS, egress rules, firewalls) and that the server answers MCP requests."
+        )
     if isinstance(exc, httpx.LocalProtocolError):
         return (
             "Failed to connect to MCP server: a request header is malformed. "
@@ -92,6 +98,9 @@ def _connection_error_message(exc: BaseException) -> str:
 
 
 if MCP_AVAILABLE:
+    from mcp.types import Tool as MCPTool
+
+    from litellm.experimental_mcp_client.client import MCPClient
     from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
         _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES,
         global_mcp_server_manager,
@@ -876,7 +885,6 @@ if MCP_AVAILABLE:
                         return (), classify_list_exception(e)
                     return tools_result, ServerListOk(tool_count=len(tools_result))
 
-                # Query all servers the user has access to
                 queried_servers: Final = tuple(
                     server
                     for server in map(global_mcp_server_manager.get_mcp_server_by_id, allowed_server_ids)
@@ -1141,12 +1149,18 @@ if MCP_AVAILABLE:
         scopes: Final[list[str] | None] = scopes_raw if isinstance(scopes_raw, list) else None
         return client_id, client_secret, scopes
 
+    async def _list_tools_within(client: MCPClient, deadline: float) -> list[MCPTool] | None:
+        with anyio.move_on_after(deadline):
+            return await client.list_tools(raise_on_error=True)
+        return None
+
     async def _execute_with_mcp_client(
         request: NewMCPServerRequest,
         operation: Callable[..., Awaitable[Mapping[str, object]]],
         mcp_auth_header: str | dict[str, str] | None = None,
         oauth2_headers: dict[str, str] | None = None,
         raw_headers: dict[str, str] | None = None,
+        timeout_seconds: float = MCP_TOOL_LISTING_TIMEOUT,
     ) -> Mapping[str, object]:
         """
         Create a temporary MCP client from *request*, run *operation*, and return the result.
@@ -1162,6 +1176,10 @@ if MCP_AVAILABLE:
             oauth2_headers: Headers extracted from the incoming request (may contain the
                 litellm API key — must NOT be forwarded for M2M servers).
             raw_headers: Raw request headers forwarded for stdio env construction.
+            timeout_seconds: Cap on OAuth discovery, connect, handshake, and *operation*
+                combined. Defaults to ``MCP_TOOL_LISTING_TIMEOUT`` (30s, below common LB
+                timeouts) so an unreachable upstream yields this endpoint's JSON error
+                instead of an opaque load-balancer 504 with an empty body.
 
         Returns:
             The dict returned by *operation*, or an error dict on failure.
@@ -1252,15 +1270,16 @@ if MCP_AVAILABLE:
                 static_headers=request.static_headers,
             )
 
-            client: Final = await global_mcp_server_manager._create_mcp_client(
-                server=server_model,
-                mcp_auth_header=mcp_auth_header,
-                extra_headers=merged_headers,
-                stdio_env=stdio_env,
-                cred_provider=preview_cred_provider,
-            )
+            with anyio.fail_after(timeout_seconds):
+                client: Final = await global_mcp_server_manager._create_mcp_client(
+                    server=server_model,
+                    mcp_auth_header=mcp_auth_header,
+                    extra_headers=merged_headers,
+                    stdio_env=stdio_env,
+                    cred_provider=preview_cred_provider,
+                )
 
-            return await operation(client)
+                return await operation(client)
 
         except (KeyboardInterrupt, SystemExit, asyncio.CancelledError):
             raise
@@ -1269,7 +1288,7 @@ if MCP_AVAILABLE:
             return {
                 "status": "error",
                 "error": True,
-                "message": _connection_error_message(e),
+                "message": _connection_error_message(e, request.url, timeout_seconds),
             }
 
     async def _preview_openapi_tools(spec_path: str) -> dict:
@@ -1422,9 +1441,7 @@ if MCP_AVAILABLE:
                 getattr(client, "timeout", MCP_CLIENT_TIMEOUT) or MCP_CLIENT_TIMEOUT,
                 MCP_TOOL_LISTING_TIMEOUT,
             )
-            list_tools_result = None  # rebind-ok: set inside the timeout scope below
-            with anyio.move_on_after(listing_deadline):
-                list_tools_result = await client.list_tools(raise_on_error=True)  # rebind-ok: fills the init above
+            list_tools_result: Final = await _list_tools_within(client, listing_deadline)
             if list_tools_result is None:
                 verbose_logger.warning(
                     "MCP tools/list preview timed out after %s seconds while paginating upstream tools",

litellm/proxy/auth/ip_address_utils.py

@@ -6,8 +6,11 @@ External callers (public IPs) only see servers with available_on_public_internet
 """
 
 import ipaddress
+import os
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Any, Final
+from urllib.parse import urlparse
 
 from fastapi import Request
 from pydantic import TypeAdapter, ValidationError
@@ -137,7 +140,7 @@ class IPAddressUtils:
     @staticmethod
     def is_request_from_trusted_proxy(
         request: Request,
-        general_settings: dict[str, Any] | None = None,
+        general_settings: Mapping[str, Any] | None = None,
     ) -> bool:
         """
         Return True if X-Forwarded-* headers on this request should be trusted.
@@ -190,6 +193,36 @@ class IPAddressUtils:
         trusted_networks: Final = IPAddressUtils.parse_trusted_proxy_networks(trusted_ranges)
         return IPAddressUtils.is_trusted_proxy(direct_ip, trusted_networks)
 
+    @staticmethod
+    def is_request_https(
+        request: Request,
+        general_settings: Mapping[str, Any] | None = None,
+    ) -> bool:
+        """
+        Whether this request's PUBLIC-facing origin is HTTPS, for deciding
+        whether a cookie set on the response should be marked ``Secure``.
+
+        litellm only sees a plain-HTTP hop whenever TLS terminates at a
+        reverse proxy, so ``request.url.scheme`` alone cannot answer this in
+        that deployment shape. Resolved from the first trusted signal:
+          1. ``PROXY_BASE_URL`` (operator-declared public origin).
+          2. ``X-Forwarded-Proto``, only when the request's direct peer is a
+             configured trusted proxy -- see ``is_request_from_trusted_proxy``.
+             An untrusted caller cannot spoof this header to strip Secure.
+          3. The request's own literal scheme (direct TLS termination, or no
+             reverse proxy in front of litellm).
+        """
+        configured_base_url: Final = os.environ.get("PROXY_BASE_URL", "").strip()
+        if configured_base_url:
+            return urlparse(configured_base_url).scheme == "https"
+
+        if IPAddressUtils.is_request_from_trusted_proxy(request, general_settings=general_settings):
+            forwarded_proto: Final = request.headers.get("X-Forwarded-Proto")
+            if forwarded_proto:
+                return forwarded_proto.split(",")[0].strip().lower() == "https"
+
+        return request.url.scheme == "https"
+
     @staticmethod
     def extract_client_ip_from_xff_hops(
         xff_header: str,

litellm/proxy/guardrails/guardrail_hooks/alice/alice.py

@@ -8,6 +8,7 @@
 import json
 import os
 from collections.abc import Mapping
+from itertools import islice
 from typing import (
     TYPE_CHECKING,
     Any,  # noqa: TID251  # **kwargs forwards verbatim to CustomGuardrail.__init__; see ruff-strict.toml
@@ -341,19 +342,18 @@ def _json_safe(
     if depth >= _MAX_DEPTH or id(value) in seen:
         return None
 
-    nested: Final = seen | {id(value)}  # mutable-ok: one-shot set literal, unioned into a frozenset immediately
+    nested: Final = seen | frozenset((id(value),))
 
     if isinstance(value, dict):
-        out: dict[str, object] = {}  # mutable-ok: bounded accumulator local to this call, never escapes as-is
-        for key, item in list(value.items())[:_MAX_ITEMS]:  # mutable-ok: list() only to slice an unordered view
-            if isinstance(key, str) and key not in strip_keys:
-                out[key] = _json_safe(item, depth + 1, nested, strip_keys)
-        return out
+        return {
+            key: _json_safe(item, depth + 1, nested, strip_keys)
+            for key, item in islice(value.items(), _MAX_ITEMS)
+            if isinstance(key, str) and key not in strip_keys
+        }
 
     if isinstance(value, (list, tuple, set, frozenset)):
         return [  # mutable-ok: return value is a one-shot list, discarded by the caller after use
-            _json_safe(item, depth + 1, nested, strip_keys)
-            for item in list(value)[:_MAX_ITEMS]  # mutable-ok: list() only to slice an unordered view
+            _json_safe(item, depth + 1, nested, strip_keys) for item in islice(value, _MAX_ITEMS)
         ]
 
     dump: Final = getattr(value, "model_dump", None)

litellm/proxy/management_endpoints/sso/saml_sso.py

@@ -36,6 +36,7 @@ from pydantic import ValidationError
 
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.dual_cache import DualCache
+from litellm.proxy.auth.ip_address_utils import IPAddressUtils
 from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role
 from litellm.proxy.utils import get_custom_url
 
@@ -131,7 +132,7 @@ class SAMLAuthHandler:
 
     @staticmethod
     def _is_https(request: Request) -> bool:
-        return SAMLAuthHandler._base_url(request).startswith("https")
+        return IPAddressUtils.is_request_https(request)
 
     @staticmethod
     def _acs_url(request: Request) -> str:

litellm/proxy/management_endpoints/ui_sso.py

@@ -92,6 +92,7 @@ from litellm.proxy.auth.auth_utils import (
     has_user_setup_sso,
 )
 from litellm.proxy.auth.handle_jwt import JWTHandler
+from litellm.proxy.auth.ip_address_utils import IPAddressUtils
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.common_utils.admin_ui_utils import (
     admin_ui_disabled,
@@ -1118,7 +1119,7 @@ async def google_login(
             request=request,
         )
         if sso_redirect is not None:
-            _persist_return_to_cookie(sso_redirect, return_to)
+            _persist_return_to_cookie(sso_redirect, return_to, request)
         return sso_redirect
 
     from fastapi.responses import HTMLResponse
@@ -1138,7 +1139,7 @@ async def google_login(
     # helper the SSO branch uses, so /login can resume the connect flow instead of dead-ending at the
     # dashboard. One implementation → the two sign-in branches cannot diverge (and the login form always
     # renders, since the helper never raises on a bad return_to).
-    _persist_return_to_cookie(form_response, return_to)
+    _persist_return_to_cookie(form_response, return_to, request)
     return form_response
 
 
@@ -2741,6 +2742,7 @@ async def _sso_return_to_redirect(
     jwt_token: str,
     redis_usage_cache,
     user_api_key_cache,
+    request: Request,
 ) -> RedirectResponse | None:
     """Resolve the post-SSO redirect for a ``return_to``, or None to fall through to the dashboard.
 
@@ -2759,7 +2761,7 @@ async def _sso_return_to_redirect(
 
     if _is_same_origin_return_path(return_to):
         redirect_response = RedirectResponse(url=return_to, status_code=303)
-        redirect_response.set_cookie(key="token", value=jwt_token)
+        set_session_token_cookie(redirect_response, request, jwt_token)
         redirect_response.delete_cookie("litellm_cp_return_to")
         return redirect_response
 
@@ -2782,7 +2784,25 @@ async def _sso_return_to_redirect(
     return None
 
 
-def _persist_return_to_cookie(response: Response, return_to: str | None) -> None:
+def set_session_token_cookie(response: Response, request: Request, jwt_token: str) -> None:
+    """Set the ``token`` session cookie shared by every sign-in path.
+
+    Not HttpOnly: the dashboard reads this cookie via ``document.cookie`` to
+    populate its own Authorization headers (see
+    ``ui/litellm-dashboard/src/utils/cookieUtils.ts``), so marking it
+    HttpOnly would break login. Secure is still required whenever the public
+    origin is HTTPS, resolved the same trust-aware way as every other
+    litellm cookie."""
+    response.set_cookie(
+        key="token",
+        value=jwt_token,
+        secure=IPAddressUtils.is_request_https(request),
+        httponly=False,
+        samesite="lax",
+    )
+
+
+def _persist_return_to_cookie(response: Response, return_to: str | None, request: Request) -> None:
     """Best-effort: persist a SAFE ``return_to`` on ``response`` as the one-shot ``litellm_cp_return_to``
     cookie so ANY sign-in path — SSO / Okta / generic OR the username/password form — can resume there
     afterwards. THIS is the single source of truth, called by every sign-in branch so they cannot
@@ -2803,6 +2823,7 @@ def _persist_return_to_cookie(response: Response, return_to: str | None) -> None
             max_age=600,
             httponly=True,
             samesite="lax",
+            secure=IPAddressUtils.is_request_https(request),
         )
 
 
@@ -3079,8 +3100,11 @@ class SSOAuthenticationHandler:
                     # incoming request is HTTP (local dev).  Without
                     # ``Secure`` the cookie is sent over plain HTTP,
                     # letting a network observer read and replay the
-                    # state value and bypass this protection.
-                    secure_flag: Final = request is None or request.url.scheme == "https"
+                    # state value and bypass this protection. Trust-aware:
+                    # honors PROXY_BASE_URL / a trusted reverse proxy's
+                    # X-Forwarded-Proto instead of only the literal scheme
+                    # litellm sees on the wire.
+                    secure_flag: Final = request is None or IPAddressUtils.is_request_https(request)
                     redirect_response.set_cookie(
                         key="litellm_oauth_state",
                         value=state_value,
@@ -3628,6 +3652,7 @@ class SSOAuthenticationHandler:
             jwt_token=jwt_token,
             redis_usage_cache=redis_usage_cache,
             user_api_key_cache=user_api_key_cache,
+            request=request,
         )
         if return_to_redirect is not None:
             return return_to_redirect
@@ -3636,7 +3661,7 @@ class SSOAuthenticationHandler:
             litellm_dashboard_ui += "?login=success"
         verbose_proxy_logger.info("Redirecting to %s", litellm_dashboard_ui)
         redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
-        redirect_response.set_cookie(key="token", value=jwt_token)
+        set_session_token_cookie(redirect_response, request, jwt_token)
         return redirect_response
 
     @staticmethod

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -1731,7 +1731,7 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict:
 
 
 def get_vertex_pass_through_handler(
-    call_type: Literal["discovery", "aiplatform"],  # noqa: UP037
+    call_type: Literal["discovery", "aiplatform"],  # noqa: UP037  # ruff reports quoted Literal values here
 ) -> BaseVertexAIPassThroughHandler:
     if call_type == "discovery":
         return VertexAIDiscoveryPassThroughHandler()
@@ -2961,7 +2961,6 @@ async def handle_gigachat_passthrough_router_model(
     """
     from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
 
-    # Detect streaming based on request body
     is_streaming: Final = request_body.get("stream", False)  # pyright: ignore[reportUnknownVariableType]  # request_body is dict[Unknown, Unknown]
 
     data: dict[str, Any] = await _read_request_body(
@@ -2997,7 +2996,6 @@ async def handle_gigachat_passthrough_router_model(
     data["json"] = request_body
     data["custom_llm_provider"] = "gigachat"
 
-    # Remove sensitive keys from data
     keys: Final = [  # mutable-ok: list of keys to remove from data
         "gigachat_auth_url",
         "gigachat_access_token",

litellm/proxy/proxy_server.py

@@ -15329,7 +15329,10 @@ async def login(request: Request):
     # authorize round-trip), mirroring the SSO callback; otherwise land on the dashboard. Gated by
     # _is_same_origin_return_path (strictly relative path) so it can never be an open redirect, and the
     # one-shot cookie is cleared after use.
-    from litellm.proxy.management_endpoints.ui_sso import _sso_return_to_redirect
+    from litellm.proxy.management_endpoints.ui_sso import (
+        _sso_return_to_redirect,
+        set_session_token_cookie,
+    )
 
     # Resume through the SAME resumer the SSO callback uses, rather than a second, narrower arm.
     # _persist_return_to_cookie stores both shapes it accepts (a relative same-origin path AND a
@@ -15346,6 +15349,7 @@ async def login(request: Request):
                 jwt_token=jwt_token,
                 redis_usage_cache=redis_usage_cache,
                 user_api_key_cache=user_api_key_cache,
+                request=request,
             )
         except Exception:  # noqa: BLE001  # resuming must NEVER block a completed sign-in
             # The symmetric half of _persist_return_to_cookie's "never raises" contract. The resumer
@@ -15360,7 +15364,7 @@ async def login(request: Request):
 
     # Create redirect response with cookie
     redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
-    redirect_response.set_cookie(key="token", value=jwt_token)
+    set_session_token_cookie(redirect_response, request, jwt_token)
     if cp_return_to:
         redirect_response.delete_cookie(key="litellm_cp_return_to")
     return redirect_response
@@ -15370,6 +15374,7 @@ async def login(request: Request):
 async def login_v2(request: Request):
     global premium_user, general_settings, master_key
     from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt
+    from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie
     from litellm.proxy.utils import get_custom_url
 
     try:
@@ -15404,7 +15409,7 @@ async def login_v2(request: Request):
             content={"redirect_url": litellm_dashboard_ui, "token": jwt_token},
             status_code=status.HTTP_200_OK,
         )
-        json_response.set_cookie(key="token", value=jwt_token)
+        set_session_token_cookie(json_response, request, jwt_token)
         return json_response
     except Exception as e:
         verbose_proxy_logger.exception("litellm.proxy.proxy_server.login_v2(): Exception occurred - %s", e)
@@ -15504,6 +15509,8 @@ async def login_v3(request: Request):
 
 @router.post("/v3/login/exchange", include_in_schema=False)  # exchange single-use opaque code for JWT
 async def login_v3_exchange(request: Request):
+    from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie
+
     try:
         if not general_settings.get("control_plane_url"):
             raise ProxyException(
@@ -15550,7 +15557,7 @@ async def login_v3_exchange(request: Request):
             },
             status_code=status.HTTP_200_OK,
         )
-        json_response.set_cookie(key="token", value=cached_data["token"])
+        set_session_token_cookie(json_response, request, cached_data["token"])
         return json_response
     except ProxyException:
         raise

litellm/proxy/rag_endpoints/endpoints.py

@@ -729,7 +729,7 @@ async def rag_query(
         # conflict so callers cannot override the store's provider or credentials.
         managed_store: Final = resolved_stores.get(retrieval_config["vector_store_id"])
         store_data: Final = (
-            await build_request_data_from_managed_vector_store(managed_store)
+            build_request_data_from_managed_vector_store(managed_store)
             if managed_store is not None
             else MappingProxyType({})
         )