Merge origin/litellm_internal_staging into litellm_python_version_ci

.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

@@ -1,9 +1,9 @@
 {
   "reportAny": {
-    "limit": 14076
+    "limit": 14074
   },
   "reportArgumentType": {
-    "limit": 2216
+    "limit": 2215
   },
   "reportAssignmentType": {
     "limit": 319
@@ -24,7 +24,7 @@
     "limit": 19
   },
   "reportExplicitAny": {
-    "limit": 4128
+    "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": 38348
+    "limit": 38330
   },
   "reportUnknownParameterType": {
-    "limit": 19624
+    "limit": 19623
   },
   "reportUnknownVariableType": {
-    "limit": 29889
+    "limit": 29860
   },
   "reportUnnecessaryCast": {
     "limit": 111
   },
   "reportUnnecessaryComparison": {
-    "limit": 692
+    "limit": 687
   },
   "reportUnnecessaryContains": {
-    "limit": 5
+    "limit": 4
   },
   "reportUnnecessaryIsInstance": {
-    "limit": 826
+    "limit": 823
   },
   "reportUntypedBaseClass": {
     "limit": 0
@@ -138,7 +138,7 @@
     "limit": 138
   },
   "reportUnusedImport": {
-    "limit": 543
+    "limit": 542
   },
   "reportUnusedVariable": {
     "limit": 137

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/audio_transcription/common_utils.rs

@@ -1,46 +0,0 @@
-use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
-use litellm_core::error::Error;
-use litellm_core::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
-use serde_json::{Map, Value};
-use std::collections::BTreeMap;
-
-pub(super) fn audio_transcription_provider_config(
-    provider: &str,
-) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
-    match provider {
-        "bedrock" => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG),
-        _ => None,
-    }
-}
-
-pub(super) fn string_headers(
-    headers: Option<Map<String, Value>>,
-) -> Result<BTreeMap<String, String>, Error> {
-    headers
-        .unwrap_or_default()
-        .into_iter()
-        .map(|(key, value)| {
-            value
-                .as_str()
-                .map(|value| (key.clone(), value.to_string()))
-                .ok_or_else(|| {
-                    Error::InvalidRequest(format!(
-                        "audio transcription extra_headers.{key} must be a string"
-                    ))
-                })
-        })
-        .collect()
-}
-
-pub(super) fn has_header(headers: &BTreeMap<String, String>, name: &str) -> bool {
-    headers.keys().any(|key| key.eq_ignore_ascii_case(name))
-}
-
-pub(super) fn truncate_error_body(body: &str) -> String {
-    let truncated: String = body.chars().take(256).collect();
-    if truncated.chars().count() == body.chars().count() {
-        truncated
-    } else {
-        format!("{truncated}... (truncated)")
-    }
-}

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

@@ -1,84 +0,0 @@
-use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
-use litellm_core::error::Error;
-use litellm_core::providers::bedrock::audio_transcription::aws_auth_config;
-use litellm_core::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
-use serde_json::Value;
-use std::time::SystemTime;
-
-use super::common_utils::truncate_error_body;
-use super::types::ProviderAudioTranscriptionRequest;
-use crate::client::http_client;
-
-pub(crate) async fn execute_audio_transcription_provider_call(
-    request: ProviderAudioTranscriptionRequest,
-) -> Result<Value, Error> {
-    let body = serde_json::to_vec(&request.body)
-        .map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?;
-    let mut request_builder = http_client().post(&request.url).body(body.clone());
-    for (key, value) in &request.upstream_headers {
-        request_builder = request_builder.header(key, value);
-    }
-    if let Some(duration) = request.timeout {
-        request_builder = request_builder.timeout(duration);
-    }
-    let response = request_builder
-        .send()
-        .await
-        .map_err(|error| Error::Network(error.to_string()))?;
-    let status = response.status();
-    let text = response
-        .text()
-        .await
-        .map_err(|error| Error::Network(error.to_string()))?;
-    if !status.is_success() {
-        return Err(Error::Http {
-            status: status.as_u16(),
-            body: truncate_error_body(&text),
-        });
-    }
-    let response_json: Value = serde_json::from_str(&text)
-        .map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?;
-    Ok(request
-        .config
-        .transform_transcription_response(&request.model, response_json)?
-        .into_json())
-}
-
-pub(crate) async fn sign_request(
-    request: &ProviderAudioTranscriptionRequest,
-    optional_params: &serde_json::Map<String, Value>,
-) -> Result<ProviderAudioTranscriptionRequest, Error> {
-    let env_lookup = environment_lookup;
-    let auth = request
-        .config
-        .auth_strategy(&request.model, optional_params, &env_lookup)?;
-    let body = serde_json::to_vec(&request.body)
-        .map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?;
-    let mut headers = super::common_utils::string_headers(None)?;
-    headers.insert("Content-Type".to_string(), "application/json".to_string());
-    headers.extend(request.upstream_headers.iter().cloned());
-    match auth {
-        AudioTranscriptionAuth::Bearer => {}
-        AudioTranscriptionAuth::AwsSigV4 { region, .. } => {
-            let credentials =
-                resolve_credentials(aws_auth_config(optional_params, &env_lookup), &env_lookup)
-                    .await?;
-            headers.extend(sign_bedrock_post(
-                &request.url,
-                &body,
-                &headers,
-                &region,
-                &credentials,
-                SystemTime::now(),
-            )?);
-        }
-    }
-    Ok(ProviderAudioTranscriptionRequest {
-        upstream_headers: headers.into_iter().collect(),
-        ..request.clone()
-    })
-}
-
-pub(super) fn environment_lookup(key: &str) -> Option<String> {
-    std::env::var(key).ok()
-}

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

@@ -1,13 +1,14 @@
-use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
+use litellm_core::audio_transcription::{
+    AudioTranscriptionRequest as CoreAudioTranscriptionRequest, ProviderAudioTranscriptionRequest,
+    prepare_audio_transcription_provider_call,
+};
 use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
 use litellm_core::error::Error;
 use serde_json::{Map, Value, json};
 use std::future::Future;
 use std::pin::Pin;
 
-use super::common_utils::{audio_transcription_provider_config, has_header, string_headers};
-use super::handler::sign_request;
-use super::types::{PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
+use super::types::PreparedAudioTranscriptionRequest;
 use crate::integrations::custom_guardrail::{
     CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
 };
@@ -88,46 +89,29 @@ impl AudioTranscriptionLifecycleHooks {
         &self,
         request: PreparedAudioTranscriptionRequest,
     ) -> Result<ProviderAudioTranscriptionRequest, Error> {
-        let config = audio_transcription_provider_config(&request.custom_llm_provider)
-            .ok_or_else(|| Error::InvalidProvider(request.custom_llm_provider.clone()))?;
-        let env_lookup = super::handler::environment_lookup;
-        let headers = string_headers(request.extra_headers)?;
-        let url = config.complete_url(
-            request.api_base.as_deref(),
-            &request.model,
-            &request.optional_params,
-            &env_lookup,
-        )?;
-        let filtered_params = config.map_transcription_params(&request.optional_params);
-        let body = config.transform_transcription_request(
-            &request.model,
-            request.audio,
-            filtered_params,
-        )?;
-        let auth = config.auth_strategy(&request.model, &request.optional_params, &env_lookup)?;
-        let mut upstream_headers = headers.into_iter().collect::<Vec<_>>();
-        if matches!(auth, AudioTranscriptionAuth::Bearer)
-            && !has_header(
-                &upstream_headers
-                    .iter()
-                    .cloned()
-                    .collect::<std::collections::BTreeMap<_, _>>(),
-                "authorization",
-            )
-            && let Some(api_key) = request.api_key.as_deref()
-        {
-            upstream_headers.push(("Authorization".to_string(), format!("Bearer {api_key}")));
-        }
-        let provider_request = ProviderAudioTranscriptionRequest {
-            model: request.model,
-            config,
-            url,
-            body: body.body,
-            upstream_headers,
-            timeout: request.timeout,
-        };
-        let provider_request = self.run_during_call_guardrails(provider_request).await?;
-        sign_request(&provider_request, &request.optional_params).await
+        let PreparedAudioTranscriptionRequest {
+            model,
+            custom_llm_provider,
+            audio,
+            api_key,
+            api_base,
+            extra_headers,
+            optional_params,
+            timeout,
+            ..
+        } = request;
+        let provider_request =
+            prepare_audio_transcription_provider_call(CoreAudioTranscriptionRequest {
+                model: &model,
+                audio,
+                api_key: api_key.as_deref(),
+                api_base: api_base.as_deref(),
+                custom_llm_provider: Some(&custom_llm_provider),
+                extra_headers,
+                optional_params,
+                timeout,
+            })?;
+        self.run_during_call_guardrails(provider_request).await
     }
 
     async fn run_during_call_guardrails(
@@ -142,10 +126,10 @@ impl AudioTranscriptionLifecycleHooks {
             .run_during_call(
                 &guardrail_context(&self.request_metadata),
                 GuardrailRequest::new(json!({
-                    "model": request.model,
-                    "custom_llm_provider": "bedrock",
-                    "url": request.url,
-                    "body": request.body,
+                    "model": request.model(),
+                    "custom_llm_provider": request.custom_llm_provider(),
+                    "url": request.url(),
+                    "body": request.body(),
                 })),
             )
             .await
@@ -158,7 +142,7 @@ impl AudioTranscriptionLifecycleHooks {
         let body = data.remove("body").ok_or_else(|| {
             Error::InvalidRequest("audio transcription guardrail removed body".to_string())
         })?;
-        Ok(ProviderAudioTranscriptionRequest { body, ..request })
+        Ok(request.with_body(body))
     }
 
     fn logging_payload(

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

@@ -1,16 +1,14 @@
 use litellm_core::Error;
+use litellm_core::audio_transcription::execute_audio_transcription_provider_call;
 use litellm_core::call_lifecycle::CallLifecycle;
 use serde_json::Value;
 
-mod common_utils;
-mod handler;
 mod hooks;
 mod prepare;
 mod types;
 
 pub use types::AudioTranscriptionRequest;
 
-use handler::execute_audio_transcription_provider_call;
 use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call};
 
 pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {

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

@@ -1,7 +1,6 @@
 use std::sync::Arc;
 use std::time::Duration;
 
-use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
 use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
 use serde_json::{Map, Value};
 
@@ -46,13 +45,3 @@ impl CallLifecycleRequest for PreparedAudioTranscriptionRequest {
         )
     }
 }
-
-#[derive(Clone)]
-pub(crate) struct ProviderAudioTranscriptionRequest {
-    pub(crate) model: String,
-    pub(crate) config: &'static dyn AudioTranscriptionProviderConfig,
-    pub(crate) url: String,
-    pub(crate) body: Value,
-    pub(crate) upstream_headers: Vec<(String, String)>,
-    pub(crate) timeout: Option<Duration>,
-}

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/client.rs

@@ -0,0 +1,14 @@
+use std::sync::OnceLock;
+use std::time::Duration;
+
+use crate::constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS;
+
+pub(super) fn http_client() -> &'static reqwest::Client {
+    static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
+    CLIENT.get_or_init(|| {
+        reqwest::Client::builder()
+            .timeout(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS))
+            .build()
+            .unwrap_or_else(|_| reqwest::Client::new())
+    })
+}

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

@@ -0,0 +1,91 @@
+use serde_json::Value;
+
+use crate::error::Error;
+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> {
+    let body = serde_json::to_vec(&request.body)
+        .map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?;
+    let headers = signed_headers(&request, &body).await?;
+    let mut request_builder = http_client().post(&request.url).body(body);
+    for (key, value) in headers {
+        request_builder = request_builder.header(key, value);
+    }
+    if let Some(duration) = request.timeout {
+        request_builder = request_builder.timeout(duration);
+    }
+    let response = http_request(request_builder)
+        .await
+        .map_err(|error| Error::Network(error.to_string()))?;
+    let status = response.status();
+    let text = response
+        .text()
+        .await
+        .map_err(|error| Error::Network(error.to_string()))?;
+    if !status.is_success() {
+        return Err(Error::Http {
+            status: status.as_u16(),
+            body: truncate_error_body(&text),
+        });
+    }
+    let response_json = serde_json::from_str(&text)
+        .map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?;
+    Ok(request
+        .config
+        .transform_transcription_response(&request.model, response_json)?
+        .into_json())
+}
+
+#[cfg(feature = "bedrock-auth")]
+async fn signed_headers(
+    request: &ProviderAudioTranscriptionRequest,
+    body: &[u8],
+) -> Result<Vec<(String, String)>, Error> {
+    use std::collections::BTreeMap;
+    use std::time::SystemTime;
+
+    use crate::audio_transcription::transformation::AudioTranscriptionAuth;
+    use crate::providers::bedrock::audio_transcription::aws_auth_config;
+    use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
+
+    let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else {
+        return Ok(request.upstream_headers.clone());
+    };
+    let env_lookup = |key: &str| std::env::var(key).ok();
+    let credentials = resolve_credentials(
+        aws_auth_config(&request.optional_params, &env_lookup),
+        &env_lookup,
+    )
+    .await?;
+    let unsigned: BTreeMap<String, String> = request.upstream_headers.iter().cloned().collect();
+    let signature = sign_bedrock_post(
+        &request.url,
+        body,
+        &unsigned,
+        region,
+        &credentials,
+        SystemTime::now(),
+    )?;
+    Ok(unsigned.into_iter().chain(signature).collect())
+}
+
+#[cfg(not(feature = "bedrock-auth"))]
+async fn signed_headers(
+    request: &ProviderAudioTranscriptionRequest,
+    _body: &[u8],
+) -> Result<Vec<(String, String)>, Error> {
+    use crate::audio_transcription::transformation::AudioTranscriptionAuth;
+
+    match request.auth {
+        AudioTranscriptionAuth::AwsSigV4 { .. } => Err(Error::Unsupported(
+            "AWS SigV4 requires the bedrock-auth feature",
+        )),
+        AudioTranscriptionAuth::Bearer => Ok(request.upstream_headers.clone()),
+    }
+}

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

@@ -1,2 +1,21 @@
+use crate::Error;
+mod client;
+mod handler;
+mod prepare;
 pub mod transformation;
 pub mod types;
+
+use serde_json::Value;
+
+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
+}
+
+#[cfg(test)]
+mod tests;

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

@@ -0,0 +1,74 @@
+use crate::error::Error;
+use crate::http_utils::{has_header, string_headers};
+#[cfg(feature = "bedrock-auth")]
+use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
+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" {
+        return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG);
+    }
+    let _ = provider;
+    None
+}
+
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+pub fn prepare_audio_transcription_provider_call(
+    request: AudioTranscriptionRequest<'_>,
+) -> Result<ProviderAudioTranscriptionRequest, Error> {
+    let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
+        .or_else(|| {
+            request
+                .custom_llm_provider
+                .map(|provider| CustomLlmProvider {
+                    model: request.model,
+                    custom_llm_provider: provider,
+                })
+        })
+        .ok_or_else(|| {
+            Error::InvalidProvider(
+                "unable to resolve custom_llm_provider for audio transcription request".to_string(),
+            )
+        })?;
+    let model = provider_info.model.to_string();
+    let config = provider_config(provider_info.custom_llm_provider)
+        .ok_or_else(|| Error::InvalidProvider(provider_info.custom_llm_provider.to_string()))?;
+    let env_lookup = |key: &str| std::env::var(key).ok();
+    let mut headers = string_headers("audio transcription", request.extra_headers)?;
+    let auth = config.auth_strategy(&model, &request.optional_params, &env_lookup)?;
+    if matches!(auth, AudioTranscriptionAuth::Bearer)
+        && !has_header(&headers, "authorization")
+        && let Some(api_key) = request.api_key
+    {
+        headers.push(("Authorization".to_string(), format!("Bearer {api_key}")));
+    }
+    if !has_header(&headers, "content-type") {
+        headers.push(("Content-Type".to_string(), "application/json".to_string()));
+    }
+    let url = config.complete_url(
+        request.api_base,
+        &model,
+        &request.optional_params,
+        &env_lookup,
+    )?;
+    let filtered_params = config.map_transcription_params(&request.optional_params);
+    let transformed =
+        config.transform_transcription_request(&model, request.audio, filtered_params)?;
+    Ok(ProviderAudioTranscriptionRequest {
+        model,
+        custom_llm_provider: provider_info.custom_llm_provider.to_string(),
+        config,
+        url,
+        body: transformed.body,
+        upstream_headers: headers,
+        auth,
+        #[cfg(feature = "bedrock-auth")]
+        optional_params: request.optional_params,
+        timeout: request.timeout,
+    })
+}

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

@@ -0,0 +1,50 @@
+use std::io::{Read, Write};
+use std::net::TcpListener;
+use std::thread;
+
+use serde_json::{Map, json};
+
+use super::audio_transcription;
+use super::types::AudioTranscriptionRequest;
+
+#[tokio::test]
+async fn bedrock_request_is_signed_and_contains_audio() {
+    let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
+    let address = listener.local_addr().expect("address");
+    let server = thread::spawn(move || {
+        let (mut stream, _) = listener.accept().expect("connection");
+        let mut request = Vec::new();
+        let mut buffer = [0_u8; 16_384];
+        let count = stream.read(&mut buffer).expect("request");
+        request.extend_from_slice(&buffer[..count]);
+        let request = String::from_utf8_lossy(&request);
+        assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse"));
+        assert!(request.contains("authorization: AWS4-HMAC-SHA256"));
+        assert!(request.contains("x-amz-date:"));
+        assert!(request.contains("\"bytes\":\"AQI=\""));
+        assert!(request.contains("Transcribe the audio. Respond with only the transcript."));
+        let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}";
+        stream.write_all(response).expect("response");
+    });
+
+    let optional_params = Map::from_iter([
+        ("aws_access_key_id".to_string(), json!("access-key")),
+        ("aws_secret_access_key".to_string(), json!("secret-key")),
+        ("aws_region_name".to_string(), json!("us-east-1")),
+    ]);
+    let api_base = format!("http://{address}");
+    let response = audio_transcription(AudioTranscriptionRequest {
+        model: "mistral.voxtral-mini-3b-2507",
+        audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}),
+        api_key: None,
+        api_base: Some(&api_base),
+        custom_llm_provider: Some("bedrock"),
+        extra_headers: None,
+        optional_params,
+        timeout: None,
+    })
+    .await
+    .expect("transcription");
+    assert_eq!(response, json!({"text": "hello"}));
+    server.join().expect("server");
+}

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/audio_transcription/types.rs

@@ -1,5 +1,56 @@
+use std::time::Duration;
+
 use serde::{Deserialize, Serialize};
-use serde_json::Value;
+use serde_json::{Map, Value};
+
+use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
+
+pub struct AudioTranscriptionRequest<'a> {
+    pub model: &'a str,
+    pub audio: Value,
+    pub api_key: Option<&'a str>,
+    pub api_base: Option<&'a str>,
+    pub custom_llm_provider: Option<&'a str>,
+    pub extra_headers: Option<Map<String, Value>>,
+    pub optional_params: Map<String, Value>,
+    pub timeout: Option<Duration>,
+}
+
+#[derive(Clone)]
+pub struct ProviderAudioTranscriptionRequest {
+    pub(super) model: String,
+    pub(super) custom_llm_provider: String,
+    pub(super) config: &'static dyn AudioTranscriptionProviderConfig,
+    pub(super) url: String,
+    pub(super) body: Value,
+    pub(super) upstream_headers: Vec<(String, String)>,
+    pub(super) auth: AudioTranscriptionAuth,
+    #[cfg(feature = "bedrock-auth")]
+    pub(super) optional_params: Map<String, Value>,
+    pub(super) timeout: Option<Duration>,
+}
+
+impl ProviderAudioTranscriptionRequest {
+    pub fn model(&self) -> &str {
+        &self.model
+    }
+
+    pub fn custom_llm_provider(&self) -> &str {
+        &self.custom_llm_provider
+    }
+
+    pub fn url(&self) -> &str {
+        &self.url
+    }
+
+    pub fn body(&self) -> &Value {
+        &self.body
+    }
+
+    pub fn with_body(self, body: Value) -> Self {
+        Self { body, ..self }
+    }
+}
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 pub struct AudioTranscriptionRequestData {

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/constants.rs

@@ -30,6 +30,8 @@ pub(crate) const CHAT_COMPLETIONS_TIMEOUT_SECS: u64 = 600;
 /// Connect timeout for chat completions provider calls, in seconds.
 pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10;
 
+pub(crate) const AUDIO_TRANSCRIPTION_TIMEOUT_SECS: u64 = 600;
+
 /// `object` field every non-streaming chat completion response carries.
 pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";
 

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::audio_transcription::{
-    AudioTranscriptionRequest, audio_transcription as run_audio_transcription,
-};
-use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
 use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
-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,453 +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,
-                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)),
+    #[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,
-            callbacks: Vec::new(),
-            guardrails: Vec::new(),
-            request_metadata: Default::default(),
-            litellm_call_id: None,
-        })
-        .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

@@ -29,7 +29,7 @@ def _dev_env_hot_reload_enabled() -> bool:
 if os.getenv("LITELLM_MODE", "DEV") == "DEV":
     _dotenv.load_dotenv(override=_dev_env_hot_reload_enabled())
 
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from typing import (
     Any,
     Callable,
@@ -424,6 +424,10 @@ anthropic_beta_headers_url: str = os.getenv(
     "LITELLM_ANTHROPIC_BETA_HEADERS_URL",
     "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/anthropic_beta_headers_config.json",
 )
+autorouter_presets_url: str = os.getenv(
+    "LITELLM_AUTOROUTER_PRESETS_URL",
+    "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/proxy/public_endpoints/autorouter_presets.json",
+)
 suppress_debug_info: bool = False
 dynamodb_table_name: Optional[str] = None
 s3_callback_params: Optional[Dict] = None
@@ -490,6 +494,7 @@ public_mcp_hub_strict_whitelist: bool = True
 public_model_groups: Optional[List[str]] = None
 public_agent_groups: Optional[List[str]] = None
 agent_search_embedding_model: Optional[str] = None
+mcp_tool_search: Optional[Mapping[str, object]] = None
 # Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]])
 # New format: { "displayName": { "url": "...", "index": 0 } }
 # Old format: { "displayName": "url" } (for backward compatibility)
@@ -1416,7 +1421,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/_logging.py

@@ -6,6 +6,7 @@ import sys
 from datetime import datetime
 from logging import Formatter
 from typing import Any, Final, TextIO
+from urllib.parse import unquote
 
 import litellm
 from litellm.constants import (
@@ -16,7 +17,11 @@ from litellm.constants import (
 from litellm.litellm_core_utils.env_utils import get_env_int
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
-from litellm.litellm_core_utils.secret_redaction import redact_string, redact_structured_value
+from litellm.litellm_core_utils.secret_redaction import (
+    redact_internal_details,
+    redact_string,
+    redact_structured_value,
+)
 
 set_verbose = False
 
@@ -88,6 +93,14 @@ def redact_secrets(value: str) -> str:
     return _redact_string(value)
 
 
+def redact_internal_details_from_client_message(value: str) -> str:
+    """Public API: redact_secrets() plus filesystem paths, internal hostnames, and an
+    embedded traceback, for a string about to leave the process in an HTTP response."""
+    if not _ENABLE_SECRET_REDACTION:
+        return value
+    return redact_internal_details(value)
+
+
 def _substituted_color_message(record: logging.LogRecord) -> str | None:
     """Render a record's ``color_message`` against its args, or None if absent.
 
@@ -146,6 +159,72 @@ class SecretRedactionFilter(logging.Filter):
 _secret_filter: Final = SecretRedactionFilter()
 
 
+_MAX_SCRUBBED_ACCESS_ARG: Final = 512
+
+_REDACTION_PLACEHOLDER: Final = "REDACTED"
+
+
+def _hides_a_credential(value: str) -> bool:
+    """Whether *value* only looks clean until it is percent-decoded."""
+    decoded: Final = unquote(value)
+    return _redact_string(decoded) != decoded
+
+
+def _drop_encoded_credential(scrubbed: str) -> str:
+    """Drop the part of a request target that only decoding shows to be a secret.
+
+    The request parser decodes query names and values, so `?k%65y=sk%2D...` is a
+    working credential that the patterns, which match literal text, do not see.
+    The decoded text is never logged back: it can carry a newline, and forging
+    log lines is not a trade worth making for a readable request target.
+    """
+    path, separator, _query = scrubbed.partition("?")
+    if _hides_a_credential(path):
+        return _REDACTION_PLACEHOLDER
+    if separator and _hides_a_credential(scrubbed):
+        return f"{path}?{_REDACTION_PLACEHOLDER}"
+    return scrubbed
+
+
+def _scrub_access_arg(value: str) -> str:
+    """Redact one access-log positional arg, bounding the scanned length.
+
+    The request target is the only input to the secret regex an unauthenticated
+    caller controls end to end, so it is cut back to a whole query parameter
+    before it is scanned; a half-parameter would be too short to match its
+    pattern and would then be logged raw.
+    """
+    if len(value) <= _MAX_SCRUBBED_ACCESS_ARG:
+        return _drop_encoded_credential(_redact_string(value))
+    head: Final = value[:_MAX_SCRUBBED_ACCESS_ARG]
+    kept: Final = head[: max(head.rfind("?"), head.rfind("&"))] if "?" in head else head
+    scrubbed: Final = _drop_encoded_credential(_redact_string(kept))
+    return f"{scrubbed}... ({len(value) - len(kept)} more chars truncated) ..."
+
+
+class AccessLogRedactionFilter(logging.Filter):
+    """Scrubs known secret/credential patterns from HTTP access-log records.
+
+    uvicorn's AccessFormatter unpacks ``record.args`` as a five-element tuple at
+    emit time, so SecretRedactionFilter cannot be reused here: it collapses the
+    record into ``record.msg`` and clears the args, and the formatter then raises.
+    """
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        if not _ENABLE_SECRET_REDACTION:
+            return True
+        if isinstance(record.args, tuple) and record.args:
+            record.args = tuple(  # rebind-ok: a Filter scrubs records in place
+                _scrub_access_arg(arg) if isinstance(arg, str) else arg for arg in record.args
+            )
+            return True
+        # No positional args means everything is in msg, where collapsing is correct.
+        return _secret_filter.filter(record)
+
+
+_access_log_filter: Final = AccessLogRedactionFilter()
+
+
 def _get_max_string_length_stdout_log() -> int:
     """Read the limit per record so a value loaded later via proxy config
     environment_variables is honored."""
@@ -553,6 +632,14 @@ _REDACTED_THIRD_PARTY_LOGGERS: Final[tuple[str, ...]] = (
     "uvicorn.error",
 )
 
+# Access loggers, which emit the full request target, so a credential passed as a
+# query parameter (e.g. `/key/info?key=`) lands on stdout verbatim. uvicorn.access
+# covers uvicorn.run, --run_gunicorn (its worker_class is UvicornWorker, so the
+# access line is still uvicorn's) and an embedding host app. --run_hypercorn and
+# --run_granian log through their own loggers in their own record shapes, and
+# both ship with access logging off.
+_REDACTED_ACCESS_LOGGERS: Final[tuple[str, ...]] = ("uvicorn.access",)
+
 
 def _redact_third_party_loggers() -> None:
     """Extend secret redaction to records litellm does not emit directly.
@@ -575,6 +662,8 @@ def _redact_third_party_loggers() -> None:
     """
     for name in _REDACTED_THIRD_PARTY_LOGGERS:
         logging.getLogger(name).addFilter(_secret_filter)
+    for name in _REDACTED_ACCESS_LOGGERS:
+        logging.getLogger(name).addFilter(_access_log_filter)
 
 
 # Call the suppression function

litellm/constants.py

@@ -1742,6 +1742,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
     "anthropic_prompt_caching_ttl",
     "max_ui_session_budget",
     "budget_rollover",
+    "mcp_tool_search",
 ]
 SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"]
 DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60))

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/anthropic_cache_control_hook.py

@@ -755,6 +755,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
     def record_gateway_injection(
         request_kwargs: Mapping[str, object],
         added: int,
+        injected_for_every_deployment: bool = False,
     ) -> None:
         """Name the deployment whose payload the gateway, not the client, put breakpoints on.
 
@@ -771,7 +772,16 @@ class AnthropicCacheControlHook(CustomPromptManagement):
 
         A pass that runs before a deployment is chosen, which is what the proxy does for
         prompt templates, injects into the payload every leg goes on to send, so it marks
-        the request for all of them rather than for one.
+        the request for all of them rather than for one. Such a pass says so with
+        ``injected_for_every_deployment`` instead of relying on the shape of
+        ``request_kwargs``: the router's prompt-management factory stamps a provisional
+        deployment's ``model_info`` into kwargs before the prompt pass runs, and billing
+        the request through any other deployment would silently drop the credit. An
+        every-deployment mark, once written, also never narrows: a later per-leg stamp
+        (the Bedrock converse tool_config one included) describes one leg of a payload
+        every leg sends, so narrowing to it would uncredit whichever leg gets billed
+        after a failover. Both losses are fail-closed under-crediting, which is why the
+        guard only protects the sentinel and per-leg marks still overwrite each other.
 
         Only what this pass actually placed counts. A ``tool_config`` point is placed by
         the Bedrock converse transform, and only when the request carries tools, so the
@@ -801,13 +811,19 @@ class AnthropicCacheControlHook(CustomPromptManagement):
             ),
             None,
         )
-        if bucket is not None:
-            model_info: Final = request_kwargs.get("model_info")
-            bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = (
-                model_info.get("id", GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT)
-                if isinstance(model_info, dict)
-                else GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT
-            )
+        if bucket is None:
+            return
+        if bucket.get(GATEWAY_INJECTED_CACHE_METADATA_KEY) == GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT:
+            return
+        if injected_for_every_deployment:
+            bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT
+            return
+        model_info: Final = request_kwargs.get("model_info")
+        bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = (
+            model_info.get("id", GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT)
+            if isinstance(model_info, dict)
+            else GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT
+        )
 
     @staticmethod
     def maybe_inject_cache_control(

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/litellm_core_utils/litellm_logging.py

@@ -901,6 +901,7 @@ class Logging(LiteLLMLoggingBaseClass):
         prompt_label: str | None = None,
         prompt_version: int | None = None,
         request_kwargs: dict[str, object] | None = None,  # mutable-ok: marker stamped into live request kwargs
+        injected_for_every_deployment: bool = False,
     ) -> tuple[str, list[AllMessageValues], dict]:
         from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
 
@@ -933,6 +934,7 @@ class Logging(LiteLLMLoggingBaseClass):
                 AnthropicCacheControlHook.record_gateway_injection(
                     request_kwargs,
                     AnthropicCacheControlHook.count_request_cache_breakpoints(messages) - breakpoints_before,
+                    injected_for_every_deployment=injected_for_every_deployment,
                 )
         self.messages = messages
         return model, messages, non_default_params
@@ -950,6 +952,7 @@ class Logging(LiteLLMLoggingBaseClass):
         prompt_label: str | None = None,
         prompt_version: int | None = None,
         request_kwargs: dict[str, object] | None = None,  # mutable-ok: marker stamped into live request kwargs
+        injected_for_every_deployment: bool = False,
     ) -> tuple[str, list[AllMessageValues], dict]:
         from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
 
@@ -985,6 +988,7 @@ class Logging(LiteLLMLoggingBaseClass):
                 AnthropicCacheControlHook.record_gateway_injection(
                     request_kwargs,
                     AnthropicCacheControlHook.count_request_cache_breakpoints(messages) - breakpoints_before,
+                    injected_for_every_deployment=injected_for_every_deployment,
                 )
         self.messages = messages
         return model, messages, non_default_params

litellm/litellm_core_utils/secret_redaction.py

@@ -30,6 +30,11 @@ def _build_secret_patterns() -> "re.Pattern[str]":
         r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}",
         # OpenAI / Anthropic sk- prefixed keys
         rf"sk-[A-Za-z0-9\-_]{{{MINIMUM_CUSTOM_KEY_LENGTH - len('sk-')},}}",
+        # Credentials passed as URL query params. Terminated by "&" like the key=
+        # and sig= patterns below, so the rest of the request line survives in an
+        # access log. Must precede the generic patterns to win at the same position.
+        r"(?<=[?&])(?:api[_-]?key|\w*(?:token|password|passwd|client_secret|secret_key|_secret))"
+        r"=[^\s&'\"]+",
         # Generic api_key / api-key / apikey (handles 'key': 'value' dict repr)
         r"(?:api[_-]?key)['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]{8,}",
         # x-api-key / api-key header values (handles 'key': 'value' dict repr)
@@ -45,8 +50,10 @@ def _build_secret_patterns() -> "re.Pattern[str]":
         # Word boundary prevents O(n^2) backtracking on long word-char runs.
         r"(?:^|(?<=\W))\w*(?:password|passwd|client_secret|secret_key|_secret)"
         r"['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+",
-        # Database connection string credentials (scheme://user:pass@host)
-        r"(?<=://)[^\s'\"]*:[^\s'\"@]+(?=@)",
+        # Database connection string credentials (scheme://user:pass@host).
+        # The user half stops at the ":" separator and both halves are length-capped,
+        # so a long attacker-supplied URL cannot backtrack quadratically.
+        r"(?<=://)[^\s'\":]{0,4096}:[^\s'\"]{1,4096}(?=@)",
         # Databricks personal access tokens
         r"dapi[0-9a-f]{32}",
         # Module-level provider keys logged as litellm.<provider>_key=<value>
@@ -67,8 +74,10 @@ def _build_secret_patterns() -> "re.Pattern[str]":
         r"""['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+""",
         # Raw JWTs (without Bearer prefix)
         r"\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]*",
-        # Azure SAS tokens in URLs
-        r"[?&]sig=[A-Za-z0-9%+/=]+",
+        # Azure SAS tokens in URLs. The delimiter is a lookbehind, like the
+        # `key=` pattern above, so the `?` or `&` survives and the redacted URL
+        # stays well formed (this string is often a request line in a log).
+        r"(?<=[?&])sig=[A-Za-z0-9%+/=]+",
         # Full JSON service-account blobs (single-line and multi-line)
         r'\{[^{}]*"type"\s*:\s*"service_account"[^{}]*(?:\{[^{}]*\}[^{}]*)*\}',
     ]
@@ -83,6 +92,27 @@ def redact_string(value: str) -> str:
     return _SECRET_RE.sub(_REDACTED, value)
 
 
+_UNIX_SYSTEM_PATH: Final = r"/(?:etc|var|opt|usr|home|root|private|Users|tmp|mnt|srv)/[^\s'\"\)\]}>,]+"
+_WINDOWS_DRIVE_PATH: Final = r"[A-Za-z]:\\[^\s'\"\)\]}>,]+"
+_PRIVATE_OR_LOOPBACK_IPV4: Final = (
+    r"\b(?:10(?:\.\d{1,3}){3}|172\.(?:1[6-9]|2\d|3[01])(?:\.\d{1,3}){2}|192\.168(?:\.\d{1,3}){2}|127(?:\.\d{1,3}){3})\b"
+)
+_INTERNAL_SUFFIX_HOSTNAME: Final = r"\b[A-Za-z0-9-]+(?:\.[A-Za-z0-9-]+)*\.(?:internal|local|corp|lan|intra|private)\b"
+_INTERNAL_DETAIL_RE: Final = re.compile(
+    "|".join((_UNIX_SYSTEM_PATH, _WINDOWS_DRIVE_PATH, _PRIVATE_OR_LOOPBACK_IPV4, _INTERNAL_SUFFIX_HOSTNAME)),
+    re.IGNORECASE,
+)
+_TRACEBACK_MARKER: Final = "Traceback (most recent call last):"
+
+
+def redact_internal_details(value: str) -> str:
+    """Drop an embedded traceback and scrub filesystem paths and internal hostnames,
+    on top of redact_string(). For client-facing messages only: server logs keep this detail."""
+    marker_index: Final = value.find(_TRACEBACK_MARKER)
+    without_traceback: Final = value[:marker_index].rstrip() if marker_index != -1 else value
+    return _INTERNAL_DETAIL_RE.sub(_REDACTED, redact_string(without_traceback))
+
+
 def redact_structured_value(key: str | None, value: str) -> str:
     """Scrub *value* as it appeared under *key* inside a structured record.
 

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))