diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index aada0fcf239..1b71232bc2e 100644 --- a/.github/workflows/test-rust.yml +++ b/.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 diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index c2dff805772..6da5fc07e80 100644 --- a/.github/workflows/test-unit.yml +++ b/.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 diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index d094c98f5ec..3f96531cf6f 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 14074 }, "reportArgumentType": { - "limit": 2216 + "limit": 2215 }, "reportAssignmentType": { "limit": 319 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 4125 + "limit": 4124 }, "reportFunctionMemberAccess": { "limit": 7 @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15306 + "limit": 15290 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,31 +99,31 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44364 + "limit": 44362 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 38350 + "limit": 38332 }, "reportUnknownParameterType": { "limit": 19625 }, "reportUnknownVariableType": { - "limit": 29877 + "limit": 29861 }, "reportUnnecessaryCast": { "limit": 111 }, "reportUnnecessaryComparison": { - "limit": 692 + "limit": 687 }, "reportUnnecessaryContains": { - "limit": 5 + "limit": 4 }, "reportUnnecessaryIsInstance": { - "limit": 826 + "limit": 823 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index f6647268624..d22484bc0e8 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/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: diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index 8d66bf872de..406f07eb792 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/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) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index dd41cf0e84b..b3dac5ca935 100644 --- a/litellm-rust/Cargo.lock +++ b/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" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index c447d915abe..a13dd4c04b0 100644 --- a/litellm-rust/Cargo.toml +++ b/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" diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 541beabe170..e3dbdf24ce6 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/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. diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index e0ce165dc93..c1fb328893b 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/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, Error> { let Some(object) = document.as_object() else { return Ok(None); diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index 815bc84363a..856d9571201 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/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 { +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +pub(crate) async fn execute_ocr_provider_call( + request: PreparedOcrRequest, + hooks: &OcrLifecycleHooks, +) -> Result { + 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()))?; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 401e26d3b29..f8c4f8fe8c5 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/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 { - 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 for OcrLifecycleHooks { +impl CallLifecycleHooks 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 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 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, diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index b59ab626fd3..d9230af1c59 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/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 { 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 } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs index 6231393c889..fedacc62760 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs +++ b/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( diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs index 8c3f0425149..85e4c408045 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs +++ b/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, diff --git a/litellm-rust/crates/ai-gateway/src/ocr/types.rs b/litellm-rust/crates/ai-gateway/src/ocr/types.rs index bde734a4dd1..95e551d79ca 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/types.rs +++ b/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, diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index ab8050734f2..389dbd49505 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/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 } diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index 30ba0da5e68..9a96b9d1140 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -1,11 +1,12 @@ use serde_json::Value; use crate::error::Error; -use crate::http_utils::truncate_error_body; +use crate::http_utils::{http_request, truncate_error_body}; use super::client::http_client; use super::types::ProviderAudioTranscriptionRequest; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn execute_audio_transcription_provider_call( request: ProviderAudioTranscriptionRequest, ) -> Result { @@ -19,8 +20,7 @@ pub async fn execute_audio_transcription_provider_call( if let Some(duration) = request.timeout { request_builder = request_builder.timeout(duration); } - let response = request_builder - .send() + let response = http_request(request_builder) .await .map_err(|error| Error::Network(error.to_string()))?; let status = response.status(); diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index b71748082bf..31b6de4b3e4 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -11,6 +11,7 @@ pub use handler::execute_audio_transcription_provider_call; pub use prepare::prepare_audio_transcription_provider_call; pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result { execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?) .await diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 6288e96b380..bbef97341a9 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -7,6 +7,7 @@ use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider} use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig}; use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> { #[cfg(feature = "bedrock-auth")] if provider == "bedrock" { @@ -16,6 +17,7 @@ fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProv None } +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub fn prepare_audio_transcription_provider_call( request: AudioTranscriptionRequest<'_>, ) -> Result { diff --git a/litellm-rust/crates/core/src/audio_transcription/transformation.rs b/litellm-rust/crates/core/src/audio_transcription/transformation.rs index 16a28fbcac0..aa9846427dc 100644 --- a/litellm-rust/crates/core/src/audio_transcription/transformation.rs +++ b/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) -> Map { params .iter() diff --git a/litellm-rust/crates/core/src/chat_completions/common_utils.rs b/litellm-rust/crates/core/src/chat_completions/common_utils.rs index ca51471eb7c..69e5f175ad5 100644 --- a/litellm-rust/crates/core/src/chat_completions/common_utils.rs +++ b/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> { diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index 7e2731442cc..96d001e2892 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/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 { + 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. diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 0d009d36d16..32dea17d202 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/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 { - 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 diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index 142b2f2aaed..3be2ba21de4 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/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, 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 { +) -> Result, 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 { + 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, diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index 2858d180e27..f8594dee447 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/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 { + prepare_provider_request(resolve_request(request)?) +} fn request<'a>( model: &'a str, diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/chat_completions/transformation.rs index a0868209305..d7b9704c46c 100644 --- a/litellm-rust/crates/core/src/chat_completions/transformation.rs +++ b/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, ) -> Option { 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, ) -> Option { @@ -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")) diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs index 35dd543a986..3238d09b6b5 100644 --- a/litellm-rust/crates/core/src/chat_completions/types.rs +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -22,6 +22,17 @@ pub struct ChatCompletionsRequest<'a> { pub timeout: Option, } +pub(super) struct ResolvedChatCompletionsRequest<'a> { + pub(super) model: String, + pub(super) config: &'static dyn ChatCompletionsProviderConfig, + pub(super) messages: Vec, + pub(super) optional_params: Map, + pub(super) api_key: Option<&'a str>, + pub(super) api_base: Option<&'a str>, + pub(super) extra_headers: Option>, + pub(super) timeout: Option, +} + pub(super) struct ProviderChatCompletionsRequest { pub(super) model: String, pub(super) config: &'static dyn ChatCompletionsProviderConfig, diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index 10661fadf96..3633130528d 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/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 { + 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 { diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 8dfdb2e361a..8f0f6652fa4 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/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> { diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 13a65d86131..61ff81bcdc8 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/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 { + 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 { + 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(); diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index ee2877e61fc..cfa8bda1104 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/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 { - execute_messages_provider_call(prepare_messages_call(request)?).await + execute_messages_provider_call(request).await } pub async fn messages_stream(request: MessagesRequest<'_>) -> Result { - execute_messages_provider_stream(prepare_messages_call(request)?).await + execute_messages_provider_stream(request).await } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 3b253ac3766..ec83d03f535 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/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 { 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>, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result, 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) } diff --git a/litellm-rust/crates/core/src/messages/transformation.rs b/litellm-rust/crates/core/src/messages/transformation.rs index 673a5728aca..a5904c085a0 100644 --- a/litellm-rust/crates/core/src/messages/transformation.rs +++ b/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, diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index 3d3c16c8cb6..ad484c8f968 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/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) -> Map { 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, ) -> Result; + #[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, + ) -> Result, 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 } diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs index 97cc48aa6f2..a7d5a8ad0cf 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs +++ b/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, ) -> Option { - 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, diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs index 8fcc0f36c7d..f31b961e78a 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs +++ b/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>, diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index 70dad0300f1..b8ca10461fb 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/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>, diff --git a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs index bb4f6afe5f9..9bf1f73a74d 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs @@ -46,10 +46,12 @@ fn optional_string<'a>(params: &'a Map, 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, diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs index ef5f44b4a14..7be3d108d44 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs +++ b/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) -> 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, ) -> Option { - 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( diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index 6a8a38204a9..9648321d7ff 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/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) -> Map Result { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 498003de149..637e5580170 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/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" diff --git a/litellm-rust/crates/python-bridge/src/constants.rs b/litellm-rust/crates/python-bridge/src/constants.rs new file mode 100644 index 00000000000..07b2836b838 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/constants.rs @@ -0,0 +1 @@ +pub(crate) const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace"; diff --git a/litellm-rust/crates/python-bridge/src/diagnostics.rs b/litellm-rust/crates/python-bridge/src/diagnostics.rs new file mode 100644 index 00000000000..cc153a89b8f --- /dev/null +++ b/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> { + 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(()) +} diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs new file mode 100644 index 00000000000..914e2e1e033 --- /dev/null +++ b/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::())?; + module.add("RustUpstreamError", py.get_type::()) +} diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs new file mode 100644 index 00000000000..f3648158cf6 --- /dev/null +++ b/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( + py: Python<'_>, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + Send + 'static, +{ + run_sync_on( + py, + pyo3_async_runtimes::tokio::get_runtime(), + future, + map_error, + ) +} + +fn run_sync_on( + py: Python<'_>, + runtime: &Runtime, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + 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( + py: Python<'_>, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + 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(result: Result, map_error: fn(Error) -> PyErr) -> PyResult { + 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(future: F) -> PyResult> +where + F: Future>, +{ + AssertUnwindSafe(future) + .catch_unwind() + .await + .map_err(panic_to_pyerr) +} + +async fn wait_for_sync_result(future: F) -> PyResult> +where + F: Future>, +{ + 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(&self, _serializer: S) -> Result + where + S: Serializer, + { + panic!("serializer panicked") + } + } + + #[pyfunction] + fn async_serialization_panic(py: Python<'_>) -> PyResult> { + run_async(py, async { Ok(PanickingOutput) }, runtime_error) + } + + #[pyfunction] + fn async_runtime_probe(py: Python<'_>) -> PyResult> { + 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>) -> 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::(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::( + py, + poll_fn(|_| -> Poll> { panic!("route future panicked") }), + runtime_error, + ) + .expect_err("panicked route should become a Python exception"); + + assert!(error.is_instance_of::(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::( + 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::(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::(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"); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/function_trace.rs b/litellm-rust/crates/python-bridge/src/function_trace.rs new file mode 100644 index 00000000000..420d237c79d --- /dev/null +++ b/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 { + Plain(T), + Traced { + response: T, + trace: Vec, + }, +} + +pub(crate) async fn trace_call( + future: impl Future>, + enabled: bool, +) -> Result, 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>>, +} + +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 { + self.events + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone() + } +} + +struct FunctionTraceLayer { + trace: FunctionTrace, +} + +impl Layer 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, + }, + ] + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 2e2624acbe1..5f36a22370a 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,142 +1,18 @@ -use std::collections::HashMap; -use std::time::Duration; +mod constants; +mod diagnostics; +mod errors; +mod execution; +pub mod function_trace; +mod marshal; +mod routes; -use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; -use litellm_core::audio_transcription::{ - AudioTranscriptionRequest, audio_transcription as run_audio_transcription, -}; -use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse}; -use litellm_core::chat_completions::{ - chat_completions as run_chat_completions, chat_completions_decline_reason, -}; -use litellm_core::error::Error; -use litellm_core::messages::messages as run_messages; -use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; -use litellm_python_interop::{from_py, release_count, release_gil, to_py}; -use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; -use pyo3::types::{PyAny, PyDict}; -use serde_json::{Map, Value}; +use pyo3::types::PyAny; +use serde_json::Value; -pyo3::create_exception!( - _native, - RustBridgeDeclined, - pyo3::exceptions::PyException, - "The route declined before calling the provider, so the host may retry on its own path." -); - -pyo3::create_exception!( - _native, - RustUpstreamError, - pyo3::exceptions::PyException, - "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." -); - -type MarshaledOcrInputs = ( - Value, - Option>, - Map, - Option, -); - -fn messages_response_to_py( - py: Python<'_>, - response: AnthropicMessagesResponse, -) -> PyResult> { - to_py(py, &response) -} - -fn chat_completions_response_to_py( - py: Python<'_>, - response: ChatCompletionsResponse, -) -> PyResult> { - 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>, -) -> PyResult> { - 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) -> Option { - 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>, -) -> PyResult> { - 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>, + #[pyo3(from_py_with = litellm_python_interop::from_py)] headers: Option, timeout_seconds: Option, ) -> PyResult> { - let headers = marshal_headers(py, headers)?; + let headers = marshal_headers(headers)?; let timeout = optional_timeout(timeout_seconds); pyo3_async_runtimes::tokio::future_into_py(py, async move { let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout) .await .map_err(core_error_to_pyerr)?; - Python::attach(|py| Py::new(py, ResponsesWebSocketConnection { inner })) + Ok(ResponsesWebSocketConnection { inner }) }) } @@ -186,445 +62,126 @@ impl ResponsesWebSocketConnection { } } -fn marshal_inputs( - py: Python<'_>, - document: Py, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult { - 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, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - 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::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, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - 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, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let audio = from_py(audio.bind(py))?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let timeout = optional_timeout(timeout_seconds); - let result = release_gil(py, || { - pyo3_async_runtimes::tokio::get_runtime().block_on(run_audio_transcription( - AudioTranscriptionRequest { - model: &model, - audio, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - }, - )) - }); - match result { - Ok(value) => to_py(py, &value), - Err(err) => Err(core_error_to_pyerr(err)), + #[test] + fn module_registration_preserves_the_public_surface() { + Python::initialize(); + Python::attach(|py| { + let module = pyo3::wrap_pymodule!(_native)(py).into_bound(py); + + let expected = [ + "RustBridgeDeclined", + "RustUpstreamError", + "ocr", + "aocr", + "transcription", + "atranscription", + "messages", + "amessages", + "chat_completions_decline", + "chat_completions", + "achat_completions", + "ResponsesWebSocketConnection", + "gil_stats", + ]; + + let public_names: Vec = module + .dict() + .keys() + .extract::>() + .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, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - optional_params: Option>, - timeout_seconds: Option, -) -> PyResult> { - let audio = from_py(audio.bind(py))?; - let extra_headers = match extra_headers { - Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), - None => None, - }; - let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; - let timeout = optional_timeout(timeout_seconds); - pyo3_async_runtimes::tokio::future_into_py(py, async move { - let value = run_audio_transcription(AudioTranscriptionRequest { - model: &model, - audio, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - optional_params, - timeout, - }) - .await - .map_err(core_error_to_pyerr)?; - Python::attach(|py| to_py(py, &value)) - }) -} - -type MarshaledMessagesInputs = (Value, Option>, Option); - -fn marshal_messages_inputs( - py: Python<'_>, - body: Py, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult { - 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, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - 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, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - 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, - Option>, - Option, -); - -fn marshal_chat_completions_inputs( - py: Python<'_>, - messages: Py, - optional_params: Option>, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult { - 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, - optional_params: Option>, - custom_llm_provider: Option, -) -> PyResult> { - 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, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - 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, - optional_params: Option>, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - 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> { - 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::())?; - module.add("RustUpstreamError", py.get_type::())?; - 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::()?; - module.add_function(wrap_pyfunction!(gil_stats, module)?)?; - #[cfg(feature = "panic-test")] - module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?; - Ok(()) -} diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs new file mode 100644 index 00000000000..a14e4b55d82 --- /dev/null +++ b/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, + pub(crate) api_base: Option, + pub(crate) custom_llm_provider: Option, + pub(crate) extra_headers: Option>, + pub(crate) timeout: Option, +} + +pub(crate) struct RouteOptionsInputs { + pub(crate) model: String, + pub(crate) api_key: Option, + pub(crate) api_base: Option, + pub(crate) custom_llm_provider: Option, + pub(crate) extra_headers: Option, + pub(crate) timeout_seconds: Option, +} + +impl RouteOptions { + pub(crate) fn from_python(inputs: RouteOptionsInputs) -> PyResult { + 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 { + 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, +) -> PyResult> { + match value { + Some(value) => object(name, value), + None => Ok(Map::new()), + } +} + +fn optional_object( + name: &'static str, + value: Option, +) -> PyResult>> { + value.map(|value| object(name, value)).transpose() +} + +fn object(name: &'static str, value: Value) -> PyResult> { + 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) -> Option { + 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) -> PyResult> { + 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() +} diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs new file mode 100644 index 00000000000..10b86132be7 --- /dev/null +++ b/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> + 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, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + extra_headers: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + optional_params: Option, + timeout_seconds: Option, + }, + prepare = prepare_transcription, + errors = core_error_to_pyerr, +} diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs new file mode 100644 index 00000000000..68b7762cb10 --- /dev/null +++ b/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> + 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, + custom_llm_provider: Option, +) -> PyResult> { + 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, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + extra_headers: Option, + timeout_seconds: Option, + }, + prepare = prepare_chat_completions, + errors = chat_completions_error_to_pyerr, + extra = [chat_completions_decline], +} diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs new file mode 100644 index 00000000000..21a7fd5a766 --- /dev/null +++ b/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> { + 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> { + 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> + 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::(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" + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs new file mode 100644 index 00000000000..2bb64a7a763 --- /dev/null +++ b/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> + 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, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + extra_headers: Option, + timeout_seconds: Option, + }, + prepare = prepare_messages, + errors = core_error_to_pyerr, +} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs new file mode 100644 index 00000000000..bf611c26d44 --- /dev/null +++ b/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) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs new file mode 100644 index 00000000000..5588c400972 --- /dev/null +++ b/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> + 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, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + extra_headers: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + optional_params: Option, + timeout_seconds: Option, + }, + prepare = prepare_ocr, + errors = core_error_to_pyerr, +} diff --git a/litellm-rust/crates/python-bridge/src/routes/runtime.rs b/litellm-rust/crates/python-bridge/src/routes/runtime.rs new file mode 100644 index 00000000000..87a0c3e0104 --- /dev/null +++ b/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( + py: Python<'_>, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + Send + 'static, +{ + run_sync_on( + py, + pyo3_async_runtimes::tokio::get_runtime(), + future, + map_error, + ) +} + +fn run_sync_on( + py: Python<'_>, + runtime: &Runtime, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + 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( + py: Python<'_>, + future: F, + map_error: fn(Error) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + F: Future> + 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(result: Result, map_error: fn(Error) -> PyErr) -> PyResult { + 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(future: F) -> PyResult> +where + F: Future>, +{ + AssertUnwindSafe(future) + .catch_unwind() + .await + .map_err(panic_to_pyerr) +} + +async fn wait_for_sync_result(future: F) -> PyResult> +where + F: Future>, +{ + 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(&self, _serializer: S) -> Result + where + S: Serializer, + { + panic!("serializer panicked") + } + } + + #[pyfunction] + fn async_serialization_panic(py: Python<'_>) -> PyResult> { + run_async(py, async { Ok(PanickingOutput) }, runtime_error) + } + + #[pyfunction] + fn async_runtime_probe(py: Python<'_>) -> PyResult> { + 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>) -> 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::(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::( + py, + poll_fn(|_| -> Poll> { panic!("route future panicked") }), + runtime_error, + ) + .expect_err("panicked route should become a Python exception"); + + assert!(error.is_instance_of::(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::( + 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::(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::(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"); + }); + } +} diff --git a/litellm-rust/crates/python-interop/src/lib.rs b/litellm-rust/crates/python-interop/src/lib.rs index df2bd260fdb..2e562bdae70 100644 --- a/litellm-rust/crates/python-interop/src/lib.rs +++ b/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}; diff --git a/litellm-rust/crates/python-interop/src/marshal.rs b/litellm-rust/crates/python-interop/src/marshal.rs index c3d0638427c..a16d1e0ae13 100644 --- a/litellm-rust/crates/python-interop/src/marshal.rs +++ b/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(pub T); + +impl<'py, T> IntoPyObject<'py> for Pythonized +where + T: Serialize, +{ + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + 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) -> PyErr { + let message = payload + .downcast_ref::() + .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(&self, _serializer: S) -> Result + where + S: Serializer, + { + panic!("serializer panicked") + } + } + + #[test] + fn pythonized_converts_on_the_attached_thread() { + Python::initialize(); + Python::attach(|py| { + let value: Vec = 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::(py)); + assert_eq!(error.to_string(), "PanicException: serializer panicked"); + }); + } +} diff --git a/litellm/__init__.py b/litellm/__init__.py index 61794dabddc..44f2e7c1f02 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1417,7 +1417,7 @@ from .skills.main import ( ) from .containers.main import * from .ocr.main import * -from .rust_bridge.ocr import use_litellm_rust +from .rust_bridge import use_litellm_rust from .rag.main import * from .sandbox.main import * from .search.main import * diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 748ef938cea..dc41c7dadc8 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/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, ) diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index 07d4f959489..e012d35b8f3 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/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, }, ) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index dc61ee38a8c..2d737bc34e7 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/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: diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index aa805ccea71..5f7ac73c919 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/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 diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index d23690976ad..6079b709bcc 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/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": }`` + 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: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 69985bcdaa3..b82903d6f87 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/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): diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 3d62b8b4784..988f81c9eb4 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/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, diff --git a/litellm/llms/azure_ai/anthropic/transformation.py b/litellm/llms/azure_ai/anthropic/transformation.py index c5053448627..864d2134a84 100644 --- a/litellm/llms/azure_ai/anthropic/transformation.py +++ b/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) diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index 64e72b819b1..db1a0fc89a3 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/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: diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 63e99c0915a..9624a721870 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/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") diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 5363c3c0366..e097805f54a 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/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)) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 2a4c38e71ea..67720451c00 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/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 ( diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 66ee5f10679..1e5329c90dd 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/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): diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index b6e93f590ca..e1f0fc9e7d3 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/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, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 6f42d42de00..0f6966b0ae2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/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 {}) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index c587146005f..65622d62af2 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/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: diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index d6b217d5746..73086ba395b 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/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) diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 2875b30232e..0a4cbd8e520 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/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 diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 8f23c5175ec..89920ebd27b 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/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( diff --git a/litellm/llms/gigachat/embedding/transformation.py b/litellm/llms/gigachat/embedding/transformation.py index 2ec8324e33c..0db4475be8f 100644 --- a/litellm/llms/gigachat/embedding/transformation.py +++ b/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( diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index a0edc6f5682..e1f73d04275 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/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), diff --git a/litellm/llms/gigachat/utils.py b/litellm/llms/gigachat/utils.py index cbb35cd1b57..ce7e848ed7f 100644 --- a/litellm/llms/gigachat/utils.py +++ b/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" diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 0e8fa294f5d..764d80c6f82 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/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) diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index c3581abfbcc..4f3c366d8c1 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/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: diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index de626b468f0..181894646e3 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/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( diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 4d774f6f165..bcd4ea43243 100644 --- a/litellm/llms/openai/common_utils.py +++ b/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, ) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 1530c154e93..5a5970fb867 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/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, diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py new file mode 100644 index 00000000000..b596adfad6f --- /dev/null +++ b/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)) diff --git a/litellm/llms/openai_like/README.md b/litellm/llms/openai_like/README.md index e9aaafe48a1..e1409b81c35 100644 --- a/litellm/llms/openai_like/README.md +++ b/litellm/llms/openai_like/README.md @@ -54,7 +54,10 @@ That's it! The provider will be automatically loaded and available. "constraints": { "temperature_max": 1.0, "temperature_min": 0.0, - "temperature_min_with_n_gt_1": 0.3 + "temperature_min_with_n_gt_1": 0.3, + // /v1/messages providers only: keep Anthropic cache_control extensions + // such as ttl instead of stripping them down to {"type": ...} + "cache_control_ttl": true }, // Optional: Special handling flags diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 11dc236064d..ac99617521c 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -1,11 +1,13 @@ from typing import Any, Final import litellm +from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.openai_like.json_loader import SimpleProviderConfig from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01" @@ -19,10 +21,17 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): ``"/v1/messages"``. The inbound Anthropic payload (system, cache_control, thinking, tools, ...) is forwarded essentially unchanged to ``{api_base}/v1/messages``, so Anthropic-only features that the - Anthropic->OpenAI translation would otherwise drop are preserved. Response - parsing and streaming are inherited from the native Anthropic config. + Anthropic->OpenAI translation would otherwise drop are preserved. The one + exception is ``cache_control``, whose Anthropic-only extensions (``ttl``) + are stripped unless the deployment opts in with + ``model_info.cache_control_ttl: true``. Response parsing and streaming are + inherited from the native Anthropic config. """ + def __init__(self, cache_control_ttl: bool = False) -> None: + super().__init__() + self._cache_control_ttl: Final = cache_control_ttl + def validate_anthropic_messages_environment( self, headers: dict[str, str], @@ -53,6 +62,35 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): def should_filter_anthropic_beta_headers(self) -> bool: return False + def supports_cache_control_ttl(self) -> bool: + return self._cache_control_ttl + + def transform_anthropic_messages_request( + self, + model: str, + messages: list[dict], # mutable-ok: matches dict-typed base signature + anthropic_messages_optional_request_params: dict, # mutable-ok: matches dict-typed base signature + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: matches dict-typed base signature + ) -> dict: # mutable-ok: matches dict-typed base signature + """ + Anthropic ignores prompt-caching hints it cannot honor, but strict + non-Anthropic implementations of the Messages API 400 the whole request + on Anthropic-only ``cache_control`` extensions (``cache_control.ttl: 1h + is not supported``), so unless the provider declares ttl support the + hints are reduced to their portable ``{"type": ...}`` core. + """ + request: Final = super().transform_anthropic_messages_request( + model=model, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + if self.supports_cache_control_ttl(): + return request + return normalize_cache_control_in_anthropic_payload(request) + def get_complete_url( self, api_base: str | None, @@ -81,7 +119,7 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): """ def __init__(self, provider: SimpleProviderConfig): - super().__init__() + super().__init__(cache_control_ttl=bool(provider.constraints.get("cache_control_ttl"))) self._provider = provider @property diff --git a/litellm/llms/valkey/vector_stores/transformation.py b/litellm/llms/valkey/vector_stores/transformation.py index 3cbfca0f1a9..b250f71cf3f 100644 --- a/litellm/llms/valkey/vector_stores/transformation.py +++ b/litellm/llms/valkey/vector_stores/transformation.py @@ -15,7 +15,10 @@ import httpx from pydantic import BaseModel, ConfigDict import litellm -from litellm.llms.base_llm.vector_store.transformation import BaseDirectVectorStoreConfig +from litellm.llms.base_llm.vector_store.transformation import ( + BaseDirectVectorStoreConfig, + VectorStoreEmbeddingExecutor, +) from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector from litellm.types.utils import EmbeddingResponse from litellm.types.vector_stores import ( @@ -213,6 +216,7 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig): vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, litellm_logging_obj: "LiteLLMLoggingObj", litellm_params: Mapping[str, object], + embedding_executor: VectorStoreEmbeddingExecutor | None = None, timeout: float | httpx.Timeout | None = None, ) -> VectorStoreSearchResponse: params: Final = _ValkeySearchParams.model_validate(litellm_params) @@ -222,10 +226,18 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig): embedding_field=params.embedding_field, text_field=params.text_field, ) - embedding_response: Final = self.embedding_fn( - model=params.require_embedding_model(), - input=[query_text], # mutable-ok: litellm.embedding's input contract is a list - **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG), + embedding_response: Final = ( + embedding_executor.embed( + params.require_embedding_model(), + query_text, + params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG, + ) + if embedding_executor is not None + else self.embedding_fn( + model=params.require_embedding_model(), + input=[query_text], # mutable-ok: the injected embedding callable requires list input + **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG), + ) ) vec_params: Final = {"vec": pack_vector(embedding_response.data[0]["embedding"])} # mutable-ok: redis-py API @@ -252,6 +264,7 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig): vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, litellm_logging_obj: "LiteLLMLoggingObj", litellm_params: Mapping[str, object], + embedding_executor: VectorStoreEmbeddingExecutor | None = None, timeout: float | httpx.Timeout | None = None, ) -> VectorStoreSearchResponse: params: Final = _ValkeySearchParams.model_validate(litellm_params) @@ -261,10 +274,18 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig): embedding_field=params.embedding_field, text_field=params.text_field, ) - embedding_response: Final = await self.aembedding_fn( - model=params.require_embedding_model(), - input=[query_text], # mutable-ok: litellm.embedding's input contract is a list - **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG), + embedding_response: Final = ( + await embedding_executor.aembed( + params.require_embedding_model(), + query_text, + params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG, + ) + if embedding_executor is not None + else await self.aembedding_fn( + model=params.require_embedding_model(), + input=[query_text], # mutable-ok: the injected embedding callable requires list input + **(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG), + ) ) vec_params: Final = {"vec": pack_vector(embedding_response.data[0]["embedding"])} # mutable-ok: redis-py API diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index ef03e61a858..7579bc8c02e 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -177,6 +177,10 @@ class VertexAIAnthropicConfig(AnthropicConfig): # Restore original model name for any other processing model = original_model + AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model( + model=original_model, optional_params=optional_params, custom_llm_provider="vertex_ai" + ) + return optional_params def transform_response( diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index b918f013700..b260ec6e06f 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -194,6 +194,12 @@ def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS +def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool: + raw_request_override: Final = prepared_request.litellm_params.get("rust") + request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None + return rust_ocr_bridge.rust_ocr_enabled(request_override=request_override) + + def _rust_bridge_optional_params( prepared_request: _PreparedOCRRequest, resolve_secret: Callable[[str], str | None], @@ -422,7 +428,7 @@ async def aocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled(): + if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): from litellm.secret_managers.main import get_secret_str rust_response: Final = await _run_rust_aocr( @@ -694,7 +700,7 @@ def ocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and rust_ocr_bridge.rust_ocr_enabled(): + if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): from litellm.secret_managers.main import get_secret_str rust_response: Final = _run_rust_ocr( diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index c4bd03fb1c3..7076683f294 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -113,10 +113,8 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): ) ) - # Compliant: Save a strong reference to prevent GC self._background_tasks.add(task) - # Remove the task from the set when it finishes to avoid memory leaks task.add_done_callback(self._background_tasks.discard) except Exception as e: # noqa: BLE001 # Safe catch-all for verbose logging verbose_logger.exception( @@ -578,7 +576,6 @@ def llm_passthrough_route( else: return response except Exception as e: - # provider_config is guaranteed non-None here due to the earlier guard assert provider_config is not None raise base_llm_http_handler._handle_error( e=e, diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index d1ef73a15cd..37474f85fe7 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -74,7 +74,13 @@ _MCP_GUARDRAIL_REJECTIONS: Final = ( ) -def _connection_error_message(exc: BaseException) -> str: +def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str: + if isinstance(exc, TimeoutError): + return ( + f"Failed to connect to MCP server: no response from {url or 'the server'} " + f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL " + "from its network (DNS, egress rules, firewalls) and that the server answers MCP requests." + ) if isinstance(exc, httpx.LocalProtocolError): return ( "Failed to connect to MCP server: a request header is malformed. " @@ -92,6 +98,9 @@ def _connection_error_message(exc: BaseException) -> str: if MCP_AVAILABLE: + from mcp.types import Tool as MCPTool + + from litellm.experimental_mcp_client.client import MCPClient from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, global_mcp_server_manager, @@ -876,7 +885,6 @@ if MCP_AVAILABLE: return (), classify_list_exception(e) return tools_result, ServerListOk(tool_count=len(tools_result)) - # Query all servers the user has access to queried_servers: Final = tuple( server for server in map(global_mcp_server_manager.get_mcp_server_by_id, allowed_server_ids) @@ -1141,12 +1149,18 @@ if MCP_AVAILABLE: scopes: Final[list[str] | None] = scopes_raw if isinstance(scopes_raw, list) else None return client_id, client_secret, scopes + async def _list_tools_within(client: MCPClient, deadline: float) -> list[MCPTool] | None: + with anyio.move_on_after(deadline): + return await client.list_tools(raise_on_error=True) + return None + async def _execute_with_mcp_client( request: NewMCPServerRequest, operation: Callable[..., Awaitable[Mapping[str, object]]], mcp_auth_header: str | dict[str, str] | None = None, oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, + timeout_seconds: float = MCP_TOOL_LISTING_TIMEOUT, ) -> Mapping[str, object]: """ Create a temporary MCP client from *request*, run *operation*, and return the result. @@ -1162,6 +1176,10 @@ if MCP_AVAILABLE: oauth2_headers: Headers extracted from the incoming request (may contain the litellm API key — must NOT be forwarded for M2M servers). raw_headers: Raw request headers forwarded for stdio env construction. + timeout_seconds: Cap on OAuth discovery, connect, handshake, and *operation* + combined. Defaults to ``MCP_TOOL_LISTING_TIMEOUT`` (30s, below common LB + timeouts) so an unreachable upstream yields this endpoint's JSON error + instead of an opaque load-balancer 504 with an empty body. Returns: The dict returned by *operation*, or an error dict on failure. @@ -1252,15 +1270,16 @@ if MCP_AVAILABLE: static_headers=request.static_headers, ) - client: Final = await global_mcp_server_manager._create_mcp_client( - server=server_model, - mcp_auth_header=mcp_auth_header, - extra_headers=merged_headers, - stdio_env=stdio_env, - cred_provider=preview_cred_provider, - ) + with anyio.fail_after(timeout_seconds): + client: Final = await global_mcp_server_manager._create_mcp_client( + server=server_model, + mcp_auth_header=mcp_auth_header, + extra_headers=merged_headers, + stdio_env=stdio_env, + cred_provider=preview_cred_provider, + ) - return await operation(client) + return await operation(client) except (KeyboardInterrupt, SystemExit, asyncio.CancelledError): raise @@ -1269,7 +1288,7 @@ if MCP_AVAILABLE: return { "status": "error", "error": True, - "message": _connection_error_message(e), + "message": _connection_error_message(e, request.url, timeout_seconds), } async def _preview_openapi_tools(spec_path: str) -> dict: @@ -1422,9 +1441,7 @@ if MCP_AVAILABLE: getattr(client, "timeout", MCP_CLIENT_TIMEOUT) or MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT, ) - list_tools_result = None # rebind-ok: set inside the timeout scope below - with anyio.move_on_after(listing_deadline): - list_tools_result = await client.list_tools(raise_on_error=True) # rebind-ok: fills the init above + list_tools_result: Final = await _list_tools_within(client, listing_deadline) if list_tools_result is None: verbose_logger.warning( "MCP tools/list preview timed out after %s seconds while paginating upstream tools", diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index 558ea54495f..c2614b85016 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -6,8 +6,11 @@ External callers (public IPs) only see servers with available_on_public_internet """ import ipaddress +import os +from collections.abc import Mapping from dataclasses import dataclass from typing import Any, Final +from urllib.parse import urlparse from fastapi import Request from pydantic import TypeAdapter, ValidationError @@ -137,7 +140,7 @@ class IPAddressUtils: @staticmethod def is_request_from_trusted_proxy( request: Request, - general_settings: dict[str, Any] | None = None, + general_settings: Mapping[str, Any] | None = None, ) -> bool: """ Return True if X-Forwarded-* headers on this request should be trusted. @@ -190,6 +193,36 @@ class IPAddressUtils: trusted_networks: Final = IPAddressUtils.parse_trusted_proxy_networks(trusted_ranges) return IPAddressUtils.is_trusted_proxy(direct_ip, trusted_networks) + @staticmethod + def is_request_https( + request: Request, + general_settings: Mapping[str, Any] | None = None, + ) -> bool: + """ + Whether this request's PUBLIC-facing origin is HTTPS, for deciding + whether a cookie set on the response should be marked ``Secure``. + + litellm only sees a plain-HTTP hop whenever TLS terminates at a + reverse proxy, so ``request.url.scheme`` alone cannot answer this in + that deployment shape. Resolved from the first trusted signal: + 1. ``PROXY_BASE_URL`` (operator-declared public origin). + 2. ``X-Forwarded-Proto``, only when the request's direct peer is a + configured trusted proxy -- see ``is_request_from_trusted_proxy``. + An untrusted caller cannot spoof this header to strip Secure. + 3. The request's own literal scheme (direct TLS termination, or no + reverse proxy in front of litellm). + """ + configured_base_url: Final = os.environ.get("PROXY_BASE_URL", "").strip() + if configured_base_url: + return urlparse(configured_base_url).scheme == "https" + + if IPAddressUtils.is_request_from_trusted_proxy(request, general_settings=general_settings): + forwarded_proto: Final = request.headers.get("X-Forwarded-Proto") + if forwarded_proto: + return forwarded_proto.split(",")[0].strip().lower() == "https" + + return request.url.scheme == "https" + @staticmethod def extract_client_ip_from_xff_hops( xff_header: str, diff --git a/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py b/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py index 27018769909..9cabac2d0fa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py +++ b/litellm/proxy/guardrails/guardrail_hooks/alice/alice.py @@ -8,6 +8,7 @@ import json import os from collections.abc import Mapping +from itertools import islice from typing import ( TYPE_CHECKING, Any, # noqa: TID251 # **kwargs forwards verbatim to CustomGuardrail.__init__; see ruff-strict.toml @@ -341,19 +342,18 @@ def _json_safe( if depth >= _MAX_DEPTH or id(value) in seen: return None - nested: Final = seen | {id(value)} # mutable-ok: one-shot set literal, unioned into a frozenset immediately + nested: Final = seen | frozenset((id(value),)) if isinstance(value, dict): - out: dict[str, object] = {} # mutable-ok: bounded accumulator local to this call, never escapes as-is - for key, item in list(value.items())[:_MAX_ITEMS]: # mutable-ok: list() only to slice an unordered view - if isinstance(key, str) and key not in strip_keys: - out[key] = _json_safe(item, depth + 1, nested, strip_keys) - return out + return { + key: _json_safe(item, depth + 1, nested, strip_keys) + for key, item in islice(value.items(), _MAX_ITEMS) + if isinstance(key, str) and key not in strip_keys + } if isinstance(value, (list, tuple, set, frozenset)): return [ # mutable-ok: return value is a one-shot list, discarded by the caller after use - _json_safe(item, depth + 1, nested, strip_keys) - for item in list(value)[:_MAX_ITEMS] # mutable-ok: list() only to slice an unordered view + _json_safe(item, depth + 1, nested, strip_keys) for item in islice(value, _MAX_ITEMS) ] dump: Final = getattr(value, "model_dump", None) diff --git a/litellm/proxy/management_endpoints/sso/saml_sso.py b/litellm/proxy/management_endpoints/sso/saml_sso.py index 3e67b211f62..466b100ea1f 100644 --- a/litellm/proxy/management_endpoints/sso/saml_sso.py +++ b/litellm/proxy/management_endpoints/sso/saml_sso.py @@ -36,6 +36,7 @@ from pydantic import ValidationError from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache +from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role from litellm.proxy.utils import get_custom_url @@ -131,7 +132,7 @@ class SAMLAuthHandler: @staticmethod def _is_https(request: Request) -> bool: - return SAMLAuthHandler._base_url(request).startswith("https") + return IPAddressUtils.is_request_https(request) @staticmethod def _acs_url(request: Request) -> str: diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 6b98d9f9a26..1feefa5725d 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -92,6 +92,7 @@ from litellm.proxy.auth.auth_utils import ( has_user_setup_sso, ) from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.admin_ui_utils import ( admin_ui_disabled, @@ -1118,7 +1119,7 @@ async def google_login( request=request, ) if sso_redirect is not None: - _persist_return_to_cookie(sso_redirect, return_to) + _persist_return_to_cookie(sso_redirect, return_to, request) return sso_redirect from fastapi.responses import HTMLResponse @@ -1138,7 +1139,7 @@ async def google_login( # helper the SSO branch uses, so /login can resume the connect flow instead of dead-ending at the # dashboard. One implementation → the two sign-in branches cannot diverge (and the login form always # renders, since the helper never raises on a bad return_to). - _persist_return_to_cookie(form_response, return_to) + _persist_return_to_cookie(form_response, return_to, request) return form_response @@ -2741,6 +2742,7 @@ async def _sso_return_to_redirect( jwt_token: str, redis_usage_cache, user_api_key_cache, + request: Request, ) -> RedirectResponse | None: """Resolve the post-SSO redirect for a ``return_to``, or None to fall through to the dashboard. @@ -2759,7 +2761,7 @@ async def _sso_return_to_redirect( if _is_same_origin_return_path(return_to): redirect_response = RedirectResponse(url=return_to, status_code=303) - redirect_response.set_cookie(key="token", value=jwt_token) + set_session_token_cookie(redirect_response, request, jwt_token) redirect_response.delete_cookie("litellm_cp_return_to") return redirect_response @@ -2782,7 +2784,25 @@ async def _sso_return_to_redirect( return None -def _persist_return_to_cookie(response: Response, return_to: str | None) -> None: +def set_session_token_cookie(response: Response, request: Request, jwt_token: str) -> None: + """Set the ``token`` session cookie shared by every sign-in path. + + Not HttpOnly: the dashboard reads this cookie via ``document.cookie`` to + populate its own Authorization headers (see + ``ui/litellm-dashboard/src/utils/cookieUtils.ts``), so marking it + HttpOnly would break login. Secure is still required whenever the public + origin is HTTPS, resolved the same trust-aware way as every other + litellm cookie.""" + response.set_cookie( + key="token", + value=jwt_token, + secure=IPAddressUtils.is_request_https(request), + httponly=False, + samesite="lax", + ) + + +def _persist_return_to_cookie(response: Response, return_to: str | None, request: Request) -> None: """Best-effort: persist a SAFE ``return_to`` on ``response`` as the one-shot ``litellm_cp_return_to`` cookie so ANY sign-in path — SSO / Okta / generic OR the username/password form — can resume there afterwards. THIS is the single source of truth, called by every sign-in branch so they cannot @@ -2803,6 +2823,7 @@ def _persist_return_to_cookie(response: Response, return_to: str | None) -> None max_age=600, httponly=True, samesite="lax", + secure=IPAddressUtils.is_request_https(request), ) @@ -3079,8 +3100,11 @@ class SSOAuthenticationHandler: # incoming request is HTTP (local dev). Without # ``Secure`` the cookie is sent over plain HTTP, # letting a network observer read and replay the - # state value and bypass this protection. - secure_flag: Final = request is None or request.url.scheme == "https" + # state value and bypass this protection. Trust-aware: + # honors PROXY_BASE_URL / a trusted reverse proxy's + # X-Forwarded-Proto instead of only the literal scheme + # litellm sees on the wire. + secure_flag: Final = request is None or IPAddressUtils.is_request_https(request) redirect_response.set_cookie( key="litellm_oauth_state", value=state_value, @@ -3628,6 +3652,7 @@ class SSOAuthenticationHandler: jwt_token=jwt_token, redis_usage_cache=redis_usage_cache, user_api_key_cache=user_api_key_cache, + request=request, ) if return_to_redirect is not None: return return_to_redirect @@ -3636,7 +3661,7 @@ class SSOAuthenticationHandler: litellm_dashboard_ui += "?login=success" verbose_proxy_logger.info("Redirecting to %s", litellm_dashboard_ui) redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303) - redirect_response.set_cookie(key="token", value=jwt_token) + set_session_token_cookie(redirect_response, request, jwt_token) return redirect_response @staticmethod diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 78d8ce296b8..b48b8d81494 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1731,7 +1731,7 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict: def get_vertex_pass_through_handler( - call_type: Literal["discovery", "aiplatform"], # noqa: UP037 + call_type: Literal["discovery", "aiplatform"], # noqa: UP037 # ruff reports quoted Literal values here ) -> BaseVertexAIPassThroughHandler: if call_type == "discovery": return VertexAIDiscoveryPassThroughHandler() @@ -2961,7 +2961,6 @@ async def handle_gigachat_passthrough_router_model( """ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - # Detect streaming based on request body is_streaming: Final = request_body.get("stream", False) # pyright: ignore[reportUnknownVariableType] # request_body is dict[Unknown, Unknown] data: dict[str, Any] = await _read_request_body( @@ -2997,7 +2996,6 @@ async def handle_gigachat_passthrough_router_model( data["json"] = request_body data["custom_llm_provider"] = "gigachat" - # Remove sensitive keys from data keys: Final = [ # mutable-ok: list of keys to remove from data "gigachat_auth_url", "gigachat_access_token", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1e0792cb8a5..27132c90e05 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15329,7 +15329,10 @@ async def login(request: Request): # authorize round-trip), mirroring the SSO callback; otherwise land on the dashboard. Gated by # _is_same_origin_return_path (strictly relative path) so it can never be an open redirect, and the # one-shot cookie is cleared after use. - from litellm.proxy.management_endpoints.ui_sso import _sso_return_to_redirect + from litellm.proxy.management_endpoints.ui_sso import ( + _sso_return_to_redirect, + set_session_token_cookie, + ) # Resume through the SAME resumer the SSO callback uses, rather than a second, narrower arm. # _persist_return_to_cookie stores both shapes it accepts (a relative same-origin path AND a @@ -15346,6 +15349,7 @@ async def login(request: Request): jwt_token=jwt_token, redis_usage_cache=redis_usage_cache, user_api_key_cache=user_api_key_cache, + request=request, ) except Exception: # noqa: BLE001 # resuming must NEVER block a completed sign-in # The symmetric half of _persist_return_to_cookie's "never raises" contract. The resumer @@ -15360,7 +15364,7 @@ async def login(request: Request): # Create redirect response with cookie redirect_response: Final = RedirectResponse(url=litellm_dashboard_ui, status_code=303) - redirect_response.set_cookie(key="token", value=jwt_token) + set_session_token_cookie(redirect_response, request, jwt_token) if cp_return_to: redirect_response.delete_cookie(key="litellm_cp_return_to") return redirect_response @@ -15370,6 +15374,7 @@ async def login(request: Request): async def login_v2(request: Request): global premium_user, general_settings, master_key from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie from litellm.proxy.utils import get_custom_url try: @@ -15404,7 +15409,7 @@ async def login_v2(request: Request): content={"redirect_url": litellm_dashboard_ui, "token": jwt_token}, status_code=status.HTTP_200_OK, ) - json_response.set_cookie(key="token", value=jwt_token) + set_session_token_cookie(json_response, request, jwt_token) return json_response except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.login_v2(): Exception occurred - %s", e) @@ -15504,6 +15509,8 @@ async def login_v3(request: Request): @router.post("/v3/login/exchange", include_in_schema=False) # exchange single-use opaque code for JWT async def login_v3_exchange(request: Request): + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie + try: if not general_settings.get("control_plane_url"): raise ProxyException( @@ -15550,7 +15557,7 @@ async def login_v3_exchange(request: Request): }, status_code=status.HTTP_200_OK, ) - json_response.set_cookie(key="token", value=cached_data["token"]) + set_session_token_cookie(json_response, request, cached_data["token"]) return json_response except ProxyException: raise diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 0ab7d99e4e4..e144ff965ae 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -729,7 +729,7 @@ async def rag_query( # conflict so callers cannot override the store's provider or credentials. managed_store: Final = resolved_stores.get(retrieval_config["vector_store_id"]) store_data: Final = ( - await build_request_data_from_managed_vector_store(managed_store) + build_request_data_from_managed_vector_store(managed_store) if managed_store is not None else MappingProxyType({}) ) diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 7d64e648e08..1feda0b0bb5 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -16,9 +16,6 @@ from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.utils import jsonify_object -from litellm.proxy.vector_store_endpoints.management_endpoints import ( - _resolve_embedding_config, -) from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, assert_user_can_access_vector_store, @@ -57,19 +54,9 @@ def reject_caller_embedding_selection_params(payload: Mapping[str, object], sour ######################################################## -async def build_request_data_from_managed_vector_store( +def build_request_data_from_managed_vector_store( vector_store: LiteLLM_ManagedVectorStore, ) -> Mapping[str, object]: - """ - Build request params (provider, credential ref, litellm_params) from an - already-resolved managed vector store. - - ``litellm_embedding_config`` is resolved here, at request-handling time, - instead of at row-creation time: the resolved api_key/api_base/api_version - lives only in the returned per-request mapping and is never persisted back - to the registry cache. Legacy rows that already carry a resolved - (cleartext) config skip the lookup and pass through unchanged. - """ top_level: Final = MappingProxyType( { key: vector_store.get(key) @@ -78,18 +65,7 @@ async def build_request_data_from_managed_vector_store( } ) litellm_params: Final = vector_store.get("litellm_params") or MappingProxyType({}) - embedding_model: Final = litellm_params.get("litellm_embedding_model") - if not embedding_model or litellm_params.get("litellm_embedding_config"): - return MappingProxyType({**top_level, **litellm_params}) - - from litellm.proxy.proxy_server import prisma_client - - resolved_config: Final = await _resolve_embedding_config( - embedding_model=embedding_model, prisma_client=prisma_client - ) - if not resolved_config: - return MappingProxyType({**top_level, **litellm_params}) - return MappingProxyType({**top_level, **litellm_params, "litellm_embedding_config": resolved_config}) + return MappingProxyType({**top_level, **litellm_params}) async def _update_request_data_with_litellm_managed_vector_store_registry( @@ -118,7 +94,7 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( vector_store=vector_store_to_run, user_api_key_dict=user_api_key_dict, ) - return {**data, **(await build_request_data_from_managed_vector_store(vector_store_to_run))} + return {**data, **build_request_data_from_managed_vector_store(vector_store_to_run)} @router.post( diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 244798ba05e..c928398a87f 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -18,11 +18,8 @@ if TYPE_CHECKING: from prisma.models import LiteLLM_ManagedVectorStoresTable as _VectorStoreRow from litellm.proxy.utils import PrismaClient - from litellm.router import Router - import litellm from litellm._logging import verbose_proxy_logger -from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import REDACTED_BY_LITELM_STRING from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker @@ -32,13 +29,10 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store -from litellm.repositories.model_repository import ModelRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ManagedVectorStoresRepository -from litellm.secret_managers.main import get_secret from litellm.types.vector_stores import ( LiteLLM_ManagedVectorStore, LiteLLM_ManagedVectorStoreListResponse, @@ -64,28 +58,6 @@ _LITELLM_PARAMS_MASKER: Final = SensitiveDataMasker() _REDACT_LITELLM_PARAMS_MAX_DEPTH: Final = 10 -# Use-time embedding-config resolution runs on every vector-store request -# whose persisted row carries only a model reference (the post-fix shape). -# Without a cache, that's one ``litellm_proxymodeltable.find_first`` per -# request — the no-DB-in-critical-path rule. Hold the resolved config in -# memory for a short TTL so a hot model name pays the DB lookup at most -# once per ``_EMBEDDING_CONFIG_CACHE_TTL`` seconds. Cleartext credentials -# only ever live in process memory (never persisted, never echoed in -# management responses), so the cache doesn't widen the disclosure surface. -_EMBEDDING_CONFIG_CACHE_TTL: Final = 60 -_EMBEDDING_CONFIG_CACHE_MAX_SIZE: Final = 256 -_embedding_config_cache: InMemoryCache | None = None - - -def _get_embedding_config_cache() -> InMemoryCache: - global _embedding_config_cache - if _embedding_config_cache is None: - _embedding_config_cache = InMemoryCache( - max_size_in_memory=_EMBEDDING_CONFIG_CACHE_MAX_SIZE, - default_ttl=_EMBEDDING_CONFIG_CACHE_TTL, - ) - return _embedding_config_cache - def _redact_sensitive_litellm_params(litellm_params: Any, _depth: int = 0) -> Any: """ @@ -155,235 +127,6 @@ async def _fetch_and_authorize_vector_store( return typed -def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, object] | None: - """ - Resolve embedding config from router's config-defined models. - - Config-defined models (from proxy_config.yaml) are stored in the router's model_list, - not in the database. This function looks up the model in the router and extracts - api_key, api_base, and api_version from the deployment's litellm_params. - - Args: - embedding_model: The embedding model string (e.g., "text-embedding-ada-002" or "azure/text-embedding-3-large") - llm_router: The LiteLLM router instance - - Returns: - Dictionary with api_key, api_base, and api_version if model found, None otherwise - """ - if not embedding_model or llm_router is None: - return None - - # Extract model name candidates - could be "text-embedding-ada-002" or "azure/text-embedding-3-large" - # Try exact match first, then try without provider prefix - model_name_candidates: Final = [embedding_model] - if "/" in embedding_model: - # If it has a provider prefix, also try without it - _, model_name = embedding_model.split("/", 1) - model_name_candidates.append(model_name) - - # Try to find model in router - for model_name in model_name_candidates: - try: - # Try to get deployment by model group name (model_name in config) - deployment = llm_router.get_deployment_by_model_group_name(model_group_name=model_name) - - if deployment is not None and deployment.litellm_params is not None: - litellm_params = deployment.litellm_params - - # Build embedding config from model params - embedding_config: dict[str, object] = {} - - # Extract api_key - api_key = getattr(litellm_params, "api_key", None) - if api_key: - # Handle os.environ/ prefix - if isinstance(api_key, str) and api_key.startswith("os.environ/"): - api_key = get_secret(api_key) - embedding_config["api_key"] = api_key - - # Extract api_base - api_base = getattr(litellm_params, "api_base", None) - if api_base: - # Handle os.environ/ prefix - if isinstance(api_base, str) and api_base.startswith("os.environ/"): - api_base = get_secret(api_base) - embedding_config["api_base"] = api_base - - # Extract api_version - api_version = getattr(litellm_params, "api_version", None) - if api_version: - embedding_config["api_version"] = api_version - - project_id = getattr(litellm_params, "project_id", None) - if project_id: - embedding_config["project_id"] = project_id - - # Only return config if we have at least api_key or api_base - if embedding_config: - verbose_proxy_logger.debug( - "Resolved embedding config from router model %s: %s", model_name, list(embedding_config.keys()) - ) - return embedding_config - except Exception as e: - verbose_proxy_logger.debug("Error resolving embedding config from router for model %s: %s", model_name, e) - continue - - return None - - -async def _resolve_embedding_config_from_db( - embedding_model: str, prisma_client: "PrismaClient" -) -> dict[str, object] | None: - """ - Resolve embedding config from database model configuration. - - If litellm_embedding_model is provided but litellm_embedding_config is not, - this function looks up the model in the database and extracts api_key, api_base, - and api_version from the model's litellm_params to build the embedding config. - - Args: - embedding_model: The embedding model string (e.g., "text-embedding-ada-002" or "azure/text-embedding-3-large") - prisma_client: The Prisma client instance - - Returns: - Dictionary with api_key, api_base, and api_version if model found, None otherwise - """ - if not embedding_model: - return None - - # Extract model name - could be "text-embedding-ada-002" or "azure/text-embedding-3-large" - # Try to find model by exact match first, then try without provider prefix - model_name_candidates: Final = [embedding_model] - if "/" in embedding_model: - # If it has a provider prefix, also try without it - _, model_name = embedding_model.split("/", 1) - model_name_candidates.append(model_name) - - # Try to find model in database - for model_name in model_name_candidates: - try: - db_model = await ModelRepository(prisma_client).table.find_first(where={"model_name": model_name}) - - if db_model and db_model.litellm_params: - # Extract litellm_params (could be dict or JSON string) - model_params = db_model.litellm_params - if isinstance(model_params, str): # pyright: ignore[reportUnnecessaryIsInstance] # prisma Json is str - model_params = json.loads(model_params) - - # Decrypt values from database (similar to how proxy_server.py does it) - # Values stored in DB are encrypted, so we need to decrypt them first - decrypted_params = {} - if isinstance(model_params, dict): - for k, v in model_params.items(): - if isinstance(v, str): - # Decrypt value - returns original value if decryption fails or no key is set - decrypted_value = decrypt_value_helper(value=v, key=k, return_original_value=True) - decrypted_params[k] = decrypted_value - else: - decrypted_params[k] = v - else: - decrypted_params = model_params - - # Build embedding config from model params - embedding_config = {} - - # Extract api_key - api_key = decrypted_params.get("api_key") - if api_key: - # Handle os.environ/ prefix (after decryption, values may be os.environ/ prefixed) - if isinstance(api_key, str) and api_key.startswith("os.environ/"): - api_key = get_secret(api_key) - embedding_config["api_key"] = api_key - - # Extract api_base - api_base = decrypted_params.get("api_base") - if api_base: - # Handle os.environ/ prefix (after decryption, values may be os.environ/ prefixed) - if isinstance(api_base, str) and api_base.startswith("os.environ/"): - api_base = get_secret(api_base) - embedding_config["api_base"] = api_base - - # Extract api_version - api_version = decrypted_params.get("api_version") - if api_version: - embedding_config["api_version"] = api_version - - # Only return config if we have at least api_key or api_base - if embedding_config: - verbose_proxy_logger.debug( - "Resolved embedding config from database model %s: %s", - model_name, - list(embedding_config.keys()), - ) - return embedding_config - except Exception as e: - verbose_proxy_logger.debug("Error resolving embedding config for model %s: %s", model_name, e) - continue - - return None - - -async def _resolve_embedding_config( - embedding_model: str, prisma_client: "PrismaClient | None", llm_router: "Router | None" = None -) -> dict[str, object] | None: - """ - Resolve embedding config from either router (config-defined) or database models. - - This function first checks the router for config-defined models, then falls back - to the database. This allows users to use models defined in either location. - - Results are cached in process memory for ``_EMBEDDING_CONFIG_CACHE_TTL`` - seconds so the request-handling path doesn't hit the database on every - vector-store call. Negative results (model not found) are intentionally - not cached to avoid blocking a freshly-added model behind the TTL. - - Args: - embedding_model: The embedding model string (e.g., "text-embedding-ada-002" or "azure/text-embedding-3-large") - prisma_client: The Prisma client instance - llm_router: The LiteLLM router instance (optional, will be imported if not provided) - - Returns: - Dictionary with api_key, api_base, and api_version if model found, None otherwise - """ - if not embedding_model: - return None - - cache: Final = _get_embedding_config_cache() - cached: Final = cache.get_cache(embedding_model) - if cached is not None: - return cached - - # Import llm_router if not provided - if llm_router is None: - try: - from litellm.proxy.proxy_server import llm_router - except ImportError: - llm_router = None - - # First try to resolve from router (config-defined models) - if llm_router is not None: - router_config = _resolve_embedding_config_from_router(embedding_model=embedding_model, llm_router=llm_router) - if router_config: - verbose_proxy_logger.debug("Resolved embedding config from router for model %s", embedding_model) - cache.set_cache(embedding_model, router_config) - return router_config - - # Fall back to database - if prisma_client is not None: - db_config: Final = await _resolve_embedding_config_from_db( - embedding_model=embedding_model, prisma_client=prisma_client - ) - if db_config: - verbose_proxy_logger.debug("Resolved embedding config from database for model %s", embedding_model) - cache.set_cache(embedding_model, db_config) - return db_config - - verbose_proxy_logger.debug( - "Could not resolve embedding config for model %s from router or database", embedding_model - ) - return None - - ######################################################## # Helper Functions ######################################################## @@ -469,10 +212,9 @@ async def create_vector_store_in_db( # (``api_key``, ``api_base``, ``api_version``) into this row. That # exposed every env-stored embedding-model credential on the # ``/vector_store/{new,info,update,list}`` responses. Keep the user's - # raw ``litellm_embedding_model`` reference; resolution now happens in - # ``build_request_data_from_managed_vector_store`` - # at request-handling time so the cleartext config exists only in - # per-request memory and never reaches the database. + # raw ``litellm_embedding_model`` reference; each search embeds the + # query through the router at request time, so the credentials stay + # on the deployment and never reach the database. if litellm_params: litellm_params_dict: Final = GenericLiteLLMParams(**litellm_params).model_dump(exclude_none=True) data_to_create["litellm_params"] = safe_dumps(litellm_params_dict) @@ -862,11 +604,9 @@ async def update_vector_store( # Handle litellm_params if provided. As with the create path, the # embedding-config auto-resolve previously persisted cleartext - # credentials into the row; resolution now happens at request- - # handling time in - # ``build_request_data_from_managed_vector_store`` - # so this row only ever stores the user-supplied - # ``litellm_embedding_model`` reference. + # credentials into the row; each search now embeds the query + # through the router at request time, so this row only ever stores + # the user-supplied ``litellm_embedding_model`` reference. if "litellm_params" in update_data: _input_litellm_params: Final[dict] = update_data.get("litellm_params", {}) or {} litellm_params_dict: Final = GenericLiteLLMParams(**_input_litellm_params).model_dump(exclude_none=True) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 5f3e88bb12f..b2d1a69e0d8 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -6,6 +6,7 @@ import json import re import uuid from collections.abc import Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -102,6 +103,15 @@ from .custom_tools import ( NamespaceNameMap: TypeAlias = Mapping[str, tuple[str, str]] NamespaceTool: TypeAlias = Mapping[str, object] ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None +ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool +NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n" + + +@dataclass(frozen=True, slots=True) +class ResponsesToolChatForm: + chat_tools: tuple[ChatToolParam, ...] + web_search_options: OpenAIWebSearchOptions | None + if TYPE_CHECKING: from openai.types.responses.response_apply_patch_tool_call import ( @@ -1771,7 +1781,7 @@ class LiteLLMCompletionResponsesConfig: tool_name: Final = str(namespace_tool.get("name") or "") raw_description: Final = str(namespace_tool.get("description") or "") description: Final = ( - f"{namespace_description}\n\n{raw_description}" + f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}{raw_description}" if nested and namespace_description and raw_description else namespace_description if nested and namespace_description @@ -1837,9 +1847,78 @@ class LiteLLMCompletionResponsesConfig: + ", ".join(sorted(conflicting_tool_names)) ) + @staticmethod + def _responses_tool_to_chat_form(tool: Mapping[str, object]) -> ResponsesToolChatForm: + tool_type: Final = tool.get("type") + if tool_type == "mcp": + return ResponsesToolChatForm(chat_tools=(cast(OpenAIMcpServerTool, tool),), web_search_options=None) + if tool_type == "web_search_preview" or tool_type == "web_search": + _search_context_size: Final[Literal["low", "medium", "high"]] = cast( + Literal["low", "medium", "high"], tool.get("search_context_size") + ) + _user_location: Final[OpenAIWebSearchUserLocation | None] = cast( + OpenAIWebSearchUserLocation | None, + tool.get("user_location") or None, + ) + return ResponsesToolChatForm( + chat_tools=(), + web_search_options=OpenAIWebSearchOptions( + search_context_size=_search_context_size, + user_location=_user_location, + ), + ) + if tool_type == "function": + typed_tool: Final = cast(FunctionToolParam, tool) + raw_parameters: Final = typed_tool.get("parameters", {}) or {} + parameters: Final = ( + {**raw_parameters} # mutable-ok: json.dumps rejects MappingProxyType + if "type" in raw_parameters + else {**raw_parameters, "type": "object"} # mutable-ok: json.dumps rejects MappingProxyType + ) + chat_completion_tool: Final[dict[str, object]] = { + "type": "function", + "function": { + "name": typed_tool.get("name") or "", + "description": typed_tool.get("description") or "", + "parameters": parameters, + "strict": typed_tool.get("strict", False) or False, + }, + } + if tool.get("cache_control"): + chat_completion_tool["cache_control"] = tool.get("cache_control") + if tool.get("defer_loading"): + chat_completion_tool["defer_loading"] = tool.get("defer_loading") + if tool.get("allowed_callers"): + chat_completion_tool["allowed_callers"] = tool.get("allowed_callers") + if tool.get("input_examples"): + chat_completion_tool["input_examples"] = tool.get("input_examples") + return ResponsesToolChatForm( + chat_tools=(cast(ChatCompletionToolParam, chat_completion_tool),), web_search_options=None + ) + if tool_type == "namespace": + return ResponsesToolChatForm( + chat_tools=LiteLLMCompletionResponsesConfig._namespace_chat_tools(tool), web_search_options=None + ) + if tool_type == "custom": + converted: Final = convert_custom_tool_to_function_tool(tool) + return ResponsesToolChatForm(chat_tools=() if converted is None else (converted,), web_search_options=None) + if tool_type in ("computer_use", "image_generation", "shell"): + verbose_logger.warning( + "Dropping Responses API tool of type '%s': it has no Chat Completions " + "equivalent and the target provider would reject the request.", + tool_type, + ) + return ResponsesToolChatForm(chat_tools=(), web_search_options=None) + return ResponsesToolChatForm(chat_tools=(cast(ChatToolParam, tool),), web_search_options=None) + + @staticmethod + def responses_tools_to_chat_forms(tools: ResponseTools) -> tuple[ResponsesToolChatForm, ...]: + LiteLLMCompletionResponsesConfig._validate_namespace_name_collisions(tools) + return tuple(LiteLLMCompletionResponsesConfig._responses_tool_to_chat_form(tool) for tool in tools or ()) + @staticmethod def transform_responses_api_tools_to_chat_completion_tools( - tools: list[FunctionToolParam | OpenAIMcpServerTool] | None, + tools: ResponseTools, ) -> tuple[ list[ChatCompletionToolParam | OpenAIMcpServerTool], OpenAIWebSearchOptions | None, @@ -1849,73 +1928,16 @@ class LiteLLMCompletionResponsesConfig: """ if tools is None: return [], None - LiteLLMCompletionResponsesConfig._validate_namespace_name_collisions(tools) - chat_completion_tools: Final[list[ChatCompletionToolParam | OpenAIMcpServerTool]] = [] - web_search_options: OpenAIWebSearchOptions | None = None - for tool in tools: - if tool.get("type") == "mcp": - chat_completion_tools.append(cast(OpenAIMcpServerTool, tool)) - elif tool.get("type") == "web_search_preview" or tool.get("type") == "web_search": - _search_context_size: Literal["low", "medium", "high"] = cast( - Literal["low", "medium", "high"], tool.get("search_context_size") - ) - _user_location: OpenAIWebSearchUserLocation | None = cast( - OpenAIWebSearchUserLocation | None, - tool.get("user_location") or None, - ) - web_search_options = OpenAIWebSearchOptions( - search_context_size=_search_context_size, - user_location=_user_location, - ) - elif tool.get("type") == "function": - typed_tool = cast(FunctionToolParam, tool) - # Ensure parameters has "type": "object" as required by providers like Anthropic - parameters = dict(typed_tool.get("parameters", {}) or {}) - if not parameters or "type" not in parameters: - parameters["type"] = "object" - chat_completion_tool: dict[str, object] = { - "type": "function", - "function": { - "name": typed_tool.get("name") or "", - "description": typed_tool.get("description") or "", - "parameters": parameters, - "strict": typed_tool.get("strict", False) or False, - }, - } - if tool.get("cache_control"): - chat_completion_tool["cache_control"] = tool.get("cache_control") - if tool.get("defer_loading"): - chat_completion_tool["defer_loading"] = tool.get("defer_loading") - if tool.get("allowed_callers"): - chat_completion_tool["allowed_callers"] = tool.get("allowed_callers") - if tool.get("input_examples"): - chat_completion_tool["input_examples"] = tool.get("input_examples") - chat_completion_tools.append(cast(ChatCompletionToolParam, chat_completion_tool)) - elif tool.get("type") == "namespace": - chat_completion_tools.extend(LiteLLMCompletionResponsesConfig._namespace_chat_tools(tool)) - elif tool.get("type") == "custom": - converted = convert_custom_tool_to_function_tool(tool) - if converted is not None: - chat_completion_tools.append(converted) - else: - _tool_type = tool.get("type") - if _tool_type in ("computer_use", "image_generation", "shell"): - # Drop unsupported Responses-API-only tool types that have no - # Chat Completions equivalent. Passing them through verbatim - # causes providers to reject the request with "'function' is a - # required property". - verbose_logger.warning( - "Dropping Responses API tool of type '%s': it has no Chat Completions " - "equivalent and the target provider would reject the request.", - _tool_type, - ) - continue - chat_completion_tools.append(cast(ChatCompletionToolParam | OpenAIMcpServerTool, tool)) - return chat_completion_tools, web_search_options + forms: Final = LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(tools) + web_search_options: Final = next( + (form.web_search_options for form in reversed(forms) if form.web_search_options is not None), + None, + ) + return [chat_tool for form in forms for chat_tool in form.chat_tools], web_search_options @staticmethod def transform_chat_completion_tool_params_to_responses_api_tools( - chat_completion_tools: list[ChatCompletionToolParam | OpenAIMcpServerTool] | None, + chat_completion_tools: Sequence[Mapping[str, object]] | None, ) -> list[dict[str, object]]: """ Transform Chat Completion tool params (e.g. from guardrail output) back to @@ -1926,9 +1948,6 @@ class LiteLLMCompletionResponsesConfig: return [] result: Final[list[dict[str, object]]] = [] for tool in chat_completion_tools: - if not isinstance(tool, dict): - result.append(tool) - continue if tool.get("type") == "function": fn = cast(_ToolFunctionDefinition, tool.get("function") or {}) parameters = dict(fn.get("parameters", {}) or {}) diff --git a/litellm/router.py b/litellm/router.py index 9353e391e76..0af514fe8a2 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -85,6 +85,10 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_credentials_in_payload, mask_sensitive_structure, ) +from litellm.llms.base_llm.vector_store.transformation import ( + RouterVectorStoreEmbeddingExecutor, + vector_store_request_metadata, +) from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.least_busy import LeastBusyLoggingHandler @@ -354,6 +358,8 @@ _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset({"model", "api_base", "api_key", "api_version"}) _ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" +_CLAUDE_CODE_SESSION_ID_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") +_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType( { @@ -810,6 +816,10 @@ class Router: self.cache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache() ) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc. + self._claude_code_session_router_cache: DualCache = DualCache( + redis_cache=redis_cache, + in_memory_cache=InMemoryCache(), + ) ### SCHEDULER ### self.scheduler = Scheduler(polling_interval=polling_interval, redis_cache=redis_cache) @@ -1131,8 +1141,8 @@ class Router: ``` and caching to just work. """ - if self.cache.redis_cache is None: - self.cache.redis_cache = cache + self.cache.attach_redis_cache(cache) + self._claude_code_session_router_cache.attach_redis_cache(cache) # Maps a routing strategy string to the attribute on `self` that holds # the default group's strategy selector for that strategy. (The selectors @@ -6475,11 +6485,24 @@ class Router: if custom_llm_provider and "custom_llm_provider" not in kwargs else MappingProxyType(kwargs) ) - if provider_kwargs.get("model"): - return self._generic_api_call_with_fallbacks(original_function=original_function, **provider_kwargs) + search_kwargs: Final = ( + MappingProxyType( + { + **provider_kwargs, + "_direct_vector_store_embedding_executor": RouterVectorStoreEmbeddingExecutor( + router=self, + metadata=self._vector_store_request_metadata(kwargs), + ), + } + ) + if call_type == "vector_store_search" + else provider_kwargs + ) + if search_kwargs.get("model"): + return self._generic_api_call_with_fallbacks(original_function=original_function, **search_kwargs) if call_type == "vector_store_search": - return original_function(**MappingProxyType({**provider_kwargs, "router": self})) - return original_function(**provider_kwargs) + return original_function(**MappingProxyType({**search_kwargs, "router": self})) + return original_function(**search_kwargs) return vector_store_sync_wrapper @@ -6652,11 +6675,22 @@ class Router: "avector_store_update", "avector_store_delete", ): + vector_store_kwargs: Final = ( + { # mutable-ok: the async routed request requires dynamic keyword arguments + **kwargs, + "_direct_vector_store_embedding_executor": RouterVectorStoreEmbeddingExecutor( + router=self, + metadata=self._vector_store_request_metadata(kwargs), + ), + } + if call_type == "avector_store_search" + else kwargs + ) return await self._init_vector_store_api_endpoints( original_function=original_function, custom_llm_provider=custom_llm_provider, call_type=call_type, - **kwargs, + **vector_store_kwargs, ) elif call_type in ("afile_delete", "afile_content"): return await self._ageneric_api_call_with_fallbacks( @@ -6692,6 +6726,10 @@ class Router: return async_wrapper + @staticmethod + def _vector_store_request_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: + return vector_store_request_metadata(kwargs) + async def _init_vector_store_api_endpoints( self, original_function: Callable, @@ -12651,6 +12689,100 @@ class Router: return None return candidates[0] + @staticmethod + def _request_header(request_kwargs: Mapping[str, object], header_name: str) -> str | None: + proxy_server_request: Final = request_kwargs.get("proxy_server_request") + if not isinstance(proxy_server_request, Mapping): + return None + headers: Final = proxy_server_request.get("headers") + if not isinstance(headers, Mapping): + return None + return next( + ( + value + for key, value in headers.items() + if isinstance(key, str) and key.lower() == header_name and isinstance(value, str) + ), + None, + ) + + def _claude_code_session_router_cache_key(self, request_kwargs: Mapping[str, object]) -> str | None: + session_id: Final = self._request_header(request_kwargs, "x-claude-code-session-id") + if session_id is None or _CLAUDE_CODE_SESSION_ID_RE.fullmatch(session_id) is None: + return None + metadata_name: Final = "litellm_metadata" if "litellm_metadata" in request_kwargs else "metadata" + metadata: Final = request_kwargs.get(metadata_name) + if not isinstance(metadata, Mapping): + return None + caller_scope: Final = metadata.get("user_api_key_hash") + if not isinstance(caller_scope, str) or not caller_scope: + return None + return f"claude_code_session_router:v1:{caller_scope}:{session_id}" + + async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: + try: + await self._claude_code_session_router_cache.async_delete_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request + verbose_router_logger.warning( + "Failed to delete Claude Code session router binding; the binding may remain until its TTL expires: %s", + e, + ) + + async def _get_claude_code_session_router_binding(self, cache_key: str) -> object: + session_cache: Final = self._claude_code_session_router_cache + try: + if session_cache.redis_cache is None: + return await session_cache.async_get_cache(key=cache_key) + return await session_cache.redis_cache.async_get_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # an optional binding must not make routing depend on Redis + verbose_router_logger.warning( + "Failed to read Claude Code session router binding; using the requested model: %s", + e, + ) + return None + + async def _resolve_claude_code_session_router( + self, + model: str, + registered_model_name: str, + request_kwargs: Mapping[str, object], + ) -> str: + if not any((self.auto_routers, self.complexity_routers, self.adaptive_routers, self.quality_routers)): + return registered_model_name + cache_key: Final = self._claude_code_session_router_cache_key(request_kwargs) + if cache_key is None or not isinstance(request_kwargs, dict): + return registered_model_name + if request_kwargs.get("fallback_depth") not in (None, 0): + return registered_model_name + + agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") + if agent_id is not None: + bound_model: Final = await self._get_claude_code_session_router_binding(cache_key) + if not isinstance(bound_model, str): + return registered_model_name + bound_registered_model: Final = self._get_model_from_alias(model=bound_model) or bound_model + if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: + await self._delete_claude_code_session_router_binding(cache_key) + return registered_model_name + await self._claude_code_session_router_cache.async_set_cache( + key=cache_key, + value=bound_model, + ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, + ) + self._stamp_or_clear_metadata_key(request_kwargs, "model_group", bound_model) + return bound_registered_model + + if self._request_header(request_kwargs, "x-app") != "cli": + return registered_model_name + if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: + return registered_model_name + await self._claude_code_session_router_cache.async_set_cache( + key=cache_key, + value=model, + ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, + ) + return registered_model_name + async def async_pre_routing_hook( self, model: str, @@ -12670,7 +12802,12 @@ class Router: the alias, since spend metadata is stamped before routing and the response carries the tier group the strategy picked. """ - registered_model_name: Final = self._get_model_from_alias(model=model) or model + requested_registered_model_name: Final = self._get_model_from_alias(model=model) or model + registered_model_name: Final = await self._resolve_claude_code_session_router( + model=model, + registered_model_name=requested_registered_model_name, + request_kwargs=request_kwargs, + ) ######################################################### # Run the routing-plugin pipeline, if any plugins are configured. @@ -13467,6 +13604,9 @@ class Router: def flush_cache(self): litellm.cache = None self.cache.flush_cache() + session_in_memory_cache: Final = self._claude_code_session_router_cache.in_memory_cache + if session_in_memory_cache is not None: + session_in_memory_cache.flush_cache() def reset(self): ## clean up on close diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index bc8df67cc28..ad8b67d5e8f 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -68,6 +68,36 @@ still resolve to a deployment in `model_list`; this configuration does not creat - abc ``` +### Heuristic v2 + +Set `classifier_type: heuristic_v2` to classify with the bundled calibrated +success-probability model instead of the hand-written weighted scorer + +```yaml +model_list: + - model_name: smart-router + litellm_params: + model: auto_router/complexity_router + complexity_router_config: + classifier_type: heuristic_v2 + tiers: + SIMPLE: luna + MEDIUM: terra + COMPLEX: sol + REASONING: sol-ultra +``` + +No classifier model call or per-model training data is required. The classifier +uses global tier quality, request-type quality, and similar-request cohorts from +the bundled UltraFeedback artifact. It estimates success at every tier, enforces +monotonic probabilities, and returns the first tier meeting the trained 0.75 +threshold. The existing complexity-router tier pool then selects and dispatches +a model from that tier + +Spend logs record `routing_decision.cause: heuristic_v2`, the detected request +type, and all four predicted probabilities. Existing `classifier_type: heuristic` +configurations keep the original weighted scorer unchanged + ### Renaming the tiers `tier_labels` puts your own vocabulary on the four tiers: @@ -244,6 +274,49 @@ except that the heuristic outcome is the one already computed rather than a seco Spend logs record `routing_decision.cause` as `heuristic_first_short_circuit` when the classifier was skipped, and `llm_classifier` when it ran, so the two are told apart per request. +### Hybrid + +`classifier_type: hybrid` also scores locally first, but it asks a different question than +`heuristic_first`. Where heuristic-first asks how CHEAP the scorer's tier is and pays for the +classifier on everything above a ceiling, hybrid asks how DECIDED the score is and pays for the +classifier only where the score lands near a tier boundary. A confident score keeps its tier at +every tier, the most expensive one included: + +```yaml +model_list: + - model_name: smart-router + litellm_params: + model: auto_router/complexity_router + complexity_router_config: + classifier_type: hybrid + hybrid_boundary_margin: 0.03 + classifier_llm_config: + model: gpt-4o-mini + tiers: + SIMPLE: gpt-4o-mini + MEDIUM: gpt-4o + COMPLEX: claude-sonnet-4 + REASONING: o1-preview +``` + +A request routes on the scorer's own tier when its score is further than `hybrid_boundary_margin` +from every active boundary. Everything else goes to the classifier: a score inside the band, where a +hair's difference would have named the adjacent tier and its model pool, and a prompt where no +dimension fired at all, which has no opinion to be confident about. `hybrid_boundary_margin` is +required for this type and rejected on the others, the same way `heuristic_first_max_tier` is +required for heuristic-first, so the two modes are told apart by the knob each one takes rather than +by a shared field that means something different per type. + +Pick the margin against the score distribution rather than by intuition. The scorer combines a small +set of discretely weighted dimensions, so achievable scores cluster on a lumpy grid instead of +spreading smoothly, and widening the margin admits whole clusters at once rather than a few more +requests. Spend logs record `routing_decision.cause` as `hybrid_short_circuit` when the classifier +was skipped and `llm_classifier` when it ran. + +Operator-defined tier sets (`tier_definitions`) are not supported here, for the same reason they are +not supported under heuristic-first: the scorer only produces the built-in tiers. Classifier failure +behaves exactly as it does under `classifier_type: llm`. + ### Reasoning Override If 2+ reasoning markers are detected in the user message, the request is promoted to the REASONING tier even when the weighted score maps lower, so complex reasoning tasks get the appropriate model. The promotion requires the score to reach `reasoning_override_min_score`, which tracks `tier_boundaries.simple_medium` unless set, so stock phrases on an otherwise trivial prompt cannot buy the top tier. Set it to `0` to promote on the markers alone. diff --git a/litellm/router_strategy/complexity_router/artifacts/ultrafeedback_tiers.json b/litellm/router_strategy/complexity_router/artifacts/ultrafeedback_tiers.json new file mode 100644 index 00000000000..4fcb599907c --- /dev/null +++ b/litellm/router_strategy/complexity_router/artifacts/ultrafeedback_tiers.json @@ -0,0 +1,4069 @@ +{ + "schema_version": 1, + "global_statistics": [ + { + "tier": 1, + "successes": 36619.0, + "observations": 45504.0 + }, + { + "tier": 2, + "successes": 59797.0, + "observations": 70062.0 + }, + { + "tier": 3, + "successes": 48604.0, + "observations": 52245.0 + }, + { + "tier": 4, + "successes": 11393.0, + "observations": 11561.0 + } + ], + "domain_statistics": [ + { + "tier": 1, + "successes": 1592.0, + "observations": 2211.0, + "request_type": "analytical_reasoning" + }, + { + "tier": 2, + "successes": 2654.0, + "observations": 3374.0, + "request_type": "analytical_reasoning" + }, + { + "tier": 3, + "successes": 2243.0, + "observations": 2481.0, + "request_type": "analytical_reasoning" + }, + { + "tier": 4, + "successes": 538.0, + "observations": 546.0, + "request_type": "analytical_reasoning" + }, + { + "tier": 1, + "successes": 750.0, + "observations": 1015.0, + "request_type": "code_generation" + }, + { + "tier": 2, + "successes": 1271.0, + "observations": 1511.0, + "request_type": "code_generation" + }, + { + "tier": 3, + "successes": 1030.0, + "observations": 1111.0, + "request_type": "code_generation" + }, + { + "tier": 4, + "successes": 233.0, + "observations": 235.0, + "request_type": "code_generation" + }, + { + "tier": 1, + "successes": 243.0, + "observations": 277.0, + "request_type": "code_understanding" + }, + { + "tier": 2, + "successes": 385.0, + "observations": 425.0, + "request_type": "code_understanding" + }, + { + "tier": 3, + "successes": 322.0, + "observations": 334.0, + "request_type": "code_understanding" + }, + { + "tier": 4, + "successes": 74.0, + "observations": 76.0, + "request_type": "code_understanding" + }, + { + "tier": 1, + "successes": 2014.0, + "observations": 2170.0, + "request_type": "factual_lookup" + }, + { + "tier": 2, + "successes": 3120.0, + "observations": 3266.0, + "request_type": "factual_lookup" + }, + { + "tier": 3, + "successes": 2612.0, + "observations": 2670.0, + "request_type": "factual_lookup" + }, + { + "tier": 4, + "successes": 540.0, + "observations": 542.0, + "request_type": "factual_lookup" + }, + { + "tier": 1, + "successes": 30571.0, + "observations": 38161.0, + "request_type": "general" + }, + { + "tier": 2, + "successes": 50037.0, + "observations": 58821.0, + "request_type": "general" + }, + { + "tier": 3, + "successes": 40460.0, + "observations": 43618.0, + "request_type": "general" + }, + { + "tier": 4, + "successes": 9565.0, + "observations": 9716.0, + "request_type": "general" + }, + { + "tier": 1, + "successes": 159.0, + "observations": 170.0, + "request_type": "technical_design" + }, + { + "tier": 2, + "successes": 282.0, + "observations": 303.0, + "request_type": "technical_design" + }, + { + "tier": 3, + "successes": 231.0, + "observations": 236.0, + "request_type": "technical_design" + }, + { + "tier": 4, + "successes": 47.0, + "observations": 47.0, + "request_type": "technical_design" + }, + { + "tier": 1, + "successes": 1290.0, + "observations": 1500.0, + "request_type": "writing" + }, + { + "tier": 2, + "successes": 2048.0, + "observations": 2362.0, + "request_type": "writing" + }, + { + "tier": 3, + "successes": 1706.0, + "observations": 1795.0, + "request_type": "writing" + }, + { + "tier": 4, + "successes": 396.0, + "observations": 399.0, + "request_type": "writing" + } + ], + "cohort_statistics": [ + { + "tier": 1, + "successes": 272.0, + "observations": 372.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 420.0, + "observations": 519.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 341.0, + "observations": 381.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 82.0, + "observations": 84.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 0.0, + "observations": 1.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 0.0, + "observations": 2.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 11.0, + "observations": 18.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 33.0, + "observations": 45.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 28.0, + "observations": 34.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "analytical_reasoning|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 131.0, + "observations": 176.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 210.0, + "observations": 274.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 187.0, + "observations": 209.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 39.0, + "observations": 41.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 9.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 11.0, + "observations": 13.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 6.0, + "observations": 7.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 10.0, + "observations": 15.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 14.0, + "observations": 16.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 15.0, + "observations": 16.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "analytical_reasoning|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 15.0, + "observations": 20.0, + "cohort": "analytical_reasoning|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 33.0, + "observations": 38.0, + "cohort": "analytical_reasoning|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 28.0, + "observations": 29.0, + "cohort": "analytical_reasoning|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "analytical_reasoning|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 2.0, + "cohort": "analytical_reasoning|long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "analytical_reasoning|long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 26.0, + "observations": 34.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 39.0, + "observations": 50.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 42.0, + "observations": 45.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 11.0, + "observations": 11.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 3.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 4.0, + "observations": 4.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 452.0, + "observations": 634.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 745.0, + "observations": 962.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 634.0, + "observations": 691.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 151.0, + "observations": 153.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 3.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 4.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 10.0, + "observations": 20.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 30.0, + "observations": 44.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 12.0, + "observations": 15.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "analytical_reasoning|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 178.0, + "observations": 270.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 304.0, + "observations": 402.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 266.0, + "observations": 295.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 73.0, + "observations": 73.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 7.0, + "observations": 18.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 17.0, + "observations": 32.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 17.0, + "observations": 27.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 7.0, + "observations": 7.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 11.0, + "observations": 17.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 17.0, + "observations": 25.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 10.0, + "observations": 16.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 6.0, + "observations": 6.0, + "cohort": "analytical_reasoning|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 58.0, + "observations": 72.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 90.0, + "observations": 103.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 72.0, + "observations": 78.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 19.0, + "observations": 19.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 37.0, + "observations": 56.0, + "cohort": "analytical_reasoning|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 95.0, + "observations": 110.0, + "cohort": "analytical_reasoning|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 73.0, + "observations": 82.0, + "cohort": "analytical_reasoning|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 20.0, + "observations": 20.0, + "cohort": "analytical_reasoning|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 69.0, + "observations": 83.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 102.0, + "observations": 110.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 93.0, + "observations": 98.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 17.0, + "observations": 17.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 6.0, + "observations": 6.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 21.0, + "observations": 32.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 45.0, + "observations": 59.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 45.0, + "observations": 48.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 9.0, + "observations": 9.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|short|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 7.0, + "observations": 9.0, + "cohort": "analytical_reasoning|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 8.0, + "observations": 12.0, + "cohort": "analytical_reasoning|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 11.0, + "observations": 11.0, + "cohort": "analytical_reasoning|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 5.0, + "observations": 9.0, + "cohort": "analytical_reasoning|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 9.0, + "observations": 11.0, + "cohort": "analytical_reasoning|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 8.0, + "observations": 8.0, + "cohort": "analytical_reasoning|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 136.0, + "observations": 176.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 221.0, + "observations": 267.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 172.0, + "observations": 194.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 37.0, + "observations": 39.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 6.0, + "observations": 11.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 10.0, + "observations": 15.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 8.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "analytical_reasoning|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 68.0, + "observations": 84.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 103.0, + "observations": 136.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 113.0, + "observations": 119.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 21.0, + "observations": 21.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 3.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 4.0, + "cohort": "analytical_reasoning|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 13.0, + "observations": 16.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 27.0, + "observations": 34.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 15.0, + "observations": 16.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 6.0, + "observations": 6.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 5.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 4.0, + "observations": 7.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 32.0, + "observations": 39.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 48.0, + "observations": 63.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 35.0, + "observations": 40.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 14.0, + "observations": 14.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 3.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 2.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 2.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "analytical_reasoning|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 28.0, + "observations": 31.0, + "cohort": "code_generation|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 41.0, + "observations": 46.0, + "cohort": "code_generation|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 25.0, + "observations": 26.0, + "cohort": "code_generation|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "code_generation|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 6.0, + "observations": 9.0, + "cohort": "code_generation|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 20.0, + "observations": 23.0, + "cohort": "code_generation|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 16.0, + "observations": 16.0, + "cohort": "code_generation|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "code_generation|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 46.0, + "observations": 60.0, + "cohort": "code_generation|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 72.0, + "observations": 91.0, + "cohort": "code_generation|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 59.0, + "observations": 64.0, + "cohort": "code_generation|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 17.0, + "observations": 17.0, + "cohort": "code_generation|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 32.0, + "observations": 49.0, + "cohort": "code_generation|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 60.0, + "observations": 76.0, + "cohort": "code_generation|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 48.0, + "observations": 52.0, + "cohort": "code_generation|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 15.0, + "observations": 15.0, + "cohort": "code_generation|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 93.0, + "observations": 121.0, + "cohort": "code_generation|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 149.0, + "observations": 170.0, + "cohort": "code_generation|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 146.0, + "observations": 151.0, + "cohort": "code_generation|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 26.0, + "observations": 26.0, + "cohort": "code_generation|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 12.0, + "observations": 16.0, + "cohort": "code_generation|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 22.0, + "observations": 25.0, + "cohort": "code_generation|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 13.0, + "observations": 15.0, + "cohort": "code_generation|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "code_generation|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 196.0, + "observations": 268.0, + "cohort": "code_generation|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 312.0, + "observations": 370.0, + "cohort": "code_generation|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 234.0, + "observations": 253.0, + "cohort": "code_generation|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 56.0, + "observations": 57.0, + "cohort": "code_generation|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 0.0, + "observations": 1.0, + "cohort": "code_generation|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_generation|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_generation|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 40.0, + "observations": 59.0, + "cohort": "code_generation|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 73.0, + "observations": 92.0, + "cohort": "code_generation|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 78.0, + "observations": 83.0, + "cohort": "code_generation|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 14.0, + "observations": 14.0, + "cohort": "code_generation|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 3.0, + "cohort": "code_generation|medium|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_generation|medium|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 106.0, + "observations": 140.0, + "cohort": "code_generation|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 220.0, + "observations": 247.0, + "cohort": "code_generation|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 167.0, + "observations": 178.0, + "cohort": "code_generation|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 43.0, + "observations": 43.0, + "cohort": "code_generation|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 4.0, + "cohort": "code_generation|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 7.0, + "observations": 9.0, + "cohort": "code_generation|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 7.0, + "cohort": "code_generation|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 147.0, + "observations": 199.0, + "cohort": "code_generation|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 233.0, + "observations": 269.0, + "cohort": "code_generation|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 189.0, + "observations": 209.0, + "cohort": "code_generation|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 34.0, + "observations": 35.0, + "cohort": "code_generation|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 7.0, + "observations": 10.0, + "cohort": "code_generation|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 18.0, + "observations": 22.0, + "cohort": "code_generation|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 10.0, + "observations": 10.0, + "cohort": "code_generation|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_generation|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 4.0, + "cohort": "code_generation|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 5.0, + "observations": 8.0, + "cohort": "code_generation|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 8.0, + "cohort": "code_generation|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_generation|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 3.0, + "observations": 5.0, + "cohort": "code_generation|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "code_generation|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_generation|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 7.0, + "observations": 9.0, + "cohort": "code_generation|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 7.0, + "observations": 8.0, + "cohort": "code_generation|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 9.0, + "observations": 9.0, + "cohort": "code_generation|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_generation|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 24.0, + "observations": 33.0, + "cohort": "code_generation|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 25.0, + "observations": 45.0, + "cohort": "code_generation|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 19.0, + "observations": 27.0, + "cohort": "code_generation|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 7.0, + "observations": 7.0, + "cohort": "code_generation|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_understanding|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_understanding|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 6.0, + "observations": 6.0, + "cohort": "code_understanding|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_understanding|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 5.0, + "cohort": "code_understanding|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 3.0, + "observations": 3.0, + "cohort": "code_understanding|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 3.0, + "cohort": "code_understanding|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 17.0, + "observations": 23.0, + "cohort": "code_understanding|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 27.0, + "observations": 37.0, + "cohort": "code_understanding|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 32.0, + "observations": 33.0, + "cohort": "code_understanding|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "code_understanding|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 24.0, + "observations": 28.0, + "cohort": "code_understanding|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 37.0, + "observations": 44.0, + "cohort": "code_understanding|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 33.0, + "observations": 33.0, + "cohort": "code_understanding|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 7.0, + "observations": 7.0, + "cohort": "code_understanding|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 8.0, + "observations": 9.0, + "cohort": "code_understanding|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 19.0, + "observations": 21.0, + "cohort": "code_understanding|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 10.0, + "observations": 10.0, + "cohort": "code_understanding|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "code_understanding|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 3.0, + "cohort": "code_understanding|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 3.0, + "cohort": "code_understanding|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 73.0, + "observations": 76.0, + "cohort": "code_understanding|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 107.0, + "observations": 111.0, + "cohort": "code_understanding|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 87.0, + "observations": 87.0, + "cohort": "code_understanding|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 22.0, + "observations": 22.0, + "cohort": "code_understanding|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 32.0, + "observations": 33.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 44.0, + "observations": 44.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 44.0, + "observations": 47.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 8.0, + "observations": 8.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 0.0, + "observations": 1.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 0.0, + "observations": 2.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|medium|code=1|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 15.0, + "observations": 16.0, + "cohort": "code_understanding|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 23.0, + "observations": 23.0, + "cohort": "code_understanding|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 12.0, + "observations": 14.0, + "cohort": "code_understanding|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "code_understanding|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 3.0, + "cohort": "code_understanding|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 36.0, + "observations": 43.0, + "cohort": "code_understanding|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 71.0, + "observations": 74.0, + "cohort": "code_understanding|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 48.0, + "observations": 50.0, + "cohort": "code_understanding|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 12.0, + "observations": 13.0, + "cohort": "code_understanding|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 5.0, + "observations": 6.0, + "cohort": "code_understanding|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 10.0, + "observations": 10.0, + "cohort": "code_understanding|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 7.0, + "cohort": "code_understanding|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 0.0, + "observations": 1.0, + "cohort": "code_understanding|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "code_understanding|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 9.0, + "observations": 11.0, + "cohort": "code_understanding|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 18.0, + "observations": 23.0, + "cohort": "code_understanding|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 11.0, + "observations": 13.0, + "cohort": "code_understanding|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 3.0, + "cohort": "code_understanding|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "code_understanding|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 13.0, + "observations": 17.0, + "cohort": "code_understanding|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 19.0, + "observations": 23.0, + "cohort": "code_understanding|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 27.0, + "observations": 28.0, + "cohort": "code_understanding|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 7.0, + "observations": 8.0, + "cohort": "code_understanding|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 49.0, + "observations": 54.0, + "cohort": "factual_lookup|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 64.0, + "observations": 75.0, + "cohort": "factual_lookup|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 80.0, + "observations": 80.0, + "cohort": "factual_lookup|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 15.0, + "observations": 15.0, + "cohort": "factual_lookup|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 9.0, + "observations": 9.0, + "cohort": "factual_lookup|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 12.0, + "observations": 14.0, + "cohort": "factual_lookup|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 9.0, + "observations": 10.0, + "cohort": "factual_lookup|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "factual_lookup|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 3.0, + "observations": 3.0, + "cohort": "factual_lookup|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 19.0, + "observations": 22.0, + "cohort": "factual_lookup|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 35.0, + "observations": 38.0, + "cohort": "factual_lookup|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 37.0, + "observations": 41.0, + "cohort": "factual_lookup|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 7.0, + "observations": 7.0, + "cohort": "factual_lookup|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 66.0, + "observations": 75.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 73.0, + "observations": 86.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 69.0, + "observations": 74.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 21.0, + "observations": 21.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 174.0, + "observations": 181.0, + "cohort": "factual_lookup|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 204.0, + "observations": 215.0, + "cohort": "factual_lookup|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 206.0, + "observations": 210.0, + "cohort": "factual_lookup|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 42.0, + "observations": 42.0, + "cohort": "factual_lookup|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 31.0, + "observations": 37.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 41.0, + "observations": 48.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 44.0, + "observations": 45.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 10.0, + "observations": 10.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 0.0, + "observations": 1.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 131.0, + "observations": 142.0, + "cohort": "factual_lookup|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 162.0, + "observations": 171.0, + "cohort": "factual_lookup|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 146.0, + "observations": 151.0, + "cohort": "factual_lookup|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 24.0, + "observations": 24.0, + "cohort": "factual_lookup|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 41.0, + "observations": 46.0, + "cohort": "factual_lookup|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 55.0, + "observations": 59.0, + "cohort": "factual_lookup|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 45.0, + "observations": 46.0, + "cohort": "factual_lookup|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 9.0, + "observations": 9.0, + "cohort": "factual_lookup|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1382.0, + "observations": 1473.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 2290.0, + "observations": 2348.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 1824.0, + "observations": 1853.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 368.0, + "observations": 370.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 5.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 12.0, + "observations": 13.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 7.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "factual_lookup|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 11.0, + "observations": 16.0, + "cohort": "factual_lookup|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 28.0, + "observations": 31.0, + "cohort": "factual_lookup|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 15.0, + "observations": 17.0, + "cohort": "factual_lookup|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "factual_lookup|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 36.0, + "observations": 39.0, + "cohort": "factual_lookup|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 63.0, + "observations": 66.0, + "cohort": "factual_lookup|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 64.0, + "observations": 66.0, + "cohort": "factual_lookup|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 13.0, + "observations": 13.0, + "cohort": "factual_lookup|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 4.0, + "cohort": "factual_lookup|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 5.0, + "observations": 6.0, + "cohort": "factual_lookup|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 5.0, + "cohort": "factual_lookup|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 27.0, + "observations": 31.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 29.0, + "observations": 37.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 26.0, + "observations": 29.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 11.0, + "observations": 11.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 0.0, + "observations": 2.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 7.0, + "observations": 8.0, + "cohort": "factual_lookup|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 17.0, + "observations": 20.0, + "cohort": "factual_lookup|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 14.0, + "observations": 14.0, + "cohort": "factual_lookup|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "factual_lookup|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 4.0, + "cohort": "factual_lookup|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 6.0, + "observations": 6.0, + "cohort": "factual_lookup|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 5.0, + "observations": 5.0, + "cohort": "factual_lookup|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "factual_lookup|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 17.0, + "observations": 21.0, + "cohort": "factual_lookup|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 19.0, + "observations": 25.0, + "cohort": "factual_lookup|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 13.0, + "observations": 13.0, + "cohort": "factual_lookup|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "factual_lookup|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2828.0, + "observations": 3552.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 4621.0, + "observations": 5551.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 3628.0, + "observations": 3931.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 911.0, + "observations": 922.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 127.0, + "observations": 406.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 250.0, + "observations": 583.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 219.0, + "observations": 396.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 98.0, + "observations": 107.0, + "cohort": "general|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 94.0, + "observations": 127.0, + "cohort": "general|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 168.0, + "observations": 196.0, + "cohort": "general|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 135.0, + "observations": 146.0, + "cohort": "general|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 34.0, + "observations": 35.0, + "cohort": "general|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 912.0, + "observations": 1219.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1417.0, + "observations": 1736.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 1163.0, + "observations": 1258.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 278.0, + "observations": 283.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 33.0, + "observations": 122.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 41.0, + "observations": 150.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 61.0, + "observations": 96.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 24.0, + "observations": 24.0, + "cohort": "general|long|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 16.0, + "observations": 24.0, + "cohort": "general|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 24.0, + "observations": 34.0, + "cohort": "general|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 18.0, + "observations": 21.0, + "cohort": "general|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "general|long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 271.0, + "observations": 348.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 488.0, + "observations": 555.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 369.0, + "observations": 386.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 67.0, + "observations": 71.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 3.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 0.0, + "observations": 3.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|long|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 9.0, + "observations": 9.0, + "cohort": "general|long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 14.0, + "observations": 17.0, + "cohort": "general|long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 9.0, + "observations": 10.0, + "cohort": "general|long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "general|long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 340.0, + "observations": 426.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 544.0, + "observations": 635.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 472.0, + "observations": 491.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 120.0, + "observations": 120.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|long|code=1|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 0.0, + "observations": 2.0, + "cohort": "general|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 0.0, + "observations": 1.0, + "cohort": "general|long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 8332.0, + "observations": 10133.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 13225.0, + "observations": 15509.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 10614.0, + "observations": 11347.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2499.0, + "observations": 2531.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 437.0, + "observations": 1294.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 863.0, + "observations": 1932.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 738.0, + "observations": 1269.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 296.0, + "observations": 325.0, + "cohort": "general|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 117.0, + "observations": 173.0, + "cohort": "general|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 183.0, + "observations": 252.0, + "cohort": "general|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 171.0, + "observations": 191.0, + "cohort": "general|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 47.0, + "observations": 52.0, + "cohort": "general|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 927.0, + "observations": 1353.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1550.0, + "observations": 2023.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 1273.0, + "observations": 1430.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 348.0, + "observations": 354.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 60.0, + "observations": 255.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 143.0, + "observations": 377.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 145.0, + "observations": 243.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 58.0, + "observations": 61.0, + "cohort": "general|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 27.0, + "observations": 47.0, + "cohort": "general|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 30.0, + "observations": 48.0, + "cohort": "general|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 30.0, + "observations": 36.0, + "cohort": "general|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "general|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 1119.0, + "observations": 1348.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1816.0, + "observations": 2024.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 1472.0, + "observations": 1552.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 301.0, + "observations": 304.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|medium|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 8.0, + "observations": 10.0, + "cohort": "general|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 9.0, + "observations": 12.0, + "cohort": "general|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 8.0, + "observations": 9.0, + "cohort": "general|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|medium|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 349.0, + "observations": 439.0, + "cohort": "general|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 549.0, + "observations": 622.0, + "cohort": "general|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 495.0, + "observations": 526.0, + "cohort": "general|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 111.0, + "observations": 113.0, + "cohort": "general|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|medium|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 4.0, + "observations": 6.0, + "cohort": "general|medium|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 4.0, + "cohort": "general|medium|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|medium|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 11635.0, + "observations": 12910.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 19249.0, + "observations": 20591.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 15415.0, + "observations": 15867.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3362.0, + "observations": 3392.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 68.0, + "observations": 155.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 113.0, + "observations": 202.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 97.0, + "observations": 124.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 31.0, + "observations": 31.0, + "cohort": "general|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 6.0, + "observations": 11.0, + "cohort": "general|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 20.0, + "observations": 21.0, + "cohort": "general|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 23.0, + "observations": 24.0, + "cohort": "general|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "general|short|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 199.0, + "observations": 270.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 334.0, + "observations": 414.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 309.0, + "observations": 337.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 75.0, + "observations": 75.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 5.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 5.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|short|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|short|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|short|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "general|short|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 629.0, + "observations": 778.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1131.0, + "observations": 1286.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 919.0, + "observations": 989.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 227.0, + "observations": 227.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 4.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 3.0, + "observations": 3.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 5.0, + "cohort": "general|short|code=1|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|short|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "general|short|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 28.0, + "observations": 37.0, + "cohort": "general|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 54.0, + "observations": 64.0, + "cohort": "general|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 41.0, + "observations": 46.0, + "cohort": "general|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 9.0, + "observations": 9.0, + "cohort": "general|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1066.0, + "observations": 1437.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1639.0, + "observations": 2007.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 1397.0, + "observations": 1507.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 333.0, + "observations": 341.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 6.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 4.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "general|very_long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 65.0, + "observations": 89.0, + "cohort": "general|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 119.0, + "observations": 139.0, + "cohort": "general|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 98.0, + "observations": 103.0, + "cohort": "general|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 21.0, + "observations": 21.0, + "cohort": "general|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 403.0, + "observations": 546.0, + "cohort": "general|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 700.0, + "observations": 875.0, + "cohort": "general|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 556.0, + "observations": 612.0, + "cohort": "general|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 133.0, + "observations": 135.0, + "cohort": "general|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 17.0, + "observations": 27.0, + "cohort": "general|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 19.0, + "observations": 33.0, + "cohort": "general|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 25.0, + "observations": 27.0, + "cohort": "general|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 9.0, + "observations": 9.0, + "cohort": "general|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 210.0, + "observations": 280.0, + "cohort": "general|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 307.0, + "observations": 378.0, + "cohort": "general|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 258.0, + "observations": 289.0, + "cohort": "general|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 64.0, + "observations": 65.0, + "cohort": "general|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 17.0, + "observations": 23.0, + "cohort": "general|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 31.0, + "observations": 39.0, + "cohort": "general|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 16.0, + "observations": 19.0, + "cohort": "general|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "general|very_long|code=1|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 202.0, + "observations": 282.0, + "cohort": "general|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 361.0, + "observations": 478.0, + "cohort": "general|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 262.0, + "observations": 310.0, + "cohort": "general|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 82.0, + "observations": 82.0, + "cohort": "general|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 5.0, + "cohort": "general|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 9.0, + "observations": 11.0, + "cohort": "general|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 4.0, + "cohort": "general|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 19.0, + "observations": 21.0, + "cohort": "technical_design|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 23.0, + "observations": 25.0, + "cohort": "technical_design|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 28.0, + "observations": 29.0, + "cohort": "technical_design|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "technical_design|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 8.0, + "observations": 9.0, + "cohort": "technical_design|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 9.0, + "observations": 9.0, + "cohort": "technical_design|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 14.0, + "observations": 15.0, + "cohort": "technical_design|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 24.0, + "observations": 25.0, + "cohort": "technical_design|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 18.0, + "observations": 18.0, + "cohort": "technical_design|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 44.0, + "observations": 45.0, + "cohort": "technical_design|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 84.0, + "observations": 86.0, + "cohort": "technical_design|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 74.0, + "observations": 74.0, + "cohort": "technical_design|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 15.0, + "observations": 15.0, + "cohort": "technical_design|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 6.0, + "observations": 6.0, + "cohort": "technical_design|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 23.0, + "observations": 24.0, + "cohort": "technical_design|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 47.0, + "observations": 53.0, + "cohort": "technical_design|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 28.0, + "observations": 29.0, + "cohort": "technical_design|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 10.0, + "observations": 10.0, + "cohort": "technical_design|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 4.0, + "cohort": "technical_design|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 7.0, + "observations": 7.0, + "cohort": "technical_design|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 8.0, + "cohort": "technical_design|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 13.0, + "observations": 13.0, + "cohort": "technical_design|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 31.0, + "observations": 31.0, + "cohort": "technical_design|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 23.0, + "observations": 23.0, + "cohort": "technical_design|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 5.0, + "observations": 5.0, + "cohort": "technical_design|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 3.0, + "observations": 4.0, + "cohort": "technical_design|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 10.0, + "observations": 15.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 26.0, + "observations": 30.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 19.0, + "observations": 21.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 5.0, + "observations": 6.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 6.0, + "observations": 7.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 6.0, + "observations": 6.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 4.0, + "observations": 6.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 8.0, + "observations": 8.0, + "cohort": "technical_design|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 5.0, + "observations": 6.0, + "cohort": "technical_design|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 5.0, + "observations": 5.0, + "cohort": "technical_design|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "technical_design|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "technical_design|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "technical_design|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 176.0, + "observations": 226.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 267.0, + "observations": 329.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 229.0, + "observations": 245.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 51.0, + "observations": 52.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 0.0, + "observations": 2.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 2.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|long|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 5.0, + "cohort": "writing|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 4.0, + "observations": 5.0, + "cohort": "writing|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 36.0, + "observations": 47.0, + "cohort": "writing|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 53.0, + "observations": 59.0, + "cohort": "writing|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 47.0, + "observations": 51.0, + "cohort": "writing|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 14.0, + "observations": 15.0, + "cohort": "writing|long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 17.0, + "observations": 21.0, + "cohort": "writing|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 27.0, + "observations": 32.0, + "cohort": "writing|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 24.0, + "observations": 24.0, + "cohort": "writing|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "writing|long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 2.0, + "cohort": "writing|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 2.0, + "cohort": "writing|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 6.0, + "observations": 6.0, + "cohort": "writing|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 541.0, + "observations": 598.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 892.0, + "observations": 997.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 728.0, + "observations": 756.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 164.0, + "observations": 165.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 7.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 7.0, + "observations": 15.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 10.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 4.0, + "observations": 4.0, + "cohort": "writing|medium|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|medium|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 25.0, + "observations": 35.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 41.0, + "observations": 63.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 39.0, + "observations": 46.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 12.0, + "observations": 12.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 2.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 4.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|medium|code=0|math=1|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|medium|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 24.0, + "observations": 27.0, + "cohort": "writing|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 49.0, + "observations": 56.0, + "cohort": "writing|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 46.0, + "observations": 49.0, + "cohort": "writing|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 8.0, + "observations": 8.0, + "cohort": "writing|medium|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 11.0, + "observations": 13.0, + "cohort": "writing|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 13.0, + "observations": 13.0, + "cohort": "writing|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 10.0, + "observations": 10.0, + "cohort": "writing|medium|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 327.0, + "observations": 353.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 508.0, + "observations": 555.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 415.0, + "observations": 432.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 92.0, + "observations": 92.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|short|code=0|math=0|mc=0|intl=1" + }, + { + "tier": 1, + "successes": 4.0, + "observations": 4.0, + "cohort": "writing|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 7.0, + "observations": 7.0, + "cohort": "writing|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 3.0, + "observations": 3.0, + "cohort": "writing|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|short|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 3.0, + "cohort": "writing|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 8.0, + "observations": 9.0, + "cohort": "writing|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 4.0, + "cohort": "writing|short|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|short|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 63.0, + "observations": 85.0, + "cohort": "writing|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 111.0, + "observations": 139.0, + "cohort": "writing|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 95.0, + "observations": 99.0, + "cohort": "writing|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 29.0, + "observations": 29.0, + "cohort": "writing|very_long|code=0|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 3.0, + "observations": 5.0, + "cohort": "writing|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 2.0, + "cohort": "writing|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 4.0, + "observations": 4.0, + "cohort": "writing|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 4, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|very_long|code=0|math=0|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 26.0, + "observations": 36.0, + "cohort": "writing|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 30.0, + "observations": 44.0, + "cohort": "writing|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 23.0, + "observations": 28.0, + "cohort": "writing|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 8.0, + "observations": 8.0, + "cohort": "writing|very_long|code=0|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 1.0, + "observations": 1.0, + "cohort": "writing|very_long|code=0|math=1|mc=1|intl=0" + }, + { + "tier": 1, + "successes": 14.0, + "observations": 14.0, + "cohort": "writing|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 14.0, + "observations": 16.0, + "cohort": "writing|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 6.0, + "observations": 7.0, + "cohort": "writing|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 4, + "successes": 3.0, + "observations": 3.0, + "cohort": "writing|very_long|code=1|math=0|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 5.0, + "observations": 5.0, + "cohort": "writing|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 2, + "successes": 7.0, + "observations": 8.0, + "cohort": "writing|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 3, + "successes": 7.0, + "observations": 7.0, + "cohort": "writing|very_long|code=1|math=1|mc=0|intl=0" + }, + { + "tier": 1, + "successes": 0.0, + "observations": 1.0, + "cohort": "writing|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 2, + "successes": 2.0, + "observations": 2.0, + "cohort": "writing|very_long|code=1|math=1|mc=1|intl=0" + }, + { + "tier": 3, + "successes": 0.0, + "observations": 1.0, + "cohort": "writing|very_long|code=1|math=1|mc=1|intl=0" + } + ], + "domain_prior_mass": 200.0, + "cohort_prior_mass": 20.0, + "routing_threshold": 0.75, + "datasets": [ + { + "name": "openbmb/UltraFeedback", + "url": "https://huggingface.co/datasets/openbmb/UltraFeedback", + "license": "MIT", + "rows": 255864, + "success_definition": "UltraFeedback overall_score >= 4" + } + ], + "success_definition": "UltraFeedback overall_score >= 4", + "split_method": "sha256(prompt): 70% train, 15% validation, 15% test" +} diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 577cee0920d..d205db90607 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -33,6 +33,11 @@ from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload from litellm.llms.base_llm.base_utils import type_to_response_format_param +from litellm.router_strategy.adaptive_router.classifier import classify_prompt +from litellm.router_strategy.complexity_router.tier_predictor import ( + TierSuccessPredictor, + resolve_tier_artifact, +) from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, ModelResponse, @@ -790,9 +795,11 @@ class ClassificationOutcome(NamedTuple): signals: tuple[str, ...] cause: Literal[ "heuristic_scorer", + "heuristic_v2", "reasoning_override", "llm_classifier", "heuristic_first_short_circuit", + "hybrid_short_circuit", "housekeeping", "classifier_plugin", "classifier_fallback", @@ -978,6 +985,11 @@ class ComplexityRouter(CustomLogger): if llm_classifier_configured else None ) + self._tier_success_predictor: TierSuccessPredictor | None = ( + TierSuccessPredictor(resolve_tier_artifact(self.config.heuristic_v2_artifact)) + if self.config.classifier_type == "heuristic_v2" + else None + ) verbose_router_logger.debug("ComplexityRouter initialized for %s with tiers: %s", model_name, self.config.tiers) @@ -1230,6 +1242,15 @@ class ComplexityRouter(CustomLogger): return tier, weighted_score, tuple(signals), "heuristic_scorer" + def _is_near_tier_boundary(self, score: float, margin: float) -> bool: + boundaries: Final = self._effective_tier_boundaries() + active_boundaries: Final = ( + boundaries["simple_medium"], + boundaries["medium_complex"], + boundaries["complex_reasoning"], + ) + return any(abs(score - boundary) <= margin for boundary in active_boundaries) + def _effective_reasoning_override_min_score(self) -> float: """The score a request must reach before the reasoning-marker override may promote it. @@ -1350,15 +1371,37 @@ class ComplexityRouter(CustomLogger): custom tier set, and classifier_fallback otherwise decides between the heuristic scorer and default_model. The outcome's `cause` reports which path actually ran. """ + if self.config.classifier_type == "heuristic_v2": + return self._classify_with_heuristic_v2(prompt) if self.config.classifier_type == "custom": return await self._classify_with_plugin(prompt, system_prompt, request_kwargs, raw_messages) if self.config.classifier_type == "heuristic_first" and self.config.classifier_llm_config is not None: return await self._classify_heuristic_first(prompt, system_prompt, request_kwargs, messages) + if self.config.classifier_type == "hybrid" and self.config.classifier_llm_config is not None: + return await self._classify_hybrid(prompt, system_prompt, request_kwargs, messages) if self.config.classifier_type != "llm" or self.config.classifier_llm_config is None: tier, score, signals, cause = self._score_and_classify(prompt, system_prompt) return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause) return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages) + def _classify_with_heuristic_v2(self, prompt: str) -> ClassificationOutcome: + predictor: Final = self._tier_success_predictor + if predictor is None: + raise ValueError("heuristic v2 predictor is not configured") + request_type: Final = classify_prompt(prompt) + prediction: Final = predictor.predict(prompt, request_type) + tier: Final = TIER_SEVERITY_ORDER[prediction.required_tier - 1] + probability_signals: Final = tuple( + f"tier-probability:{candidate.value.lower()}={prediction.probabilities[index]:.6f}" + for index, candidate in enumerate(TIER_SEVERITY_ORDER, start=1) + ) + return ClassificationOutcome( + tier=tier, + score=None, + signals=(f"request-type:{request_type.value}", *probability_signals), + cause="heuristic_v2", + ) + async def _classify_heuristic_first( self, prompt: str, @@ -1387,6 +1430,29 @@ class ComplexityRouter(CustomLogger): return ClassificationOutcome(tier=tier, score=score, signals=signals, cause="heuristic_first_short_circuit") return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, scored=scored) + async def _classify_hybrid( + self, + prompt: str, + system_prompt: str | None, + request_kwargs: dict[str, Any] | None, # mutable-ok: handed to _classify_with_llm as-is + messages: Sequence[Mapping[str, object]] | None, + ) -> ClassificationOutcome: + """Score locally, and only pay for the classifier when the score sits near a tier boundary. + + Where heuristic_first asks how CHEAP the scorer's tier is, this asks how DECIDED it is, so a + confident score keeps its tier at every tier including the most expensive one. Two things make + a score undecided: landing within hybrid_boundary_margin of an active boundary, where a + hair's difference in score would have named the adjacent tier and its model pool, and firing + no dimension at all, which scores 0.0 and lands SIMPLE by default rather than by evidence. + """ + tier, score, signals, cause = self._score_and_classify(prompt, system_prompt) + scored: Final = ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause) + margin: Final = self.config.hybrid_boundary_margin + decided: Final = margin is not None and bool(signals) and not self._is_near_tier_boundary(score, margin) + if decided: + return ClassificationOutcome(tier=tier, score=score, signals=signals, cause="hybrid_short_circuit") + return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, scored=scored) + async def _llm_classifier_outcome( self, prompt: str, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 70aeecb31c6..0ae0db63fad 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -14,6 +14,8 @@ from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_seriali from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin +from .tier_predictor import TrainedTierArtifact + class ComplexityTier(str, Enum): """Complexity tiers for routing decisions.""" @@ -41,7 +43,7 @@ DEFAULT_CLASSIFICATION_RUBRIC: Final[ClassificationRubric] = ClassificationRubri # The classifier_type values that can call classifier_llm_config.model. Every consumer asking # "is the classifier model a real dependency of this router" resolves it here, including the ones # that only hold the raw config mapping and cannot reach ComplexityRouterConfig.uses_llm_classifier. -LLM_CLASSIFIER_TYPES: Final[frozenset[str]] = frozenset({"llm", "heuristic_first"}) +LLM_CLASSIFIER_TYPES: Final[frozenset[str]] = frozenset({"llm", "heuristic_first", "hybrid"}) TIER_SEVERITY_ORDER: Final[tuple[ComplexityTier, ...]] = ( @@ -625,17 +627,28 @@ class ComplexityRouterConfig(BaseModel): ) # Classifier strategy - classifier_type: Literal["heuristic", "llm", "custom", "heuristic_first"] = Field( + classifier_type: Literal["heuristic", "heuristic_v2", "llm", "custom", "heuristic_first", "hybrid"] = Field( default="heuristic", description=( - "Classification strategy: local regex/keyword scoring, an LLM call, a custom classifier " - "plugin, or 'heuristic_first', which scores locally and only pays for the LLM classifier " - "when the local scorer does not confidently land a cheap tier" + "Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, " + "an LLM call, a custom classifier plugin, 'heuristic_first', which scores locally and only pays " + "for the LLM classifier when the local scorer does not confidently land a cheap tier, or 'hybrid', " + "which trusts the local scorer everywhere except when its score lands near a tier boundary" + ), + ) + heuristic_v2_artifact: TrainedTierArtifact | Literal["ultrafeedback"] = Field( + default="ultrafeedback", + description=( + "Success-probability artifact used by classifier_type 'heuristic_v2'. The bundled " + "UltraFeedback artifact is selected by default; an inline trained artifact may replace it" ), ) classifier_llm_config: ClassifierLLMConfig | None = Field( default=None, - description="Configuration for the LLM classifier; required when classifier_type is 'llm' or 'heuristic_first'", + description=( + "Configuration for the LLM classifier; required when classifier_type is 'llm', " + "'heuristic_first' or 'hybrid'" + ), ) heuristic_first_max_tier: str | None = Field( default=None, @@ -650,6 +663,19 @@ class ComplexityRouterConfig(BaseModel): "may not name the highest one, since that would make the LLM classifier unreachable." ), ) + hybrid_boundary_margin: float | None = Field( + default=None, + ge=0, + le=1, + description=( + "How close to a tier boundary a heuristic score has to land before the LLM classifier breaks the " + "tie; required when classifier_type is 'hybrid' and rejected otherwise. Everything further than " + "this from every active boundary routes on the scorer's own tier with no classifier call, at any " + "tier, which is what separates 'hybrid' from 'heuristic_first' and its cheap-tier ceiling. A " + "prompt where no dimension fired still goes to the classifier, since the scorer has no opinion " + "to be near a boundary with. 0 escalates only scores sitting exactly on a boundary." + ), + ) classifier_plugin: ClassifierPlugin | None = Field( default=None, description=( @@ -1126,6 +1152,23 @@ class ComplexityRouterConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_hybrid_boundary_margin(self) -> "ComplexityRouterConfig": + if self.classifier_type != "hybrid": + if self.hybrid_boundary_margin is not None: + raise ValueError( + f"hybrid_boundary_margin is set but classifier_type is {self.classifier_type!r}; " + "the scorer would never consult the classifier on a near-boundary score. Set " + "classifier_type 'hybrid' or remove hybrid_boundary_margin" + ) + return self + if self.hybrid_boundary_margin is None: + raise ValueError( + "hybrid_boundary_margin is required when classifier_type is 'hybrid': without a margin no " + "score is ever near enough to a boundary to escalate, which is classifier_type 'heuristic'" + ) + return self + @field_validator("fallback_tier") @classmethod def _reject_blank_optional_text(cls, value: str | None) -> str | None: @@ -1248,10 +1291,10 @@ class ComplexityRouterConfig(BaseModel): ) if duplicated: raise ValueError(f"tier_definitions names must be unique (case-insensitive): {', '.join(duplicated)}") - if self.classifier_type in ("heuristic", "heuristic_first"): + if self.classifier_type in ("heuristic", "heuristic_v2", "heuristic_first", "hybrid"): raise ValueError( "tier_definitions requires classifier_type 'llm' or 'custom': the heuristic scorer only " - "produces the built-in tiers" + "produces the four built-in tiers, as does heuristic_v2" ) conflicts: Final = self._tier_definition_conflicts() if conflicts: diff --git a/litellm/router_strategy/complexity_router/tier_predictor.py b/litellm/router_strategy/complexity_router/tier_predictor.py new file mode 100644 index 00000000000..764f6e6ad56 --- /dev/null +++ b/litellm/router_strategy/complexity_router/tier_predictor.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import re +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal + +from pydantic import BaseModel, Field, model_validator + +from litellm.types.router import RequestType + + +class TierGlobalStatistic(BaseModel): + tier: int = Field(ge=1, le=4) + successes: float = Field(ge=0.0) + observations: float = Field(gt=0.0) + + @model_validator(mode="after") + def _successes_do_not_exceed_observations(self) -> TierGlobalStatistic: + if self.successes > self.observations: + raise ValueError("successes cannot exceed observations") + return self + + +class TierDomainStatistic(TierGlobalStatistic): + request_type: RequestType + + +class TierCohortStatistic(TierGlobalStatistic): + cohort: str = Field(min_length=1) + + +class TierDataset(BaseModel): + name: str = Field(min_length=1) + url: str = Field(min_length=1) + license: str = Field(min_length=1) + rows: int = Field(gt=0) + success_definition: str = Field(default="quality score meets the dataset success threshold", min_length=1) + + +class TrainedTierArtifact(BaseModel): + schema_version: Literal[1] = 1 + global_statistics: tuple[TierGlobalStatistic, ...] + domain_statistics: tuple[TierDomainStatistic, ...] = () + cohort_statistics: tuple[TierCohortStatistic, ...] = () + domain_prior_mass: float = Field(default=200.0, gt=0.0) + cohort_prior_mass: float = Field(default=20.0, gt=0.0) + routing_threshold: float = Field(default=0.75, ge=0.0, le=1.0) + datasets: tuple[TierDataset, ...] = () + success_definition: str = Field(default="quality score meets the dataset success threshold", min_length=1) + split_method: str = Field(default="sha256(prompt): 70% train, 15% validation, 15% test", min_length=1) + + @model_validator(mode="after") + def _statistics_are_unique(self) -> TrainedTierArtifact: + global_tiers: Final = tuple(stat.tier for stat in self.global_statistics) + if frozenset(global_tiers) != frozenset((1, 2, 3, 4)) or len(global_tiers) != 4: + raise ValueError("global statistics must contain each tier exactly once") + domain_keys: Final = tuple((stat.request_type, stat.tier) for stat in self.domain_statistics) + if len(domain_keys) != len(frozenset(domain_keys)): + raise ValueError("domain statistics must contain unique request_type and tier pairs") + cohort_keys: Final = tuple((stat.cohort, stat.tier) for stat in self.cohort_statistics) + if len(cohort_keys) != len(frozenset(cohort_keys)): + raise ValueError("cohort statistics must contain unique cohort and tier pairs") + return self + + +_CODE_PATTERN: Final = re.compile( + r"```|\b(def|class|function|python|javascript|typescript|sql|code)\b", + re.IGNORECASE, +) +_MATH_PATTERN: Final = re.compile( + r"\b(solve|calculate|equation|probability|theorem|proof|integral)\b|[$=]", + re.IGNORECASE, +) +_MULTIPLE_CHOICE_PATTERN: Final = re.compile(r"(?:^|\s)[A-D][.)]\s") +_TIERS: Final = (1, 2, 3, 4) +_BUILTIN_ARTIFACTS: Final = MappingProxyType({"ultrafeedback": "ultrafeedback_tiers.json"}) + + +def resolve_tier_artifact(artifact: TrainedTierArtifact | str) -> TrainedTierArtifact: + if isinstance(artifact, TrainedTierArtifact): + return artifact + filename: Final = _BUILTIN_ARTIFACTS.get(artifact) + if filename is None: + raise ValueError(f"unknown complexity router tier artifact: {artifact}") + path: Final = Path(__file__).with_name("artifacts") / filename + return TrainedTierArtifact.model_validate_json(path.read_text()) + + +def similarity_cohort(prompt: str, request_type: RequestType) -> str: + length: Final = len(prompt) + length_bucket: Final = ( + "short" if length < 200 else "medium" if length < 800 else "long" if length < 2000 else "very_long" + ) + code: Final = int(bool(_CODE_PATTERN.search(prompt))) + math: Final = int(bool(_MATH_PATTERN.search(prompt))) + multiple_choice: Final = int(bool(_MULTIPLE_CHOICE_PATTERN.search(prompt))) + non_ascii: Final = int(sum(ord(character) > 127 for character in prompt) / max(1, length) > 0.1) + return f"{request_type.value}|{length_bucket}|code={code}|math={math}|mc={multiple_choice}|intl={non_ascii}" + + +@dataclass(frozen=True, slots=True) +class TierPrediction: + probabilities: Mapping[int, float] + required_tier: int + + +class TierSuccessPredictor: + def __init__(self, artifact: TrainedTierArtifact) -> None: + self._artifact = artifact + self._global: Mapping[int, TierGlobalStatistic] = MappingProxyType( + {stat.tier: stat for stat in artifact.global_statistics} + ) + self._domain: Mapping[tuple[RequestType, int], TierDomainStatistic] = MappingProxyType( + {(stat.request_type, stat.tier): stat for stat in artifact.domain_statistics} + ) + self._cohort: Mapping[tuple[str, int], TierCohortStatistic] = MappingProxyType( + {(stat.cohort, stat.tier): stat for stat in artifact.cohort_statistics} + ) + + @property + def routing_threshold(self) -> float: + return self._artifact.routing_threshold + + def predict(self, prompt: str, request_type: RequestType) -> TierPrediction: + cohort: Final = similarity_cohort(prompt, request_type) + raw: Final = tuple(self._probability(tier, request_type, cohort) for tier in _TIERS) + monotonic: Final = tuple(max(raw[:index]) for index in range(1, len(raw) + 1)) + probabilities: Final[Mapping[int, float]] = MappingProxyType( + {int(tier): probability for tier, probability in zip(_TIERS, monotonic)} + ) + required_tier: Final = next( + (tier for tier in _TIERS if probabilities[tier] >= self._artifact.routing_threshold), + 4, + ) + return TierPrediction(probabilities=probabilities, required_tier=required_tier) + + def _probability(self, tier: int, request_type: RequestType, cohort: str) -> float: + global_stat: Final = self._global[tier] + global_mean: Final = (global_stat.successes + 1.0) / (global_stat.observations + 2.0) + domain_stat: Final = self._domain.get((request_type, tier)) + domain_mean: Final = self._posterior_mean(domain_stat, self._artifact.domain_prior_mass, global_mean) + cohort_stat: Final = self._cohort.get((cohort, tier)) + return self._posterior_mean(cohort_stat, self._artifact.cohort_prior_mass, domain_mean) + + @staticmethod + def _posterior_mean( + statistic: TierGlobalStatistic | None, + prior_mass: float, + prior_mean: float, + ) -> float: + if statistic is None: + return prior_mean + return (statistic.successes + prior_mass * prior_mean) / (statistic.observations + prior_mass) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 3d37ca216a7..f7855cb38ff 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any, Final import litellm from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, @@ -231,8 +232,6 @@ def record_pre_routing_selection(request_kwargs: Mapping[str, Any] | None, selec on /v1/messages the top-level ``metadata`` dict is the provider's own request field, so a blanket write would forward the tier stamp upstream. """ - from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs - if request_kwargs is None: return bucket: Final = request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs)) @@ -267,10 +266,13 @@ def get_pre_routing_selection(kwargs: Mapping[str, Any]) -> str | None: def fallback_lookup_groups(kwargs: Mapping[str, Any], model_group: str | None) -> tuple[str, ...]: """ Ordered keys for resolving a fallback chain: the tier a pre-routing hook selected wins, - and the requested group still resolves when no tier-keyed chain exists, so configs keyed - on the router name (the documented contract) keep working behind auto-routers. + then the routed group, then the requested group. The routed group differs when Claude Code + session affinity remaps a subagent's concrete model to its bound router. """ - ordered: Final = (get_pre_routing_selection(kwargs), model_group) + metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) + routed_group_value: Final = metadata.get("model_group") if isinstance(metadata, Mapping) else None + routed_group: Final = routed_group_value if isinstance(routed_group_value, str) else None + ordered: Final = (get_pre_routing_selection(kwargs), routed_group, model_group) return tuple(dict.fromkeys(group for group in ordered if group)) @@ -470,10 +472,11 @@ async def run_async_fallback( attempted: Final = ( carried_targets if isinstance(carried_targets, AttemptedFallbackTargets) else AttemptedFallbackTargets() ) - attempted.record(original_model_group) + failed_model_group: Final = get_pre_routing_selection(kwargs) or original_model_group + attempted.record(failed_model_group) for mg in fallback_model_group: - if mg == original_model_group: + if mg == failed_model_group: continue if same_model_group_only and _get_fallback_target_model_group(mg) != original_model_group: verbose_router_logger.info( diff --git a/litellm/rust_bridge/__init__.py b/litellm/rust_bridge/__init__.py index 3da5b98449b..e6d8ffef48c 100644 --- a/litellm/rust_bridge/__init__.py +++ b/litellm/rust_bridge/__init__.py @@ -1,9 +1,9 @@ """LiteLLM Rust bridge package.""" +from litellm.rust_bridge.configuration import use_litellm_rust from litellm.rust_bridge.loader import ( get_native_bridge, native_bridge_available, ) -from litellm.rust_bridge.ocr import use_litellm_rust __all__ = ["get_native_bridge", "native_bridge_available", "use_litellm_rust"] diff --git a/litellm/rust_bridge/bindings.py b/litellm/rust_bridge/bindings.py new file mode 100644 index 00000000000..d16f150a2aa --- /dev/null +++ b/litellm/rust_bridge/bindings.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import Final, Generic, TypeVar + +from litellm.rust_bridge.loader import get_native_bridge + +BindingT = TypeVar("BindingT") + + +class _Unset: + pass + + +_UNSET: Final = _Unset() + + +class NativeBinding(Generic[BindingT]): + """Resolve one native attribute with an explicit, resettable test override.""" + + def __init__(self, attribute: str, *, validate: Callable[[object], BindingT | None]) -> None: + self._attribute: Final = attribute + self._validate: Final = validate + self._override: BindingT | None | _Unset = _UNSET + + def load(self) -> BindingT | None: + if not isinstance(self._override, _Unset): + return self._override + native: Final = get_native_bridge() + if native is None: + return None + return self._validate(getattr(native, self._attribute, None)) + + def override(self, value: BindingT | None) -> None: + self._override = value + + def reset(self) -> None: + self._override = _UNSET + + +def native_exception_types() -> tuple[type[BaseException], type[BaseException]] | None: + native: Final = get_native_bridge() + if native is None: + return None + declined: Final = getattr(native, "RustBridgeDeclined", None) + upstream: Final = getattr(native, "RustUpstreamError", None) + if not isinstance(declined, type) or not isinstance(upstream, type): + return None + return declined, upstream diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index acda3086051..c599667ab17 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -13,7 +13,6 @@ retrying it there would bill the customer for the same work twice. from __future__ import annotations import json -import os from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Protocol @@ -27,6 +26,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo convert_to_model_response_object, ) from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned +from litellm.rust_bridge.configuration import rust_enabled from litellm.rust_bridge.loader import get_native_bridge from litellm.rust_bridge.timeouts import timeout_to_seconds from litellm.types.utils import ModelResponse @@ -44,8 +44,6 @@ _LITELLM_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) RUST_RESPONSE_HEADER: Final = "x-litellm-rust" -_TRUTHY_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"}) - class RustChatCompletions(Protocol): def __call__( @@ -181,10 +179,6 @@ def load_rust_achat_completions() -> RustAchatCompletions | None: return loaded -def _env_enables_rust() -> bool: - return os.getenv("LITELLM_RUST", "").strip().lower() in _TRUTHY_ENV_VALUES - - def _load_rust_decline() -> RustChatCompletionsDecline | None: if _STATE.decline is not None: return _STATE.decline @@ -253,8 +247,8 @@ def rust_chat_completions_accepts( return False if stream: return False - opted_in: Final = litellm_params is not None and litellm_params.get("rust") is True - if not opted_in and not _env_enables_rust(): + request_override: Final = litellm_params.get("rust") if litellm_params is not None else None + if not rust_enabled(request_override=request_override if isinstance(request_override, bool) else None): return False if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params): verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path") diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py new file mode 100644 index 00000000000..d54b15f060c --- /dev/null +++ b/litellm/rust_bridge/configuration.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import os +import warnings +from typing import TYPE_CHECKING, Final + +if TYPE_CHECKING: + from litellm.rust_bridge.messages import RustAmessages, RustMessages + from litellm.rust_bridge.ocr import RustAocr, RustOcr + from litellm.rust_bridge.responses_websocket import RustResponsesWebSocketConnection + from litellm.rust_bridge.transcription import RustAtranscription, RustTranscription + +DEFAULT_RUST_ENABLED: Final = False +_TRUE_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"}) +_GLOBAL_ENV_NAME: Final = "LITELLM_RUST" +_LEGACY_OCR_ENV_NAME: Final = "LITELLM_USE_RUST_OCR" + + +class _Unset: + pass + + +_UNSET: Final = _Unset() + + +class _RustConfiguration: + def __init__(self) -> None: + self.override: bool | None = None + + +_CONFIGURATION: Final = _RustConfiguration() + + +def _parse_env_bool(value: str | None) -> bool | None: + if value is None: + return None + return value.strip().lower() in _TRUE_ENV_VALUES + + +def resolve_rust_enabled( + *, + request_override: bool | None, + process_override: bool | None, + environment_override: bool | None, + legacy_ocr_override: bool | None = None, + release_default: bool = DEFAULT_RUST_ENABLED, +) -> bool: + if request_override is not None: + return request_override + if process_override is not None: + return process_override + if environment_override is not None: + return environment_override + if legacy_ocr_override is not None: + return legacy_ocr_override + return release_default + + +def rust_enabled(*, request_override: bool | None = None) -> bool: + if request_override is not None: + return request_override + process_override: Final = _CONFIGURATION.override + if process_override is not None: + return process_override + return resolve_rust_enabled( + request_override=None, + process_override=None, + environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)), + ) + + +def rust_ocr_enabled(*, request_override: bool | None = None) -> bool: + if request_override is not None: + return request_override + process_override: Final = _CONFIGURATION.override + if process_override is not None: + return process_override + global_override: Final = _parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)) + legacy_override: Final = None if global_override is not None else _parse_env_bool(os.getenv(_LEGACY_OCR_ENV_NAME)) + if legacy_override is not None: + warnings.warn( + f"{_LEGACY_OCR_ENV_NAME} is deprecated; use {_GLOBAL_ENV_NAME} instead", + DeprecationWarning, + stacklevel=2, + ) + return resolve_rust_enabled( + request_override=None, + process_override=None, + environment_override=global_override, + legacy_ocr_override=legacy_override, + ) + + +def reset_rust_configuration() -> None: + _CONFIGURATION.override = None + + +def use_litellm_rust( + enabled: bool = True, + *, + ocr: RustOcr | None | _Unset = _UNSET, + aocr: RustAocr | None | _Unset = _UNSET, + messages: RustMessages | None | _Unset = _UNSET, + amessages: RustAmessages | None | _Unset = _UNSET, + responses_websocket: type[RustResponsesWebSocketConnection] | None | _Unset = _UNSET, + transcription: RustTranscription | None | _Unset = _UNSET, + atranscription: RustAtranscription | None | _Unset = _UNSET, +) -> None: + """Set the process override for optional Rust paths. + + Rust-only paths, including Bedrock transcription, are not controlled by this switch. + """ + _CONFIGURATION.override = enabled + bindings: Final = (ocr, aocr, messages, amessages, responses_websocket, transcription, atranscription) + if all(isinstance(binding, _Unset) for binding in bindings): + return + warnings.warn( + "Injecting Rust bridge implementations through use_litellm_rust() is deprecated; " + "use the internal bridge setters in tests", + DeprecationWarning, + stacklevel=2, + ) + + if not isinstance(ocr, _Unset) or not isinstance(aocr, _Unset): + from litellm.rust_bridge.ocr import set_rust_ocr + + if not isinstance(ocr, _Unset): + set_rust_ocr(ocr=ocr) + if not isinstance(aocr, _Unset): + set_rust_ocr(aocr=aocr) + if not isinstance(messages, _Unset) or not isinstance(amessages, _Unset): + from litellm.rust_bridge.messages import set_rust_messages + + if not isinstance(messages, _Unset): + set_rust_messages(messages=messages) + if not isinstance(amessages, _Unset): + set_rust_messages(amessages=amessages) + if not isinstance(responses_websocket, _Unset): + from litellm.rust_bridge.responses_websocket import set_rust_responses_websocket + + set_rust_responses_websocket(connection=responses_websocket) + if not isinstance(transcription, _Unset) or not isinstance(atranscription, _Unset): + from litellm.rust_bridge.transcription import configure_rust_transcription + + if not isinstance(transcription, _Unset): + configure_rust_transcription(transcription=transcription) + if not isinstance(atranscription, _Unset): + configure_rust_transcription(atranscription=atranscription) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 82297d35170..b5b0a35a498 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -2,16 +2,16 @@ from __future__ import annotations -import os from collections.abc import Awaitable -from typing import TYPE_CHECKING, Any, Final, Protocol, cast +from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables import httpx +from litellm.rust_bridge import configuration as _configuration from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds -if TYPE_CHECKING: - from litellm.rust_bridge.messages import RustAmessages, RustMessages +rust_ocr_enabled = _configuration.rust_ocr_enabled +use_litellm_rust = _configuration.use_litellm_rust class RustOcr(Protocol): @@ -51,69 +51,20 @@ class _Unset: _UNSET: Final[_Unset] = _Unset() -def _env_enables_rust_ocr() -> bool: - return os.getenv("LITELLM_USE_RUST_OCR", "").strip().lower() in { - "1", - "true", - "yes", - "on", - } - - -_rust_ocr_enabled = _env_enables_rust_ocr() _rust_ocr_impl: RustOcr | None = None _rust_aocr_impl: RustAocr | None = None -def use_litellm_rust( - enabled: bool = True, +def set_rust_ocr( *, ocr: RustOcr | None | _Unset = _UNSET, aocr: RustAocr | None | _Unset = _UNSET, - messages: RustMessages | None | _Unset = _UNSET, - amessages: RustAmessages | None | _Unset = _UNSET, - responses_websocket: Any | None | _Unset = _UNSET, - transcription: Any | None | _Unset = _UNSET, - atranscription: Any | None | _Unset = _UNSET, ) -> None: - global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl - configuring_ocr: Final = not isinstance(ocr, _Unset) or not isinstance(aocr, _Unset) - configuring_messages: Final = not isinstance(messages, _Unset) or not isinstance(amessages, _Unset) - configuring_responses_websocket: Final = not isinstance(responses_websocket, _Unset) - configuring_transcription: Final = not isinstance(transcription, _Unset) or not isinstance(atranscription, _Unset) - if configuring_ocr or (not configuring_messages and not configuring_responses_websocket): - _rust_ocr_enabled = enabled + global _rust_ocr_impl, _rust_aocr_impl if not isinstance(ocr, _Unset): _rust_ocr_impl = ocr if not isinstance(aocr, _Unset): _rust_aocr_impl = aocr - if configuring_transcription: - from litellm.rust_bridge.transcription import configure_rust_transcription - - configure_rust_transcription( - enabled=enabled, - transcription=transcription, - atranscription=atranscription, - ) - if not configuring_messages and not configuring_responses_websocket: - return - if configuring_messages: - from litellm.rust_bridge.messages import set_rust_messages - - if not isinstance(messages, _Unset) and not isinstance(amessages, _Unset): - set_rust_messages(messages=messages, amessages=amessages) - elif not isinstance(messages, _Unset): - set_rust_messages(messages=messages) - else: - set_rust_messages(amessages=amessages) - if configuring_responses_websocket: - from litellm.rust_bridge.responses_websocket import set_rust_responses_websocket - - set_rust_responses_websocket(connection=responses_websocket) - - -def rust_ocr_enabled() -> bool: - return _rust_ocr_enabled def load_rust_ocr() -> RustOcr | None: diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py new file mode 100644 index 00000000000..00f06c046a2 --- /dev/null +++ b/litellm/rust_bridge/runtime.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from enum import Enum +from typing import Final, Generic, NoReturn, TypeAlias, TypeVar + +from litellm.exceptions import APIError +from litellm.rust_bridge.bindings import native_exception_types + +NativeT = TypeVar("NativeT") +ResultT = TypeVar("ResultT") + + +class FallbackMode(Enum): + PYTHON = "python" + RUST_REQUIRED = "rust_required" + + +@dataclass(frozen=True, slots=True) +class RustHandled(Generic[ResultT]): + value: ResultT + + +@dataclass(frozen=True, slots=True) +class RustDeclined: + reason: str + + +@dataclass(frozen=True, slots=True) +class RustUnavailable: + pass + + +RustAttempt: TypeAlias = RustHandled[ResultT] | RustDeclined | RustUnavailable + + +@dataclass(frozen=True, slots=True) +class BridgeErrorContext: + route: str + provider: str + model: str + + +def invoke( + *, + native_call: Callable[[], NativeT] | None, + fallback: Callable[[], ResultT], + adapt: Callable[[NativeT], ResultT], + mode: FallbackMode, + context: BridgeErrorContext, +) -> ResultT: + result: Final = attempt(native_call=native_call, adapt=adapt, context=context) + if isinstance(result, RustHandled): + return result.value + if mode is FallbackMode.PYTHON: + return fallback() + _raise_required(result, context) + + +async def ainvoke( + *, + native_call: Callable[[], Awaitable[NativeT]] | None, + fallback: Callable[[], Awaitable[ResultT]], + adapt: Callable[[NativeT], ResultT], + mode: FallbackMode, + context: BridgeErrorContext, +) -> ResultT: + result: Final = await aattempt(native_call=native_call, adapt=adapt, context=context) + if isinstance(result, RustHandled): + return result.value + if mode is FallbackMode.PYTHON: + return await fallback() + _raise_required(result, context) + + +def attempt( + *, + native_call: Callable[[], NativeT] | None, + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, +) -> RustAttempt[ResultT]: + if native_call is None: + return RustUnavailable() + exceptions: Final = native_exception_types() + if exceptions is None: + return RustHandled(adapt(native_call())) + declined, upstream = exceptions + try: + value: Final = native_call() + except declined as error: + return RustDeclined(reason=_decline_reason(error)) + except upstream as error: + _raise_upstream(error, context) + return RustHandled(adapt(value)) + + +async def aattempt( + *, + native_call: Callable[[], Awaitable[NativeT]] | None, + adapt: Callable[[NativeT], ResultT], + context: BridgeErrorContext, +) -> RustAttempt[ResultT]: + if native_call is None: + return RustUnavailable() + exceptions: Final = native_exception_types() + if exceptions is None: + return RustHandled(adapt(await native_call())) + declined, upstream = exceptions + try: + value: Final = await native_call() + except declined as error: + return RustDeclined(reason=_decline_reason(error)) + except upstream as error: + _raise_upstream(error, context) + return RustHandled(adapt(value)) + + +def call(operation: Callable[[], ResultT], context: BridgeErrorContext) -> ResultT: + exceptions: Final = native_exception_types() + if exceptions is None: + return operation() + upstream: Final = exceptions[1] + try: + return operation() + except upstream as error: + _raise_upstream(error, context) + + +async def acall(operation: Callable[[], Awaitable[ResultT]], context: BridgeErrorContext) -> ResultT: + exceptions: Final = native_exception_types() + if exceptions is None: + return await operation() + upstream: Final = exceptions[1] + try: + return await operation() + except upstream as error: + _raise_upstream(error, context) + + +def _decline_reason(error: BaseException) -> str: + reason: Final[object] = error.args[0] if error.args else str(error) + return reason if isinstance(reason, str) else str(reason) + + +def _raise_required( + result: RustDeclined | RustUnavailable, + context: BridgeErrorContext, +) -> NoReturn: + raise RuntimeError(f"Rust {context.route} bridge {_required_reason(result)}") + + +def _required_reason(result: RustDeclined | RustUnavailable) -> str: + match result: + case RustUnavailable(): + return "is unavailable" + case RustDeclined(reason=reason): + return f"declined the request: {reason}" + + +def _raise_upstream(error: BaseException, context: BridgeErrorContext) -> NoReturn: + args: Final[tuple[object, ...]] = error.args + status_value: Final = args[0] if args else 0 + message_value: Final = args[1] if len(args) > 1 else str(error) + status: Final = status_value if isinstance(status_value, int) else 0 + message: Final = message_value if isinstance(message_value, str) else str(message_value) + raise APIError( + status_code=status or 500, + message=f"litellm rust {context.route}: {message}", + llm_provider=context.provider, + model=context.model, + ) from error + + +def identity(value: ResultT) -> ResultT: + return value + + +async def async_none() -> None: + return None diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index 903781b2ccd..a76e6cf1187 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -4,8 +4,10 @@ https://docs.cohere.com/reference/rerank """ -from pydantic import BaseModel, PrivateAttr -from typing_extensions import Required, TypedDict +from typing import Literal + +from pydantic import BaseModel, ConfigDict, PrivateAttr +from typing_extensions import ReadOnly, Required, TypedDict class RerankRequest(BaseModel): @@ -21,6 +23,18 @@ class RerankRequest(BaseModel): # (e.g. hosted vLLM / Qwen3-Reranker, DeepInfra). Omitted from the outgoing # request when None, so this is fully backward-compatible. instruction: str | None = None + truncate_prompt_tokens: int | None = None + truncation_side: Literal["left", "right"] | None = None + max_tokens_per_query: int | None = None + + +class HostedVLLMRerankTruncationParams(BaseModel): + model_config = ConfigDict(frozen=True) + + truncate_prompt_tokens: int | None = None + truncation_side: Literal["left", "right"] | None = None + max_tokens_per_query: int | None = None + max_tokens_per_doc: int | None = None class OptionalRerankParams(TypedDict, total=False): @@ -32,6 +46,9 @@ class OptionalRerankParams(TypedDict, total=False): max_chunks_per_doc: int | None max_tokens_per_doc: int | None instruction: str | None + truncate_prompt_tokens: ReadOnly[int | None] + truncation_side: ReadOnly[Literal["left", "right"] | None] + max_tokens_per_query: ReadOnly[int | None] class RerankBilledUnits(TypedDict, total=False): diff --git a/litellm/types/router.py b/litellm/types/router.py index 2a5f264cee3..4f4df1a8d2e 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -305,6 +305,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): """ custom_llm_provider: str | None = None + rust: bool | None = None tpm: int | None = None rpm: int | None = None itpm: int | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5783a39b30c..ee6f09e05dc 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2839,6 +2839,7 @@ class StandardLoggingRoutingDecisionTierBoundaries(TypedDict): RoutingDecisionCause = Literal[ "heuristic_scorer", + "heuristic_v2", # The scorer found 2+ reasoning markers and forced REASONING regardless of score. # A distinct cause rather than a marker inside `signals`, because it is the fact # that tells a reader the score did NOT choose the tier; encoding it as free text @@ -2851,6 +2852,7 @@ RoutingDecisionCause = Literal[ # scorer, and from "classifier_fallback", which is the scorer running because a call failed: # only this cause means an LLM classifier was configured, reachable, and deliberately skipped. "heuristic_first_short_circuit", + "hybrid_short_circuit", # The operator's classifier plugin (classifier_type 'custom') decided the tier. "classifier_plugin", # The LLM classifier or classifier plugin failed on a router with an operator-defined diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index cd576755f5f..636bdd4b52e 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -15,6 +15,11 @@ import litellm from litellm.constants import request_timeout from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.vector_store.transformation import ( + BaseQueryEmbeddingVectorStoreConfig, + VectorStoreEmbeddingExecutor, + vector_store_request_metadata, +) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( @@ -38,6 +43,16 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# +def _direct_vector_store_embedding_executor( + value: object, router: "Router | None", request_kwargs: Mapping[str, object] +) -> VectorStoreEmbeddingExecutor: + if value is not None and not isinstance(value, VectorStoreEmbeddingExecutor): + raise TypeError("Invalid direct vector store embedding executor") + return BaseQueryEmbeddingVectorStoreConfig.query_embedding_executor( + value, router, vector_store_request_metadata(request_kwargs) + ) + + def mock_vector_store_search_response( mock_results: list[VectorStoreSearchResult] | None = None, ): @@ -289,7 +304,12 @@ async def asearch( """ Async: Search a vector store for relevant chunks based on a query and file attributes filter. """ - local_vars: Final = locals() + embedding_executor: Final = _direct_vector_store_embedding_executor( + kwargs.pop("_direct_vector_store_embedding_executor", None), router, kwargs + ) + local_vars: Final = { # mutable-ok: exception logging requires a sanitized mutable snapshot + key: value for key, value in locals().items() if key != "embedding_executor" + } try: loop: Final = asyncio.get_event_loop() @@ -312,6 +332,7 @@ async def asearch( extra_body=extra_body, timeout=timeout, custom_llm_provider=custom_llm_provider, + _direct_vector_store_embedding_executor=embedding_executor, router=router, **kwargs, ) @@ -369,12 +390,16 @@ def search( Returns: VectorStoreSearchResponse containing the search results. """ - local_vars: Final = locals() + embedding_executor: Final = _direct_vector_store_embedding_executor( + kwargs.pop("_direct_vector_store_embedding_executor", None), router, kwargs + ) + local_vars: Final = { # mutable-ok: exception logging requires a sanitized mutable snapshot + key: value for key, value in locals().items() if key != "embedding_executor" + } try: litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) _is_async: Final = kwargs.pop("asearch", False) is True - # pull credentials from registry if available if litellm.vector_store_registry is not None and vector_store_id is not None: try: @@ -451,6 +476,7 @@ def search( custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, logging_obj=litellm_logging_obj, + embedding_executor=embedding_executor, extra_headers=extra_headers, extra_body=extra_body, timeout=timeout or request_timeout, diff --git a/pyproject.toml b/pyproject.toml index 60162544612..d0e5723d1cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -278,7 +278,10 @@ bindings = "pyo3" features = ["extension-module"] profile = "release" editable-profile = "dev" -include = ["litellm/proxy/_experimental/out/**"] +include = [ + "litellm/proxy/_experimental/out/**", + "litellm/router_strategy/complexity_router/artifacts/*.json", +] exclude = [ "litellm/proxy/enterprise", "litellm/proxy/enterprise/**", diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index ae91b711e13..4fcf650a8bc 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 307 + "limit": 304 }, "ASYNC230": { "limit": 11 @@ -156,7 +156,7 @@ "limit": 215 }, "PLW0603": { - "limit": 191 + "limit": 190 }, "PLW1508": { "limit": 190 @@ -177,7 +177,7 @@ "limit": 8 }, "RUF019": { - "limit": 31 + "limit": 27 }, "RUF046": { "limit": 4 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1073 + "limit": 1071 }, "TRY002": { "limit": 524 diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index a5e00799519..0af29f069c6 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -82,6 +82,11 @@ ignored_function_names = [ "_invalidate_access_groups_cache", # Tested indirectly via set_model_list, upsert_model etc. (test files lack "router" in name) "has_buffered_provider_output", # Property, so its reads in test_router.py are never an ast.Call "_resolved_provider", # Tested via get_pattern in test_pattern_match_deployments.py (file lacks "router" in name) + "_request_header", # Tested through Claude Code session routing in test_router.py + "_claude_code_session_router_cache_key", # Tested through Claude Code session routing in test_router.py + "_delete_claude_code_session_router_binding", # Tested through Redis cleanup failure in test_router.py + "_resolve_claude_code_session_router", # Tested through Claude Code session routing in test_router.py + "_get_claude_code_session_router_binding", # Tested through the two-worker session routing test in test_router.py ] diff --git a/tests/e2e/ui/constants.ts b/tests/e2e/ui/constants.ts index bb33c90ddf3..3a62252915c 100644 --- a/tests/e2e/ui/constants.ts +++ b/tests/e2e/ui/constants.ts @@ -30,6 +30,7 @@ export const E2E_PROXY_ADMIN_EMAIL = "admin@test.local"; export const E2E_INTERNAL_USER_ID = "e2e-internal-user"; export const E2E_INTERNAL_USER_EMAIL = "internal@test.local"; export const E2E_TEAM_ADMIN_USER_ID = "e2e-team-admin"; +export const E2E_SEEDED_USER_PASSWORD = "E2e-Test-Pass-2026!"; // Key aliases for seeded test keys (match seed.sql) export const E2E_UPDATE_LIMITS_KEY_ALIAS = "e2eUpdateLimitsKey"; diff --git a/tests/e2e/ui/fixtures/seed.sql b/tests/e2e/ui/fixtures/seed.sql index e77b4a16b3d..00ea668ed8f 100644 --- a/tests/e2e/ui/fixtures/seed.sql +++ b/tests/e2e/ui/fixtures/seed.sql @@ -24,18 +24,18 @@ INSERT INTO "LiteLLM_OrganizationTable" ( 'e2e-proxy-admin', 'e2e-proxy-admin' ); --- 4. Users (password hash is scrypt of "test") +-- 4. Users (password hash is scrypt of E2E_SEEDED_USER_PASSWORD from constants.ts) INSERT INTO "LiteLLM_UserTable" ("user_id", "user_email", "user_role", "teams", "password") VALUES - ('e2e-proxy-admin', 'admin@test.local', 'proxy_admin', '{"e2e-team-crud"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-admin-viewer', 'adminviewer@test.local', 'proxy_admin_viewer', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-internal-user', 'internal@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-org","e2e-team-keygen"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-internal-viewer', 'viewer@test.local', 'internal_user_viewer', '{"e2e-team-crud"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-team-admin', 'teamadmin@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-delete"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-invitable-user', 'invitable@test.local', 'internal_user', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-internal-noteam', 'noteam@test.local', 'internal_user', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-invitable-by-team-admin', 'invitable-team@test.local', 'internal_user', '{}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'), - ('e2e-removable-member', 'removable@test.local', 'internal_user', '{"e2e-team-crud"}', 'scrypt:MU5CcTAi6rVK1HfY1rVPEWq6r4sxg837eq9dG4n5Q6BhDJ44442+seC6LAhLEAYr'); + ('e2e-proxy-admin', 'admin@test.local', 'proxy_admin', '{"e2e-team-crud"}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-admin-viewer', 'adminviewer@test.local', 'proxy_admin_viewer', '{}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-internal-user', 'internal@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-org","e2e-team-keygen"}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-internal-viewer', 'viewer@test.local', 'internal_user_viewer', '{"e2e-team-crud"}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-team-admin', 'teamadmin@test.local', 'internal_user', '{"e2e-team-crud","e2e-team-delete"}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-invitable-user', 'invitable@test.local', 'internal_user', '{}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-internal-noteam', 'noteam@test.local', 'internal_user', '{}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-invitable-by-team-admin', 'invitable-team@test.local', 'internal_user', '{}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'), + ('e2e-removable-member', 'removable@test.local', 'internal_user', '{"e2e-team-crud"}', 'scrypt:KdnTJwPb3gswdqSznPE5CC6apeFIMycd6BG7yRWndZa3QZPcVs37y7jvrQCaPUNq'); -- 5. Teams (members_with_roles is required JSON) INSERT INTO "LiteLLM_TeamTable" ( diff --git a/tests/e2e/ui/fixtures/users.ts b/tests/e2e/ui/fixtures/users.ts index 79ee237f334..0457361d9d4 100644 --- a/tests/e2e/ui/fixtures/users.ts +++ b/tests/e2e/ui/fixtures/users.ts @@ -1,6 +1,7 @@ import { ADMIN_STORAGE_PATH, ADMIN_VIEWER_STORAGE_PATH, + E2E_SEEDED_USER_PASSWORD, INTERNAL_USER_STORAGE_PATH, INTERNAL_VIEWER_STORAGE_PATH, TEAM_ADMIN_STORAGE_PATH, @@ -23,22 +24,22 @@ export const users: Record { // Log in via the form as the no-team seeded user. await page.goto("/ui/login"); await page.getByPlaceholder("Enter your username").fill("noteam@test.local"); - await page.getByPlaceholder("Enter your password").fill("test"); + await page.getByPlaceholder("Enter your password").fill(E2E_SEEDED_USER_PASSWORD); await page.getByRole("button", { name: "Login", exact: true }).click(); await expect(page.getByRole("complementary").getByText("Virtual Keys")).toBeVisible({ timeout: 30_000 }); expect(new URL(page.url()).pathname).not.toMatch(/\/connect$/); diff --git a/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts b/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts index 5a8bc84cc13..73263c844fa 100644 --- a/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/secondAdmin.spec.ts @@ -10,7 +10,7 @@ test.describe("Second proxy admin", () => { test("an invited admin can log in, mint a key, and call a model with it", async ({ page, browser, request }) => { const suffix = Date.now(); const email = `second-admin-${suffix}@test.local`; - const password = "e2e-second-admin-password"; + const password = "E2e-Second-Admin-Pass-1!"; const auth = { Authorization: `Bearer ${masterKey()}` }; const inviteAdminUser = async (): Promise => { diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 0c362db8853..98045725177 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -71,6 +71,48 @@ def setup_vector_store_registry(): ) +@pytest.mark.asyncio +async def test_vector_store_hook_routes_search_through_proxy_router( + setup_vector_store_registry, +): + proxy_router = Mock() + proxy_router.avector_store_search = AsyncMock( + return_value=VectorStoreSearchResponse( + object="vector_store.search_results.page", + search_query="what is litellm?", + data=[ + VectorStoreSearchResult( + score=1.0, + content=[VectorStoreResultContent(text="routed context", type="text")], + ) + ], + ) + ) + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_params": {"metadata": {"user_api_key_team_id": "team-a"}} + } + + with patch("litellm.proxy.proxy_server.llm_router", proxy_router): + _, messages, _ = await VectorStorePreCallHook().async_get_chat_completion_prompt( + model="chat-model", + messages=[{"role": "user", "content": "what is litellm?"}], + non_default_params={"vector_store_ids": ["T37J8R4WTM"]}, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + litellm_logging_obj=logging_obj, + ) + + proxy_router.avector_store_search.assert_awaited_once_with( + vector_store_id="T37J8R4WTM", + query="what is litellm?", + custom_llm_provider="bedrock", + metadata={"user_api_key_team_id": "team-a"}, + ) + assert messages[0]["content"] == "Context:\n\nrouted context\n\n" + + @pytest.mark.asyncio async def test_e2e_bedrock_knowledgebase_retrieval_with_completion( setup_vector_store_registry, diff --git a/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py index 75dacbaf08e..2cc9914c9b3 100644 --- a/tests/router_unit_tests/test_router_embedding_integration.py +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -5,17 +5,218 @@ These tests simulate real-world scenarios where headers and configuration need to be properly propagated through the router to the LLM API. """ -from unittest.mock import MagicMock, patch, AsyncMock +import json +from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx - +import litellm from litellm import Router +from litellm.llms.base_llm.vector_store.transformation import ( + LiteLLMVectorStoreEmbeddingExecutor, + RouterVectorStoreEmbeddingExecutor, +) + +QUERY_VECTOR = [0.5, -0.25, 0.125] +OPENAI_EMBEDDINGS_URL = "https://api.openai.com/v1/embeddings" +STORE_EMBEDDINGS_URL = "https://embedding.example/v1/embeddings" + + +def _mock_embedding_route(respx_mock: respx.MockRouter, url: str) -> respx.Route: + return respx_mock.post(url).mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": QUERY_VECTOR}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + ) + + +def _sent(route: respx.Route, index: int) -> tuple[str, str, list[str]]: + request = route.calls[index].request + body = json.loads(request.read()) + return request.headers["authorization"], body["model"], body["input"] + + +def _alias_router() -> Router: + return Router( + model_list=[ + { + "model_name": "team-alias", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "deployment-key", + }, + } + ] + ) class TestRouterEmbeddingIntegration: """Integration tests for embedding with router configuration.""" + def test_vector_store_request_metadata_prefers_litellm_metadata(self): + assert Router._vector_store_request_metadata( + { + "litellm_metadata": {"user_api_key_team_id": "team-a"}, + "metadata": {"user_api_key_team_id": "team-b"}, + } + ) == {"user_api_key_team_id": "team-a"} + + assert Router._vector_store_request_metadata({"metadata": {"user_api_key_team_id": "team-b"}}) == { + "user_api_key_team_id": "team-b" + } + assert Router._vector_store_request_metadata({}) == {} + + def test_sync_vector_store_wrapper_injects_router_embedding_executor(self): + router = Router(model_list=[]) + original = MagicMock(return_value="searched") + wrapped = router.factory_function(original, call_type="vector_store_search") + + assert ( + wrapped( + vector_store_id="store", + query="query", + custom_llm_provider="valkey", + metadata={"user_api_key_team_id": "team-a"}, + ) + == "searched" + ) + + call_kwargs = original.call_args.kwargs + assert call_kwargs["custom_llm_provider"] == "valkey" + executor = call_kwargs["_direct_vector_store_embedding_executor"] + assert isinstance(executor, RouterVectorStoreEmbeddingExecutor) + assert executor.metadata == {"user_api_key_team_id": "team-a"} + + def test_sync_vector_store_wrapper_preserves_model_routing(self): + router = Router(model_list=[]) + original = MagicMock() + wrapped = router.factory_function(original, call_type="vector_store_search") + + with patch.object(router, "_generic_api_call_with_fallbacks", return_value="routed") as fallback: + assert wrapped(model="vector-alias", vector_store_id="store", query="query") == "routed" + + assert fallback.call_args.kwargs["model"] == "vector-alias" + assert fallback.call_args.kwargs["original_function"] is original + assert isinstance( + fallback.call_args.kwargs["_direct_vector_store_embedding_executor"], + RouterVectorStoreEmbeddingExecutor, + ) + + @pytest.mark.asyncio + async def test_vector_store_embedding_executors_cover_sdk_and_router_paths( + self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + openai_route = _mock_embedding_route(respx_mock, OPENAI_EMBEDDINGS_URL) + store_route = _mock_embedding_route(respx_mock, STORE_EMBEDDINGS_URL) + sdk_executor = LiteLLMVectorStoreEmbeddingExecutor() + + sync_response = sdk_executor.embed("openai/text-embedding-3-small", "sync", {"api_key": "explicit"}) + async_response = await sdk_executor.aembed("openai/text-embedding-3-small", "async", {"api_key": "explicit"}) + + assert sync_response.data[0]["embedding"] == QUERY_VECTOR + assert async_response.data[0]["embedding"] == QUERY_VECTOR + assert _sent(openai_route, 0) == ("Bearer explicit", "text-embedding-3-small", ["sync"]) + assert _sent(openai_route, 1) == ("Bearer explicit", "text-embedding-3-small", ["async"]) + + explicit_config = { + "api_base": "https://embedding.example/v1", + "api_key": "store-key", + "metadata": { + "configured": True, + "user_api_key_team_id": "untrusted-team", + }, + "model": "untrusted-model", + } + mock_router = MagicMock() + mock_router.embedding.return_value = sync_response + router_executor = RouterVectorStoreEmbeddingExecutor( + router=mock_router, + metadata={"user_api_key_team_id": "team-a"}, + ) + assert router_executor.embed("team-alias", "query", explicit_config) is sync_response + mock_router.embedding.assert_called_once_with( + model="team-alias", + input=["query"], + api_base="https://embedding.example/v1", + api_key="store-key", + metadata={"configured": True, "user_api_key_team_id": "team-a"}, + ) + + alias_executor = RouterVectorStoreEmbeddingExecutor( + router=_alias_router(), + metadata={"user_api_key_team_id": "team-a"}, + ) + sync_alias = alias_executor.embed("team-alias", "sync query", explicit_config) + async_alias = await alias_executor.aembed("team-alias", "async query", explicit_config) + + assert sync_alias.data[0]["embedding"] == QUERY_VECTOR + assert async_alias.data[0]["embedding"] == QUERY_VECTOR + assert openai_route.call_count == 2 + assert _sent(store_route, 0) == ("Bearer store-key", "text-embedding-3-small", ["sync query"]) + assert _sent(store_route, 1) == ("Bearer store-key", "text-embedding-3-small", ["async query"]) + + @pytest.mark.asyncio + async def test_router_executor_falls_back_to_sdk_for_models_the_router_does_not_serve( + self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + store_route = _mock_embedding_route(respx_mock, STORE_EMBEDDINGS_URL) + executor = RouterVectorStoreEmbeddingExecutor( + router=_alias_router(), + metadata={"user_api_key_team_id": "team-a"}, + ) + inline_config = {"api_base": "https://embedding.example/v1", "api_key": "store-key"} + + sync_response = executor.embed("openai/text-embedding-3-large", "sync query", inline_config) + async_response = await executor.aembed("openai/text-embedding-3-large", "async query", inline_config) + + assert sync_response.data[0]["embedding"] == QUERY_VECTOR + assert async_response.data[0]["embedding"] == QUERY_VECTOR + assert _sent(store_route, 0) == ("Bearer store-key", "text-embedding-3-large", ["sync query"]) + assert _sent(store_route, 1) == ("Bearer store-key", "text-embedding-3-large", ["async query"]) + + @pytest.mark.asyncio + async def test_router_executor_rejects_unserved_models_without_explicit_config( + self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + openai_route = _mock_embedding_route(respx_mock, OPENAI_EMBEDDINGS_URL) + executor = RouterVectorStoreEmbeddingExecutor( + router=_alias_router(), + metadata={"user_api_key_team_id": "team-a"}, + ) + + with pytest.raises(litellm.BadRequestError): + executor.embed("openai/text-embedding-3-large", "sync query", {}) + with pytest.raises(litellm.BadRequestError): + await executor.aembed("openai/text-embedding-3-large", "async query", {}) + + assert openai_route.call_count == 0 + + def test_router_executor_routes_deployment_model_names_through_the_router( + self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + openai_route = _mock_embedding_route(respx_mock, OPENAI_EMBEDDINGS_URL) + executor = RouterVectorStoreEmbeddingExecutor(router=_alias_router(), metadata={}) + + response = executor.embed("openai/text-embedding-3-small", "query", {}) + + assert response.data[0]["embedding"] == QUERY_VECTOR + assert _sent(openai_route, 0) == ("Bearer deployment-key", "text-embedding-3-small", ["query"]) + def test_embedding_with_deployment_specific_headers(self): """ Test that deployment-specific headers are propagated. @@ -122,9 +323,7 @@ class TestRouterEmbeddingIntegration: router = Router( model_list=model_list, - default_litellm_params={ - "metadata": {"environment": "test", "service": "embedding-service"} - }, + default_litellm_params={"metadata": {"environment": "test", "service": "embedding-service"}}, ) with patch("litellm.embedding") as mock_embedding: @@ -240,9 +439,7 @@ class TestRouterEmbeddingIntegration: # Make multiple calls and verify headers are always present for i in range(5): with patch("litellm.embedding") as mock_embedding: - mock_embedding.return_value = MagicMock( - data=[{"embedding": [0.1, 0.2]}] - ) + mock_embedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2]}]) router.embedding(model="shared-embedding-model", input=[f"test {i}"]) @@ -327,9 +524,7 @@ class TestRouterEmbeddingIntegration: router = Router( model_list=model_list, - default_litellm_params={ - "headers": {"X-Custom-Azure-Header": "azure-value"} - }, + default_litellm_params={"headers": {"X-Custom-Azure-Header": "azure-value"}}, ) with patch("litellm.embedding") as mock_embedding: diff --git a/tests/sdk_function_trace/README.md b/tests/sdk_function_trace/README.md new file mode 100644 index 00000000000..d3a3b654aea --- /dev/null +++ b/tests/sdk_function_trace/README.md @@ -0,0 +1,30 @@ +# SDK function tracing + +The compare runner executes the same SDK calls through the Python engine and the Rust native bridge against a local HTTP provider fixture, then prints their pipeline trees side by side. Matching calls align on the same row in green; Python-only calls are blue, Rust-only calls yellow, and reordered calls red. Gaps preserve execution order and each column retains its own nesting. A comparison column labels every row even without color. Colors are enabled in terminals unless `NO_COLOR` is set. A difference summary follows (shared step order, python-only steps, rust-only steps). Each invocation must issue exactly one HTTP request. It requires the LiteLLM Python dependencies and the native extension built with tracing support + +From the repository root, using the project's Python environment: + +```bash +uv run python -m tests.sdk_function_trace.compare +uv run python -m tests.sdk_function_trace.compare --route ocr +uv run python -m tests.sdk_function_trace.compare --route ocr --sync +uv run python -m tests.sdk_function_trace.compare --route all --both --check +``` + +Calls default to async; use `--sync` for synchronous calls or `--both` for the complete matrix. Python sync Messages raises `not implemented for sync calls`; only that exact failure is marked `SKIP`, and the runner still executes Rust sync Messages and subsequent routes. Bedrock transcription has no independent Python provider implementation: its Python trace covers SDK dispatch into Rust + +Both engines are projected onto a shared per-route step table (`steps.py`): canonical names such as `transform_ocr_request` map Python functions (`MistralOCRConfig.transform_ocr_request`) and Rust spans (`transform_ocr_request`) to the same label. Only the first occurrence of each step is kept. Python indentation uses each event's actual frame ancestors and the nearest already displayed ancestor, so returned helpers and coroutine resumptions do not create false parents. Rust indentation uses instrumented span ancestry. Unmatched Rust span names pass through unchanged. `--full` prints every captured runtime event; validation still uses projected steps + +Every report checks required stage presence and dependency order. Provider lookup must precede request transformation, which must precede HTTP, followed by response transformation. The handler must precede HTTP; parameter mapping and supported-parameter checks must precede request transformation. Environment validation and URL construction, where mapped, must precede HTTP. Python transcription is checked only through native dispatch. `--check` also requires identical canonical step sequences for comparable routes and exits nonzero for missing, extra, or reordered steps, or an unexpected call failure, after finishing all selected cases + +Individual stage checks are separate from cross-language `step parity`. Passing stage checks cannot override a failing step comparison. Bedrock transcription and Python sync Messages report `UNAVAILABLE` for cross-language parity because they lack an independent Python execution to compare. Absolute nesting depth is not a cross-language gate: async Python Messages dispatches its handler onto another thread. See `route-comparison.md` for the audited matrix and remaining contract limitations + +The Python runner uses the existing `profile_python` / `sys.setprofile` collector, selecting executed code under the installed `litellm` source directory instead of maintaining a function-name allowlist. It prints source locations and qualified function names, including repeated calls. Coroutine resumptions are counted once per invocation. It profiles the current thread and threads created during the call, including the fresh async executor. Existing worker threads are not retroactively profiled; background Python calls may appear, and indentation follows selected Python stack ancestors within each thread + +The Rust runner calls the compiled PyO3 SDK entrypoints with `trace=True`. The existing `FunctionTrace` subscriber collects `#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]` spans for the route entrypoint, preparation, provider lookup, HTTP handler, and selected provider transformations. The shared `http_request` helper instruments the existing Rust send operation without changing clients, timeouts, signing, or error mapping. Function names come from the actual functions. `WithSubscriber` attaches the collector to each future across async polls. Arguments and provider payloads are not recorded in trace events. Uninstrumented functions do not appear; this is scoped instrumentation, not an exhaustive native call graph + +Tracing is opt-in: native calls without `trace=True` keep their original response shape. Traced calls return `{"response": ..., "trace": [{"function": ..., "depth": ...}]}`. The runners print only trace events. Missing native support or empty traces fail instead of falling back to source searching. The old `--repo`, `--signatures`, and `--calls` options are removed + +`profile_python(functions)` still supports direct function references for focused parity checks. `assert_function_trace_parity` compares selected Python events with Rust events supplied by an executable scenario. Successful stage checks prove the declared pipeline ran in a valid dependency order for this fixture; they do not assert identical function contracts, request bodies, responses, streaming behavior, or live-provider correctness + +Build the extension with `maturin develop` in the project's virtual environment. Then run either command above to get the executed function order diff --git a/tests/sdk_function_trace/__init__.py b/tests/sdk_function_trace/__init__.py new file mode 100644 index 00000000000..da62b8041f6 --- /dev/null +++ b/tests/sdk_function_trace/__init__.py @@ -0,0 +1,13 @@ +from tests.sdk_function_trace.harness import ( + TraceScenario, + TraceStep, + assert_function_trace_parity, +) +from tests.sdk_function_trace.profiler import FunctionTraceEvent + +__all__ = [ + "FunctionTraceEvent", + "TraceScenario", + "TraceStep", + "assert_function_trace_parity", +] diff --git a/tests/sdk_function_trace/compare.py b/tests/sdk_function_trace/compare.py new file mode 100644 index 00000000000..941c1b6e067 --- /dev/null +++ b/tests/sdk_function_trace/compare.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import argparse +import os +import sys +from typing import Final + +from tests.sdk_function_trace.fixtures import ROUTES +from tests.sdk_function_trace.report import compare, render + + +def _run(route: str, asynchronous: bool, *, full: bool, colorize: bool) -> bool: + comparison: Final = compare(route, asynchronous=asynchronous) + sys.stdout.write(render(comparison, full=full, colorize=colorize)) + return comparison.passed + + +def main() -> None: + parser: Final = argparse.ArgumentParser(description="Compare Python and Rust SDK pipeline steps per route") + parser.add_argument("--route", choices=("all", *ROUTES), default="all") + mode: Final = parser.add_mutually_exclusive_group() + mode.add_argument("--async", dest="asynchronous", action="store_true", default=True) + mode.add_argument("--sync", dest="asynchronous", action="store_false") + mode.add_argument("--both", action="store_true", help="run async and sync for every selected route") + parser.add_argument( + "--check", action="store_true", help="exit nonzero for missing, extra, or reordered comparable steps" + ) + parser.add_argument( + "--full", action="store_true", help="print every captured runtime event instead of pipeline steps" + ) + args: Final = parser.parse_args() + os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True") + colorize: Final = sys.stdout.isatty() and "NO_COLOR" not in os.environ + results: Final = tuple( + _run(selected, selected_mode, full=args.full, colorize=colorize) + for selected in ROUTES + if args.route in ("all", selected) + for selected_mode in ((True, False) if args.both else (args.asynchronous,)) + ) + if args.check and not all(results): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tests/sdk_function_trace/fixtures.py b/tests/sdk_function_trace/fixtures.py new file mode 100644 index 00000000000..47bbe839627 --- /dev/null +++ b/tests/sdk_function_trace/fixtures.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +import base64 +import io +import json +import wave +from collections.abc import Callable +from dataclasses import dataclass +from typing import Final, Protocol, cast + +from tests.sdk_function_trace.mock_provider import MockProviderResponse +from tests.sdk_function_trace.steps import Engine + +ANTHROPIC_MODEL: Final = "claude-sonnet-5" +OCR_MODEL: Final = "mistral-ocr-latest" +AUDIO_MODEL: Final = "mistral.voxtral-mini-3b-2507" + + +class SdkCall(Protocol): + def __call__(self, **kwargs: object) -> object: ... + + +@dataclass(frozen=True, slots=True) +class Fixture: + kwargs: dict[str, object] + provider_response: MockProviderResponse + + +@dataclass(frozen=True, slots=True) +class RouteSpec: + label: str + python_entrypoints: tuple[str, str] + rust_entrypoints: tuple[str, str] + fixture: Callable[[Engine], Fixture] + + +@dataclass(frozen=True, slots=True) +class Invocation: + function: SdkCall + kwargs: dict[str, object] + provider_response: MockProviderResponse + label: str + + +def audio_bytes() -> bytes: + with io.BytesIO() as buffer: + with wave.open(buffer, "wb") as audio: + audio.setnchannels(1) + audio.setsampwidth(2) + audio.setframerate(16000) + audio.writeframes(b"\x00\x00" * 1600) + return buffer.getvalue() + + +def _anthropic_message_response() -> MockProviderResponse: + body: Final = { + "id": "msg_trace", + "type": "message", + "role": "assistant", + "model": ANTHROPIC_MODEL, + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 3}, + } + return MockProviderResponse(200, (("content-type", "application/json"),), json.dumps(body).encode()) + + +def _conversation() -> dict[str, object]: + return {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16} + + +def _ocr_fixture(engine: Engine) -> Fixture: + return Fixture( + kwargs={ + "model": f"mistral/{OCR_MODEL}", + "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, + **({"optional_params": {"pages": [0]}} if engine == "rust" else {"pages": [0]}), + }, + provider_response=MockProviderResponse( + 200, + (("content-type", "application/json"),), + json.dumps( + { + "pages": [{"index": 0, "markdown": "hello"}], + "model": OCR_MODEL, + "usage_info": {"pages_processed": 1}, + } + ).encode(), + ), + ) + + +def _chat_completions_fixture(engine: Engine) -> Fixture: + conversation: Final = _conversation() + payload: Final = ( + {"messages": conversation["messages"], "optional_params": {"max_tokens": 16}} + if engine == "rust" + else conversation + ) + return Fixture( + kwargs={"model": f"anthropic/{ANTHROPIC_MODEL}", **payload}, + provider_response=_anthropic_message_response(), + ) + + +def _messages_fixture(engine: Engine) -> Fixture: + conversation: Final = _conversation() + payload: Final = {"body": {**conversation, "model": ANTHROPIC_MODEL}} if engine == "rust" else conversation + return Fixture( + kwargs={"model": f"anthropic/{ANTHROPIC_MODEL}", **payload}, + provider_response=_anthropic_message_response(), + ) + + +def _transcription_fixture(engine: Engine) -> Fixture: + credentials: Final = { + "aws_access_key_id": "test-access", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", + } + payload: Final = ( + { + "audio": {"data": base64.b64encode(audio_bytes()).decode(), "format": "wav"}, + "optional_params": credentials, + } + if engine == "rust" + else {"file": ("sample.wav", audio_bytes(), "audio/wav"), **credentials} + ) + return Fixture( + kwargs={"model": f"bedrock/{AUDIO_MODEL}", **payload}, + provider_response=MockProviderResponse( + 200, + (("content-type", "application/json"),), + json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 2, "outputTokens": 3, "totalTokens": 5}, + } + ).encode(), + ), + ) + + +ROUTE_SPECS: Final[dict[str, RouteSpec]] = { + "chat_completions": RouteSpec( + label="anthropic", + python_entrypoints=("completion", "acompletion"), + rust_entrypoints=("chat_completions", "achat_completions"), + fixture=_chat_completions_fixture, + ), + "audio_transcription": RouteSpec( + label="bedrock (Rust-only provider; Python trace covers SDK dispatch)", + python_entrypoints=("transcription", "atranscription"), + rust_entrypoints=("transcription", "atranscription"), + fixture=_transcription_fixture, + ), + "messages": RouteSpec( + label="anthropic", + python_entrypoints=("create", "acreate"), + rust_entrypoints=("messages", "amessages"), + fixture=_messages_fixture, + ), + "ocr": RouteSpec( + label="mistral", + python_entrypoints=("ocr", "aocr"), + rust_entrypoints=("ocr", "aocr"), + fixture=_ocr_fixture, + ), +} + +ROUTES: Final = tuple(ROUTE_SPECS) + + +def sdk_invocation(route: str, *, engine: Engine, asynchronous: bool) -> Invocation: + import litellm + from litellm.anthropic_interface import messages as sdk_messages + from litellm.rust_bridge import get_native_bridge + + rust: Final = engine == "rust" + bridge: Final = get_native_bridge() if rust else None + if rust and bridge is None: + raise RuntimeError("Build the native extension first: maturin develop") + spec: Final = ROUTE_SPECS.get(route) + if spec is None: + raise ValueError(f"Unknown route: {route}") + fixture: Final = spec.fixture(engine) + owner: Final = bridge if rust else (sdk_messages if route == "messages" else litellm) + entrypoint: Final = (spec.rust_entrypoints if rust else spec.python_entrypoints)[int(asynchronous)] + return Invocation( + function=cast(SdkCall, getattr(owner, entrypoint)), + kwargs={ + **fixture.kwargs, + "api_key": "test-key", + **({"trace": True, "timeout_seconds": 5} if rust else {"timeout": 5}), + }, + provider_response=fixture.provider_response, + label=spec.label, + ) diff --git a/tests/sdk_function_trace/harness.py b/tests/sdk_function_trace/harness.py new file mode 100644 index 00000000000..8f707402449 --- /dev/null +++ b/tests/sdk_function_trace/harness.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from types import FunctionType +from typing import Final, cast + +from tests.sdk_function_trace.profiler import FunctionTraceEvent, profile_python + + +@dataclass(frozen=True, slots=True) +class TraceStep: + function: FunctionType + depth: int + + +@dataclass(frozen=True, slots=True) +class TraceScenario: + steps: tuple[TraceStep, ...] + invoke_python: Callable[[], object] + invoke_rust: Callable[[], Sequence[FunctionTraceEvent]] + + +def assert_function_trace_parity(scenario: TraceScenario) -> None: + expected: Final = tuple( + FunctionTraceEvent(function=step.function.__name__, depth=step.depth) for step in scenario.steps + ) + functions: Final = cast(tuple[FunctionType, ...], tuple(step.function for step in scenario.steps)) + with profile_python(functions) as profiler: + scenario.invoke_python() + python_trace: Final = tuple(profiler.events) + rust_trace: Final = tuple(scenario.invoke_rust()) + + if python_trace != expected: + raise AssertionError(f"Python function trace differs: {python_trace!r} != {expected!r}") + if rust_trace != expected: + raise AssertionError(f"Rust function trace differs: {rust_trace!r} != {expected!r}") + if python_trace != rust_trace: + raise AssertionError(f"Python and Rust function traces differ: {python_trace!r} != {rust_trace!r}") diff --git a/tests/sdk_function_trace/mock_provider.py b/tests/sdk_function_trace/mock_provider.py new file mode 100644 index 00000000000..37eca665586 --- /dev/null +++ b/tests/sdk_function_trace/mock_provider.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Lock, Thread +from typing import Final, cast + + +@dataclass(frozen=True, slots=True) +class MockProviderResponse: + status_code: int + headers: tuple[tuple[str, str], ...] + body: bytes + + +class _MockProviderServer(ThreadingHTTPServer): + def __init__(self, response: MockProviderResponse) -> None: + super().__init__(("127.0.0.1", 0), _MockProviderHandler) + self.response: Final = response + self._request_count = 0 + self._request_count_lock: Final = Lock() + + def record_request(self) -> None: + with self._request_count_lock: + self._request_count += 1 + + @property + def request_count(self) -> int: + with self._request_count_lock: + return self._request_count + + +class _MockProviderHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + content_length: Final = int(self.headers.get("content-length", "0")) + self.rfile.read(content_length) + server: Final = cast(_MockProviderServer, self.server) + server.record_request() + self.send_response(server.response.status_code) + for name, value in server.response.headers: + self.send_header(name, value) + self.send_header("content-length", str(len(server.response.body))) + self.end_headers() + self.wfile.write(server.response.body) + + def log_message(self, format: str, *args: object) -> None: # noqa: A002 # matches BaseHTTPRequestHandler + pass + + +@contextmanager +def mock_provider(response: MockProviderResponse) -> Generator[str]: + server: Final = _MockProviderServer(response) + thread: Final = Thread(target=server.serve_forever, daemon=True) + thread.start() + host, port = cast(tuple[str, int], server.server_address) + try: + yield f"http://{host}:{port}" + finally: + server.shutdown() + server.server_close() + thread.join() + if server.request_count != 1: + raise AssertionError(f"expected one provider request, received {server.request_count}") diff --git a/tests/sdk_function_trace/ocr-comparison.md b/tests/sdk_function_trace/ocr-comparison.md new file mode 100644 index 00000000000..d252480e218 --- /dev/null +++ b/tests/sdk_function_trace/ocr-comparison.md @@ -0,0 +1,59 @@ +# OCR Python and Rust comparison + +Audited implementation revision: `edcba483b2`. The implementations do not match in function contracts, call structure, or all tested response behavior. This audit changes the source listing coverage, not OCR runtime behavior + +Run both source listings from the repository root: + +```bash +python3 tests/sdk_function_trace/list_python_steps.py --route ocr --signatures --calls +uv run tests/sdk_function_trace/list_rust_steps.py --route ocr --signatures --calls +``` + +Both cover Mistral, Azure AI Mistral, Azure Document Intelligence, Vertex Mistral, and Vertex DeepSeek. Listings show declarations and source call sites, not executed traces + +## Function contracts + +Comparing Python `BaseOCRConfig` with Rust `OcrProviderConfig`, omitting `self` and language-specific ownership details: + +| Python | Rust | Difference | +| --- | --- | --- | +| `get_supported_ocr_params(model)` | `supported_ocr_params()` | Name and model argument | +| `get_api_key_env_var()` | No corresponding method | Missing contract | +| `map_ocr_params(non_default_params, optional_params, model)` | `map_ocr_params(non_default_params)` | Missing accumulator and model | +| `validate_environment(headers, model, api_key, api_base, litellm_params, **kwargs)` | Separate auth/key/header helpers | Different contract | +| `get_complete_url(api_base, model, optional_params, litellm_params, **kwargs)` | `complete_url(api_base, model, optional_params, env_lookup)` | Name and context | +| `transform_ocr_request(model, document, optional_params, headers, **kwargs)` | `transform_ocr_request(model, document, optional_params)` | Missing headers and extra context | +| `async_transform_ocr_request(...)` | No corresponding method | Missing async override | +| `transform_ocr_response(model, raw_response, logging_obj, **kwargs)` | `transform_ocr_response(model, response_json)` | Missing HTTP metadata, logging and extra context | +| `async_transform_ocr_response(...)` | No corresponding method | Missing async override | +| `get_error_class(error_message, status_code, headers)` | Central Rust error mapping | Different contract | + +Python's default mapper returns the supplied `optional_params`; Rust's filters `non_default_params`. Provider overrides must also be compared + +Python maps parameters during SDK preparation, before HTTP-handler environment validation and URL construction. Rust resolves auth and URL before mapping parameters in `prepare_provider_request`. Python has async provider transforms; both native entrypoints execute the same Rust async route using synchronous transform hooks, with polling and document downloading in gateway helpers + +The native bindings also accept `optional_params` and `timeout_seconds`, while the Python SDK accepts `**kwargs` and `timeout`. Public SDK calls with Rust enabled still execute Python preparation before entering Rust, so matching SDK responses would not prove matching standalone Rust steps + +## Runtime results + +Built the native extension from the audited source using `cargo build -p litellm-python-bridge --features extension-module --offline`. Supplied that build's functions through `use_litellm_rust` dependency injection. Ran public `litellm.ocr` and `litellm.aocr` with Rust disabled and enabled against identical local HTTP response fixtures, requiring one request per invocation + +Successful `model_dump()` results and failure exception classes were compared. These checks cover Mistral response outcomes only, not request equality, error messages, live providers, or every execution branch + +| Mistral response fixture | Sync | Async | Observation | +| --- | --- | --- | --- | +| Valid page/model/usage | Match | Match | Same normalized response | +| Model omitted | Match | Match | Both use the requested model | +| `model: null` | Different | Different | Python rejects; Rust uses the requested model | +| `pages: null` | Different | Different | Python rejects; Rust returns an empty array | +| Invalid page element | Match | Match | Both reject during response validation | + +Six of ten fixture/mode comparisons match, four differ. Rust's Mistral response transform conflates missing values with explicit nulls through `as_array`/`as_str` fallbacks. Python preserves explicit nulls into response validation, which rejects them + +## Other provider gaps found in source + +Azure Document Intelligence's Python configuration supports `pages`, `features`, and `req_format`; Rust lists only `pages`. Python normalizes parameters before URL construction; Rust normalizes pages during URL construction + +Python preserves Azure `content`, `tables`, and `keyValuePairs`, and supports retaining the native operation payload. Rust's `OcrResponseData` has no corresponding fields, and its Azure transform does not preserve those values + +Azure and Vertex async document transforms and Azure polling also use different helper contracts. Their runtime equivalence was not tested in this audit diff --git a/tests/sdk_function_trace/profiler.py b/tests/sdk_function_trace/profiler.py new file mode 100644 index 00000000000..c71c74ab0d3 --- /dev/null +++ b/tests/sdk_function_trace/profiler.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import sys +import threading +from collections.abc import Generator, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from types import CodeType, FrameType, FunctionType +from typing import Final + + +@dataclass(frozen=True, slots=True) +class FunctionTraceEvent: + function: str + depth: int + ancestors: tuple[str, ...] | None = None + + +class PythonProfiler: + def __init__(self, functions: Sequence[FunctionType], source_root: Path | None = None) -> None: + self._source_root: Final = str(source_root.resolve()) + "/" if source_root is not None else None + self._names_by_code: Final = {function.__code__: function.__name__ for function in functions} + self._seen_frames: Final[set[FrameType]] = set() + self.events: Final[list[FunctionTraceEvent]] = [] + + def __call__(self, frame: FrameType, event: str, _arg: object) -> None: + if event != "call" or frame in self._seen_frames: + return + function_name: Final = self.function_name(frame.f_code) + if function_name is None: + return + ancestors: Final = tuple( + name for ancestor in _frame_ancestors(frame) if (name := self.function_name(ancestor.f_code)) is not None + ) + self._seen_frames.add(frame) + self.events.append( + FunctionTraceEvent( + function=function_name, + depth=len(ancestors), + ancestors=ancestors if self._source_root is not None else None, + ) + ) + + def function_name(self, code: CodeType) -> str | None: + if self._source_root is None: + return self._names_by_code.get(code) + if not code.co_filename.startswith(self._source_root): + return None + relative: Final = code.co_filename.removeprefix(self._source_root) + return f"{relative}:{code.co_firstlineno} {getattr(code, 'co_qualname', code.co_name)}" + + +def _frame_ancestors(frame: FrameType) -> Generator[FrameType]: + ancestor: Final = frame.f_back + if ancestor is not None: + yield ancestor + yield from _frame_ancestors(ancestor) + + +@contextmanager +def profile_python( + functions: Sequence[FunctionType] = (), *, source_root: Path | None = None, threads: bool = False +) -> Generator[PythonProfiler]: + profiler: Final = PythonProfiler(functions, source_root) + previous_thread: Final = threading.getprofile() + if threads: + threading.setprofile(profiler) + previous: Final = sys.getprofile() + sys.setprofile(profiler) + try: + yield profiler + finally: + sys.setprofile(previous) + if threads: + threading.setprofile(previous_thread) diff --git a/tests/sdk_function_trace/report.py b/tests/sdk_function_trace/report.py new file mode 100644 index 00000000000..9b654e571f8 --- /dev/null +++ b/tests/sdk_function_trace/report.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +from tests.sdk_function_trace.fixtures import ROUTE_SPECS +from tests.sdk_function_trace.profiler import FunctionTraceEvent +from tests.sdk_function_trace.runtime import ( + TraceDiff, + TraceFailed, + TraceOk, + TraceRun, + TraceSkipped, + attempt_trace, + trace_diff, +) +from tests.sdk_function_trace.steps import Engine, pipeline_issues, pipeline_steps +from tests.sdk_function_trace.table import format_trace_table + +_PYTHON_ONLY_COLOR: Final = "\033[34m" +_RUST_ONLY_COLOR: Final = "\033[33m" +_RESET: Final = "\033[0m" + +_ENGINE_COLOR: Final[dict[Engine, str]] = {"python": _PYTHON_ONLY_COLOR, "rust": _RUST_ONLY_COLOR} + + +@dataclass(frozen=True, slots=True) +class EngineReport: + engine: Engine + run: TraceRun + events: tuple[FunctionTraceEvent, ...] + steps: tuple[FunctionTraceEvent, ...] + issues: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class Comparison: + route: str + label: str + asynchronous: bool + engines: tuple[EngineReport, ...] + diff: TraceDiff + + @property + def comparable(self) -> bool: + return self.route != "audio_transcription" and all(isinstance(report.run, TraceOk) for report in self.engines) + + @property + def passed(self) -> bool: + return ( + (not self.comparable or self.diff.matches) + and not any(report.issues for report in self.engines) + and all(not isinstance(report.run, TraceFailed) for report in self.engines) + ) + + +def _events(run: TraceRun) -> tuple[FunctionTraceEvent, ...]: + match run: + case TraceOk(events=events): + return events + case TraceSkipped() | TraceFailed(): + return () + + +def _engine_report(route: str, engine: Engine, run: TraceRun) -> EngineReport: + events: Final = _events(run) + steps: Final = pipeline_steps(route, engine, events) + issues: Final = pipeline_issues(route, engine, steps) if isinstance(run, TraceOk) else () + return EngineReport(engine=engine, run=run, events=events, steps=steps, issues=issues) + + +def compare(route: str, *, asynchronous: bool) -> Comparison: + runs: Final = { + engine: attempt_trace(route, engine=engine, asynchronous=asynchronous) for engine in ("python", "rust") + } + engines: Final = tuple(_engine_report(route, engine, run) for engine, run in runs.items()) + return Comparison( + route=route, + label=ROUTE_SPECS[route].label, + asynchronous=asynchronous, + engines=engines, + diff=trace_diff(engines[0].steps, engines[1].steps), + ) + + +def _tree_line(event: FunctionTraceEvent, only: frozenset[str], marker: str, color: str, *, colorize: bool) -> str: + line: Final = f"{' ' * event.depth}{event.function}" + (f" {marker}" if event.function in only else "") + return f"{color}{line}{_RESET}\n" if colorize and event.function in only else f"{line}\n" + + +def _tree_lines( + events: tuple[FunctionTraceEvent, ...], + only: frozenset[str], + marker: str, + color: str, + *, + colorize: bool, +) -> tuple[str, ...]: + return tuple(_tree_line(event, only, marker, color, colorize=colorize) for event in events) + + +def _engine_lines( + report: EngineReport, diff: TraceDiff, *, comparable: bool, full: bool, colorize: bool +) -> tuple[str, ...]: + match report.run: + case TraceSkipped(reason=reason): + return (f"{report.engine}: SKIP ({reason})\n\n",) + case TraceFailed(reason=reason): + return (f"{report.engine}: FAIL ({reason})\n\n",) + case TraceOk(): + shown: Final = report.events if full else report.steps + only: Final = ( + () if full or not comparable else (diff.python_only if report.engine == "python" else diff.rust_only) + ) + return ( + f"{report.engine} ({len(shown)} steps)\n\n", + *_tree_lines( + shown, + frozenset(only), + f"<- {report.engine} only", + _ENGINE_COLOR[report.engine], + colorize=colorize, + ), + "\n", + ) + + +def _parity_lines(comparison: Comparison) -> tuple[str, ...]: + if not comparison.comparable: + if comparison.route == "audio_transcription": + return ("step parity: UNAVAILABLE (Bedrock transcription has no independent Python implementation)\n",) + return ("step parity: UNAVAILABLE (both engines must complete)\n",) + diff: Final = comparison.diff + order: Final = "the same" if diff.shared_order_matches else "a different" + return ( + "diff\n\n", + f"shared steps appear in {order} order\n", + f"python-only: {', '.join(diff.python_only) or 'none'}\n", + f"rust-only: {', '.join(diff.rust_only) or 'none'}\n\n", + f"step parity: {'PASS' if diff.matches else 'FAIL'}\n", + ) + + +def _stage_lines(comparison: Comparison) -> tuple[str, ...]: + return tuple( + f"{report.engine} " + f"{'SDK dispatch only' if comparison.route == 'audio_transcription' and report.engine == 'python' else 'pipeline'}: " + f"{'FAIL: ' + '; '.join(report.issues) if report.issues else 'PASS'}\n" + for report in comparison.engines + if isinstance(report.run, TraceOk) + ) + + +def render(comparison: Comparison, *, full: bool, colorize: bool) -> str: + mode: Final = "async" if comparison.asynchronous else "sync" + traces: Final = ( + (format_trace_table(comparison.engines[0].steps, comparison.engines[1].steps, colorize=colorize) + "\n\n",) + if not full and all(isinstance(report.run, TraceOk) for report in comparison.engines) + else tuple( + line + for report in comparison.engines + for line in _engine_lines( + report, comparison.diff, comparable=comparison.comparable, full=full, colorize=colorize + ) + ) + ) + return "".join( + ( + f"route: {comparison.route} provider: {comparison.label} mode: {mode}\n\n", + *traces, + *_parity_lines(comparison), + *_stage_lines(comparison), + "Each successful invocation issued exactly one local provider request\n\n", + ) + ) diff --git a/tests/sdk_function_trace/route-comparison.md b/tests/sdk_function_trace/route-comparison.md new file mode 100644 index 00000000000..009d3544d05 --- /dev/null +++ b/tests/sdk_function_trace/route-comparison.md @@ -0,0 +1,26 @@ +# SDK route trace audit + +Run the four native HTTP route families in both modes from the repository root: + +```bash +uv run python -m tests.sdk_function_trace.compare --route all --both --check +``` + +The local fixture matrix on 2026-09-02 completed 15 successful engine invocations and one expected skip. Every successful invocation issued exactly one local HTTP request. All five comparable route/mode pairs have identical canonical steps in the same order, with no Python-only or Rust-only steps + +| Route | Python async | Python sync | Rust async | Rust sync | +| --- | --- | --- | --- | --- | +| Chat completions, Anthropic | Pass | Pass | Pass | Pass | +| Messages, Anthropic | Pass | Unsupported, skipped | Pass | Pass | +| OCR, Mistral | Pass | Pass | Pass | Pass | +| Audio transcription, Bedrock | Dispatch only | Dispatch only | Pass | Pass | + +The same canonical step sequence ran in sync and async for each engine with both modes available. Bedrock transcription's Python SDK delegates to Rust, so its two successful calls do not establish independent provider parity. Realtime and Responses WebSockets are outside this HTTP fixture runner + +Chat and OCR also have identical projected nesting in both modes. Async Messages has the same helper nesting beneath its handler, but Python starts that handler on a worker thread, so it appears as a second root. The comparison preserves this physical thread boundary and checks step order independently of absolute depth + +Rust now resolves chat providers and supported parameters before entering its handler. Chat and Messages validate the environment and transform requests inside their handlers. Messages builds the final URL after transformation. OCR resolves its config and maps supported parameters during preparation, then validates credentials, builds the URL, and transforms the request inside its handler. Its during-call guardrails still run before HTTP, within the provider-call lifecycle phase + +The environment hooks execute credential and header validation. Chat's supported-parameter hooks return OpenAI names paired with provider names and feed the existing request acceptance checks. The direct Rust API still accepts provider-mapped parameters, and its supported subset is smaller than Python's. Matching the pipeline does not establish identical parameter contracts + +`--check` now fails if either comparable engine has missing, extra, or reordered canonical steps, even if its individual stage checks pass. Bedrock transcription and sync Messages report `UNAVAILABLE` for cross-language parity; native execution is still checked. Passing establishes step coverage and order for one non-streaming fixture per route, not complete request, response, error, or provider parity. The previously recorded OCR response gaps remain in `ocr-comparison.md` diff --git a/tests/sdk_function_trace/runtime.py b/tests/sdk_function_trace/runtime.py new file mode 100644 index 00000000000..d5bf15694bc --- /dev/null +++ b/tests/sdk_function_trace/runtime.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import asyncio +import os +from collections.abc import Awaitable, Generator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final, cast +from unittest.mock import patch + +from pydantic import BaseModel, ConfigDict + +from tests.sdk_function_trace.fixtures import Invocation, sdk_invocation +from tests.sdk_function_trace.mock_provider import mock_provider +from tests.sdk_function_trace.profiler import FunctionTraceEvent, profile_python +from tests.sdk_function_trace.steps import Engine + + +class TraceEventPayload(BaseModel): + model_config = ConfigDict(strict=True, extra="forbid") + function: str + depth: int + + +class TraceResponsePayload(BaseModel): + model_config = ConfigDict(strict=True, extra="forbid") + response: object + trace: tuple[TraceEventPayload, ...] | list[TraceEventPayload] + + +@contextmanager +def _python_engine() -> Generator[None]: + from litellm.rust_bridge import ocr as ocr_bridge + + previous_ocr: Final = ocr_bridge.rust_ocr_enabled() + with patch.dict(os.environ, {"LITELLM_RUST": "false"}): + ocr_bridge.use_litellm_rust(False) + try: + yield + finally: + ocr_bridge.use_litellm_rust(previous_ocr) + + +def _invoke(case: Invocation, api_base: str, *, asynchronous: bool) -> object: + async def invoke_async() -> object: + return await cast("Awaitable[object]", case.function(**case.kwargs, api_base=api_base)) + + if asynchronous: + return asyncio.run(invoke_async()) + return case.function(**case.kwargs, api_base=api_base) + + +def collect(case: Invocation, api_base: str, *, engine: Engine, asynchronous: bool) -> tuple[FunctionTraceEvent, ...]: + import litellm + + if engine == "rust": + payload: Final = TraceResponsePayload.model_validate(_invoke(case, api_base, asynchronous=asynchronous)) + return tuple(FunctionTraceEvent(event.function, event.depth) for event in payload.trace) + with profile_python(source_root=Path(litellm.__file__).parent, threads=True) as profiler: + _invoke(case, api_base, asynchronous=asynchronous) + return tuple(profiler.events) + + +def run_trace(route: str, *, engine: Engine, asynchronous: bool = False) -> tuple[FunctionTraceEvent, ...]: + case: Final = sdk_invocation(route, engine=engine, asynchronous=asynchronous) + with _python_engine(), mock_provider(case.provider_response) as api_base: + events: Final = collect(case, api_base, engine=engine, asynchronous=asynchronous) + if not events: + raise RuntimeError(f"No runtime events for {route}; rebuild the native extension with tracing support") + return events + + +@dataclass(frozen=True, slots=True) +class TraceOk: + events: tuple[FunctionTraceEvent, ...] + + +@dataclass(frozen=True, slots=True) +class TraceSkipped: + reason: str + + +@dataclass(frozen=True, slots=True) +class TraceFailed: + reason: str + + +TraceRun = TraceOk | TraceSkipped | TraceFailed + + +def attempt_trace(route: str, *, engine: Engine, asynchronous: bool) -> TraceRun: + try: + return TraceOk(run_trace(route, engine=engine, asynchronous=asynchronous)) + except Exception as error: + reason: Final = f"{type(error).__name__}: {error}" + if ( + route == "messages" + and engine == "python" + and not asynchronous + and isinstance(error, ValueError) + and str(error) == "anthropic_messages_handler is not implemented for sync calls" + ): + return TraceSkipped(reason) + return TraceFailed(reason) + + +@dataclass(frozen=True, slots=True) +class TraceDiff: + python_only: tuple[str, ...] + rust_only: tuple[str, ...] + shared_order_matches: bool + + @property + def matches(self) -> bool: + return not self.python_only and not self.rust_only and self.shared_order_matches + + +def trace_diff(python: tuple[FunctionTraceEvent, ...], rust: tuple[FunctionTraceEvent, ...]) -> TraceDiff: + python_names: Final = {event.function for event in python} + rust_names: Final = {event.function for event in rust} + shared_python: Final = tuple(event.function for event in python if event.function in rust_names) + shared_rust: Final = tuple(event.function for event in rust if event.function in python_names) + return TraceDiff( + python_only=tuple(event.function for event in python if event.function not in rust_names), + rust_only=tuple(event.function for event in rust if event.function not in python_names), + shared_order_matches=bool(shared_python) and shared_python == shared_rust, + ) diff --git a/tests/sdk_function_trace/steps.py b/tests/sdk_function_trace/steps.py new file mode 100644 index 00000000000..bb50d4ebe57 --- /dev/null +++ b/tests/sdk_function_trace/steps.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import re +from collections.abc import Sequence +from dataclasses import dataclass +from functools import reduce +from typing import Final, Literal + +from tests.sdk_function_trace.profiler import FunctionTraceEvent + +Engine = Literal["python", "rust"] + + +@dataclass(frozen=True, slots=True) +class Step: + name: str + python: re.Pattern[str] | None + rust: str | None + + +def _step(name: str, python: str | None = None, rust: str | None = None) -> Step: + return Step(name, re.compile(python) if python is not None else None, rust) + + +_POST: Final = r"AsyncHTTPHandler\.post$|HTTPHandler\.post$" + +STEPS: Final[dict[str, tuple[Step, ...]]] = { + "ocr": ( + _step("ocr", r"ocr/main\.py:\d+ a?ocr$", "ocr"), + _step("prepare_ocr_call", r"ocr/main\.py:\d+ _prepare_ocr_request$", "prepare_ocr_call"), + _step("get_provider_ocr_config", r"ProviderConfigManager\.get_provider_ocr_config$", "ocr_provider_config"), + _step("supported_ocr_params", r"get_supported_ocr_params$", "supported_ocr_params"), + _step("map_ocr_params", r"(? tuple[str, ...]: + names: Final = tuple(event.function for event in events) + required: Final = tuple(step.name for step in STEPS[route] if getattr(step, engine) is not None) + missing: Final = tuple(f"missing {name}" for name in required if name not in names) + provider: Final = next(name for name in required if name.startswith("get_provider_")) + handler: Final = next(name for name in required if name.startswith("execute_")) + dispatch_only: Final = route == "audio_transcription" and engine == "python" + request: Final = next( + (name for name in required if name.startswith("transform_") and name.endswith("request")), handler + ) + response: Final = next( + (name for name in required if name.startswith("transform_") and name.endswith("response")), handler + ) + phases: Final = ( + (route, "map_transcription_params", provider, handler) + if dispatch_only + else (route, provider, request, "http_request", response) + ) + extra_edges: Final = ( + () + if dispatch_only + else ( + (handler, "http_request"), + *((name, request) for name in required if name.startswith(("map_", "supported_"))), + *((name, "http_request") for name in ("validate_environment", "complete_url") if name in required), + ) + ) + edges: Final = (*zip(phases, phases[1:]), *extra_edges) + return missing + tuple( + f"{before} must precede {after}" + for before, after in edges + if before in names and after in names and names.index(before) >= names.index(after) + ) + + +def _canonical_name(route: str, engine: Engine, function: str) -> str | None: + for step in STEPS[route]: + if engine == "python": + if step.python is not None and step.python.search(function): + return step.name + elif step.rust is not None and function == step.rust: + return step.name + return function if engine == "rust" else None + + +@dataclass(frozen=True, slots=True) +class _Projection: + shown: tuple[FunctionTraceEvent, ...] = () + stack: tuple[tuple[int, int], ...] = () + seen: frozenset[str] = frozenset() + + +def _project(route: str, engine: Engine, state: _Projection, event: FunctionTraceEvent) -> _Projection: + stack: Final = tuple(pair for pair in state.stack if event.depth > pair[0]) + name: Final = _canonical_name(route, engine, event.function) + if name is None or name in state.seen: + return _Projection(state.shown, stack, state.seen) + depth: Final = ( + next( + ( + kept.depth + 1 + for ancestor in event.ancestors + for kept in state.shown + if kept.function == _canonical_name(route, engine, ancestor) + ), + 0, + ) + if event.ancestors is not None + else stack[-1][1] + 1 + if stack + else 0 + ) + return _Projection( + state.shown + (FunctionTraceEvent(function=name, depth=depth),), + stack + ((event.depth, depth),), + state.seen | {name}, + ) + + +def pipeline_steps(route: str, engine: Engine, events: Sequence[FunctionTraceEvent]) -> tuple[FunctionTraceEvent, ...]: + projection: Final = reduce(lambda state, event: _project(route, engine, state, event), events, _Projection()) + return projection.shown diff --git a/tests/sdk_function_trace/table.py b/tests/sdk_function_trace/table.py new file mode 100644 index 00000000000..2124d7e3faf --- /dev/null +++ b/tests/sdk_function_trace/table.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +from collections.abc import Iterator +from difflib import SequenceMatcher +from typing import Final + +from tests.sdk_function_trace.profiler import FunctionTraceEvent + + +def _aligned_rows( + python: tuple[FunctionTraceEvent, ...], rust: tuple[FunctionTraceEvent, ...] +) -> Iterator[tuple[FunctionTraceEvent | None, FunctionTraceEvent | None]]: + matcher: Final = SequenceMatcher( + a=tuple(event.function for event in python), + b=tuple(event.function for event in rust), + autojunk=False, + ) + for tag, python_start, python_end, rust_start, rust_end in matcher.get_opcodes(): + if tag == "equal": + yield from zip(python[python_start:python_end], rust[rust_start:rust_end]) + else: + yield from ((event, None) for event in python[python_start:python_end]) + yield from ((None, event) for event in rust[rust_start:rust_end]) + + +def _label(event: FunctionTraceEvent | None) -> str: + return f"{' ' * event.depth}{event.function}" if event is not None else "" + + +def _status( + python: FunctionTraceEvent | None, + rust: FunctionTraceEvent | None, + python_names: frozenset[str], + rust_names: frozenset[str], +) -> tuple[str, str]: + if python is not None and rust is not None: + return "match", "\033[32m" + if python is not None: + return ("reordered", "\033[31m") if python.function in rust_names else ("python only", "\033[34m") + if rust is not None: + return ("reordered", "\033[31m") if rust.function in python_names else ("rust only", "\033[33m") + return "", "" + + +def format_trace_table( + python: tuple[FunctionTraceEvent, ...], + rust: tuple[FunctionTraceEvent, ...], + *, + colorize: bool, +) -> str: + python_header: Final = f"python ({len(python)} steps)" + rust_header: Final = f"rust ({len(rust)} steps)" + python_width: Final = max(len(python_header), *(len(_label(event)) for event in python), 0) + rust_width: Final = max(len(rust_header), *(len(_label(event)) for event in rust), 0) + python_names: Final = frozenset(event.function for event in python) + rust_names: Final = frozenset(event.function for event in rust) + border: Final = f"+-{'-' * python_width}-+-{'-' * rust_width}-+-------------+" + rows: Final = tuple( + f"{color}{line}\033[0m" if colorize else line + for left, right in _aligned_rows(python, rust) + for status, color in (_status(left, right, python_names, rust_names),) + for line in (f"| {_label(left):<{python_width}} | {_label(right):<{rust_width}} | {status:<11} |",) + ) + return "\n".join( + ( + border, + f"| {python_header:<{python_width}} | {rust_header:<{rust_width}} | {'comparison':<11} |", + border, + *rows, + border, + ) + ) diff --git a/tests/sdk_function_trace/test_mock_provider.py b/tests/sdk_function_trace/test_mock_provider.py new file mode 100644 index 00000000000..88d7d5392d0 --- /dev/null +++ b/tests/sdk_function_trace/test_mock_provider.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from contextlib import ExitStack +from typing import Final +from urllib.error import HTTPError +from urllib.request import Request, urlopen + +import pytest + +from tests.sdk_function_trace.mock_provider import MockProviderResponse, mock_provider + + +def test_mock_provider_preserves_error_response() -> None: + response: Final = MockProviderResponse(429, (("retry-after", "2"),), b'{"error":"rate limited"}') + with mock_provider(response) as api_base: + with pytest.raises(HTTPError) as error: + urlopen(Request(api_base, data=b"{}"), timeout=5) + with error.value as received: + assert received.code == 429 + assert received.headers["retry-after"] == "2" + assert received.read() == response.body + + +@pytest.mark.parametrize("request_count", [0, 2]) +def test_mock_provider_rejects_missing_or_duplicate_requests(request_count: int) -> None: + response: Final = MockProviderResponse(200, (), b"{}") + with ExitStack() as stack: + api_base: Final = stack.enter_context(mock_provider(response)) + for _ in range(request_count): + with urlopen(Request(api_base, data=b"{}"), timeout=5) as received: + assert received.read() == response.body + with pytest.raises(AssertionError, match=f"expected one provider request, received {request_count}"): + stack.close() diff --git a/tests/sdk_function_trace/test_profiler.py b/tests/sdk_function_trace/test_profiler.py new file mode 100644 index 00000000000..10a266fb1e8 --- /dev/null +++ b/tests/sdk_function_trace/test_profiler.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +import asyncio +import sys +from pathlib import Path +from types import FunctionType +from typing import Final, cast + +import pytest + +from tests.sdk_function_trace import ( + FunctionTraceEvent, + TraceScenario, + TraceStep, + assert_function_trace_parity, +) +from tests.sdk_function_trace.profiler import profile_python + + +class First: + @staticmethod + def run() -> None: + return None + + +class Second: + @staticmethod + def run() -> None: + return None + + +def test_profiler_matches_code_objects_and_keeps_repeated_calls() -> None: + with profile_python((First.run,)) as profiler: + Second.run() + First.run() + First.run() + + assert profiler.events == [ + FunctionTraceEvent(function="run", depth=0), + FunctionTraceEvent(function="run", depth=0), + ] + + +def test_profiler_records_selected_function_nesting_depth() -> None: + class Nested: + @staticmethod + def run() -> None: + First.run() + + with profile_python((Nested.run, First.run)) as profiler: + Nested.run() + + assert profiler.events == [ + FunctionTraceEvent(function="run", depth=0), + FunctionTraceEvent(function="run", depth=1), + ] + + +def test_profiler_restores_previous_profiler_after_failure() -> None: + previous: Final = sys.getprofile() + + with profile_python((First.run,)) as outer: + with pytest.raises(RuntimeError, match="stop"): + with profile_python((Second.run,)): + raise RuntimeError("stop") + assert sys.getprofile() is outer + First.run() + + assert sys.getprofile() is previous + assert outer.events == [FunctionTraceEvent(function="run", depth=0)] + + +def test_profiler_does_not_count_coroutine_resumption_as_another_call() -> None: + async def suspended() -> None: + await asyncio.sleep(0) + First.run() + await asyncio.sleep(0) + + with profile_python((suspended, First.run)) as profiler: + asyncio.run(suspended()) + + assert profiler.events == [ + FunctionTraceEvent(function="suspended", depth=0), + FunctionTraceEvent(function="run", depth=1), + ] + + +def test_source_profiler_records_real_frame_ancestry() -> None: + def outer() -> None: + First.run() + + with profile_python(source_root=Path(__file__).parent) as profiler: + outer() + Second.run() + + outer_event, first_event, second_event = ( + event for event in profiler.events if event.function.startswith("test_profiler.py:") + ) + assert first_event.ancestors is not None + assert outer_event.function in first_event.ancestors + assert second_event.ancestors is not None + assert outer_event.function not in second_event.ancestors + + +@pytest.mark.parametrize( + "rust_trace", + [ + (), + (FunctionTraceEvent(function="renamed", depth=0),), + (FunctionTraceEvent(function="run", depth=1),), + (FunctionTraceEvent(function="run", depth=0),) * 2, + ], + ids=["missing", "renamed", "wrong-depth", "extra-call"], +) +def test_harness_rejects_rust_function_trace_drift(rust_trace: tuple[FunctionTraceEvent, ...]) -> None: + with pytest.raises(AssertionError, match="Rust function trace differs"): + assert_function_trace_parity( + TraceScenario( + steps=(TraceStep(cast(FunctionType, First.run), depth=0),), + invoke_python=First.run, + invoke_rust=lambda: rust_trace, + ) + ) + + +def test_harness_rejects_python_function_trace_drift() -> None: + with pytest.raises(AssertionError, match="Python function trace differs"): + assert_function_trace_parity( + TraceScenario( + steps=(TraceStep(cast(FunctionType, First.run), depth=0),), + invoke_python=Second.run, + invoke_rust=lambda: (FunctionTraceEvent(function="run", depth=0),), + ) + ) + + +def test_harness_accepts_matching_traces() -> None: + assert_function_trace_parity( + TraceScenario( + steps=(TraceStep(cast(FunctionType, First.run), depth=0),), + invoke_python=First.run, + invoke_rust=lambda: (FunctionTraceEvent(function="run", depth=0),), + ) + ) + + +def test_harness_rejects_reordered_calls() -> None: + def begin() -> None: + return None + + def finish() -> None: + return None + + with pytest.raises(AssertionError, match="Rust function trace differs"): + assert_function_trace_parity( + TraceScenario( + steps=( + TraceStep(cast(FunctionType, begin), depth=0), + TraceStep(cast(FunctionType, finish), depth=0), + ), + invoke_python=lambda: (begin(), finish()), + invoke_rust=lambda: ( + FunctionTraceEvent(function="finish", depth=0), + FunctionTraceEvent(function="begin", depth=0), + ), + ) + ) diff --git a/tests/sdk_function_trace/test_runtime.py b/tests/sdk_function_trace/test_runtime.py new file mode 100644 index 00000000000..015cba55083 --- /dev/null +++ b/tests/sdk_function_trace/test_runtime.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +from typing import Final + +import pytest + +from tests.sdk_function_trace.runtime import ( + TraceFailed, + TraceSkipped, + attempt_trace, + run_trace, + trace_diff, +) +from tests.sdk_function_trace.steps import pipeline_issues, pipeline_steps + + +def test_sync_messages_records_the_known_python_limitation() -> None: + result: Final = attempt_trace("messages", engine="python", asynchronous=False) + + assert isinstance(result, TraceSkipped) + assert result.reason == "ValueError: anthropic_messages_handler is not implemented for sync calls" + + +def test_unexpected_call_failure_is_not_skipped() -> None: + result: Final = attempt_trace("unknown", engine="python", asynchronous=False) + + assert isinstance(result, TraceFailed) + assert result.reason == "ValueError: Unknown route: unknown" + + +@pytest.mark.parametrize( + ("route", "asynchronous"), + (("chat_completions", False), ("chat_completions", True), ("messages", True), ("ocr", False), ("ocr", True)), +) +def test_compiled_routes_match_python_steps(route: str, asynchronous: bool) -> None: + from litellm.rust_bridge import get_native_bridge + + if get_native_bridge() is None: + pytest.skip("build the native bridge to run executed route parity") + python: Final = pipeline_steps(route, "python", run_trace(route, engine="python", asynchronous=asynchronous)) + rust: Final = pipeline_steps(route, "rust", run_trace(route, engine="rust", asynchronous=asynchronous)) + + assert pipeline_issues(route, "python", python) == () + assert pipeline_issues(route, "rust", rust) == () + assert trace_diff(python, rust).matches + if route != "messages": + assert python == rust diff --git a/tests/sdk_function_trace/test_steps.py b/tests/sdk_function_trace/test_steps.py new file mode 100644 index 00000000000..b5432951187 --- /dev/null +++ b/tests/sdk_function_trace/test_steps.py @@ -0,0 +1,244 @@ +from __future__ import annotations + +from typing import Final + +import pytest + +from tests.sdk_function_trace.profiler import FunctionTraceEvent +from tests.sdk_function_trace.runtime import trace_diff +from tests.sdk_function_trace.steps import pipeline_issues, pipeline_steps + + +def test_python_ocr_projection_keeps_pipeline_and_drops_noise() -> None: + events: Final = ( + FunctionTraceEvent("utils.py:1747 client..wrapper_async", 0), + FunctionTraceEvent("ocr/main.py:331 aocr", 1), + FunctionTraceEvent("ocr/main.py:70 _prepare_ocr_request", 2), + FunctionTraceEvent("litellm_core_utils/get_llm_provider_logic.py:142 get_llm_provider", 3), + FunctionTraceEvent("utils.py:9303 ProviderConfigManager.get_provider_ocr_config", 3), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:34 MistralOCRConfig.get_supported_ocr_params", 4), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:72 MistralOCRConfig.map_ocr_params", 4), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:34 MistralOCRConfig.get_supported_ocr_params", 5), + FunctionTraceEvent("llms/custom_httpx/llm_http_handler.py:1705 BaseLLMHTTPHandler.async_ocr", 2), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:94 MistralOCRConfig.validate_environment", 4), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:124 MistralOCRConfig.get_complete_url", 4), + FunctionTraceEvent("llms/base_llm/ocr/transformation.py:209 BaseOCRConfig.async_transform_ocr_request", 5), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:149 MistralOCRConfig.transform_ocr_request", 6), + FunctionTraceEvent("llms/custom_httpx/http_handler.py:654 AsyncHTTPHandler.post", 6), + FunctionTraceEvent("llms/base_llm/ocr/transformation.py:255 BaseOCRConfig.async_transform_ocr_response", 4), + FunctionTraceEvent("llms/mistral/ocr/transformation.py:200 MistralOCRConfig.transform_ocr_response", 5), + FunctionTraceEvent("cost_calculator.py:1874 ocr_cost", 6), + ) + + assert pipeline_steps("ocr", "python", events) == ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("prepare_ocr_call", 1), + FunctionTraceEvent("get_provider_ocr_config", 2), + FunctionTraceEvent("supported_ocr_params", 3), + FunctionTraceEvent("map_ocr_params", 3), + FunctionTraceEvent("execute_ocr_provider_call", 1), + FunctionTraceEvent("validate_environment", 2), + FunctionTraceEvent("complete_url", 2), + FunctionTraceEvent("transform_ocr_request", 3), + FunctionTraceEvent("http_request", 3), + FunctionTraceEvent("transform_ocr_response", 2), + ) + + +def test_rust_ocr_projection_reuses_step_names_and_keeps_unknown_spans() -> None: + events: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("prepare_ocr_call", 1), + FunctionTraceEvent("map_ocr_params", 2), + FunctionTraceEvent("supported_ocr_params", 3), + FunctionTraceEvent("map_ocr_params", 2), + FunctionTraceEvent("transform_ocr_request", 2), + FunctionTraceEvent("execute_ocr_provider_call", 2), + FunctionTraceEvent("transform_ocr_response", 3), + FunctionTraceEvent("new_uninstrumented_span", 3), + ) + + assert pipeline_steps("ocr", "rust", events) == ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("prepare_ocr_call", 1), + FunctionTraceEvent("map_ocr_params", 2), + FunctionTraceEvent("supported_ocr_params", 3), + FunctionTraceEvent("transform_ocr_request", 2), + FunctionTraceEvent("execute_ocr_provider_call", 2), + FunctionTraceEvent("transform_ocr_response", 3), + FunctionTraceEvent("new_uninstrumented_span", 3), + ) + + +def test_projection_resets_depth_on_thread_root() -> None: + events: Final = ( + FunctionTraceEvent("main.py:387 acompletion", 1), + FunctionTraceEvent("llms/anthropic/chat/handler.py:255 AnthropicChatCompletion.acompletion_function", 2), + FunctionTraceEvent( + "llms/anthropic/experimental_pass_through/messages/handler.py:416 anthropic_messages_handler", 0 + ), + FunctionTraceEvent( + "llms/anthropic/experimental_pass_through/messages/transformation.py:575" + " AnthropicMessagesConfig.transform_anthropic_messages_request", + 4, + ), + ) + + assert pipeline_steps("chat_completions", "python", events) == ( + FunctionTraceEvent("chat_completions", 0), + FunctionTraceEvent("execute_chat_completions_provider_call", 1), + ) + assert pipeline_steps("messages", "python", events) == ( + FunctionTraceEvent("execute_messages_provider_call", 0), + FunctionTraceEvent("transform_request", 1), + ) + + +@pytest.mark.parametrize("function", ("completion", "completion_function", "acompletion_function")) +def test_chat_projection_includes_sync_and_async_handlers(function: str) -> None: + events: Final = (FunctionTraceEvent(f"llms/anthropic/chat/handler.py:100 AnthropicChatCompletion.{function}", 0),) + + assert pipeline_steps("chat_completions", "python", events) == ( + FunctionTraceEvent("execute_chat_completions_provider_call", 0), + ) + + +def test_trace_diff_reports_no_difference_for_identical_steps() -> None: + steps: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("transform_ocr_request", 1), + ) + + diff: Final = trace_diff(steps, steps) + + assert diff.python_only == () + assert diff.rust_only == () + assert diff.shared_order_matches + assert diff.matches + + +def test_trace_diff_reports_exclusive_steps_and_reordered_shared_steps() -> None: + python: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("supported_ocr_params", 1), + FunctionTraceEvent("map_ocr_params", 1), + FunctionTraceEvent("http_request", 2), + ) + rust: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("map_ocr_params", 1), + FunctionTraceEvent("supported_ocr_params", 2), + FunctionTraceEvent("transform_ocr_response", 2), + ) + + diff: Final = trace_diff(python, rust) + + assert diff.python_only == ("http_request",) + assert diff.rust_only == ("transform_ocr_response",) + assert not diff.shared_order_matches + assert not diff.matches + + +def test_trace_diff_does_not_claim_empty_or_disjoint_traces_match() -> None: + assert not trace_diff((), ()).shared_order_matches + assert not trace_diff((FunctionTraceEvent("ocr", 0),), (FunctionTraceEvent("messages", 0),)).shared_order_matches + + +def test_projection_uses_actual_ancestors_after_coroutine_resumption() -> None: + entrypoint: Final = "main.py:387 acompletion" + handler: Final = "llms/anthropic/chat/handler.py:255 AnthropicChatCompletion.acompletion_function" + events: Final = ( + FunctionTraceEvent(entrypoint, 0, ()), + FunctionTraceEvent(handler, 1, (entrypoint,)), + FunctionTraceEvent("utils.py:100 unrelated_worker", 0, ()), + FunctionTraceEvent("llms/anthropic/chat/transformation.py:100 transform_response", 1, (handler,)), + ) + + assert pipeline_steps("chat_completions", "python", events) == ( + FunctionTraceEvent("chat_completions", 0), + FunctionTraceEvent("execute_chat_completions_provider_call", 1), + FunctionTraceEvent("transform_response", 2), + ) + + +def test_projection_does_not_nest_siblings_under_a_returned_config_lookup() -> None: + events: Final = ( + FunctionTraceEvent("main.py:387 completion", 0), + FunctionTraceEvent("utils.py:100 ProviderConfigManager.get_provider_chat_config", 1), + FunctionTraceEvent("utils.py:200 unrelated_helper", 1), + FunctionTraceEvent("llms/anthropic/chat/transformation.py:100 transform_request", 2), + ) + + assert pipeline_steps("chat_completions", "python", events) == ( + FunctionTraceEvent("chat_completions", 0), + FunctionTraceEvent("get_provider_chat_config", 1), + FunctionTraceEvent("transform_request", 1), + ) + + +CHAT_RUST_STEPS: Final = ( + "chat_completions", + "get_provider_chat_config", + "supported_openai_params", + "execute_chat_completions_provider_call", + "validate_environment", + "transform_request", + "http_request", + "transform_response", +) + + +@pytest.mark.parametrize("missing", CHAT_RUST_STEPS) +def test_pipeline_check_rejects_missing_stages(missing: str) -> None: + steps: Final = tuple(FunctionTraceEvent(name, 0) for name in CHAT_RUST_STEPS if name != missing) + + assert f"missing {missing}" in pipeline_issues("chat_completions", "rust", steps) + + +def test_pipeline_check_rejects_http_before_request_transformation() -> None: + steps: Final = tuple( + FunctionTraceEvent(name, 0) + for name in ( + "chat_completions", + "get_provider_chat_config", + "supported_openai_params", + "execute_chat_completions_provider_call", + "validate_environment", + "http_request", + "transform_request", + "transform_response", + ) + ) + + assert "transform_request must precede http_request" in pipeline_issues("chat_completions", "rust", steps) + + +def test_step_parity_rejects_different_handler_boundaries_even_with_valid_stages() -> None: + rust: Final = tuple(FunctionTraceEvent(name, 0) for name in CHAT_RUST_STEPS) + python: Final = tuple( + FunctionTraceEvent(name, 0) + for name in ( + "chat_completions", + "get_provider_chat_config", + "supported_openai_params", + "validate_environment", + "transform_request", + "execute_chat_completions_provider_call", + "http_request", + "transform_response", + ) + ) + + assert not trace_diff(python, rust).shared_order_matches + assert not trace_diff(python, rust).matches + assert pipeline_issues("chat_completions", "python", python) == () + assert pipeline_issues("chat_completions", "rust", rust) == () + + +def test_step_parity_rejects_an_exclusive_helper_with_matching_shared_order() -> None: + rust: Final = tuple(FunctionTraceEvent(name, 0) for name in CHAT_RUST_STEPS) + python: Final = (*rust, FunctionTraceEvent("unmatched_helper", 0)) + diff: Final = trace_diff(python, rust) + + assert diff.shared_order_matches + assert not diff.matches diff --git a/tests/sdk_function_trace/test_table.py b/tests/sdk_function_trace/test_table.py new file mode 100644 index 00000000000..c2341a391a9 --- /dev/null +++ b/tests/sdk_function_trace/test_table.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import re +from typing import Final + +from tests.sdk_function_trace.profiler import FunctionTraceEvent +from tests.sdk_function_trace.table import format_trace_table + + +def test_table_aligns_matches_after_missing_steps_and_preserves_indentation() -> None: + python: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("python_helper", 1), + FunctionTraceEvent("http_request", 2), + ) + rust: Final = ( + FunctionTraceEvent("ocr", 0), + FunctionTraceEvent("rust_helper", 1), + FunctionTraceEvent("http_request", 1), + ) + output: Final = format_trace_table(python, rust, colorize=False) + rows: Final = tuple(line.split("|")[1:-1] for line in output.splitlines() if line.startswith("|")) + + assert tuple(tuple(cell.strip() for cell in row) for row in rows) == ( + ("python (3 steps)", "rust (3 steps)", "comparison"), + ("ocr", "ocr", "match"), + ("python_helper", "", "python only"), + ("", "rust_helper", "rust only"), + ("http_request", "http_request", "match"), + ) + assert rows[-1][0].startswith(" http_request") + assert rows[-1][1].startswith(" http_request") + assert len({len(line) for line in output.splitlines()}) == 1 + assert "\033[" not in output + + +def test_table_marks_reordered_calls_and_keeps_both_execution_orders() -> None: + python: Final = tuple(FunctionTraceEvent(name, 0) for name in ("ocr", "map", "validate", "http")) + rust: Final = tuple(FunctionTraceEvent(name, 0) for name in ("ocr", "validate", "map", "http")) + output: Final = format_trace_table(python, rust, colorize=True) + plain: Final = re.sub(r"\033\[[0-9;]*m", "", output) + rows: Final = tuple(line.split("|")[1:-1] for line in plain.splitlines() if line.startswith("|"))[1:] + + assert tuple(row[0].strip() for row in rows if row[0].strip()) == tuple(event.function for event in python) + assert tuple(row[1].strip() for row in rows if row[1].strip()) == tuple(event.function for event in rust) + assert plain.count("reordered") == 2 + assert output.count("\033[31m") == 2 + assert "only" not in output + + +def test_table_colors_match_and_exclusive_rows_without_changing_alignment() -> None: + python: Final = (FunctionTraceEvent("ocr", 0), FunctionTraceEvent("python_helper", 1)) + rust: Final = (FunctionTraceEvent("ocr", 0), FunctionTraceEvent("rust_helper", 1)) + colored: Final = format_trace_table(python, rust, colorize=True) + + assert re.sub(r"\033\[[0-9;]*m", "", colored) == format_trace_table(python, rust, colorize=False) + assert next(line for line in colored.splitlines() if "match" in line).startswith("\033[32m") + assert next(line for line in colored.splitlines() if "python only" in line).startswith("\033[34m") + assert next(line for line in colored.splitlines() if "rust only" in line).startswith("\033[33m") + + +def test_table_handles_empty_traces() -> None: + output: Final = format_trace_table((), (), colorize=False) + + assert "python (0 steps)" in output + assert "rust (0 steps)" in output + assert "match" not in output diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index fbd7e36e298..293f75b7592 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -8,6 +8,7 @@ import pytest import litellm from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.rust_bridge import configuration from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -109,10 +110,12 @@ class RaisingAsyncMessages: @pytest.fixture(autouse=True) def _reset_rust_flag(): - litellm.use_litellm_rust(False, messages=None, amessages=None) + rust_messages.set_rust_messages(messages=None, amessages=None) + configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield - litellm.use_litellm_rust(False, messages=None, amessages=None) + rust_messages.set_rust_messages(messages=None, amessages=None) + configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -122,17 +125,6 @@ def test_load_rust_messages_returns_injected_impl(): assert rust_messages.load_rust_messages() is bridge -def test_configuring_messages_does_not_enable_ocr(): - from litellm.rust_bridge.ocr import rust_ocr_enabled - - litellm.use_litellm_rust(False) - assert rust_ocr_enabled() is False - - litellm.use_litellm_rust(True, messages=RecordingMessages()) - - assert rust_ocr_enabled() is False - - def test_bare_use_litellm_rust_still_toggles_ocr(): from litellm.rust_bridge.ocr import rust_ocr_enabled @@ -264,7 +256,7 @@ async def test_gate_falls_back_to_python_when_bridge_raises(): @pytest.mark.asyncio async def test_gate_skips_rust_when_flag_absent(): bridge = ExplodingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + rust_messages.set_rust_messages(amessages=bridge) response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure")) @@ -272,6 +264,18 @@ async def test_gate_skips_rust_when_flag_absent(): assert bridge.calls == 0 +@pytest.mark.asyncio +async def test_gate_uses_process_enable_without_request_override(): + bridge = RecordingAsyncMessages() + rust_messages.set_rust_messages(amessages=bridge) + litellm.use_litellm_rust(True) + + response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure")) + + assert response is not None + assert bridge.calls[0]["custom_llm_provider"] == "azure_ai" + + @pytest.mark.asyncio async def test_gate_skips_rust_when_flag_false(): bridge = ExplodingAsyncMessages() @@ -305,7 +309,7 @@ async def test_gate_invokes_rust_for_native_anthropic_provider(): @pytest.mark.asyncio async def test_gate_invokes_rust_when_env_var_set(monkeypatch): bridge = RecordingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + rust_messages.set_rust_messages(amessages=bridge) monkeypatch.setenv("LITELLM_RUST", "1") response = await _gate( @@ -320,7 +324,7 @@ async def test_gate_invokes_rust_when_env_var_set(monkeypatch): @pytest.mark.asyncio async def test_gate_env_var_falsey_does_not_enable(monkeypatch): bridge = ExplodingAsyncMessages() - litellm.use_litellm_rust(True, amessages=bridge) + rust_messages.set_rust_messages(amessages=bridge) monkeypatch.setenv("LITELLM_RUST", "0") response = await _gate( diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 0f9f8259bef..8ea8db5fb65 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -3127,6 +3127,31 @@ def test_reasoning_effort_accepts_dict_shape_for_non_adaptive_model( ) +@pytest.mark.parametrize( + "model,budget_tokens,expected", + [ + ("claude-opus-4-8", 4096, ({"type": "adaptive"}, {"effort": "high"})), + ("claude-opus-4-7", 24000, ({"type": "adaptive"}, {"effort": "xhigh"})), + ("claude-opus-4-6", 4096, ({"type": "enabled", "budget_tokens": 4096}, None)), + ("claude-sonnet-4-5-20250929", 4096, ({"type": "enabled", "budget_tokens": 4096}, None)), + ], +) +def test_legacy_thinking_translated_to_adaptive_on_adaptive_only_models(model, budget_tokens, expected): + """Adaptive-only models reject thinking={type: enabled} with a 400, so the + legacy shape must be upgraded to adaptive + output_config.effort on + /chat/completions too, while models that accept it keep the caller's budget.""" + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": budget_tokens}, "max_tokens": 64000}, + optional_params={}, + model=model, + drop_params=False, + ) + + assert (result["thinking"], result.get("output_config")) == expected + + @pytest.mark.parametrize( "bad_value", [ diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py new file mode 100644 index 00000000000..4b95b36fec3 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py @@ -0,0 +1,79 @@ +""" +Regression tests for issue #34692. + +ollama_chat streams tool_calls in a mid-stream chunk while its final +(``done: true``) chunk carries only ``done_reason: "stop"``. The provider +iterator must remember the earlier tool_calls and stamp +``finish_reason="tool_calls"`` on the final chunk, so the Anthropic +``/v1/messages`` bridge emits ``stop_reason: "tool_use"``. Before the fix the +bridge emitted ``stop_reason: "end_turn"`` and Anthropic tool-runners +(Claude Code, ``messages.stream``) silently dropped the tool call. +""" + +import pytest + +from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( + AnthropicStreamWrapper, +) +from litellm.llms.ollama.chat.transformation import ( + OllamaChatCompletionResponseIterator, +) +from litellm.types.utils import ModelResponseStream + +_OLLAMA_TOOL_CHUNK = { + "model": "qwen3:8b", + "message": { + "role": "assistant", + "content": "", + "tool_calls": [{"function": {"name": "get_weather", "arguments": {"city": "San Francisco"}}}], + }, + "done": False, +} +_OLLAMA_DONE_CHUNK = { + "model": "qwen3:8b", + "message": {"role": "assistant", "content": ""}, + "done": True, + "done_reason": "stop", + "prompt_eval_count": 100, + "eval_count": 20, +} + + +def _ollama_streamed_chunks() -> list[ModelResponseStream]: + iterator = OllamaChatCompletionResponseIterator(streaming_response=iter([]), sync_stream=True) + return [iterator.chunk_parser(_OLLAMA_TOOL_CHUNK), iterator.chunk_parser(_OLLAMA_DONE_CHUNK)] + + +class _AsyncStream: + def __init__(self, items: list[ModelResponseStream]): + self._it = iter(items) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._it) + except StopIteration: + raise StopAsyncIteration + + +def _assert_tool_use_stop_reason(events: list[dict]) -> None: + block_types = [e["content_block"]["type"] for e in events if e.get("type") == "content_block_start"] + assert "tool_use" in block_types, f"no tool_use content block opened: {events}" + message_deltas = [e for e in events if e.get("type") == "message_delta"] + assert message_deltas, f"no message_delta emitted: {events}" + assert message_deltas[-1]["delta"]["stop_reason"] == "tool_use", ( + f"expected stop_reason 'tool_use', got: {message_deltas[-1]}" + ) + + +def test_ollama_mid_stream_tool_call_yields_tool_use_stop_reason_sync(): + wrapper = AnthropicStreamWrapper(completion_stream=iter(_ollama_streamed_chunks()), model="qwen3:8b") + _assert_tool_use_stop_reason(list(wrapper)) + + +@pytest.mark.asyncio +async def test_ollama_mid_stream_tool_call_yields_tool_use_stop_reason_async(): + wrapper = AnthropicStreamWrapper(completion_stream=_AsyncStream(_ollama_streamed_chunks()), model="qwen3:8b") + _assert_tool_use_stop_reason([event async for event in wrapper]) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index ad4c3d6bfbb..e819433c269 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -296,21 +296,15 @@ async def test_bedrock_converse_budget_tokens_preserved(): mock_acompletion.assert_called_once() call_kwargs = mock_acompletion.call_args.kwargs - print( - "acompletion call kwargs: ", json.dumps(call_kwargs, indent=4, default=str) - ) + print("acompletion call kwargs: ", json.dumps(call_kwargs, indent=4, default=str)) # Verify thinking parameter is passed through with budget_tokens preserved thinking_param = call_kwargs.get("thinking") - assert ( - thinking_param is not None - ), "thinking parameter should be passed to acompletion" - assert ( - thinking_param.get("type") == "enabled" - ), "thinking.type should be 'enabled'" - assert ( - thinking_param.get("budget_tokens") == 1024 - ), f"thinking.budget_tokens should be 1024, but got {thinking_param.get('budget_tokens')}" + assert thinking_param is not None, "thinking parameter should be passed to acompletion" + assert thinking_param.get("type") == "enabled", "thinking.type should be 'enabled'" + assert thinking_param.get("budget_tokens") == 1024, ( + f"thinking.budget_tokens should be 1024, but got {thinking_param.get('budget_tokens')}" + ) def test_openai_model_with_thinking_converts_to_reasoning(): @@ -342,23 +336,18 @@ def test_openai_model_with_thinking_converts_to_reasoning(): call_kwargs = mock_responses.call_args.kwargs # Verify reasoning is set (converted from thinking) - assert ( - "reasoning" in call_kwargs - ), "reasoning should be passed to litellm.responses" + assert "reasoning" in call_kwargs, "reasoning should be passed to litellm.responses" # budget_tokens=1024 -> effort="low" (at the LOW budget threshold) # reasoning_auto_summary is False by default, so no summary key expected_reasoning = {"effort": "low"} assert call_kwargs["reasoning"] == expected_reasoning, ( - f"reasoning should be {expected_reasoning} for budget_tokens=1024, " - f"got {call_kwargs.get('reasoning')}" + f"reasoning should be {expected_reasoning} for budget_tokens=1024, got {call_kwargs.get('reasoning')}" ) assert "summary" not in call_kwargs["reasoning"] # Verify thinking is NOT passed directly to the Responses API - assert ( - "thinking" not in call_kwargs - ), "thinking should NOT be passed directly to litellm.responses" + assert "thinking" not in call_kwargs, "thinking should NOT be passed directly to litellm.responses" class TestThinkingParameterTransformation: @@ -411,9 +400,7 @@ class TestThinkingParameterTransformation: thinking=thinking, model="openai/gpt-5.2", ) - assert result == { - "reasoning_effort": {"effort": "high", "summary": "detailed"} - } + assert result == {"reasoning_effort": {"effort": "high", "summary": "detailed"}} finally: litellm.reasoning_auto_summary = original @@ -611,9 +598,9 @@ class TestThinkingSummaryPreservation: mock_responses.assert_called_once() call_kwargs = mock_responses.call_args.kwargs reasoning = call_kwargs["reasoning"] - assert ( - reasoning["summary"] == "concise" - ), f"Expected summary='concise', got summary='{reasoning.get('summary')}'" + assert reasoning["summary"] == "concise", ( + f"Expected summary='concise', got summary='{reasoning.get('summary')}'" + ) def test_responses_adapter_preserves_summary(self): """translate_thinking_to_reasoning should include summary when user provides it.""" @@ -622,9 +609,7 @@ class TestThinkingSummaryPreservation: ) thinking = {"type": "enabled", "budget_tokens": 5000, "summary": "concise"} - result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( - thinking - ) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) assert result == {"effort": "high", "summary": "concise"} def test_responses_adapter_no_summary_by_default(self): @@ -638,11 +623,7 @@ class TestThinkingSummaryPreservation: try: litellm.reasoning_auto_summary = False thinking = {"type": "enabled", "budget_tokens": 5000} - result = ( - LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning( - thinking - ) - ) + result = LiteLLMAnthropicToResponsesAPIAdapter.translate_thinking_to_reasoning(thinking) assert result == {"effort": "high"} assert result is not None and "summary" not in result finally: @@ -659,9 +640,7 @@ class TestThinkingSummaryPreservation: thinking=thinking, model="openai/gpt-5.2", ) - assert result == { - "reasoning_effort": {"effort": "high", "summary": "concise"} - } + assert result == {"reasoning_effort": {"effort": "high", "summary": "concise"}} def test_translate_thinking_for_model_disabled_stays_plain_string_when_auto_summary_enabled(self): """Disabled thinking must stay a plain string even when reasoning_auto_summary is on.""" @@ -807,9 +786,7 @@ def test_presanitized_flag_not_leaked_to_provider_params(): def fake_base_handler(*args, **kwargs): captured.update(kwargs) - captured["optional"] = kwargs.get( - "anthropic_messages_optional_request_params", {} - ) + captured["optional"] = kwargs.get("anthropic_messages_optional_request_params", {}) return "stub" with patch.object( @@ -974,6 +951,38 @@ def test_gate_passthrough_skipped_when_only_chat_completions_supported(monkeypat assert "config" not in captured +@pytest.mark.parametrize( + "model_info, expected_ttl_support", + [ + ({"supported_endpoints": ["/v1/messages"]}, False), + ({"supported_endpoints": ["/v1/messages"], "cache_control_ttl": True}, True), + ({"supported_endpoints": ["/v1/messages"], "cache_control_ttl": "yes"}, False), + ], +) +def test_gate_passthrough_forwards_cache_control_ttl_only_when_deployment_opts_in( + monkeypatch, model_info, expected_ttl_support +): + """The passthrough config strips cache_control.ttl unless the deployment sets + model_info.cache_control_ttl to exactly true.""" + from litellm.llms.anthropic.experimental_pass_through.messages.handler import ( + anthropic_messages_handler, + ) + + captured, _ = _gate_stubs(monkeypatch) + + result = anthropic_messages_handler( + max_tokens=100, + messages=[{"role": "user", "content": "Hello"}], + model="openai/some-model", + api_key="sk-test", + api_base="https://host/v1", + model_info=model_info, + ) + + assert result == "native-passthrough" + assert captured["config"].supports_cache_control_ttl() is expected_ttl_support + + def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_flag(): """Regional and provider-prefixed Claude 4.8+/5 entries carry ``supports_mid_conversation_system``, but the bare first-party keys @@ -987,9 +996,7 @@ def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_sys import litellm - cost_map_path = os.path.join( - os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json" - ) + cost_map_path = os.path.join(os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json") with open(cost_map_path) as f: cost_map = json.load(f) rules = cost_map["fallback_generalizations"]["rules"] @@ -1028,9 +1035,7 @@ def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_sys ("perplexity/sonar", "sonar", "https://api.perplexity.ai/chat/completions"), ], ) -async def test_messages_strips_provider_prefix_exactly_once( - requested_model, expected_wire_model, expected_url -): +async def test_messages_strips_provider_prefix_exactly_once(requested_model, expected_wire_model, expected_url): """ BerriAI/litellm#37716: only the leading provider segment may be stripped on the way upstream. diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py index 9dac914ca4d..9e2bfb08852 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py @@ -362,8 +362,8 @@ class TestAzureAnthropicConfig: ) assert "xhigh" in str(exc_info.value) - def test_extra_body_promotion_does_not_clobber_top_level(self): - """Top-level ``optional_params`` wins over duplicates in ``extra_body``.""" + def test_extra_body_promotion_overrides_mapped_top_level(self): + """The caller's ``extra_body`` wins over a mapped top-level duplicate, like the native ``anthropic`` passthrough.""" config = AzureAnthropicConfig() messages = [{"role": "user", "content": "Hello"}] @@ -383,7 +383,31 @@ class TestAzureAnthropicConfig: headers=headers, ) - assert result["output_config"] == {"effort": "low"} + assert result["output_config"] == {"effort": "high"} + + def test_legacy_thinking_upgrade_keeps_caller_effort_from_extra_body(self, local_model_cost_map): + config = AzureAnthropicConfig() + + mapped = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 1024}, "max_tokens": 100}, + optional_params={}, + model="claude-opus-4-8", + drop_params=False, + ) + assert mapped["thinking"] == {"type": "adaptive"} + assert mapped["output_config"] == {"effort": "low"} + + result = config.transform_request( + model="claude-opus-4-8", + messages=[{"role": "user", "content": "Hello"}], + optional_params={**mapped, "extra_body": {"output_config": {"effort": "high"}}}, + litellm_params={"api_key": "test-key"}, + headers={"api-key": "test-key", "anthropic-version": "2023-06-01"}, + ) + + assert result["thinking"] == {"type": "adaptive"} + assert result["output_config"] == {"effort": "high"} + assert "extra_body" not in result def test_context_management_mixed_edits_beta_headers(self): """Test that context_management with both compact and other edits adds both beta headers""" diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 41d82e4f960..0d7573a2536 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -671,3 +671,46 @@ def test_bedrock_chat_invoke_fable_5_1_response_format_avoids_forced_tool_choice assert "output_format" not in result assert "tools" in result assert "tool_choice" not in result + + +@pytest.mark.parametrize("model", ["us.anthropic.claude-sonnet-5", "us.anthropic.claude-fable-5-1"]) +def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_thinking(local_model_cost_map, model): + result = AmazonAnthropicClaudeConfig().map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": {"type": "object", "properties": {"result": {"type": "string"}}}, + }, + }, + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "max_tokens": 8192, + }, + optional_params={}, + model=model, + drop_params=False, + ) + + assert "tools" in result + assert result["thinking"] == {"type": "adaptive"} + assert result["output_config"] == {"effort": "high"} + + +def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking(local_model_cost_map): + """Regression: the tool-based ``response_format`` path swaps in a Claude 3 stub + model before the shared Anthropic mapping, which hid the adaptive-only model + from the legacy ``thinking`` upgrade and left ``type=enabled`` on the wire.""" + result = AmazonAnthropicClaudeConfig().map_openai_params( + non_default_params={ + "response_format": {"type": "json_object"}, + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "max_tokens": 8192, + }, + optional_params={}, + model="us.anthropic.claude-fable-5-1", + drop_params=False, + ) + + assert result["thinking"] == {"type": "adaptive"} + assert result["output_config"] == {"effort": "high"} diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 70f3153ed7e..cb05cdb9451 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -979,6 +979,34 @@ def test_config_blocks_do_not_leak_into_inference_config(): assert data["serviceTier"] == {"type": "priority"} +@pytest.mark.parametrize( + "model", + [ + "anthropic.claude-opus-4-8", + "us.anthropic.claude-opus-4-8", + "amazon.nova-pro-v1:0", + "us.meta.llama4-maverick-17b-instruct-v1:0", + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456", + "arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.amazon.nova-pro-v1:0", + ], +) +def test_client_metadata_stripped_from_converse_request(model): + data = AmazonConverseConfig()._transform_request_helper( + model=model, + system_content_blocks=[], + optional_params={ + "maxTokens": 16, + "anthropic_beta": ["computer-use-2025-01-24"], + "client_metadata": {"originator": "codex_cli_rs"}, + }, + messages=None, + ) + + fields = data["additionalModelRequestFields"] + assert "client_metadata" not in fields + assert fields["anthropic_beta"] == ["computer-use-2025-01-24"] + + def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost @@ -6347,6 +6375,82 @@ def test_adaptive_thinking_passes_through_on_46_plus_converse(model): assert optional_params.get("thinking") == {"type": "adaptive"} +@pytest.mark.parametrize( + "model,budget_tokens,expected_effort", + [ + ("anthropic.claude-opus-4-8", 4096, "high"), + ("us.anthropic.claude-opus-4-8", 2000, "low"), + ("global.anthropic.claude-opus-4-8", 12000, "xhigh"), + ("us.anthropic.claude-opus-4-7", 3000, "medium"), + ("anthropic.claude-fable-5", 4096, "high"), + ], +) +def test_legacy_thinking_translated_to_adaptive_on_adaptive_only_converse(model, budget_tokens, expected_effort): + """Adaptive-only models (4.7+, 5 families) reject thinking={type: enabled} + with a 400 on Bedrock Converse, so the legacy shape from callers like Claude + Code must be upgraded to thinking={type: adaptive} + output_config.effort + derived from budget_tokens, matching the /v1/messages passthrough.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": budget_tokens}, "max_tokens": 64000}, + optional_params={}, + model=model, + drop_params=False, + ) + request = config.transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert request["additionalModelRequestFields"]["thinking"] == {"type": "adaptive"} + assert request["additionalModelRequestFields"]["output_config"] == {"effort": expected_effort} + + +def test_legacy_thinking_translation_keeps_caller_output_config_effort_converse(): + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={ + "output_config": {"effort": "low"}, + "thinking": {"type": "enabled", "budget_tokens": 12000}, + "max_tokens": 64000, + }, + optional_params={}, + model="anthropic.claude-opus-4-8", + drop_params=False, + ) + + assert optional_params["thinking"] == {"type": "adaptive"} + assert optional_params["output_config"] == {"effort": "low"} + + +@pytest.mark.parametrize( + "model", + [ + "us.anthropic.claude-opus-4-6", + "anthropic.claude-3-5-sonnet-20241022-v2:0", + ], +) +def test_legacy_thinking_forwarded_verbatim_when_model_accepts_it_converse(model): + """The 4.6 family and pre-adaptive models accept thinking={type: enabled} + natively, so the caller's budget_tokens cap must keep applying.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 4096}, "max_tokens": 8192}, + optional_params={}, + model=model, + drop_params=False, + ) + + assert optional_params["thinking"] == {"type": "enabled", "budget_tokens": 4096} + assert "output_config" not in optional_params + + def test_adaptive_thinking_dropped_when_max_tokens_too_small_converse(): """When max_tokens can't fit even the minimum thinking budget, the raw adaptive block must be dropped entirely rather than translated, so the diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index d34517f61f6..09be2118001 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -128,6 +128,12 @@ def test_mantle_messages_url_construction(): _VPC_ENDPOINT = "https://vpce-0a1b2c3d.bedrock-mantle.us-gov-west-1.vpce.amazonaws.com" +@pytest.fixture(autouse=True) +def no_ambient_mantle_api_base(monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + + + def test_mantle_chat_url_honors_api_base_host(): config = AmazonMantleConfig() url = config.get_complete_url( @@ -193,6 +199,48 @@ def test_mantle_messages_url_honors_aws_bedrock_runtime_endpoint(): assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" +_ENV_ENDPOINT = "https://bedrock-mantle.us-east-1.api.aws.internal.example.com" + + +@pytest.mark.parametrize("config_cls", [AmazonMantleConfig, AmazonMantleMessagesConfig]) +@pytest.mark.parametrize( + "env_value", + [_ENV_ENDPOINT, f"{_ENV_ENDPOINT}/", f"{_ENV_ENDPOINT}/v1", f"{_ENV_ENDPOINT}/openai/v1"], +) +def test_mantle_url_honors_bedrock_mantle_api_base_env(monkeypatch, config_cls, env_value): + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", env_value) + url = config_cls().get_complete_url( + api_base=None, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params={"aws_region_name": "us-east-1"}, + litellm_params={}, + ) + assert url == f"{_ENV_ENDPOINT}/anthropic/v1/messages" + + +@pytest.mark.parametrize("config_cls", [AmazonMantleConfig, AmazonMantleMessagesConfig]) +@pytest.mark.parametrize( + ("api_base", "optional_params"), + [ + (_VPC_ENDPOINT, {"aws_region_name": "us-gov-west-1"}), + (None, {"aws_region_name": "us-gov-west-1", "aws_bedrock_runtime_endpoint": _VPC_ENDPOINT}), + ], +) +def test_mantle_url_explicit_endpoint_beats_bedrock_mantle_api_base_env( + monkeypatch, config_cls, api_base, optional_params +): + monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", _ENV_ENDPOINT) + url = config_cls().get_complete_url( + api_base=api_base, + api_key=None, + model="mantle/anthropic.claude-mythos-preview", + optional_params=optional_params, + litellm_params={}, + ) + assert url == f"{_VPC_ENDPOINT}/anthropic/v1/messages" + + def test_mantle_transform_request_strips_prefix_and_adds_model(): config = AmazonMantleConfig() request = config.transform_request( diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index cd775abf136..56b111f294e 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -486,6 +486,71 @@ class TestBedrockMantleChatAuth: assert "/us-east-2/bedrock/aws4_request" in authorization assert requests[0]["url"].startswith("https://bedrock-mantle.us-east-2.api.aws") + def test_completion_per_request_role_reaches_signer_and_not_the_body(self, monkeypatch): + from unittest.mock import MagicMock, Mock + + from botocore.credentials import Credentials + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.utils import ModelResponse + + for var in ("BEDROCK_MANTLE_API_KEY", "AWS_BEARER_TOKEN_BEDROCK", "BEDROCK_MANTLE_API_BASE"): + monkeypatch.delenv(var, raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + return_value=Credentials( + access_key="ASIAEXAMPLE", + secret_key="YXNzdW1lZC1yb2xlLXNlY3JldC1hc3N1bWVk", + token="assumed-session-token", + ) + ) + url = "https://bedrock-mantle.us-east-1.api.aws/openai/v1/chat/completions" + client = HTTPHandler(client=httpx.Client()) + client.post = Mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1733529600, + "model": "google.gemma-4-31b", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + request=httpx.Request("POST", url), + ) + ) + + BaseLLMHTTPHandler().completion( + model="google.gemma-4-31b", + messages=[{"role": "user", "content": "hello"}], + api_base=None, + custom_llm_provider="bedrock_mantle", + model_response=ModelResponse(), + encoding=None, + logging_obj=Mock(), + optional_params={}, + timeout=10, + litellm_params={ + "aws_role_name": "arn:aws:iam::000000000000:role/attributed-role", + "aws_session_name": "user-123", + "aws_region_name": "us-east-1", + }, + acompletion=False, + client=client, + provider_config=BedrockMantleChatConfig(aws_signer=signer), + ) + + credential_kwargs = signer.get_credentials.call_args.kwargs + assert credential_kwargs["aws_role_name"] == "arn:aws:iam::000000000000:role/attributed-role" + assert credential_kwargs["aws_session_name"] == "user-123" + sent = client.post.call_args.kwargs + assert sent["headers"]["Authorization"].startswith("AWS4-HMAC-SHA256") + assert not [key for key in json.loads(sent["data"]) if key.startswith("aws_")] + class TestBedrockMantleProjectHeader: def test_validate_environment_sets_openai_project_header(self): diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 16d57437043..9e64bfafa54 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1314,3 +1314,236 @@ async def test_finalizer_on_live_loop_disposes_foreign_loop_session_without_sche assert AsyncHTTPHandler._finalizer_close_tasks == baseline_tasks assert session.closed + + +@pytest.fixture +def forward_proxy_server(): + """Plain HTTP forward proxy that records the absolute URIs it is asked to fetch.""" + from http.server import BaseHTTPRequestHandler, HTTPServer + from socketserver import ThreadingMixIn + + seen_uris: list[str] = [] + + class RecordingProxyHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self): + seen_uris.append(self.path) + self.send_response(200) + self.send_header("Content-Length", "9") + self.end_headers() + self.wfile.write(b"via-proxy") + + def log_message(self, format, *args): + pass + + class ThreadedServer(ThreadingMixIn, HTTPServer): + daemon_threads = True + + server = ThreadedServer(("127.0.0.1", 0), RecordingProxyHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}", seen_uris + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +# `.invalid` never resolves (RFC 6761), so the only way this request can succeed is through the proxy +_PROXY_ONLY_UPSTREAM_URL = "http://upstream.invalid/v1/models" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("disable_aiohttp_transport", [True, False]) +@pytest.mark.parametrize("force_ipv4", [True, False]) +async def test_async_handler_honours_proxy_env_for_every_transport( + forward_proxy_server, monkeypatch: pytest.MonkeyPatch, disable_aiohttp_transport: bool, force_ipv4: bool +): + proxy_url, seen_uris = forward_proxy_server + monkeypatch.setenv("HTTP_PROXY", proxy_url) + monkeypatch.delenv("NO_PROXY", raising=False) + monkeypatch.delenv("no_proxy", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", disable_aiohttp_transport) + monkeypatch.setattr(litellm, "force_ipv4", force_ipv4) + + handler = AsyncHTTPHandler() + try: + response = await handler.get(_PROXY_ONLY_UPSTREAM_URL) + finally: + await handler.close() + + assert response.text == "via-proxy" + assert seen_uris == [_PROXY_ONLY_UPSTREAM_URL] + + +@pytest.mark.parametrize("force_ipv4", [True, False]) +def test_sync_handler_honours_proxy_env(forward_proxy_server, monkeypatch: pytest.MonkeyPatch, force_ipv4: bool): + proxy_url, seen_uris = forward_proxy_server + monkeypatch.setenv("HTTP_PROXY", proxy_url) + monkeypatch.delenv("NO_PROXY", raising=False) + monkeypatch.delenv("no_proxy", raising=False) + monkeypatch.setattr(litellm, "force_ipv4", force_ipv4) + + handler = HTTPHandler() + try: + response = handler.get(_PROXY_ONLY_UPSTREAM_URL) + finally: + handler.close() + + assert response.text == "via-proxy" + assert seen_uris == [_PROXY_ONLY_UPSTREAM_URL] + + +@pytest.mark.asyncio +async def test_force_ipv4_httpx_transport_honours_no_proxy(keepalive_server, monkeypatch: pytest.MonkeyPatch): + """NO_PROXY hosts must still go direct when the proxy mounts are supplied by litellm instead of httpx.""" + monkeypatch.setenv("HTTP_PROXY", "http://proxy.invalid:3128") + monkeypatch.setenv("NO_PROXY", "127.0.0.1") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "force_ipv4", True) + + handler = AsyncHTTPHandler() + try: + response = await handler.get(keepalive_server) + finally: + await handler.close() + + assert response.text == "ok" + + +@pytest.fixture +def private_ca_tls_upstream(tmp_path: pathlib.Path): + """HTTPS server behind a CONNECT proxy, both on localhost; the server's cert is signed by a test-only CA.""" + import datetime + import select + import socket + from http.server import BaseHTTPRequestHandler, HTTPServer + from socketserver import ThreadingMixIn + + from cryptography import x509 + from cryptography.hazmat.primitives import hashes, serialization + from cryptography.hazmat.primitives.asymmetric import rsa + from cryptography.x509.oid import NameOID + + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "upstream.invalid")]) + now = datetime.datetime.now(datetime.timezone.utc) + cert = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - datetime.timedelta(minutes=1)) + .not_valid_after(now + datetime.timedelta(hours=1)) + .add_extension(x509.SubjectAlternativeName([x509.DNSName("upstream.invalid")]), critical=False) + .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True) + .sign(key, hashes.SHA256()) + ) + ca_pem = tmp_path / "ca.pem" + ca_pem.write_bytes(cert.public_bytes(serialization.Encoding.PEM)) + key_pem = tmp_path / "key.pem" + key_pem.write_bytes( + key.private_bytes( + serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption() + ) + ) + + class OkTlsHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self): + self.send_response(200) + self.send_header("Content-Length", "6") + self.end_headers() + self.wfile.write(b"ok-tls") + + def log_message(self, format, *args): + pass + + class ThreadedServer(ThreadingMixIn, HTTPServer): + daemon_threads = True + + tls_server = ThreadedServer(("127.0.0.1", 0), OkTlsHandler) + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(str(ca_pem), str(key_pem)) + tls_server.socket = server_ctx.wrap_socket(tls_server.socket, server_side=True) + tls_port = tls_server.server_port + + class ConnectProxyHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_CONNECT(self): + upstream = socket.create_connection(("127.0.0.1", tls_port)) + self.send_response(200, "Connection established") + self.end_headers() + sockets = [self.connection, upstream] + while True: + readable, _, _ = select.select(sockets, [], [], 5) + if not readable: + break + for src in readable: + data = src.recv(65536) + if not data: + upstream.close() + return + (upstream if src is self.connection else self.connection).sendall(data) + + def log_message(self, format, *args): + pass + + proxy_server = ThreadedServer(("127.0.0.1", 0), ConnectProxyHandler) + threads = [ + threading.Thread(target=tls_server.serve_forever, daemon=True), + threading.Thread(target=proxy_server.serve_forever, daemon=True), + ] + for thread in threads: + thread.start() + try: + yield f"http://127.0.0.1:{proxy_server.server_port}", str(ca_pem) + finally: + for server in (proxy_server, tls_server): + server.shutdown() + server.server_close() + for thread in threads: + thread.join(timeout=5) + + +@pytest.mark.asyncio +async def test_force_ipv4_https_proxy_mount_uses_handler_ca_bundle( + private_ca_tls_upstream, monkeypatch: pytest.MonkeyPatch +): + proxy_url, ca_pem = private_ca_tls_upstream + monkeypatch.setenv("HTTPS_PROXY", proxy_url) + monkeypatch.delenv("NO_PROXY", raising=False) + monkeypatch.delenv("no_proxy", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "force_ipv4", True) + + handler = AsyncHTTPHandler(ssl_verify=ca_pem) + try: + response = await handler.get("https://upstream.invalid/v1/models") + finally: + await handler.close() + + assert response.text == "ok-tls" + + +def test_sync_force_ipv4_https_proxy_mount_uses_handler_ca_bundle( + private_ca_tls_upstream, monkeypatch: pytest.MonkeyPatch +): + proxy_url, ca_pem = private_ca_tls_upstream + monkeypatch.setenv("HTTPS_PROXY", proxy_url) + monkeypatch.delenv("NO_PROXY", raising=False) + monkeypatch.delenv("no_proxy", raising=False) + monkeypatch.setattr(litellm, "force_ipv4", True) + + handler = HTTPHandler(ssl_verify=ca_pem) + try: + response = handler.get("https://upstream.invalid/v1/models") + finally: + handler.close() + + assert response.text == "ok-tls" diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 26f841c1146..1d583c16ad7 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2295,6 +2295,25 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques assert retry_authorization != first_attempt_headers["Authorization"] +def test_aws_signing_overrides_only_fills_missing_credentials(): + from litellm.llms.custom_httpx.llm_http_handler import _aws_signing_overrides + + overrides = _aws_signing_overrides( + {"temperature": 0.2, "aws_region_name": "us-west-2"}, + { + "aws_role_name": "arn:aws:iam::000000000000:role/attributed", + "aws_session_name": "user-123", + "aws_region_name": "us-east-1", + "api_key": "not-an-aws-param", + }, + ) + + assert dict(overrides) == { + "aws_role_name": "arn:aws:iam::000000000000:role/attributed", + "aws_session_name": "user-123", + } + + class TestServerFulfilledToolsInRequest: """_server_fulfilled_tools_in_request gates the buffered (non-leaking) streaming mode for server-fulfilled tools like headroom_retrieve.""" diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index 41fb2589655..71661cc532b 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -422,6 +422,27 @@ def test_databricks_config_probes_capabilities_under_databricks_namespace(): assert DatabricksConfig().custom_llm_provider == "databricks" +@pytest.mark.parametrize( + "model, expected_thinking, expected_output_config", + [ + ("databricks-claude-opus-4-8", {"type": "adaptive"}, {"effort": "high"}), + ("databricks-claude-opus-4-6", {"type": "enabled", "budget_tokens": 4096}, None), + ], + ids=["adaptive_only_upgrades_to_adaptive", "legacy_capable_forwards_verbatim"], +) +def test_map_openai_params_upgrades_legacy_thinking_on_adaptive_only_claude( + model, expected_thinking, expected_output_config +): + mapped = DatabricksConfig().map_openai_params( + non_default_params={"thinking": {"type": "enabled", "budget_tokens": 4096}}, + optional_params={}, + model=model, + drop_params=False, + ) + assert mapped["thinking"] == expected_thinking + assert mapped.get("output_config") == expected_output_config + + def _streaming_chunk(usage=None, choices=None): base = { "id": "chatcmpl-test", diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index e6e6aa946d5..9a62fcf6f0f 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -1,8 +1,14 @@ +import json import os import sys +from typing import Final +from unittest.mock import MagicMock, patch +import httpx import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig from litellm.rerank_api.rerank_utils import get_optional_rerank_params from litellm.types.rerank import ( @@ -87,9 +93,7 @@ class TestHostedVLLMRerankTransform: assert "instruction" not in body def test_map_cohere_rerank_params_raises_on_max_chunks_per_doc(self): - with pytest.raises( - ValueError, match="Hosted VLLM does not support max_chunks_per_doc" - ): + with pytest.raises(ValueError, match="Hosted VLLM does not support max_chunks_per_doc"): self.config.map_cohere_rerank_params( non_default_params=None, model=self.model, @@ -104,12 +108,10 @@ class TestHostedVLLMRerankTransform: url = self.config.get_complete_url(base, self.model) assert url == "https://api.example.com/rerank" # Already ends with /rerank - url2 = self.config.get_complete_url( - "https://api.example.com/rerank", self.model - ) + url2 = self.config.get_complete_url("https://api.example.com/rerank", self.model) assert url2 == "https://api.example.com/rerank" # Raises if api_base is None - with pytest.raises(ValueError, match='api_base must be provided for Hosted VLLM rerank'): + with pytest.raises(ValueError, match="api_base must be provided for Hosted VLLM rerank"): self.config.get_complete_url(None, self.model) def test_transform_response(self): @@ -173,3 +175,121 @@ class TestGetOptionalRerankParamsInstruction: documents=["doc1", "doc2"], ) assert "instruction" not in params + + +class TestHostedVLLMRerankTruncationParams: + def setup_method(self): + self.config = HostedVLLMRerankConfig() + self.model = "hosted-vllm-model" + + def test_map_cohere_rerank_params_forwards_vllm_truncation_params(self): + params: Final = self.config.map_cohere_rerank_params( + non_default_params={ + "truncate_prompt_tokens": 512, + "truncation_side": "left", + "max_tokens_per_query": 64, + "metadata": {"user_api_key": "sk-test"}, + }, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + max_tokens_per_doc=128, + ) + assert params["truncate_prompt_tokens"] == 512 + assert params["truncation_side"] == "left" + assert params["max_tokens_per_query"] == 64 + assert params["max_tokens_per_doc"] == 128 + assert "metadata" not in params + + @pytest.mark.parametrize( + "bad_params", + [{"truncation_side": "middle"}, {"truncate_prompt_tokens": "lots"}, {"max_tokens_per_query": -1.5}], + ) + def test_map_cohere_rerank_params_rejects_invalid_truncation_params_as_400(self, bad_params: dict[str, object]): + with pytest.raises(litellm.UnsupportedParamsError) as raised: + self.config.map_cohere_rerank_params( + non_default_params=dict(bad_params), + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + assert raised.value.status_code == 400 + assert next(iter(bad_params)) in str(raised.value) + + def test_map_cohere_rerank_params_omits_truncation_params_when_absent(self): + params: Final = self.config.map_cohere_rerank_params( + non_default_params={"metadata": {"user_api_key": "sk-test"}}, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + ) + body: Final = self.config.transform_rerank_request(model=self.model, optional_rerank_params=params, headers={}) + truncation_keys: Final = { + "truncate_prompt_tokens", + "truncation_side", + "max_tokens_per_query", + "max_tokens_per_doc", + } + assert not truncation_keys & body.keys() + assert body == { + "model": self.model, + "query": "test query", + "documents": ["doc1", "doc2"], + "return_documents": True, + } + + def test_transform_request_forwards_truncation_params(self): + body: Final = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={ + "query": "test query", + "documents": ["doc1", "doc2"], + "truncate_prompt_tokens": 512, + "truncation_side": "left", + "max_tokens_per_query": 64, + "max_tokens_per_doc": 128, + }, + headers={}, + ) + assert body["truncate_prompt_tokens"] == 512 + assert body["truncation_side"] == "left" + assert body["max_tokens_per_query"] == 64 + assert body["max_tokens_per_doc"] == 128 + + def test_transform_request_omits_truncation_params_when_absent(self): + body: Final = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={"query": "test query", "documents": ["doc1", "doc2"]}, + headers={}, + ) + assert "truncate_prompt_tokens" not in body + assert "truncation_side" not in body + assert "max_tokens_per_query" not in body + assert "max_tokens_per_doc" not in body + + def test_rerank_sends_truncate_prompt_tokens_to_vllm(self): + client: Final = HTTPHandler() + mock_response: Final = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "id": "score-1", + "results": [{"index": 0, "relevance_score": 0.5}], + "usage": {"total_tokens": 512}, + } + with patch.object(client, "post", return_value=mock_response) as mock_post: + litellm.rerank( + model="hosted_vllm/BAAI/bge-reranker-base", + api_base="http://vllm.local:8000", + query="List all the unique case ids", + documents=["a document longer than the reranker context window"], + truncate_prompt_tokens=512, + truncation_side="left", + client=client, + ) + sent_body: Final = json.loads(mock_post.call_args.kwargs["data"]) + assert mock_post.call_args.kwargs["url"] == "http://vllm.local:8000/rerank" + assert sent_body["truncate_prompt_tokens"] == 512 + assert sent_body["truncation_side"] == "left" diff --git a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py b/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py index 8f3dbf7b0d9..25f9645faa0 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py @@ -615,6 +615,46 @@ class TestOllamaFinishReasonLength: result.choices[0].finish_reason == "stop" ), f"Expected 'stop' for natural finish, got '{result.choices[0].finish_reason}'" + def test_finish_reason_tool_calls_streamed_before_done_chunk(self): + """Streaming: tool_calls arriving mid-stream (not on the done chunk) must + still produce finish_reason='tool_calls' on the final chunk. + + Regression test for https://github.com/BerriAI/litellm/issues/34692: + Ollama emits tool_calls in an earlier chunk and the done chunk carries + none, which left finish_reason at 'stop' and made the Anthropic + /v1/messages bridge emit stop_reason 'end_turn' instead of 'tool_use'. + """ + iterator = OllamaChatCompletionResponseIterator( + streaming_response=iter([]), + sync_stream=True, + ) + + tool_chunk = { + "model": "qwen3:8b", + "message": { + "role": "assistant", + "content": "", + "tool_calls": [ + {"function": {"name": "get_weather", "arguments": {"city": "San Francisco"}}} + ], + }, + "done": False, + } + done_chunk = { + "model": "qwen3:8b", + "message": {"role": "assistant", "content": ""}, + "done": True, + "done_reason": "stop", + } + + tool_result = iterator.chunk_parser(tool_chunk) + assert tool_result.choices[0].delta.tool_calls is not None + + done_result = iterator.chunk_parser(done_chunk) + assert ( + done_result.choices[0].finish_reason == "tool_calls" + ), f"Expected 'tool_calls' when tool_calls were streamed earlier, got '{done_result.choices[0].finish_reason}'" + class TestOllamaReasoningContentStreaming: """Test that reasoning_content is properly extracted from all thinking chunks.""" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 315b6948bd8..d071ef78c2d 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -5,6 +5,8 @@ Tests the handler's ability to process input/output for the Responses API with guardrail transformations. """ +import copy +from collections.abc import Callable from typing import Any, List, Literal, Optional, Tuple from unittest.mock import AsyncMock, MagicMock @@ -19,6 +21,10 @@ from litellm.llms import get_guardrail_translation_mapping from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) +from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.responses.main import GenericResponseOutputItem, OutputText from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs @@ -1287,14 +1293,14 @@ class TestOpenAIResponsesHandlerToolInjection: """A tool a guardrail injects must survive the write-back to Responses format.""" def test_merge_keeps_guardrail_appended_tool(self): - """_merge_tools_after_guardrail must not drop the extra appended tool.""" - handler = OpenAIResponsesHandler() + """merge_guardrailed_tools must not drop the extra appended tool.""" original = [{"type": "function", "name": "a"}] - remapped = [ - {"type": "function", "name": "a"}, - {"type": "function", "name": "b"}, + groups = [form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original)] + guardrailed = [ + *groups[0], + {"type": "function", "function": {"name": "b", "description": "", "parameters": {"type": "object"}}}, ] - merged = handler._merge_tools_after_guardrail(original, remapped) + merged = merge_guardrailed_tools(original, groups, guardrailed) assert [t["name"] for t in merged] == ["a", "b"] @pytest.mark.asyncio @@ -1323,6 +1329,194 @@ class TestOpenAIResponsesHandlerToolInjection: assert "injected_tool" in names +class ToolEditingGuardrail(CustomGuardrail): + """Guardrail that rewrites the flattened chat tools it was handed through ``edit``""" + + def __init__(self, edit: Callable[[list[dict]], list[dict]], **kwargs): + super().__init__(**kwargs) + self.edit = edit + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Any | None = None, + ) -> GenericGuardrailAPIInputs: + inputs["tools"] = self.edit(list(inputs.get("tools") or [])) + return inputs + + +def _codex_request(input_value): + """A Responses API request shaped like what the Codex CLI sends when an MCP server is configured""" + return { + "model": "gpt-5.3-codex", + "input": input_value, + "tools": [ + { + "type": "function", + "name": "get_weather", + "description": "Weather lookup", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + "strict": False, + }, + { + "type": "namespace", + "name": "mcp__confluence", + "description": "Confluence tools", + "tools": [ + { + "type": "function", + "name": "confluence_get_page", + "description": "Get a page", + "parameters": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}, + "strict": False, + }, + { + "type": "function", + "name": "confluence_search", + "description": "Search pages", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + "strict": False, + }, + ], + }, + { + "type": "custom", + "name": "apply_patch", + "description": "Apply a patch", + "format": {"type": "grammar", "syntax": "lark", "definition": 'start: "x"'}, + }, + {"type": "web_search"}, + ], + } + + +def _tool_named(tools, name): + return next(tool for tool in tools if tool.get("name") == name) + + +class TestOpenAIResponsesHandlerNamespaceTools: + """Codex sends MCP tools as ``namespace`` tools; a guardrail must never flatten them (GH #39183)""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "input_value", + ["hi", [{"role": "user", "content": "hi", "type": "message"}]], + ids=["string_input", "list_input"], + ) + async def test_pass_through_guardrail_leaves_tools_untouched(self, input_value): + data = _codex_request(input_value) + expected_tools = copy.deepcopy(data["tools"]) + + result = await OpenAIResponsesHandler().process_input_messages( + data, MockPassThroughGuardrail(guardrail_name="test") + ) + + assert result["tools"] == expected_tools + + @pytest.mark.asyncio + async def test_appending_guardrail_keeps_namespace_and_adds_tool(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + + result = await OpenAIResponsesHandler().process_input_messages( + data, ToolAppendingGuardrail(guardrail_name="test") + ) + + assert result["tools"][:-1] == expected_tools + assert result["tools"][-1]["type"] == "function" + assert result["tools"][-1]["name"] == "injected_tool" + + @pytest.mark.asyncio + async def test_dropping_one_member_prunes_only_that_member(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + guardrail = ToolEditingGuardrail( + edit=lambda tools: [t for t in tools if t["function"]["name"] != "mcp__confluence__confluence_search"], + guardrail_name="test", + ) + + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + namespace = _tool_named(result["tools"], "mcp__confluence") + assert [member["name"] for member in namespace["tools"]] == ["confluence_get_page"] + assert namespace["tools"][0] == expected_tools[1]["tools"][0] + assert [t for t in result["tools"] if t is not namespace] == [expected_tools[0], *expected_tools[2:]] + + @pytest.mark.asyncio + async def test_editing_a_member_lands_on_that_member_without_the_namespace_prefix(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + + def redact_search(tools): + for tool in tools: + if tool["function"]["name"] == "mcp__confluence__confluence_search": + tool["function"]["description"] = "Confluence tools\n\nREDACTED" + return tools + + result = await OpenAIResponsesHandler().process_input_messages( + data, ToolEditingGuardrail(edit=redact_search, guardrail_name="test") + ) + + namespace = _tool_named(result["tools"], "mcp__confluence") + assert namespace["tools"][0] == expected_tools[1]["tools"][0] + assert namespace["tools"][1] == {**expected_tools[1]["tools"][1], "description": "REDACTED"} + assert {k: v for k, v in namespace.items() if k != "tools"} == { + k: v for k, v in expected_tools[1].items() if k != "tools" + } + + @pytest.mark.asyncio + async def test_dropping_every_member_drops_the_namespace(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + guardrail = ToolEditingGuardrail( + edit=lambda tools: [t for t in tools if not t["function"]["name"].startswith("mcp__confluence__")], + guardrail_name="test", + ) + + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert result["tools"] == [expected_tools[0], *expected_tools[2:]] + + @pytest.mark.asyncio + async def test_edited_top_level_function_is_rewritten_in_place(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + + def rename_weather(tools): + for tool in tools: + if tool["function"]["name"] == "get_weather": + tool["function"]["description"] = "Weather lookup (guarded)" + return tools + + result = await OpenAIResponsesHandler().process_input_messages( + data, ToolEditingGuardrail(edit=rename_weather, guardrail_name="test") + ) + + assert result["tools"][0] == {**expected_tools[0], "description": "Weather lookup (guarded)"} + assert result["tools"][1:] == expected_tools[1:] + + +class TestOpenAIResponsesHandlerMalformedTools: + @pytest.mark.asyncio + async def test_request_tools_that_are_not_a_list_never_reach_the_guardrail(self): + handler = OpenAIResponsesHandler() + seen: list[list[dict]] = [] + + def record(tools): + seen.append(tools) + return tools + + guardrail = ToolEditingGuardrail(edit=record, guardrail_name="test") + data = {"input": "hi", "tools": {"type": "function", "name": "get_weather"}} + + result = await handler.process_input_messages(data, guardrail) + + assert seen == [[]] + assert result["input"] == "hi" + + class TestBuildBlockSseChunks: """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE events""" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py new file mode 100644 index 00000000000..9c236d81f51 --- /dev/null +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py @@ -0,0 +1,198 @@ +""" +Unit tests for merge_guardrailed_tools, which writes guardrail-returned chat tools back onto the +Responses API request tools they were flattened from +""" + +import copy + +from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) +from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GuardrailToolParam + + +def _groups(tools): + return [form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(tools)] + + +def _flat(groups): + return [chat_tool for group in groups for chat_tool in group] + + +def _function(name, description=""): + return {"type": "function", "name": name, "description": description, "parameters": {"type": "object"}} + + +def test_unchanged_tools_come_back_as_the_original_objects(): + original = [ + _function("a"), + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("x"), _function("y")]}, + {"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}, + {"type": "web_search"}, + ] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, _flat(groups)) + + assert list(merged) == original + assert all(merged_tool is original_tool for merged_tool, original_tool in zip(merged, original)) + + +def test_guardrail_reordering_unchanged_tools_keeps_request_order(): + original = [_function("a"), {"type": "namespace", "name": "ns", "tools": [_function("x")]}, {"type": "web_search"}] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, list(reversed(_flat(groups)))) + + assert list(merged) == original + + +def test_duplicate_function_names_are_matched_by_ordinal(): + original = [_function("dup", "first"), _function("dup", "second")] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, _flat(groups)[:1]) + + assert list(merged) == [original[0]] + + +def test_interleaved_duplicate_names_keep_their_own_ordinals(): + original = [ + _function("dup", "a"), + _function("other", "x"), + _function("dup", "b"), + _function("dup", "c"), + _function("other", "y"), + ] + groups = _groups(original) + flat = _flat(groups) + edited = {**flat[3], "function": {**flat[3]["function"], "description": "changed"}} + + merged = merge_guardrailed_tools(original, groups, [*flat[:3], edited, flat[4]]) + + assert list(merged) == [*original[:3], {**_function("dup", "changed"), "strict": False}, original[4]] + assert all(merged[position] is original[position] for position in (0, 1, 2, 4)) + + +def test_edited_mcp_tool_is_rewritten(): + original = [{"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}] + groups = _groups(original) + edited = [{**groups[0][0], "allowed_tools": ["read_wiki_structure"]}] + + merged = merge_guardrailed_tools(original, groups, edited) + + assert list(merged) == edited + + +def test_injected_tool_lands_after_the_request_tools_when_request_had_none(): + injected = {"type": "function", "function": {"name": "b", "description": "d", "parameters": {"type": "object"}}} + + merged = merge_guardrailed_tools([], [], [injected]) + + assert list(merged) == [ + {"type": "function", "name": "b", "description": "d", "parameters": {"type": "object"}, "strict": False} + ] + + +def test_empty_guardrail_output_keeps_only_tools_never_sent_to_the_guardrail(): + original = [_function("a"), {"type": "web_search"}, {"type": "namespace", "name": "ns", "tools": [_function("x")]}] + + merged = merge_guardrailed_tools(original, _groups(original), []) + + assert list(merged) == [{"type": "web_search"}] + + +def test_member_edit_strips_only_the_namespace_description_prefix(): + original = [{"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("x", "X doc")]}] + groups = _groups(original) + assert groups[0][0]["function"]["description"] == "NS\n\nX doc" + edited = [{**groups[0][0], "function": {**groups[0][0]["function"], "description": "NS\n\nX doc (guarded)"}}] + + merged = merge_guardrailed_tools(original, groups, edited) + + assert list(merged) == [ + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("x", "X doc (guarded)")]} + ] + + +def test_namespace_keeps_a_non_function_member_when_a_function_member_is_edited(): + custom_member = {"type": "custom", "name": "grep", "description": "Grep", "format": {"type": "text"}} + original = [ + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("read", "Read"), custom_member]} + ] + groups = _groups(original) + edited = copy.deepcopy(_flat(groups)) + edited[0]["function"]["description"] = "NS\n\nEDITED" + + merged = merge_guardrailed_tools(original, groups, edited) + + assert len(merged) == 1 + assert [member["name"] for member in merged[0]["tools"]] == ["read", "grep"] + assert merged[0]["tools"][0]["description"] == "EDITED" + assert merged[0]["tools"][1] == custom_member + + +def test_namespace_keeps_its_non_function_members_when_every_function_member_is_dropped(): + custom_member = {"type": "custom", "name": "grep", "description": "Grep", "format": {"type": "text"}} + original = [ + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("read"), custom_member]}, + _function("a"), + ] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, [groups[1][0]]) + + assert list(merged) == [{"type": "namespace", "name": "ns", "description": "NS", "tools": [custom_member]}, _function("a")] + + +def test_member_extras_edited_by_the_guardrail_land_on_that_member(): + original = [{"type": "namespace", "name": "ns", "tools": [_function("read")]}] + groups = _groups(original) + edited = copy.deepcopy(_flat(groups)) + edited[0]["cache_control"] = {"type": "ephemeral"} + + merged = merge_guardrailed_tools(original, groups, edited) + + assert merged[0]["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert merged[0]["tools"][0]["name"] == "read" + + +def test_guardrail_output_is_read_once(): + original = [_function("a"), {"type": "namespace", "name": "ns", "tools": [_function("x")]}] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, (chat_tool for chat_tool in _flat(groups))) + + assert list(merged) == original + + +def test_pydantic_guardrail_tools_round_trip_like_dicts(): + original = [_function("a"), {"type": "namespace", "name": "ns", "tools": [_function("x", "X doc")]}] + groups = _groups(original) + models = [GuardrailToolParam.model_validate(chat_tool) for chat_tool in _flat(groups)] + + merged = merge_guardrailed_tools(original, groups, models) + + assert list(merged) == original + assert all(merged_tool is original_tool for merged_tool, original_tool in zip(merged, original)) + + +def test_pydantic_guardrail_edit_lands_on_the_member(): + original = [{"type": "namespace", "name": "ns", "tools": [_function("x", "X doc")]}] + groups = _groups(original) + edited = copy.deepcopy(_flat(groups)) + edited[0]["function"]["description"] = "EDITED" + + merged = merge_guardrailed_tools(original, groups, [GuardrailToolParam.model_validate(edited[0])]) + + assert list(merged) == [{"type": "namespace", "name": "ns", "tools": [_function("x", "EDITED")]}] + + +def test_non_object_guardrail_items_are_dropped(): + original = [_function("a")] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, [*_flat(groups), "junk", None]) + + assert list(merged) == original diff --git a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index 9a6a039a470..67a56fdcd79 100644 --- a/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -318,3 +318,203 @@ def test_json_provider_messages_config_probes_capabilities_under_provider_slug() ) assert JSONProviderAnthropicMessagesConfig(provider).custom_llm_provider == "exampleprovider" assert OpenAILikeAnthropicMessagesConfig().custom_llm_provider == "anthropic" + + +def _cache_control_request_params() -> tuple[list, dict]: + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "write a regex for a US phone number", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + ] + optional_params = { + "max_tokens": 256, + "system": [ + { + "type": "text", + "text": "You are Claude Code.", + "cache_control": {"type": "ephemeral", "ttl": "5m"}, + } + ], + "tools": [ + { + "name": "lookup", + "input_schema": {"type": "object"}, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + return messages, optional_params + + +def test_request_strips_cache_control_ttl_everywhere(config): + """Regression: Claude Code always sends ``cache_control: {type: ephemeral, + ttl: 1h}``, and strict non-Anthropic /v1/messages validators 400 the whole + request on the ttl extension (``cache_control.ttl: 1h is not supported``).""" + messages, optional_params = _cache_control_request_params() + + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["system"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert messages[0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +def test_request_defaults_missing_cache_control_type_and_drops_non_dict(config): + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "a", "cache_control": {"ttl": "1h"}}, + {"type": "text", "text": "b", "cache_control": None}, + ], + } + ], + anthropic_messages_optional_request_params={"max_tokens": 64}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + blocks = payload["messages"][0]["content"] + assert blocks[0]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in blocks[1] + + +def test_native_anthropic_config_keeps_cache_control_ttl(): + """Anthropic itself accepts ttl, so the normalization must stay scoped to + the OpenAI-like passthrough and never reach the native Anthropic path.""" + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + messages, optional_params = _cache_control_request_params() + payload = AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-sonnet-4-20250514", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + assert payload["system"][0]["cache_control"] == {"type": "ephemeral", "ttl": "5m"} + + +def test_deployment_opt_in_keeps_cache_control_ttl(): + config = OpenAILikeAnthropicMessagesConfig(cache_control_ttl=True) + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + } + ], + anthropic_messages_optional_request_params={"max_tokens": 16}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert payload["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +def test_json_provider_constraint_opts_into_cache_control_ttl(): + from litellm.llms.openai_like.json_loader import SimpleProviderConfig + from litellm.llms.openai_like.messages.transformation import ( + JSONProviderAnthropicMessagesConfig, + ) + + base_data = {"base_url": "https://api.example.com/v1", "api_key_env": "EXAMPLE_API_KEY"} + strict = JSONProviderAnthropicMessagesConfig(SimpleProviderConfig(slug="strictprov", data=base_data)) + lenient = JSONProviderAnthropicMessagesConfig( + SimpleProviderConfig(slug="lenientprov", data={**base_data, "constraints": {"cache_control_ttl": True}}) + ) + + def transform(provider_config): + messages, optional_params = _cache_control_request_params() + return provider_config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert transform(strict)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"} + assert transform(lenient)["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +def test_request_strips_ttl_only_where_the_messages_api_defines_cache_control(config): + """Regression: the sanitizer must only touch ``cache_control`` where the + Messages API defines it (request, system, tools, content blocks, tool_result + content), never application data such as ``tool_use.input`` or a tool's + ``input_schema`` that happens to contain a ``cache_control`` key.""" + tool_input = {"cache_control": {"type": "ephemeral", "ttl": "1h"}, "query": "x"} + input_schema = { + "type": "object", + "properties": {"cache_control": {"type": "string", "ttl": "1h"}}, + } + messages = [ + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": tool_input}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + "content": [ + {"type": "text", "text": "result", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ], + }, + {"type": "text", "text": "plain string content stays", "cache_control": {"ttl": "1h"}}, + ], + }, + {"role": "user", "content": "a plain string message"}, + ] + optional_params = { + "max_tokens": 64, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + "tools": [ + { + "name": "lookup", + "input_schema": input_schema, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + + payload = config.transform_anthropic_messages_request( + model="some-model", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert payload["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["tools"][0]["input_schema"] == input_schema + assert payload["messages"][0]["content"][0]["input"] == tool_input + tool_result = payload["messages"][1]["content"][0] + assert tool_result["cache_control"] == {"type": "ephemeral"} + assert tool_result["content"][0]["cache_control"] == {"type": "ephemeral"} + assert payload["messages"][1]["content"][1]["cache_control"] == {"type": "ephemeral"} + assert payload["messages"][2] == {"role": "user", "content": "a plain string message"} diff --git a/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py b/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py index a2ee2c2bdb1..aa114f128c5 100644 --- a/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py +++ b/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py @@ -67,20 +67,52 @@ class FakeAsyncEmbeddingFn(FakeEmbeddingFn): return SimpleNamespace(data=[{"embedding": self.embedding}]) +class FakeEmbeddingExecutor: + def __init__(self, embedding): + self.embedding = embedding + self.captured = None + + def embed(self, model, query, configuration): + self.captured = (model, query, configuration) + return SimpleNamespace(data=[{"embedding": self.embedding}]) + + async def aembed(self, model, query, configuration): + self.captured = (model, query, configuration) + return SimpleNamespace(data=[{"embedding": self.embedding}]) + + def _doc(doc_id, distance, **fields): return SimpleNamespace(id=doc_id, vector_distance=str(distance), **fields) -def _search(config, client=None, query="what is litellm", optional_params=None, litellm_params=None): +def _search(config, client=None, query="what is litellm", optional_params=None, litellm_params=None, executor=None): return config.execute_search_vector_store_request( vector_store_id="my_index", query=query, vector_store_search_optional_params=optional_params or {}, litellm_logging_obj=MagicMock(), litellm_params={"litellm_embedding_model": "openai/text-embedding-3-small", **(litellm_params or {})}, + embedding_executor=executor, ) +def test_sync_search_uses_request_embedding_executor_without_overwriting_explicit_config(): + executor = FakeEmbeddingExecutor([0.1, 0.2]) + config = ValkeyVectorStoreConfig(sync_client=FakeRedis()) + embedding_config = {"api_key": "store-specific-key", "aws_region_name": "us-west-2"} + + _search( + config, + litellm_params={ + "litellm_embedding_model": "team-embedding-alias", + "litellm_embedding_config": embedding_config, + }, + executor=executor, + ) + + assert executor.captured == ("team-embedding-alias", "what is litellm", embedding_config) + + def test_sync_search_builds_knn_query_with_packed_vector(): embedding_fn = FakeEmbeddingFn([0.1, 0.2, 0.3]) client = FakeRedis() diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 9419f88a981..a57672cfbfb 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -752,3 +752,26 @@ def test_vertex_ai_fable_5_1_response_format_uses_native_output_format(local_mod assert "output_format" in result_params assert "tool_choice" not in result_params assert "tools" not in result_params + + +def test_vertex_ai_anthropic_tool_based_response_format_still_upgrades_legacy_thinking(local_model_cost_map): + result_params = VertexAIAnthropicConfig().map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": {"type": "object", "properties": {"result": {"type": "string"}}}, + }, + }, + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "max_tokens": 8192, + }, + optional_params={}, + model="claude-opus-4-8", + drop_params=False, + ) + + assert "tools" in result_params + assert result_params["thinking"] == {"type": "adaptive"} + assert result_params["output_config"] == {"effort": "high"} diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index acad249a2bb..0764aec7185 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -1,7 +1,7 @@ """Tests for the optional Rust-backed OCR path.""" -import importlib import builtins +import importlib import types from typing import Any @@ -10,6 +10,7 @@ import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge import configuration # `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr` # function onto `litellm.ocr` and shadows the submodule, so import the modules @@ -214,10 +215,12 @@ def build_prepared_request( @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" - rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) + rust_bridge.set_rust_ocr(ocr=None, aocr=None) + configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield - rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) + rust_bridge.set_rust_ocr(ocr=None, aocr=None) + configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -247,7 +250,14 @@ def test_use_litellm_rust_toggles_flag(): def test_env_var_enables_rust_ocr(monkeypatch): monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") - assert rust_bridge._env_enables_rust_ocr() is True + with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): + assert rust_bridge.rust_ocr_enabled() is True + + +def test_explicit_false_overrides_process_enable(): + litellm.use_litellm_rust(True) + + assert ocr_main._rust_ocr_enabled(build_prepared_request(litellm_params={"rust": False})) is False def test_load_rust_ocr_returns_injected_impl(): @@ -471,9 +481,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): ocr_main._run_rust_ocr( prepared_request=build_prepared_request(api_key=None, timeout=None), - resolve_api_key=lambda name: ( - "sk-from-vault" if name == "MISTRAL_API_KEY" else None - ), + resolve_api_key=lambda name: "sk-from-vault" if name == "MISTRAL_API_KEY" else None, ) assert bridge.calls[0]["api_key"] == "sk-from-vault" @@ -580,9 +588,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): api_base=None, timeout=None, ), - resolve_api_key=lambda name: ( - "https://azure.example.com" if name == "AZURE_AI_API_BASE" else None - ), + resolve_api_key=lambda name: "https://azure.example.com" if name == "AZURE_AI_API_BASE" else None, ) assert bridge.calls[0]["api_base"] == "https://azure.example.com" @@ -600,9 +606,7 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): timeout=None, ), resolve_api_key=lambda name: ( - "https://document-intelligence.example.com" - if name == "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT" - else None + "https://document-intelligence.example.com" if name == "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT" else None ), ) @@ -815,9 +819,6 @@ def test_ocr_provider_configs_expose_api_key_env_vars(): assert BaseOCRConfig().get_api_key_env_var() is None assert MistralOCRConfig().get_api_key_env_var() == "MISTRAL_API_KEY" assert AzureAIOCRConfig().get_api_key_env_var() == "AZURE_AI_API_KEY" - assert ( - AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() - == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" - ) + assert AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" assert VertexAIOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" assert VertexAIDeepSeekOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 0480bbc40a7..d441c05090b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1,4 +1,5 @@ import asyncio +import inspect import json import sys from datetime import datetime @@ -13,6 +14,7 @@ import pytest from fastapi import HTTPException from starlette.requests import Request +from litellm.constants import MCP_TOOL_LISTING_TIMEOUT from litellm.proxy._experimental.mcp_server import rest_endpoints from litellm.proxy._experimental.mcp_server.auth import ( user_api_key_auth_mcp as auth_mcp, @@ -109,6 +111,71 @@ class TestExecuteWithMcpClient: assert result["status"] == "error" assert "stack_trace" not in result + @pytest.mark.asyncio + async def test_timeout_caps_hanging_operation_and_names_url(self, monkeypatch): + async def fake_create_client(*args, **kwargs): + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + ) + + async def hanging_operation(client): + await asyncio.Event().wait() + + payload = NewMCPServerRequest( + server_name="example", + url="https://mcp.example.com/mcp/", + auth_type=MCPAuth.none, + ) + + result = await asyncio.wait_for( + rest_endpoints._execute_with_mcp_client(payload, hanging_operation, timeout_seconds=0.05), + timeout=5, + ) + + assert result["error"] is True + assert "https://mcp.example.com/mcp/" in result["message"] + + @pytest.mark.asyncio + async def test_timeout_covers_client_creation(self, monkeypatch): + async def hanging_create_client(*args, **kwargs): + await asyncio.Event().wait() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + hanging_create_client, + ) + + async def unreached_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="example", + url="https://mcp.example.com/mcp/", + auth_type=MCPAuth.none, + ) + + result = await asyncio.wait_for( + rest_endpoints._execute_with_mcp_client(payload, unreached_operation, timeout_seconds=0.05), + timeout=5, + ) + + assert result["error"] is True + assert "https://mcp.example.com/mcp/" in result["message"] + + def test_timeout_defaults_to_tool_listing_timeout(self): + default = inspect.signature(rest_endpoints._execute_with_mcp_client).parameters["timeout_seconds"].default + assert default == MCP_TOOL_LISTING_TIMEOUT + + def test_connection_error_message_timeout_names_url_and_budget(self): + message = rest_endpoints._connection_error_message(TimeoutError(), "https://api.example.com/mcp/", 30.0) + assert "https://api.example.com/mcp/" in message + assert "30s" in message + @pytest.mark.asyncio async def test_forwards_static_headers(self, monkeypatch): """Ensure static_headers are forwarded to the MCP client during test calls. @@ -3168,17 +3235,21 @@ class TestConnectionErrorMessage: secret = "Bearer sk-super-secret-token" exc = httpx.LocalProtocolError(f"Illegal header value b' {secret}'") - message = rest_endpoints._connection_error_message(exc) + message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "header" in message.lower() assert secret not in message def test_connect_error_points_at_reachability(self): - message = rest_endpoints._connection_error_message(httpx.ConnectError("All connection attempts failed")) + message = rest_endpoints._connection_error_message( + httpx.ConnectError("All connection attempts failed"), "https://example.com", 30.0 + ) assert "unreachable" in message.lower() def test_timeout_error_message(self): - message = rest_endpoints._connection_error_message(httpx.ConnectTimeout("timed out")) + message = rest_endpoints._connection_error_message( + httpx.ConnectTimeout("timed out"), "https://example.com", 30.0 + ) assert "unreachable" in message.lower() def test_http_status_error_includes_status_code(self): @@ -3188,11 +3259,11 @@ class TestConnectionErrorMessage: request=httpx.Request("POST", "http://x/"), response=response, ) - message = rest_endpoints._connection_error_message(exc) + message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0) assert "503" in message def test_unknown_error_falls_back_to_generic(self): - message = rest_endpoints._connection_error_message(RuntimeError("weird")) + message = rest_endpoints._connection_error_message(RuntimeError("weird"), "https://example.com", 30.0) assert "weird" not in message assert "proxy logs" in message.lower() diff --git a/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py b/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py index e0585ab04f1..4c219760762 100644 --- a/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py +++ b/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py @@ -591,3 +591,100 @@ class TestFilterServerIdsByIpWithInfo: ) assert allowed == [] assert blocked == 2 + + +def _make_scheme_request( + scheme: str, client_host: str = "203.0.113.5", headers: dict[str, str] | None = None +) -> Request: + request = MagicMock(spec=Request) + request.client = MagicMock() + request.client.host = client_host + request.headers = headers or {} + request.url = MagicMock() + request.url.scheme = scheme + return request + + +class TestIsRequestHttps: + """Regression tests for the cookie Secure trust-boundary resolution. + + litellm only sees a plain-HTTP hop when TLS terminates at a reverse + proxy, so a cookie's Secure attribute must not be derived from the + literal request scheme alone. It must also not blindly trust a + client-spoofable X-Forwarded-Proto header with no trust boundary. + """ + + def test_direct_https_is_secure(self, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = _make_scheme_request("https") + assert IPAddressUtils.is_request_https(request, general_settings={}) is True + + def test_direct_http_is_not_secure(self, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = _make_scheme_request("http") + assert IPAddressUtils.is_request_https(request, general_settings={}) is False + + def test_spoofed_forwarded_proto_without_trusted_proxy_config_is_ignored( + self, monkeypatch + ): + # Regression: an internal HTTP hop with an attacker-supplied + # X-Forwarded-Proto: https must NOT flip Secure on, because no + # trust boundary (use_x_forwarded_for + mcp_trusted_proxy_ranges) + # is configured. Blindly trusting this header is itself a + # vulnerability. + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = _make_scheme_request( + "http", headers={"X-Forwarded-Proto": "https"} + ) + assert IPAddressUtils.is_request_https(request, general_settings={}) is False + + def test_forwarded_proto_honored_only_from_trusted_proxy(self, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = _make_scheme_request( + "http", + client_host="10.0.0.5", + headers={"X-Forwarded-Proto": "https"}, + ) + general_settings = { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + } + assert IPAddressUtils.is_request_https(request, general_settings=general_settings) is True + + def test_forwarded_proto_http_from_trusted_proxy_is_not_secure(self, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = _make_scheme_request( + "https", + client_host="10.0.0.5", + headers={"X-Forwarded-Proto": "http"}, + ) + general_settings = { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + } + assert IPAddressUtils.is_request_https(request, general_settings=general_settings) is False + + def test_untrusted_direct_peer_falls_back_to_literal_scheme(self, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = _make_scheme_request( + "http", + client_host="203.0.113.5", + headers={"X-Forwarded-Proto": "https"}, + ) + general_settings = { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + } + assert IPAddressUtils.is_request_https(request, general_settings=general_settings) is False + + def test_proxy_base_url_https_overrides_literal_http_scheme(self, monkeypatch): + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + request = _make_scheme_request("http") + assert IPAddressUtils.is_request_https(request, general_settings={}) is True + + def test_proxy_base_url_http_overrides_literal_https_scheme(self, monkeypatch): + # An explicit operator-configured plain-http public origin wins over + # the literal connection scheme, same as the https direction above. + monkeypatch.setenv("PROXY_BASE_URL", "http://litellm.internal") + request = _make_scheme_request("https") + assert IPAddressUtils.is_request_https(request, general_settings={}) is False diff --git a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py index 57decc7d458..36d6414ba9f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_saml_sso.py @@ -642,3 +642,77 @@ async def test_read_acs_post_data_rejects_oversized_stream_without_content_lengt with pytest.raises(HTTPException) as exc: await SAMLAuthHandler.read_acs_post_data(cast(Request, request)) assert exc.value.status_code == 413 + + +def _fake_request_with_scheme(scheme, headers=None, client_host="203.0.113.5"): + """A fuller fake Request than ``_fake_request``: adds ``url``, ``headers`` and + ``client``, which ``IPAddressUtils.is_request_https`` reads directly instead of + going through ``PROXY_BASE_URL``.""" + return type( + "Req", + (), + { + "base_url": URL(f"{scheme}://proxy.example.com/"), + "url": URL(f"{scheme}://proxy.example.com/sso/saml/login"), + "query_params": {}, + "cookies": {}, + "headers": headers or {}, + "client": type("Client", (), {"host": client_host})(), + }, + )() + + +class TestSAMLAuthnCookieSecureFlag: + """Regression tests for the litellm_saml_authn cookie's Secure attribute. + litellm only sees a plain-HTTP hop whenever TLS terminates at a reverse + proxy, so Secure must not be derived from the literal request scheme alone.""" + + @pytest.mark.asyncio + async def test_secure_over_direct_https(self, saml_env, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + cache = DualCache() + request = _fake_request_with_scheme("https") + redirect = await SAMLAuthHandler.build_login_redirect(request, cache) + cookie = redirect.headers["set-cookie"] + assert "Secure" in cookie + assert "SameSite=none" in cookie + + @pytest.mark.asyncio + async def test_not_secure_over_direct_http(self, saml_env, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + cache = DualCache() + request = _fake_request_with_scheme("http") + redirect = await SAMLAuthHandler.build_login_redirect(request, cache) + cookie = redirect.headers["set-cookie"] + assert "Secure" not in cookie + assert "SameSite=lax" in cookie + + @pytest.mark.asyncio + async def test_secure_behind_trusted_tls_terminating_proxy(self, saml_env, monkeypatch): + """THE regression: TLS terminates at a reverse proxy, litellm only sees a + plain-HTTP hop, but the cookie must still be marked Secure when the operator + has configured a trusted proxy reporting X-Forwarded-Proto: https.""" + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": ["10.0.0.0/8"]}, + ) + cache = DualCache() + request = _fake_request_with_scheme( + "http", headers={"X-Forwarded-Proto": "https"}, client_host="10.0.0.5" + ) + redirect = await SAMLAuthHandler.build_login_redirect(request, cache) + cookie = redirect.headers["set-cookie"] + assert "Secure" in cookie + + @pytest.mark.asyncio + async def test_untrusted_spoofed_forwarded_proto_is_ignored(self, saml_env, monkeypatch): + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + cache = DualCache() + request = _fake_request_with_scheme( + "http", headers={"X-Forwarded-Proto": "https"}, client_host="203.0.113.5" + ) + redirect = await SAMLAuthHandler.build_login_redirect(request, cache) + cookie = redirect.headers["set-cookie"] + assert "Secure" not in cookie diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index e648bd09734..dd8c752a868 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -7604,6 +7604,112 @@ class TestPKCEStateCookieBinding: assert cookie_str is not None assert "Secure" not in cookie_str + @pytest.mark.asyncio + async def test_redirect_response_sets_secure_flag_behind_trusted_tls_terminating_proxy( + self, monkeypatch + ): + """Regression: litellm sees a plain-HTTP hop when TLS terminates at a reverse + proxy. The Secure flag must still be set when the direct peer is a configured + trusted proxy and it reports X-Forwarded-Proto: https -- but NOT from an + unconfigured/untrusted caller spoofing the same header (see the sibling test + below).""" + from fastapi.responses import RedirectResponse + + from litellm.proxy.management_endpoints.ui_sso import ( + SSOAuthenticationHandler, + ) + + mock_redirect = RedirectResponse( + url="http://idp.internal/authorize?state=behind-proxy-state" + ) + mock_generic_sso = MagicMock() + mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) + mock_generic_sso.__exit__ = MagicMock(return_value=None) + mock_generic_sso.get_login_redirect = AsyncMock(return_value=mock_redirect) + + proxied_request = MagicMock(spec=Request) + proxied_request.url.scheme = "http" + proxied_request.headers = {"X-Forwarded-Proto": "https"} + proxied_request.client = MagicMock() + proxied_request.client.host = "10.0.0.5" + + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": ["10.0.0.0/8"]}, + ) + + with patch.dict( + os.environ, + { + "GENERIC_CLIENT_STATE": "behind-proxy-state", + "GENERIC_CLIENT_USE_PKCE": "true", + }, + ): + response = await SSOAuthenticationHandler.get_generic_sso_redirect_response( + generic_sso=mock_generic_sso, + state=None, + generic_authorization_endpoint="http://idp.internal/authorize", + request=proxied_request, + ) + + cookie_headers = response.headers.getlist("set-cookie") + cookie_str = next( + (c for c in cookie_headers if "litellm_oauth_state=" in c), None + ) + assert cookie_str is not None + assert "Secure" in cookie_str + + @pytest.mark.asyncio + async def test_redirect_response_ignores_spoofed_forwarded_proto_without_trust_config( + self, monkeypatch + ): + """The same X-Forwarded-Proto: https header must NOT flip Secure on when no + trusted-proxy config is present -- honoring it unconditionally would let any + client spoof the header and would not itself be the vulnerability the ticket + warns against.""" + from fastapi.responses import RedirectResponse + + from litellm.proxy.management_endpoints.ui_sso import ( + SSOAuthenticationHandler, + ) + + mock_redirect = RedirectResponse( + url="http://idp.internal/authorize?state=spoofed-state" + ) + mock_generic_sso = MagicMock() + mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) + mock_generic_sso.__exit__ = MagicMock(return_value=None) + mock_generic_sso.get_login_redirect = AsyncMock(return_value=mock_redirect) + + spoofed_request = MagicMock(spec=Request) + spoofed_request.url.scheme = "http" + spoofed_request.headers = {"X-Forwarded-Proto": "https"} + spoofed_request.client = MagicMock() + spoofed_request.client.host = "203.0.113.5" + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + with patch.dict( + os.environ, + { + "GENERIC_CLIENT_STATE": "spoofed-state", + "GENERIC_CLIENT_USE_PKCE": "true", + }, + ): + response = await SSOAuthenticationHandler.get_generic_sso_redirect_response( + generic_sso=mock_generic_sso, + state=None, + generic_authorization_endpoint="http://idp.internal/authorize", + request=spoofed_request, + ) + + cookie_headers = response.headers.getlist("set-cookie") + cookie_str = next( + (c for c in cookie_headers if "litellm_oauth_state=" in c), None + ) + assert cookie_str is not None + assert "Secure" not in cookie_str + @pytest.mark.asyncio async def test_pkce_callback_rejects_missing_cookie(self): """When PKCE is enabled and a code_verifier is in the cache, the @@ -8586,6 +8692,24 @@ class TestSameOriginReturnPath: assert _is_same_origin_return_path("") is False +def _make_https_request() -> Request: + request = MagicMock(spec=Request) + request.url.scheme = "https" + request.headers = {} + request.client = MagicMock() + request.client.host = "203.0.113.5" + return request + + +def _make_http_request() -> Request: + request = MagicMock(spec=Request) + request.url.scheme = "http" + request.headers = {} + request.client = MagicMock() + request.client.host = "203.0.113.5" + return request + + class TestPersistReturnToCookieSharedHelper: """The single shared return_to helper used by EVERY sign-in branch (SSO / Okta / generic AND the username/password form). It must be best-effort and NEVER raise — a bad return_to can never block @@ -8603,7 +8727,7 @@ class TestPersistReturnToCookieSharedHelper: monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) resp = Response() - _persist_return_to_cookie(resp, "/mcp/authorize?client_id=llm_dcrc_abc") + _persist_return_to_cookie(resp, "/mcp/authorize?client_id=llm_dcrc_abc", _make_https_request()) assert "litellm_cp_return_to=" in self._cookie(resp) def test_bad_absolute_with_control_plane_configured_does_not_raise_and_is_not_stored(self, monkeypatch): @@ -8617,7 +8741,7 @@ class TestPersistReturnToCookieSharedHelper: "litellm.proxy.proxy_server.general_settings", {"control_plane_url": "https://cp.example.com"} ) resp = Response() - _persist_return_to_cookie(resp, "https://evil.example.com/steal") # must not raise + _persist_return_to_cookie(resp, "https://evil.example.com/steal", _make_https_request()) # must not raise assert "litellm_cp_return_to=" not in self._cookie(resp) def test_none_return_to_is_a_noop(self): @@ -8626,7 +8750,7 @@ class TestPersistReturnToCookieSharedHelper: from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie resp = Response() - _persist_return_to_cookie(resp, None) + _persist_return_to_cookie(resp, None, _make_https_request()) assert "litellm_cp_return_to=" not in self._cookie(resp) def test_control_plane_matching_absolute_is_stored(self, monkeypatch): @@ -8638,5 +8762,126 @@ class TestPersistReturnToCookieSharedHelper: "litellm.proxy.proxy_server.general_settings", {"control_plane_url": "https://cp.example.com"} ) resp = Response() - _persist_return_to_cookie(resp, "https://cp.example.com/ui?page=models") + _persist_return_to_cookie(resp, "https://cp.example.com/ui?page=models", _make_https_request()) assert "litellm_cp_return_to=" in self._cookie(resp) + + def test_cookie_is_secure_and_httponly_over_https(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + resp = Response() + _persist_return_to_cookie(resp, "/mcp/authorize", _make_https_request()) + cookie = self._cookie(resp) + assert "Secure" in cookie + assert "HttpOnly" in cookie + assert "SameSite=lax" in cookie + + def test_cookie_is_not_secure_over_plain_http_direct(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + resp = Response() + _persist_return_to_cookie(resp, "/mcp/authorize", _make_http_request()) + assert "Secure" not in self._cookie(resp) + + def test_cookie_is_secure_behind_trusted_tls_terminating_proxy(self, monkeypatch): + """Regression for the reported bug: TLS terminates at a reverse proxy, litellm only + sees a plain-HTTP hop, but a trusted X-Forwarded-Proto: https must still mark the + cookie Secure.""" + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie + + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": ["10.0.0.0/8"]}, + ) + resp = Response() + request = _make_http_request() + request.client.host = "10.0.0.5" + request.headers = {"X-Forwarded-Proto": "https"} + _persist_return_to_cookie(resp, "/mcp/authorize", request) + assert "Secure" in self._cookie(resp) + + +class TestSessionTokenCookie: + """Regression tests for the ``token`` session cookie set by every sign-in path + (username/password login, SSO callback, the CLI /v2, /v3 login exchange helpers). + It was previously set with no Secure/HttpOnly/SameSite attributes at all -- always + sent over plain HTTP and readable by any script on the page. HttpOnly must stay off + deliberately: the dashboard reads this cookie via document.cookie.""" + + @staticmethod + def _cookie(resp) -> str: + return resp.headers.get("set-cookie", "") + + def test_secure_over_direct_https(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + resp = Response() + set_session_token_cookie(resp, _make_https_request(), "jwt-token-value") + cookie = self._cookie(resp) + assert "token=jwt-token-value" in cookie + assert "Secure" in cookie + assert "SameSite=lax" in cookie + assert "HttpOnly" not in cookie + + def test_not_secure_over_direct_http(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + resp = Response() + set_session_token_cookie(resp, _make_http_request(), "jwt-token-value") + assert "Secure" not in self._cookie(resp) + + def test_secure_behind_trusted_tls_terminating_proxy(self, monkeypatch): + """THE regression: TLS terminates at a reverse proxy, litellm only sees a + plain-HTTP hop, but the session cookie must still be marked Secure when the + operator has configured a trusted proxy that reports X-Forwarded-Proto: https.""" + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": ["10.0.0.0/8"]}, + ) + request = _make_http_request() + request.client.host = "10.0.0.5" + request.headers = {"X-Forwarded-Proto": "https"} + resp = Response() + set_session_token_cookie(resp, request, "jwt-token-value") + assert "Secure" in self._cookie(resp) + + def test_untrusted_spoofed_forwarded_proto_is_ignored(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + request = _make_http_request() + request.headers = {"X-Forwarded-Proto": "https"} + resp = Response() + set_session_token_cookie(resp, request, "jwt-token-value") + assert "Secure" not in self._cookie(resp) + + def test_proxy_base_url_https_overrides_literal_http_scheme(self, monkeypatch): + from fastapi import Response + + from litellm.proxy.management_endpoints.ui_sso import set_session_token_cookie + + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + resp = Response() + set_session_token_cookie(resp, _make_http_request(), "jwt-token-value") + assert "Secure" in self._cookie(resp) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 5f692f8f109..4ed6a468371 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -148,6 +148,72 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): assert mock_jwt_encode.call_args.kwargs == {"algorithm": "HS256"} +def _mock_login_v2_deps(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + AsyncMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.create_ui_token_object", + MagicMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token")) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") + monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + + +def test_login_v2_sets_secure_cookie_over_direct_https(monkeypatch): + """Regression: the token cookie previously carried no Secure/HttpOnly/SameSite + attributes at all, so it was always sent over plain HTTP.""" + _mock_login_v2_deps(monkeypatch) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + client = TestClient(app, base_url="https://testserver") + response = client.post("/v2/login", json={"username": "alice", "password": "secret"}) + + assert response.status_code == 200 + cookie = response.headers.get("set-cookie") + assert "Secure" in cookie + assert "HttpOnly" not in cookie # deliberate: the dashboard reads this cookie via JS + assert "samesite=lax" in cookie.lower() + + +def test_login_v2_does_not_set_secure_cookie_over_direct_http(monkeypatch): + _mock_login_v2_deps(monkeypatch) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + client = TestClient(app, base_url="http://testserver") + response = client.post("/v2/login", json={"username": "alice", "password": "secret"}) + + assert response.status_code == 200 + assert "Secure" not in response.headers.get("set-cookie") + + +def test_login_v2_sets_secure_cookie_behind_trusted_tls_terminating_proxy(monkeypatch): + """THE regression: litellm only sees a plain-HTTP hop when TLS terminates at a + reverse proxy, but the token cookie must still be Secure when the direct peer is + a configured trusted proxy reporting X-Forwarded-Proto: https.""" + _mock_login_v2_deps(monkeypatch) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": ["10.0.0.0/8"]}, + ) + + client = TestClient(app, base_url="http://testserver", client=("10.0.0.5", 50000)) + response = client.post( + "/v2/login", + json={"username": "alice", "password": "secret"}, + headers={"X-Forwarded-Proto": "https"}, + ) + + assert response.status_code == 200 + assert "Secure" in response.headers.get("set-cookie") + + def test_login_v2_returns_json_on_proxy_exception(monkeypatch): """Test that /v2/login returns JSON error when ProxyException is raised""" from litellm.proxy._types import ProxyErrorTypes, ProxyException @@ -356,6 +422,51 @@ def test_login_v3_exchange_happy_path(monkeypatch): assert exchange_response.cookies.get("token") == "signed-token" +def test_login_v3_exchange_sets_secure_cookie_behind_trusted_tls_terminating_proxy(monkeypatch): + """Regression: /v3/login/exchange's token cookie must be Secure behind a trusted + TLS-terminating reverse proxy even though litellm only sees a plain-HTTP hop.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + AsyncMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.create_ui_token_object", + MagicMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token")) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + { + "control_plane_url": "https://cp.example.com", + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + }, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_config = MagicMock() + mock_config.worker_registry = [] + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config) + monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") + monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + + client = TestClient(app, base_url="http://testserver", client=("10.0.0.5", 50000)) + + login_response = client.post("/v3/login", json={"username": "alice", "password": "secret"}) + code = login_response.json()["code"] + + exchange_response = client.post( + "/v3/login/exchange", + json={"code": code}, + headers={"X-Forwarded-Proto": "https"}, + ) + assert exchange_response.status_code == 200 + assert "Secure" in exchange_response.headers.get("set-cookie") + + def test_login_v3_exchange_single_use(monkeypatch): """Code can only be redeemed once.""" mock_prisma_client = MagicMock() diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 45a0221c8a6..1abbbe91e97 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -2,29 +2,24 @@ from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import Request - - -from fastapi import HTTPException +from fastapi import HTTPException, Request import litellm from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) +from litellm.llms.base_llm.vector_store.transformation import ( + LiteLLMVectorStoreEmbeddingExecutor, + RouterVectorStoreEmbeddingExecutor, +) from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.vector_store_endpoints.endpoints import ( _update_request_data_with_litellm_managed_vector_store_registry, index_create, index_list, ) -from litellm.proxy.vector_store_files_endpoints.endpoints import ( - _update_request_data_with_model_routing_hint, -) from litellm.proxy.vector_store_endpoints.management_endpoints import ( _check_vector_store_access, - _resolve_embedding_config, - _resolve_embedding_config_from_db, - _resolve_embedding_config_from_router, create_vector_store_in_db, new_vector_store, ) @@ -33,8 +28,12 @@ from litellm.proxy.vector_store_endpoints.utils import ( is_allowed_to_call_vector_store_endpoint, is_allowed_to_call_vector_store_files_endpoint, ) +from litellm.proxy.vector_store_files_endpoints.endpoints import ( + _update_request_data_with_model_routing_hint, +) +from litellm.types.utils import EmbeddingResponse, LlmProviders from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse -from litellm.types.utils import LlmProviders +from litellm.vector_stores.main import _direct_vector_store_embedding_executor def _serialize_litellm_params(litellm_params): @@ -51,17 +50,113 @@ def _serialize_litellm_params(litellm_params): return json.dumps(litellm_params or {}) -@pytest.fixture(autouse=True) -def _reset_embedding_config_cache(): - """The use-time embedding-config resolver caches results in process - memory across calls. Reset it before every test so the resolver - actually exercises the router/DB path under test instead of returning - a value cached by an earlier test.""" - from litellm.proxy.vector_store_endpoints import management_endpoints +def test_direct_vector_store_embedding_executor_rejects_invalid_value(): + with pytest.raises(TypeError, match="Invalid direct vector store embedding executor"): + _direct_vector_store_embedding_executor(object(), None, {}) - management_endpoints._embedding_config_cache = None - yield - management_endpoints._embedding_config_cache = None + +def test_router_vector_store_search_injects_executor_and_request_metadata(): + router = litellm.Router(model_list=[]) + original = MagicMock(return_value="searched") + wrapped = router.factory_function(original, call_type="vector_store_search") + + assert ( + wrapped( + vector_store_id="store", + query="query", + custom_llm_provider="valkey", + litellm_metadata={"user_api_key_team_id": "team-a"}, + ) + == "searched" + ) + + call_kwargs = original.call_args.kwargs + assert call_kwargs["custom_llm_provider"] == "valkey" + executor = call_kwargs["_direct_vector_store_embedding_executor"] + assert isinstance(executor, RouterVectorStoreEmbeddingExecutor) + assert executor.metadata == {"user_api_key_team_id": "team-a"} + assert litellm.Router._vector_store_request_metadata({"metadata": {"user_api_key_team_id": "team-b"}}) == { + "user_api_key_team_id": "team-b" + } + assert litellm.Router._vector_store_request_metadata({}) == {} + + with patch.object( # test-quality-ok: fallback dispatch is the boundary this wrapper delegates to + router, "_generic_api_call_with_fallbacks", return_value="routed" + ) as fallback: + assert wrapped(model="vector-alias", vector_store_id="store", query="query") == "routed" + assert fallback.call_args.kwargs["model"] == "vector-alias" + assert fallback.call_args.kwargs["original_function"] is original + + create_original = MagicMock(return_value="created") + wrapped_create = router.factory_function(create_original, call_type="vector_store_create") + assert wrapped_create(name="store") == "created" + create_original.assert_called_once_with(name="store") + with patch.object( # test-quality-ok: fallback dispatch is the boundary this wrapper delegates to + router, "_generic_api_call_with_fallbacks", return_value="created-through-router" + ) as fallback: + assert wrapped_create(model="vector-alias", name="store") == "created-through-router" + fallback.assert_called_once_with(original_function=create_original, model="vector-alias", name="store") + + +@pytest.mark.asyncio +async def test_vector_store_embedding_executors_preserve_explicit_configuration(): + response = EmbeddingResponse(data=[{"embedding": [0.1], "index": 0, "object": "embedding"}]) + sdk_executor = LiteLLMVectorStoreEmbeddingExecutor() + + with ( + patch( # test-quality-ok: isolates SDK dispatch from external embedding providers + "litellm.embedding", return_value=response + ) as embedding, + patch( # test-quality-ok: isolates async SDK dispatch from external embedding providers + "litellm.aembedding", new=AsyncMock(return_value=response) + ) as aembedding, + ): + assert sdk_executor.embed("openai/model", "sync", {"api_key": "explicit"}) is response + assert await sdk_executor.aembed("openai/model", "async", {"api_key": "explicit"}) is response + + embedding.assert_called_once_with(model="openai/model", input=["sync"], api_key="explicit") + aembedding.assert_awaited_once_with(model="openai/model", input=["async"], api_key="explicit") + + mock_router = MagicMock() + mock_router.embedding.return_value = response + mock_router.aembedding = AsyncMock(return_value=response) + router_executor = RouterVectorStoreEmbeddingExecutor( + router=mock_router, + metadata={"user_api_key_team_id": "team-a"}, + ) + + assert router_executor.embed("team-alias", "query", {}) is response + mock_router.embedding.assert_called_once_with( + model="team-alias", + input=["query"], + metadata={"user_api_key_team_id": "team-a"}, + ) + + with ( + patch( # test-quality-ok: verifies explicit store configuration at the SDK boundary + "litellm.embedding", return_value=response + ) as explicit_embedding, + patch( # test-quality-ok: verifies async explicit store configuration at the SDK boundary + "litellm.aembedding", new=AsyncMock(return_value=response) + ) as explicit_aembedding, + ): + assert router_executor.embed("openai/model", "query", {"api_key": "store-key"}) is response + assert await router_executor.aembed("openai/model", "query", {"api_key": "store-key"}) is response + + explicit_embedding.assert_not_called() + explicit_aembedding.assert_not_awaited() + assert mock_router.embedding.call_args.kwargs == { + "model": "openai/model", + "input": ["query"], + "api_key": "store-key", + "metadata": {"user_api_key_team_id": "team-a"}, + } + mock_router.aembedding.assert_awaited_once_with( + model="openai/model", + input=["query"], + api_key="store-key", + metadata={"user_api_key_team_id": "team-a"}, + ) @pytest.mark.asyncio @@ -82,10 +177,11 @@ async def test_router_avector_store_search_passes_correct_args(): } # Call router's avector_store_search - result = await router.avector_store_search( + await router.avector_store_search( vector_store_id="test_store_id", query="test query", custom_llm_provider="bedrock", + metadata={"user_api_key_team_id": "team-a"}, ) # Verify the internal method was called with correct args @@ -96,6 +192,38 @@ async def test_router_avector_store_search_passes_correct_args(): assert call_args[1]["vector_store_id"] == "test_store_id" assert call_args[1]["query"] == "test query" assert call_args[1]["custom_llm_provider"] == "bedrock" + executor = call_args[1]["_direct_vector_store_embedding_executor"] + assert isinstance(executor, RouterVectorStoreEmbeddingExecutor) + assert executor.metadata["user_api_key_team_id"] == "team-a" + + +@pytest.mark.asyncio +async def test_vector_store_embedding_executor_uses_team_scoped_router_deployment(): + router = litellm.Router( + model_list=[ + { + "model_name": "shared-embedding", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "team-a-key"}, + "model_info": {"team_id": "team-a", "team_public_model_name": "shared-embedding"}, + }, + { + "model_name": "shared-embedding", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "team-b-key"}, + "model_info": {"team_id": "team-b", "team_public_model_name": "shared-embedding"}, + }, + ] + ) + executor = RouterVectorStoreEmbeddingExecutor( + router=router, + metadata={"user_api_key_team_id": "team-b"}, + ) + response = EmbeddingResponse(data=[{"embedding": [0.1], "index": 0, "object": "embedding"}]) + + with patch("litellm.aembedding", new=AsyncMock(return_value=response)) as mock_aembedding: + result = await executor.aembed("shared-embedding", "query", {}) + + assert result is response + assert mock_aembedding.await_args.kwargs["api_key"] == "team-b-key" @pytest.mark.asyncio @@ -502,91 +630,30 @@ async def test_update_request_data_with_litellm_managed_vector_store_registry(): @pytest.mark.asyncio -async def test_update_request_data_resolves_embedding_config_at_use_time(): - """When the persisted vector store row carries only a - ``litellm_embedding_model`` reference (the new behaviour after - moving the auto-resolve out of write time), the request-handling - layer must resolve the embedding config so the downstream embed - call still has ``api_key`` / ``api_base`` / ``api_version``. The - resolved config lives in this per-request data dict only — never - persisted.""" - mock_vector_store: LiteLLM_ManagedVectorStore = { +async def test_managed_vector_store_keeps_embedding_reference_and_explicit_config(): + explicit_config = {"api_key": "store-specific-key", "api_base": "https://embedding.example"} + managed_vector_store: LiteLLM_ManagedVectorStore = { "vector_store_id": "test_store", - "custom_llm_provider": "azure_ai", + "custom_llm_provider": "valkey", "litellm_params": { - "litellm_embedding_model": "azure/text-embedding-3-large", - # Note: no litellm_embedding_config persisted + "litellm_embedding_model": "team-embedding-alias", + "litellm_embedding_config": explicit_config, }, } - mock_registry = MagicMock() - mock_registry.get_litellm_managed_vector_store_from_registry.return_value = ( - mock_vector_store - ) + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = managed_vector_store - resolved = { - "api_key": "use-time-resolved-key", - "api_base": "https://my-azure.example", - "api_version": "2024-09-01", - } - - with ( - patch.object(litellm, "vector_store_registry", mock_registry), - patch( - "litellm.proxy.vector_store_endpoints.endpoints._resolve_embedding_config", - new=AsyncMock(return_value=resolved), - ), - ): + with patch.object(litellm, "vector_store_registry", mock_registry): result = await _update_request_data_with_litellm_managed_vector_store_registry( - data={}, vector_store_id="test_store" + data={}, + vector_store_id="test_store", ) - assert result["litellm_embedding_model"] == "azure/text-embedding-3-large" - assert result["litellm_embedding_config"] == resolved + assert result["litellm_embedding_model"] == "team-embedding-alias" + assert result["litellm_embedding_config"] == explicit_config + assert managed_vector_store["litellm_params"]["litellm_embedding_config"] == explicit_config -@pytest.mark.asyncio -async def test_update_request_data_passes_through_legacy_embedding_config(): - """A vector store row created by an older proxy version may already - carry a fully-resolved ``litellm_embedding_config`` in its persisted - ``litellm_params`` (the very leak this PR closes). Those legacy rows - must still work — the use-time resolver skips re-resolution when - the config is already present so the embed call keeps succeeding.""" - legacy_config = { - "api_key": "legacy-cleartext-key", - "api_base": "https://legacy-azure.example", - "api_version": "2024-01-01", - } - mock_vector_store: LiteLLM_ManagedVectorStore = { - "vector_store_id": "legacy_store", - "custom_llm_provider": "azure_ai", - "litellm_params": { - "litellm_embedding_model": "azure/text-embedding-3-large", - "litellm_embedding_config": legacy_config, - }, - } - - mock_registry = MagicMock() - mock_registry.get_litellm_managed_vector_store_from_registry.return_value = ( - mock_vector_store - ) - - resolve_mock = AsyncMock() - - with ( - patch.object(litellm, "vector_store_registry", mock_registry), - patch( - "litellm.proxy.vector_store_endpoints.endpoints._resolve_embedding_config", - new=resolve_mock, - ), - ): - result = await _update_request_data_with_litellm_managed_vector_store_registry( - data={}, vector_store_id="legacy_store" - ) - - assert result["litellm_embedding_config"] == legacy_config - resolve_mock.assert_not_awaited() - class TestCheckVectorStorePermission: """Test suite for check_vector_store_permission function.""" @@ -2003,57 +2070,7 @@ async def test_vector_store_update_and_list_synchronization(): @pytest.mark.asyncio -async def test_resolve_embedding_config_from_db(): - """Test that _resolve_embedding_config_from_db correctly resolves embedding config from database.""" - mock_prisma_client = MagicMock() - - # Mock database model with litellm_params - mock_db_model = MagicMock() - mock_db_model.litellm_params = { - "api_key": "test-api-key", - "api_base": "https://api.openai.com", - "api_version": "2024-01-01", - } - - mock_prisma_client.db.litellm_proxymodeltable.find_first = AsyncMock( - return_value=mock_db_model - ) - - with patch( - "litellm.proxy.vector_store_endpoints.management_endpoints.decrypt_value_helper", - side_effect=lambda value, key, return_original_value: value, - ): - result = await _resolve_embedding_config_from_db( - embedding_model="text-embedding-ada-002", prisma_client=mock_prisma_client - ) - - assert result is not None - assert result["api_key"] == "test-api-key" - assert result["api_base"] == "https://api.openai.com" - assert result["api_version"] == "2024-01-01" - mock_prisma_client.db.litellm_proxymodeltable.find_first.assert_called_once_with( - where={"model_name": "text-embedding-ada-002"} - ) - - # Test with empty embedding_model - result_empty = await _resolve_embedding_config_from_db( - embedding_model="", prisma_client=mock_prisma_client - ) - assert result_empty is None - - # Test with model not found - mock_prisma_client.db.litellm_proxymodeltable.find_first = AsyncMock( - return_value=None - ) - result_not_found = await _resolve_embedding_config_from_db( - embedding_model="non-existent-model", prisma_client=mock_prisma_client - ) - assert result_not_found is None - - -@pytest.mark.asyncio -async def test_new_vector_store_auto_resolves_embedding_config(): - """Test that new_vector_store auto-resolves embedding config when embedding_model is provided but config is not.""" +async def test_new_vector_store_persists_embedding_reference_without_credentials(): import json from litellm.types.vector_stores import LiteLLM_ManagedVectorStore @@ -2070,14 +2087,6 @@ async def test_new_vector_store_auto_resolves_embedding_config(): }, } - # Mock database model lookup for embedding config resolution - mock_db_model = MagicMock() - mock_db_model.litellm_params = { - "api_key": "resolved-api-key", - "api_base": "https://api.openai.com", - "api_version": "2024-01-01", - } - # Mock user API key mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.user_role = None @@ -2088,10 +2097,6 @@ async def test_new_vector_store_auto_resolves_embedding_config(): mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock( return_value=None # Vector store doesn't exist yet ) - mock_prisma_client.db.litellm_proxymodeltable.find_first = AsyncMock( - return_value=mock_db_model - ) - # Track what was passed to create captured_create_data = {} @@ -2112,261 +2117,21 @@ async def test_new_vector_store_auto_resolves_embedding_config(): mock_registry = MagicMock() mock_registry.add_vector_store_to_registry = MagicMock() - # Mock router to return None (so it falls back to DB resolution) - mock_router = MagicMock() - mock_router.get_deployment_by_model_group_name.return_value = None - with ( patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), - patch("litellm.proxy.proxy_server.llm_router", mock_router), - patch( - "litellm.proxy.vector_store_endpoints.management_endpoints.decrypt_value_helper", - side_effect=lambda value, key, return_original_value: value, - ), patch.object(litellm, "vector_store_registry", mock_registry), ): - result = await new_vector_store( - vector_store=vector_store_data, user_api_key_dict=mock_user_api_key - ) + result = await new_vector_store(vector_store=vector_store_data, user_api_key_dict=mock_user_api_key) assert result["status"] == "success" - # Auto-resolve no longer happens at create time — the persisted row - # carries only the model reference, never the resolved cleartext - # credential. Resolution now happens at request-handling time inside - # ``_update_request_data_with_litellm_managed_vector_store_registry``, - # where the resolved config lives in per-request memory and is never - # written to the database. litellm_params_json = captured_create_data.get("litellm_params") assert litellm_params_json is not None litellm_params_dict = json.loads(litellm_params_json) assert "litellm_embedding_config" not in litellm_params_dict assert litellm_params_dict["litellm_embedding_model"] == "text-embedding-ada-002" - # The response must also not echo a cleartext credential — even on - # the create response, where redaction guards against caller-supplied - # cleartext or pre-existing rows that were created by an earlier - # proxy version. response_vs = result["vector_store"] - assert "resolved-api-key" not in _serialize_litellm_params( - response_vs.get("litellm_params") - ) - - -def test_resolve_embedding_config_from_router(): - """Test that _resolve_embedding_config_from_router correctly extracts credentials from config-defined models.""" - from litellm.types.router import Deployment, LiteLLM_Params - - # Create a mock router with a model - mock_router = MagicMock() - - # Create a mock deployment with litellm_params - mock_litellm_params = MagicMock(spec=LiteLLM_Params) - mock_litellm_params.api_key = "config-api-key" - mock_litellm_params.api_base = "https://config-api-base.com" - mock_litellm_params.api_version = "2024-02-01" - - mock_deployment = MagicMock(spec=Deployment) - mock_deployment.litellm_params = mock_litellm_params - - mock_router.get_deployment_by_model_group_name.return_value = mock_deployment - - # Test resolution - result = _resolve_embedding_config_from_router( - embedding_model="text-embedding-ada-002", llm_router=mock_router - ) - - assert result is not None - assert result["api_key"] == "config-api-key" - assert result["api_base"] == "https://config-api-base.com" - assert result["api_version"] == "2024-02-01" - - mock_router.get_deployment_by_model_group_name.assert_called_once_with( - model_group_name="text-embedding-ada-002" - ) - - -def test_resolve_embedding_config_from_router_with_provider_prefix(): - """Test that _resolve_embedding_config_from_router handles provider prefixes like 'azure/model-name'.""" - from litellm.types.router import Deployment, LiteLLM_Params - - # Create a mock router - mock_router = MagicMock() - - # Create a mock deployment - mock_litellm_params = MagicMock(spec=LiteLLM_Params) - mock_litellm_params.api_key = "azure-api-key" - mock_litellm_params.api_base = "https://azure-endpoint.openai.azure.com" - mock_litellm_params.api_version = "2024-02-15" - - mock_deployment = MagicMock(spec=Deployment) - mock_deployment.litellm_params = mock_litellm_params - - # First call with full name returns None, second call with stripped name returns deployment - mock_router.get_deployment_by_model_group_name.side_effect = [None, mock_deployment] - - result = _resolve_embedding_config_from_router( - embedding_model="azure/text-embedding-3-large", llm_router=mock_router - ) - - assert result is not None - assert result["api_key"] == "azure-api-key" - assert result["api_base"] == "https://azure-endpoint.openai.azure.com" - assert result["api_version"] == "2024-02-15" - - # Should have tried both the full name and stripped name - assert mock_router.get_deployment_by_model_group_name.call_count == 2 - - -def test_resolve_embedding_config_from_router_returns_none_when_not_found(): - """Test that _resolve_embedding_config_from_router returns None when model is not in router.""" - mock_router = MagicMock() - mock_router.get_deployment_by_model_group_name.return_value = None - - result = _resolve_embedding_config_from_router( - embedding_model="nonexistent-model", llm_router=mock_router - ) - - assert result is None - - -def test_resolve_embedding_config_from_router_handles_os_environ(): - """Test that _resolve_embedding_config_from_router handles os.environ/ prefixed values.""" - from litellm.types.router import Deployment, LiteLLM_Params - - mock_router = MagicMock() - - mock_litellm_params = MagicMock(spec=LiteLLM_Params) - mock_litellm_params.api_key = "os.environ/OPENAI_API_KEY" - mock_litellm_params.api_base = "https://direct-url.com" - mock_litellm_params.api_version = None - - mock_deployment = MagicMock(spec=Deployment) - mock_deployment.litellm_params = mock_litellm_params - - mock_router.get_deployment_by_model_group_name.return_value = mock_deployment - - with patch( - "litellm.proxy.vector_store_endpoints.management_endpoints.get_secret", - return_value="resolved-from-env", - ) as mock_get_secret: - result = _resolve_embedding_config_from_router( - embedding_model="text-embedding-ada-002", llm_router=mock_router - ) - - assert result is not None - assert result["api_key"] == "resolved-from-env" - assert result["api_base"] == "https://direct-url.com" - assert "api_version" not in result - - mock_get_secret.assert_called_once_with("os.environ/OPENAI_API_KEY") - - -@pytest.mark.asyncio -async def test_resolve_embedding_config_tries_router_then_db(): - """Test that _resolve_embedding_config tries router first, then falls back to DB.""" - from litellm.types.router import Deployment, LiteLLM_Params - - mock_prisma_client = MagicMock() - mock_router = MagicMock() - - # Router has the model - mock_litellm_params = MagicMock(spec=LiteLLM_Params) - mock_litellm_params.api_key = "router-api-key" - mock_litellm_params.api_base = "https://router-api-base.com" - mock_litellm_params.api_version = None - - mock_deployment = MagicMock(spec=Deployment) - mock_deployment.litellm_params = mock_litellm_params - - mock_router.get_deployment_by_model_group_name.return_value = mock_deployment - - # DB should NOT be called since router has the model - mock_prisma_client.db.litellm_proxymodeltable.find_first = AsyncMock() - - result = await _resolve_embedding_config( - embedding_model="text-embedding-ada-002", - prisma_client=mock_prisma_client, - llm_router=mock_router, - ) - - assert result is not None - assert result["api_key"] == "router-api-key" - - # DB should NOT have been called since router found the model - mock_prisma_client.db.litellm_proxymodeltable.find_first.assert_not_called() - - -@pytest.mark.asyncio -async def test_resolve_embedding_config_caches_result(): - """The first lookup should hit the router/DB; subsequent lookups for - the same model name should return the cached value without touching - the router or the database.""" - from litellm.types.router import Deployment, LiteLLM_Params - - mock_prisma_client = MagicMock() - mock_router = MagicMock() - - mock_litellm_params = MagicMock(spec=LiteLLM_Params) - mock_litellm_params.api_key = "router-api-key" - mock_litellm_params.api_base = "https://router-api-base.com" - mock_litellm_params.api_version = None - - mock_deployment = MagicMock(spec=Deployment) - mock_deployment.litellm_params = mock_litellm_params - mock_router.get_deployment_by_model_group_name.return_value = mock_deployment - - first = await _resolve_embedding_config( - embedding_model="cached-model", - prisma_client=mock_prisma_client, - llm_router=mock_router, - ) - assert first is not None - assert mock_router.get_deployment_by_model_group_name.call_count == 1 - - second = await _resolve_embedding_config( - embedding_model="cached-model", - prisma_client=mock_prisma_client, - llm_router=mock_router, - ) - assert second == first - # Router (and by extension the DB) was not consulted again. - assert mock_router.get_deployment_by_model_group_name.call_count == 1 - - -@pytest.mark.asyncio -async def test_resolve_embedding_config_falls_back_to_db(): - """Test that _resolve_embedding_config falls back to DB when router doesn't have the model.""" - mock_prisma_client = MagicMock() - mock_router = MagicMock() - - # Router doesn't have the model - mock_router.get_deployment_by_model_group_name.return_value = None - - # DB has the model - mock_db_model = MagicMock() - mock_db_model.litellm_params = { - "api_key": "db-api-key", - "api_base": "https://db-api-base.com", - } - mock_prisma_client.db.litellm_proxymodeltable.find_first = AsyncMock( - return_value=mock_db_model - ) - - with patch( - "litellm.proxy.vector_store_endpoints.management_endpoints.decrypt_value_helper", - side_effect=lambda value, key, return_original_value: value, - ): - result = await _resolve_embedding_config( - embedding_model="text-embedding-ada-002", - prisma_client=mock_prisma_client, - llm_router=mock_router, - ) - - assert result is not None - assert result["api_key"] == "db-api-key" - - # DB should have been called since router didn't find the model - mock_prisma_client.db.litellm_proxymodeltable.find_first.assert_called() + assert "api_key" not in _serialize_litellm_params(response_vs.get("litellm_params")) @pytest.mark.asyncio @@ -2425,9 +2190,7 @@ async def test_new_vector_store_auto_resolves_from_router(): } return mock_created_vector_store - mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock( - side_effect=mock_create - ) + mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock(side_effect=mock_create) mock_registry = MagicMock() mock_registry.add_vector_store_to_registry = MagicMock() diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index b2b8eb5da80..2068f10ea2d 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1928,6 +1928,19 @@ class TestToolTransformation: assert result_tool["function"]["parameters"]["type"] == "object" assert "properties" in result_tool["function"]["parameters"] + def test_transform_function_tools_parameters_keep_client_key_order(self): + tools = [ + {"type": "function", "name": "a", "parameters": {"properties": {"arg": {"type": "string"}}, "required": ["arg"]}}, + {"type": "function", "name": "b", "parameters": {"type": "object", "properties": {}}}, + ] + + result_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=tools + ) + + assert list(result_tools[0]["function"]["parameters"]) == ["properties", "required", "type"] + assert list(result_tools[1]["function"]["parameters"]) == ["type", "properties"] + def test_transform_function_tools_empty_parameters(self): """Test that empty parameters get 'type': 'object' added""" function_tool = { diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index c9a5b988be6..1233ddf1785 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -3,7 +3,7 @@ from __future__ import annotations import pytest from litellm.llms.custom_httpx.llm_http_handler import _rust_responses_websocket_enabled -from litellm.rust_bridge import responses_websocket +from litellm.rust_bridge import configuration, responses_websocket from litellm.types.router import GenericLiteLLMParams @@ -39,12 +39,33 @@ class _FakeNativeBridge: return _FakeNativeConnection() +@pytest.fixture(autouse=True) +def reset_responses_websocket(): + responses_websocket.set_rust_responses_websocket(connection=None) + configuration.reset_rust_configuration() + yield + responses_websocket.set_rust_responses_websocket(connection=None) + configuration.reset_rust_configuration() + + def test_rust_websocket_bridge_is_disabled_without_flag() -> None: assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams()) assert not _rust_responses_websocket_enabled("anthropic", GenericLiteLLMParams(rust=True)) assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=True)) +def test_explicit_false_overrides_process_enable() -> None: + configuration.use_litellm_rust(True) + + assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=False)) + + +def test_process_enable_applies_without_request_override() -> None: + configuration.use_litellm_rust(True) + + assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams()) + + @pytest.mark.asyncio async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None: adapter = responses_websocket._ConnectionAdapter(_ClosedNativeConnection()) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 1ec8be88c9b..941d78085e4 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -12,7 +12,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from pydantic import ValidationError - import litellm from litellm import Router from litellm._logging import verbose_router_logger @@ -34,10 +33,14 @@ from litellm.router_strategy.complexity_router.config import ( DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, DEFAULT_COMPLEXITY_CONFIG, DEFAULT_TECHNICAL_KEYWORDS, + ClassificationRubric, ClassifierLLMConfig, ComplexityRouterConfig, ComplexityTier, - ClassificationRubric, +) +from litellm.router_strategy.complexity_router.tier_predictor import ( + TierGlobalStatistic, + TrainedTierArtifact, ) from litellm.types.router import ( Deployment, @@ -46,6 +49,16 @@ from litellm.types.router import ( ) +def _heuristic_v2_artifact() -> TrainedTierArtifact: + return TrainedTierArtifact( + global_statistics=tuple( + TierGlobalStatistic(tier=tier, successes=successes, observations=100) + for tier, successes in enumerate((10, 20, 90, 99), start=1) + ), + routing_threshold=0.8, + ) + + @pytest.fixture def mock_router_instance(): """Create a mock LiteLLM Router instance.""" @@ -1696,6 +1709,59 @@ class TestLLMClassifier: assert outcome.cause == "heuristic_scorer" assert outcome.score is not None + @pytest.mark.asyncio + async def test_heuristic_v2_routes_directly_to_predicted_builtin_tier(self, mock_router_instance): + router = ComplexityRouter( + model_name="tier-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "classifier_type": "heuristic_v2", + "heuristic_v2_artifact": _heuristic_v2_artifact(), + "tiers": { + "SIMPLE": "simple-model", + "MEDIUM": "medium-model", + "COMPLEX": "complex-model", + "REASONING": "reasoning-model", + }, + }, + ) + + response = await router.async_pre_routing_hook( + model="tier-router", + request_kwargs={}, + messages=[{"role": "user", "content": "Handle this new request"}], + ) + + assert response is not None + assert response.model == "complex-model" + assert response.routing_decision["tier"] == "COMPLEX" + assert response.routing_decision["cause"] == "heuristic_v2" + assert response.routing_decision["signals"] == [ + "request-type:general", + "tier-probability:simple=0.107843", + "tier-probability:medium=0.205882", + "tier-probability:complex=0.892157", + "tier-probability:reasoning=0.980392", + ] + + def test_heuristic_v2_needs_no_classifier_model(self): + config = ComplexityRouterConfig(classifier_type="heuristic_v2") + + assert config.classifier_llm_config is None + assert config.heuristic_v2_artifact == "ultrafeedback" + + def test_heuristic_v2_rejects_custom_tier_definitions(self): + with pytest.raises(ValidationError, match="as does heuristic_v2"): + ComplexityRouterConfig( + classifier_type="heuristic_v2", + tier_definitions=( + {"name": "low", "description": "easy work"}, + {"name": "high", "description": "hard work"}, + ), + tiers={"low": "cheap", "high": "expensive"}, + fallback_tier="high", + ) + @pytest.mark.asyncio async def test_aclassify_llm_success_routes_by_llm_verdict(self, llm_complexity_router, mock_router_instance): """A well-formed structured LLM response should decide the tier directly. @@ -10023,6 +10089,207 @@ class TestHeuristicFirst: assert outcome.cause == "default_model_fallback" +# Scores 0.175 with one signal, so it sits 0.025 from simple_medium: the pair of tiers either side of +# that boundary are different model pools, and a hair's difference in score picks the other one. +NEAR_BOUNDARY_PROMPT = ( + "design a distributed cache with consistent hashing, then explain the failure modes step by step" +) + +# Scores 0.075 with signals, the far side of any margin under 0.075: the scorer is decided here. +CLEAR_OF_BOUNDARY_PROMPT = "explain step by step how consistent hashing rebalances keys" + + +def _hybrid_router(mock_router_instance, **config_overrides): + config = { + "tiers": dict(HEURISTIC_FIRST_TIERS), + "tier_boundaries": dict(HEURISTIC_FIRST_BOUNDARIES), + "classifier_type": "hybrid", + "hybrid_boundary_margin": 0.03, + "classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 400}, + **config_overrides, + } + return ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=config, + ) + + +class TestHybridConfig: + """Config validation for classifier_type='hybrid'.""" + + @pytest.mark.parametrize( + "overrides, expected", + [ + ({"classifier_llm_config": None}, "classifier_llm_config is required"), + ({"hybrid_boundary_margin": None}, "hybrid_boundary_margin is required"), + ({"hybrid_boundary_margin": -0.01}, "greater than or equal to 0"), + ({"hybrid_boundary_margin": 1.01}, "less than or equal to 1"), + ], + ) + def test_rejects_incoherent_config(self, overrides, expected): + config = { + "tiers": dict(HEURISTIC_FIRST_TIERS), + "classifier_type": "hybrid", + "hybrid_boundary_margin": 0.03, + "classifier_llm_config": {"model": "haiku-classifier"}, + **overrides, + } + with pytest.raises(ValidationError, match=expected): + ComplexityRouterConfig(**config) + + @pytest.mark.parametrize("classifier_type", ["heuristic", "llm", "custom", "heuristic_first"]) + def test_margin_rejected_on_every_other_classifier_type(self, classifier_type): + """A margin on a router that never compares a score to a boundary is a silent no-op, so it is + refused rather than accepted and ignored. heuristic_first is in this list on purpose: its + ceiling is a different question from proximity, and accepting both on one router would make + two modes out of one classifier_type.""" + config: dict[str, object] = { + "tiers": dict(HEURISTIC_FIRST_TIERS), + "classifier_type": classifier_type, + "hybrid_boundary_margin": 0.03, + } + if classifier_type in ("llm", "heuristic_first"): + config["classifier_llm_config"] = {"model": "haiku-classifier"} + if classifier_type == "heuristic_first": + config["heuristic_first_max_tier"] = "SIMPLE" + if classifier_type == "custom": + config["classifier_plugin"] = _FixedTierClassifier("SIMPLE") + with pytest.raises(ValidationError, match="hybrid_boundary_margin is set but classifier_type"): + ComplexityRouterConfig(**config) + + def test_the_cheap_tier_ceiling_is_rejected_here(self): + """The two modes are told apart by which knob they take, so the ceiling is refused on hybrid + exactly as the margin is refused on heuristic_first.""" + with pytest.raises(ValidationError, match="heuristic_first_max_tier is set but classifier_type"): + ComplexityRouterConfig( + tiers=dict(HEURISTIC_FIRST_TIERS), + classifier_type="hybrid", + hybrid_boundary_margin=0.03, + heuristic_first_max_tier="SIMPLE", + classifier_llm_config={"model": "haiku-classifier"}, + ) + + def test_custom_tier_set_is_rejected(self): + """The scorer only emits the four built-in tiers, so it cannot judge proximity on a replaced set.""" + with pytest.raises(ValidationError, match="tier_definitions requires classifier_type"): + ComplexityRouterConfig( + classifier_type="hybrid", + hybrid_boundary_margin=0.03, + classifier_llm_config={"model": "haiku-classifier"}, + tier_definitions=[{"name": "lo", "description": "x"}, {"name": "hi", "description": "y"}], + tiers={"lo": "gpt-4o-mini", "hi": "gpt-4o"}, + ) + + def test_classifier_model_is_a_dependency(self): + config = ComplexityRouterConfig( + tiers=dict(HEURISTIC_FIRST_TIERS), + classifier_type="hybrid", + hybrid_boundary_margin=0.03, + classifier_llm_config={"model": "haiku-classifier"}, + ) + assert config.uses_llm_classifier is True + + +class TestHybrid: + """Behavior of the hybrid chain: the scorer keeps its tier unless the score is near a boundary.""" + + @pytest.mark.asyncio + async def test_near_boundary_prompt_escalates(self, mock_router_instance): + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}')) + router = _hybrid_router(mock_router_instance) + + _tier, score, signals, _cause = router._score_and_classify(NEAR_BOUNDARY_PROMPT) + assert signals and abs(score - HEURISTIC_FIRST_BOUNDARIES["simple_medium"]) < 0.03 + + outcome = await router.aclassify(NEAR_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_awaited_once() + assert outcome.tier == ComplexityTier.COMPLEX + assert outcome.cause == "llm_classifier" + + @pytest.mark.asyncio + async def test_score_clear_of_every_boundary_keeps_the_heuristic_tier(self, mock_router_instance): + mock_router_instance.acompletion = AsyncMock() + router = _hybrid_router(mock_router_instance) + outcome = await router.aclassify(CLEAR_OF_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_not_called() + assert outcome.tier == ComplexityTier.SIMPLE + assert outcome.cause == "hybrid_short_circuit" + + @pytest.mark.asyncio + async def test_an_expensive_tier_short_circuits_too(self, mock_router_instance): + """This is the whole difference from heuristic_first, which would have escalated this by tier + alone. Hybrid asks whether the score is DECIDED, not whether the tier is cheap.""" + mock_router_instance.acompletion = AsyncMock() + router = _hybrid_router( + mock_router_instance, + tier_boundaries={"simple_medium": -0.9, "medium_complex": -0.8, "complex_reasoning": -0.7}, + ) + + tier, _score, signals, _cause = router._score_and_classify(CLEAR_OF_BOUNDARY_PROMPT) + assert (tier, bool(signals)) == (ComplexityTier.REASONING, True) + + outcome = await router.aclassify(CLEAR_OF_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_not_called() + assert outcome.tier == ComplexityTier.REASONING + assert outcome.cause == "hybrid_short_circuit" + + @pytest.mark.asyncio + async def test_widening_the_margin_escalates_what_a_narrow_one_kept(self, mock_router_instance): + """The margin is the knob: the same prompt short-circuits at 0.03 and escalates at 0.08.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "MEDIUM"}')) + router = _hybrid_router(mock_router_instance, hybrid_boundary_margin=0.08) + outcome = await router.aclassify(CLEAR_OF_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_awaited_once() + assert outcome.cause == "llm_classifier" + + @pytest.mark.asyncio + async def test_a_zero_margin_escalates_only_an_exact_boundary_score(self, mock_router_instance): + """0 is a real margin, not an off switch: a score sitting exactly on the line still escalates. + + The boundary is spelled as the scorer's own accumulated float rather than the 0.075 it prints + as, because the comparison is on raw floats: a boundary written 0.075 sits 1.4e-17 away from + this score and a zero margin correctly declines to call that exact.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "MEDIUM"}')) + on_the_line = 0.07499999999999998 + router = _hybrid_router( + mock_router_instance, + tier_boundaries={"simple_medium": on_the_line, "medium_complex": 0.35, "complex_reasoning": 0.60}, + hybrid_boundary_margin=0, + ) + + _tier, score, _signals, _cause = router._score_and_classify(CLEAR_OF_BOUNDARY_PROMPT) + assert score == on_the_line + + outcome = await router.aclassify(CLEAR_OF_BOUNDARY_PROMPT) + mock_router_instance.acompletion.assert_awaited_once() + assert outcome.cause == "llm_classifier" + + @pytest.mark.asyncio + async def test_no_signal_prompt_escalates_however_far_from_a_boundary(self, mock_router_instance): + """The scorer with no opinion has no tier to be confident about, so proximity cannot save it.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}')) + router = _hybrid_router(mock_router_instance) + + tier, score, signals, _cause = router._score_and_classify(NO_SIGNAL_PROMPT) + assert (tier, score, signals) == (ComplexityTier.SIMPLE, 0.0, ()) + + outcome = await router.aclassify(NO_SIGNAL_PROMPT) + mock_router_instance.acompletion.assert_awaited_once() + assert outcome.cause == "llm_classifier" + + @pytest.mark.asyncio + async def test_classifier_failure_falls_back_to_the_scorer(self, mock_router_instance): + mock_router_instance.acompletion = AsyncMock(side_effect=RuntimeError("classifier exploded")) + router = _hybrid_router(mock_router_instance) + expected_tier, expected_score, expected_signals, _cause = router._score_and_classify(NEAR_BOUNDARY_PROMPT) + + outcome = await router.aclassify(NEAR_BOUNDARY_PROMPT) + + assert (outcome.tier, outcome.score, outcome.signals) == (expected_tier, expected_score, expected_signals) + assert outcome.cause == "heuristic_scorer" + + def _windowed_router(*deployments: tuple) -> Router: """Real Router; each deployment is (group, provider_model, declared window or None). None means no declared override on a model the cost map does not know: unresolvable.""" @@ -10268,7 +10535,8 @@ class TestContextWindowEscalation: litellm_router_instance=_windowed_router(_SMALL, _BIG), complexity_router_config=_tier_config(session_affinity=True), ) - session_kwargs = lambda: {"metadata": {"session_id": "s-1", "user_api_key_hash": "k-1"}} # noqa: E731 + def session_kwargs() -> dict[str, object]: + return {"metadata": {"session_id": "s-1", "user_api_key_hash": "k-1"}} first = await router.async_pre_routing_hook( model="test-router", request_kwargs=session_kwargs(), messages=_OVERSIZED_TURNS @@ -10291,7 +10559,8 @@ class TestContextWindowEscalation: litellm_router_instance=_windowed_router(_SMALL, _BIG), complexity_router_config=_tier_config(session_affinity=True), ) - session_kwargs = lambda: {"metadata": {"session_id": "s-2", "user_api_key_hash": "k-2"}} # noqa: E731 + def session_kwargs() -> dict[str, object]: + return {"metadata": {"session_id": "s-2", "user_api_key_hash": "k-2"}} pinned = await router.async_pre_routing_hook( model="test-router", request_kwargs=session_kwargs(), messages=[{"role": "user", "content": "ok continue"}] diff --git a/tests/test_litellm/router_strategy/test_complexity_tier_predictor.py b/tests/test_litellm/router_strategy/test_complexity_tier_predictor.py new file mode 100644 index 00000000000..5bbe0fb5669 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_complexity_tier_predictor.py @@ -0,0 +1,91 @@ +from typing import Final + +import pytest + +from litellm.router_strategy.complexity_router.tier_predictor import ( + TierCohortStatistic, + TierDomainStatistic, + TierGlobalStatistic, + TierSuccessPredictor, + TrainedTierArtifact, + resolve_tier_artifact, + similarity_cohort, +) +from litellm.types.router import RequestType + + +def _artifact( + global_successes: tuple[float, float, float, float] = (4.0, 5.0, 6.0, 7.0), + threshold: float = 0.75, + domain_statistics: tuple[TierDomainStatistic, ...] = (), + cohort_statistics: tuple[TierCohortStatistic, ...] = (), +) -> TrainedTierArtifact: + return TrainedTierArtifact( + global_statistics=tuple( + TierGlobalStatistic(tier=tier, successes=successes, observations=10.0) + for tier, successes in enumerate(global_successes, start=1) + ), + domain_statistics=domain_statistics, + cohort_statistics=cohort_statistics, + domain_prior_mass=10.0, + cohort_prior_mass=10.0, + routing_threshold=threshold, + ) + + +def test_predictions_are_monotonic_across_tiers() -> None: + predictor: Final = TierSuccessPredictor(_artifact(global_successes=(9.0, 2.0, 7.0, 6.0))) + + prediction: Final = predictor.predict("hello", RequestType.GENERAL) + + probabilities: Final = tuple(prediction.probabilities.values()) + assert probabilities == tuple(sorted(probabilities)) + + +def test_domain_and_cohort_statistics_back_off_hierarchically() -> None: + matching_cohort: Final = similarity_cohort("hello", RequestType.GENERAL) + artifact: Final = _artifact( + global_successes=(1.0, 5.0, 6.0, 7.0), + domain_statistics=( + TierDomainStatistic( + tier=1, + request_type=RequestType.GENERAL, + successes=10.0, + observations=10.0, + ), + ), + cohort_statistics=( + TierCohortStatistic( + tier=1, + cohort=matching_cohort, + successes=0.0, + observations=10.0, + ), + ), + ) + predictor: Final = TierSuccessPredictor(artifact) + + cohort_probability: Final = predictor.predict("hello", RequestType.GENERAL).probabilities[1] + domain_probability: Final = predictor.predict("hello " * 100, RequestType.GENERAL).probabilities[1] + global_probability: Final = predictor.predict("hello", RequestType.WRITING).probabilities[1] + + assert cohort_probability == pytest.approx(7.0 / 24.0) + assert domain_probability == pytest.approx(7.0 / 12.0) + assert global_probability == pytest.approx(1.0 / 6.0) + + +def test_selects_first_tier_above_probability_threshold() -> None: + predictor: Final = TierSuccessPredictor(_artifact(global_successes=(4.0, 6.0, 8.0, 9.0), threshold=0.7)) + + prediction: Final = predictor.predict("hello", RequestType.GENERAL) + + assert prediction.required_tier == 3 + + +def test_builtin_ultrafeedback_artifact_is_loadable() -> None: + artifact: Final = resolve_tier_artifact("ultrafeedback") + + assert artifact.routing_threshold == 0.75 + assert artifact.domain_prior_mass == 200.0 + assert artifact.cohort_prior_mass == 20.0 + assert artifact.datasets[0].license == "MIT" diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 894b2d9e74f..a4965c49f07 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -614,6 +614,27 @@ async def test_run_async_fallback_forwards_attempted_model_groups_to_nested_call ) +@pytest.mark.asyncio +async def test_run_async_fallback_can_target_the_requested_group_when_a_pre_router_replaced_it(): + """The requested group was never called when a pre-router selected a tier, so a + tier fallback may legitimately target that originally requested group.""" + router = RecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["requested-model"], + original_model_group="requested-model", + original_exception=RuntimeError("selected tier failed"), + max_fallbacks=3, + fallback_depth=0, + model="requested-model", + metadata={"pre_routing_selected_model": "selected-tier"}, + ) + + assert router.received_kwargs["model"] == "requested-model" + assert router.received_kwargs["attempted_targets"].keys == frozenset({"selected-tier", "requested-model"}) + + @pytest.mark.asyncio @pytest.mark.parametrize( "entry", @@ -1199,6 +1220,28 @@ class TestOrderedFallbackLookupGroups: assert fallback_lookup_groups({}, "smart-router") == ("smart-router",) assert fallback_lookup_groups({}, None) == () + def test_session_remap_keeps_the_bound_router_between_tier_and_requested_group(self): + from litellm.router_utils.fallback_event_handlers import ( + PRE_ROUTING_SELECTED_MODEL_KEY, + fallback_lookup_groups, + ) + + kwargs = { + "litellm_metadata": { + PRE_ROUTING_SELECTED_MODEL_KEY: "tier1", + "model_group": "smart-router", + } + } + + assert fallback_lookup_groups(kwargs, "requested-model") == ( + "tier1", + "smart-router", + "requested-model", + ) + assert fallback_lookup_groups({"metadata": {"model_group": []}}, "requested-model") == ( + "requested-model", + ) + def test_first_resolving_group_wins_and_generic_idx_survives_a_miss(self): from litellm.router_utils.fallback_event_handlers import ( get_fallback_model_group_for_lookup_groups, diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py new file mode 100644 index 00000000000..0ce49d51aed --- /dev/null +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -0,0 +1,349 @@ +from __future__ import annotations + +import asyncio +import importlib.util +import json +import os +import signal +import subprocess +import sys +import tempfile +import threading +import zipfile +from http.client import HTTPMessage +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from socket import socket as Socket +from typing import Final + +REQUEST_STARTED: Final = threading.Event() +REQUEST_CANCELLED: Final = threading.Event() + +ANTHROPIC_RESPONSE: Final = ( + b'{"id":"msg_native","type":"message","role":"assistant",' + b'"model":"claude-sonnet-4-5","content":[{"type":"text","text":"native-message"}],' + b'"stop_reason":"end_turn","stop_sequence":null,' + b'"usage":{"input_tokens":2,"output_tokens":3}}' +) + + +class NativeRouteHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + content_length: Final = int(self.headers.get("content-length", "0")) + body: Final = json.loads(self.rfile.read(content_length)) + route: Final = self.headers.get("x-test-route") + outcome: Final = self.headers.get("x-test-outcome") + assert_native_request(route, outcome, self.path, self.headers, body) + if outcome == "hang": + REQUEST_STARTED.set() + self.connection.settimeout(5) + if connection_was_cancelled(self.connection): + REQUEST_CANCELLED.set() + return + + status: Final = 429 if outcome == "429" else 200 + response_body: Final = native_response(status, route) + + self.send_response(status) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(response_body))) + self.send_header("connection", "close") + self.end_headers() + self.wfile.write(response_body) + + def log_message(self, _message_format: str, *_args: object) -> None: + pass + + +def connection_was_cancelled(connection: Socket) -> bool: + try: + return connection.recv(1) == b"" + except TimeoutError: + return False + except OSError: + return True + + +def assert_native_request( + route: str | None, + outcome: str | None, + path: str, + headers: HTTPMessage, + body: object, +) -> None: + if route not in {"ocr", "transcription", "messages", "chat_completions"}: + raise AssertionError(f"unexpected route marker: {route!r}") + if outcome not in {"success", "429", "hang"}: + raise AssertionError(f"unexpected outcome marker: {outcome!r}") + if not isinstance(body, dict): + raise TypeError(f"{route} sent {type(body).__name__}, expected a JSON object") + if route == "ocr": + assert path == "/v1/ocr" + assert headers.get("authorization") == "Bearer sk-native" + assert body["model"] == "mistral-ocr-latest" + assert body["document"]["document_url"] == "https://example.com/document.pdf" + assert body["include_image_base64"] is True + return + if route == "transcription": + assert path == "/model/mistral.voxtral-mini-3b-2507/converse" + assert headers.get("authorization", "").startswith("AWS4-HMAC-SHA256 ") + assert headers.get("x-amz-date") + assert body["messages"][0]["content"][0]["audio"]["source"]["bytes"] == "AQI=" + assert "The audio language is en" in body["messages"][0]["content"][1]["text"] + return + assert path == "/v1/messages" + assert headers.get("x-api-key") == "sk-native" + assert body["model"] == "claude-sonnet-4-5" + if route == "messages": + assert body["max_tokens"] == 16 + assert body["messages"][0]["content"] == "hello-from-messages" + return + assert body["max_tokens"] == 17 + assert body["messages"][0]["content"] == [{"type": "text", "text": "hello-from-chat"}] + + +def native_response(status: int, route: str | None) -> bytes: + if status == 429: + return b'{"error":"native-rate-limit"}' + if route == "ocr": + return b'{"pages":[{"index":0,"markdown":"native-ocr"}]}' + if route == "transcription": + return b'{"output":{"message":{"content":[{"text":"native-transcription"}]}}}' + return ANTHROPIC_RESPONSE + + +def load_native(native_path: Path) -> object: + module_spec: Final = importlib.util.spec_from_file_location("litellm.rust_bridge._native", native_path) + if module_spec is None or module_spec.loader is None: + raise RuntimeError("cannot create native extension import specification") + native_module: Final = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(native_module) + return native_module + + +def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]: + common: Final = { + "api_base": api_base, + "extra_headers": {"x-test-outcome": outcome, "x-test-route": route}, + "timeout_seconds": 3.0, + } + if route == "ocr": + return common | { + "model": "mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, + "api_key": "sk-native", + "custom_llm_provider": "mistral", + "optional_params": {"include_image_base64": True}, + } + if route == "transcription": + return common | { + "model": "mistral.voxtral-mini-3b-2507", + "audio": {"data": "AQI=", "format": "wav", "filename": "audio.wav"}, + "custom_llm_provider": "bedrock", + "optional_params": { + "aws_access_key_id": "native-access-key", + "aws_secret_access_key": "native-secret-key", + "aws_region_name": "us-east-1", + "language": "en", + }, + } + if route == "messages": + return common | { + "model": "claude-sonnet-4-5", + "body": { + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hello-from-messages"}], + }, + "api_key": "sk-native", + "custom_llm_provider": "anthropic", + } + if route == "chat_completions": + return common | { + "model": "anthropic/claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hello-from-chat"}], + "optional_params": {"max_tokens": 17}, + "api_key": "sk-native", + } + raise AssertionError(f"unknown route: {route}") + + +def assert_success(route: str, response: object) -> None: + if not isinstance(response, dict): + raise TypeError(f"{route} returned {type(response).__name__}, expected dict") + actual: Final = success_value(route, response) + expected: Final = ( + "native-ocr" if route == "ocr" else "native-transcription" if route == "transcription" else "native-message" + ) + if actual != expected: + raise AssertionError(f"{route} returned {actual!r}, expected {expected!r}") + + +def assert_traced_success(route: str, response: object) -> None: + if not isinstance(response, dict): + raise TypeError(f"{route} returned {type(response).__name__}, expected a traced dict") + assert_success(route, response["response"]) + expected_function: Final = "audio_transcription" if route == "transcription" else route + assert response["trace"][0] == {"function": expected_function, "depth": 0} + + +def success_value(route: str, response: dict[object, object]) -> object: + if route == "ocr": + return response["pages"][0]["markdown"] + if route == "transcription": + return response["text"] + if route == "messages": + return response["content"][0]["text"] + return response["choices"][0]["message"]["content"] + + +def assert_rate_limit(native: object, route: str, error: BaseException) -> None: + if route == "chat_completions": + upstream_error: Final = native.RustUpstreamError + if not isinstance(error, upstream_error) or error.args[0] != 429: + raise AssertionError(f"{route} returned the wrong 429 error: {error!r}") + return + if not isinstance(error, RuntimeError) or "429" not in str(error): + raise AssertionError(f"{route} returned the wrong 429 error: {error!r}") + + +def exercise_sync(native: object, api_base: str) -> None: + for route in ("ocr", "transcription", "messages", "chat_completions"): + function: Final = getattr(native, route) + assert_success(route, function(**route_kwargs(route, api_base, "success"))) + assert_traced_success(route, function(**route_kwargs(route, api_base, "success"), trace=True)) + try: + function(**route_kwargs(route, api_base, "429")) + except (RuntimeError, native.RustUpstreamError) as error: + assert_rate_limit(native, route, error) + else: + raise AssertionError(f"{route} accepted a 429 response") + + +async def exercise_async(native: object, api_base: str) -> None: + for route in ("ocr", "transcription", "messages", "chat_completions"): + function: Final = getattr(native, f"a{route}") + assert_success(route, await function(**route_kwargs(route, api_base, "success"))) + assert_traced_success(route, await function(**route_kwargs(route, api_base, "success"), trace=True)) + try: + await function(**route_kwargs(route, api_base, "429")) + except (RuntimeError, native.RustUpstreamError) as error: + assert_rate_limit(native, route, error) + else: + raise AssertionError(f"a{route} accepted a 429 response") + + +async def exercise_async_concurrency(native: object, api_base: str) -> None: + responses: Final = await asyncio.wait_for( + asyncio.gather( + *( + native.amessages(**route_kwargs("messages", api_base, "success")) + for _ in range(32) + ) + ), + timeout=15, + ) + for response in responses: + assert_success("messages", response) + + +def exercise_routes(native_path: Path, api_base: str) -> object: + native: Final = load_native(native_path) + exercise_sync(native, api_base) + asyncio.run(exercise_async(native, api_base)) + asyncio.run(exercise_async_concurrency(native, api_base)) + return native + + +def exercise_signal(native: object, api_base: str) -> int: + try: + native.messages( + **route_kwargs("messages", api_base, "hang"), + ) + except KeyboardInterrupt: + sys.stdout.write("KeyboardInterrupt\n") + sys.stdout.flush() + sys.stdin.read(1) + return 0 + raise AssertionError("sync native route ignored SIGINT") + + +def verify_sigint(native_path: Path, api_base: str) -> None: + REQUEST_STARTED.clear() + REQUEST_CANCELLED.clear() + process: Final = subprocess.Popen( + (sys.executable, __file__, "child", str(native_path), api_base), + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + if not REQUEST_STARTED.wait(30): + process.kill() + stdout, stderr = process.communicate(timeout=5) + raise AssertionError( + f"native route matrix did not reach the hanging upstream\nstdout:\n{stdout}\nstderr:\n{stderr}" + ) + os.kill(process.pid, signal.SIGINT) + if not REQUEST_CANCELLED.wait(5): + raise AssertionError("interrupted native route did not cancel its upstream future") + if process.poll() is not None: + raise AssertionError("signal child exited before cancellation was observed") + stdout, stderr = process.communicate(input="\n", timeout=5) + if process.returncode != 0 or stdout != "KeyboardInterrupt\n": + raise AssertionError( + f"signal child exited with status {process.returncode}\nstdout:\n{stdout}\nstderr:\n{stderr}" + ) + finally: + if process.poll() is None: + process.kill() + process.wait(timeout=5) + + +def verify_wheel(wheel: Path) -> int: + with tempfile.TemporaryDirectory() as temporary_directory, zipfile.ZipFile(wheel) as archive: + wheel_root: Final = Path(temporary_directory) + for member in archive.infolist(): + target: Final = wheel_root / member.filename + if member.is_dir(): + target.mkdir(parents=True, exist_ok=True) + else: + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(archive.read(member)) + native_members: Final = tuple( + member + for member in archive.infolist() + if member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so") + ) + if len(native_members) != 1: + raise AssertionError(f"expected one native extension, found {len(native_members)}") + native_path: Final = wheel_root / native_members[0].filename + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), NativeRouteHandler) + server_thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + server_thread.start() + api_base: Final = f"http://127.0.0.1:{server.server_address[1]}" + try: + verify_sigint(native_path, api_base) + finally: + server.shutdown() + server.server_close() + server_thread.join(timeout=5) + return 0 + + +def main() -> int: + if len(sys.argv) == 2: + return verify_wheel(Path(sys.argv[1])) + if len(sys.argv) == 4 and sys.argv[1] == "child": + native: Final = exercise_routes(Path(sys.argv[2]), sys.argv[3]) + return exercise_signal(native, sys.argv[3]) + sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n") + return 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_litellm/rust_bridge/test_bindings.py b/tests/test_litellm/rust_bridge/test_bindings.py new file mode 100644 index 00000000000..88036a5a556 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_bindings.py @@ -0,0 +1,35 @@ +from types import SimpleNamespace +from typing import Final + +import pytest + +from litellm.rust_bridge import bindings + + +def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None: + native = SimpleNamespace(route=lambda: "native") + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + binding: bindings.NativeBinding[object] = bindings.NativeBinding("route", validate=lambda value: value) + + assert binding.load() is native.route + + binding.override(None) + assert binding.load() is None + + replacement = object() + binding.override(replacement) + assert binding.load() is replacement + + binding.reset() + assert binding.load() is native.route + + +@pytest.mark.parametrize(("value", "expected"), ((3, 3), ("invalid", None), (None, None))) +def test_binding_validates_native_attribute( + monkeypatch: pytest.MonkeyPatch, value: object, expected: int | None +) -> None: + native: Final = SimpleNamespace(route=value) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + binding: Final = bindings.NativeBinding("route", validate=lambda item: item if isinstance(item, int) else None) + + assert binding.load() == expected diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/test_chat_completions.py index 47cb66932b7..03921133c77 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/test_chat_completions.py @@ -11,6 +11,7 @@ import pytest import litellm from litellm.rust_bridge import chat_completions as bridge +from litellm.rust_bridge import configuration from litellm.types.utils import ModelResponse RUST_RESPONSE = { @@ -68,13 +69,11 @@ def _hide_native_bridge(monkeypatch): @pytest.fixture(autouse=True) def reset_bridge(): """Every test starts with no injected callables, and leaves none behind.""" - bridge.set_rust_chat_completions( - chat_completions=None, achat_completions=None, decline=None - ) + bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) + configuration.reset_rust_configuration() yield - bridge.set_rust_chat_completions( - chat_completions=None, achat_completions=None, decline=None - ) + bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None) + configuration.reset_rust_configuration() class _RecordingDecline: @@ -138,6 +137,18 @@ class TestGate: assert gate.calls[0]["model"] == "claude-sonnet-4-5" assert gate.calls[0]["custom_llm_provider"] == "anthropic" + def test_explicit_false_overrides_process_enable(self): + bridge.set_rust_chat_completions(decline=_RecordingDecline()) + configuration.use_litellm_rust(True) + + assert _accepts(litellm_params={"rust": False}) is False + + def test_process_enable_applies_without_request_override(self): + bridge.set_rust_chat_completions(decline=_RecordingDecline()) + configuration.use_litellm_rust(True) + + assert _accepts(litellm_params={}) is True + def test_the_env_var_opts_in_without_a_per_model_flag(self, monkeypatch): monkeypatch.setenv("LITELLM_RUST", "true") bridge.set_rust_chat_completions(decline=_RecordingDecline()) @@ -253,9 +264,7 @@ class TestSyncCall: assert result.usage.completion_tokens == 4 assert result.usage.total_tokens == 15 assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} - assert result.id == original_id, ( - "the rust path must keep the chatcmpl id litellm already minted" - ) + assert result.id == original_id, "the rust path must keep the chatcmpl id litellm already minted" def test_passes_the_timeout_through_as_seconds(self): native = _RecordingCall() @@ -269,9 +278,7 @@ class TestSyncCall: def test_falls_back_when_the_core_declines_before_calling_the_provider(self, monkeypatch): _fake_native_bridge(monkeypatch) - bridge.set_rust_chat_completions( - chat_completions=_RecordingCall(error=_FakeDeclined("streaming")) - ) + bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeDeclined("streaming"))) assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None @@ -290,13 +297,9 @@ class TestAsyncCall: assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None @pytest.mark.asyncio - async def test_falls_back_when_the_core_declines_before_calling_the_provider( - self, monkeypatch - ): + async def test_falls_back_when_the_core_declines_before_calling_the_provider(self, monkeypatch): _fake_native_bridge(monkeypatch) - bridge.set_rust_chat_completions( - achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming")) - ) + bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming"))) assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None @@ -310,25 +313,19 @@ class TestAsyncFallbackWrapper: ran.append(True) return "python" - result = await bridge.achat_completions_or_fallback( - **_call_kwargs(ModelResponse()), python_fallback=fallback - ) + result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) assert result.choices[0].message.content == "hello from rust" assert ran == [] @pytest.mark.asyncio async def test_runs_the_fallback_when_the_core_declines(self, monkeypatch): _fake_native_bridge(monkeypatch) - bridge.set_rust_chat_completions( - achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming")) - ) + bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming"))) async def fallback(): return "python" - result = await bridge.achat_completions_or_fallback( - **_call_kwargs(ModelResponse()), python_fallback=fallback - ) + result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) assert result == "python" @pytest.mark.asyncio @@ -338,9 +335,7 @@ class TestAsyncFallbackWrapper: async def fallback(): return "python" - result = await bridge.achat_completions_or_fallback( - **_call_kwargs(ModelResponse()), python_fallback=fallback - ) + result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) assert result == "python" @@ -353,17 +348,13 @@ class TestFailureClassification: _fake_native_bridge(monkeypatch) def test_a_decline_falls_back_because_nothing_was_sent(self): - bridge.set_rust_chat_completions( - chat_completions=_RecordingCall(error=_FakeDeclined("streaming")) - ) + bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeDeclined("streaming"))) assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None def test_an_upstream_failure_is_surfaced_with_its_status(self): from litellm.exceptions import APIError - bridge.set_rust_chat_completions( - chat_completions=_RecordingCall(error=_FakeUpstream(429, "429: rate limited")) - ) + bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(429, "429: rate limited"))) with pytest.raises(APIError) as raised: bridge.chat_completions(**_call_kwargs(ModelResponse())) assert raised.value.status_code == 429 @@ -372,17 +363,13 @@ class TestFailureClassification: def test_a_transport_failure_with_no_response_surfaces_as_a_500(self): from litellm.exceptions import APIError - bridge.set_rust_chat_completions( - chat_completions=_RecordingCall(error=_FakeUpstream(0, "connection reset")) - ) + bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(0, "connection reset"))) with pytest.raises(APIError) as raised: bridge.chat_completions(**_call_kwargs(ModelResponse())) assert raised.value.status_code == 500 def test_an_unrecognized_error_is_not_swallowed(self): - bridge.set_rust_chat_completions( - chat_completions=_RecordingCall(error=RuntimeError("something else")) - ) + bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=RuntimeError("something else"))) with pytest.raises(RuntimeError): bridge.chat_completions(**_call_kwargs(ModelResponse())) @@ -390,9 +377,7 @@ class TestFailureClassification: async def test_the_async_wrapper_does_not_fall_back_on_an_upstream_failure(self): from litellm.exceptions import APIError - bridge.set_rust_chat_completions( - achat_completions=_RecordingAsyncCall(error=_FakeUpstream(500, "500: boom")) - ) + bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeUpstream(500, "500: boom"))) ran = [] async def fallback(): @@ -400,9 +385,7 @@ class TestFailureClassification: return "python" with pytest.raises(APIError): - await bridge.achat_completions_or_fallback( - **_call_kwargs(ModelResponse()), python_fallback=fallback - ) + await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) assert ran == [], "a request the provider already served must not be re-issued" @pytest.mark.asyncio @@ -414,7 +397,5 @@ class TestFailureClassification: async def fallback(): return "python" - result = await bridge.achat_completions_or_fallback( - **_call_kwargs(ModelResponse()), python_fallback=fallback - ) + result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback) assert result == "python" diff --git a/tests/test_litellm/rust_bridge/test_configuration.py b/tests/test_litellm/rust_bridge/test_configuration.py new file mode 100644 index 00000000000..1c81c1fb624 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_configuration.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +import os +import subprocess +import sys +from collections.abc import Generator +from concurrent.futures import ThreadPoolExecutor +from typing import Final + +import pytest + +from litellm.rust_bridge import configuration +from litellm.rust_bridge import ocr as rust_ocr + + +class _OcrBridge: + def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + return {} + + +@pytest.fixture(autouse=True) +def _isolated_configuration( # pyright: ignore[reportUnusedFunction] # pytest discovers fixtures dynamically + monkeypatch: pytest.MonkeyPatch, +) -> Generator[None]: + configuration.reset_rust_configuration() + monkeypatch.delenv("LITELLM_RUST", raising=False) + monkeypatch.delenv("LITELLM_USE_RUST_OCR", raising=False) + rust_ocr.set_rust_ocr(ocr=None, aocr=None) + yield + configuration.reset_rust_configuration() + rust_ocr.set_rust_ocr(ocr=None, aocr=None) + + +@pytest.mark.parametrize( + ("request_override", "process", "environment", "legacy_ocr", "release_default", "expected"), + ( + (False, True, True, True, True, False), + (True, False, False, False, False, True), + (None, False, True, True, True, False), + (None, True, False, False, False, True), + (None, None, False, True, True, False), + (None, None, True, False, False, True), + (None, None, None, False, True, False), + (None, None, None, True, False, True), + (None, None, None, None, False, False), + (None, None, None, None, True, True), + ), +) +def test_resolution_precedence( + request_override: bool | None, + process: bool | None, + environment: bool | None, + legacy_ocr: bool | None, + release_default: bool, + expected: bool, +) -> None: + assert ( + configuration.resolve_rust_enabled( + request_override=request_override, + process_override=process, + environment_override=environment, + legacy_ocr_override=legacy_ocr, + release_default=release_default, + ) + is expected + ) + + +def test_release_default_remains_disabled() -> None: + assert configuration.DEFAULT_RUST_ENABLED is False + assert configuration.rust_enabled() is False + + +def test_process_override_wins_over_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_RUST", "0") + configuration.use_litellm_rust(True) + + assert configuration.rust_enabled() is True + assert configuration.rust_enabled(request_override=False) is False + + +def test_global_environment_accepts_explicit_false(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_RUST", "off") + + assert configuration.rust_enabled() is False + + +@pytest.mark.parametrize("value", ("", " ", "sometimes", "2")) +def test_invalid_environment_value_disables_rust(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + monkeypatch.setenv("LITELLM_RUST", value) + monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") + + assert configuration.rust_enabled() is False + assert configuration.rust_ocr_enabled() is False + + +@pytest.mark.parametrize("value", ("", " ", "sometimes", "2")) +def test_invalid_legacy_environment_value_disables_ocr(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + monkeypatch.setenv("LITELLM_USE_RUST_OCR", value) + + with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): + assert configuration.rust_ocr_enabled() is False + + +def test_process_override_and_reset_apply_to_existing_threads(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_RUST", "1") + + with ThreadPoolExecutor(max_workers=1) as executor: + assert executor.submit(configuration.rust_enabled).result() is True + configuration.use_litellm_rust(False) + assert executor.submit(configuration.rust_enabled).result() is False + assert executor.submit(configuration.rust_ocr_enabled).result() is False + configuration.reset_rust_configuration() + assert executor.submit(configuration.rust_enabled).result() is True + assert executor.submit(configuration.rust_ocr_enabled).result() is True + + +def test_explicit_override_precedes_invalid_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_RUST", "sometimes") + + assert configuration.rust_enabled(request_override=False) is False + configuration.use_litellm_rust(True) + assert configuration.rust_enabled() is True + + +def test_legacy_ocr_environment_is_deprecated_and_ocr_only(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") + + with pytest.warns(DeprecationWarning, match="LITELLM_USE_RUST_OCR is deprecated"): + assert configuration.rust_ocr_enabled() is True + assert configuration.rust_enabled() is False + + +def test_global_environment_precedes_legacy_ocr_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_RUST", "0") + monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") + + assert configuration.rust_ocr_enabled() is False + + +def test_deprecated_public_injection_delegates_to_internal_binding() -> None: + bridge: Final = _OcrBridge() + + with pytest.warns(DeprecationWarning, match="Injecting Rust bridge implementations"): + configuration.use_litellm_rust(True, ocr=bridge) + + assert rust_ocr.load_rust_ocr() is bridge + + +@pytest.mark.parametrize(("value", "expected"), (("1", "True"), ("0", "False"))) +def test_environment_controls_startup(value: str, expected: str) -> None: + environment: Final = {**os.environ, "LITELLM_RUST": value} + result: Final = subprocess.run( + ( + sys.executable, + "-c", + "from litellm.rust_bridge.configuration import rust_enabled; print(rust_enabled())", + ), + check=True, + capture_output=True, + text=True, + env=environment, + ) + + assert result.stdout.strip() == expected diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py new file mode 100644 index 00000000000..b0fa510069b --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from litellm.exceptions import APIError +from litellm.rust_bridge import bindings, runtime + + +class RustBridgeDeclined(Exception): + pass + + +class RustUpstreamError(Exception): + pass + + +@pytest.fixture(autouse=True) +def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> None: + native = SimpleNamespace( + RustBridgeDeclined=RustBridgeDeclined, + RustUpstreamError=RustUpstreamError, + ) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + + +def context() -> runtime.BridgeErrorContext: + return runtime.BridgeErrorContext(route="messages", provider="anthropic", model="model") + + +def test_invoke_tags_native_decline_before_running_fallback() -> None: + calls: list[str] = [] + + def decline() -> object: + calls.append("rust") + raise RustBridgeDeclined("unsupported") + + value = runtime.invoke( + native_call=decline, + fallback=lambda: calls.append("python") or "fallback", + adapt=str, + mode=runtime.FallbackMode.PYTHON, + context=context(), + ) + + assert value == "fallback" + assert calls == ["rust", "python"] + + +def test_invoke_translates_upstream_without_fallback() -> None: + def fail() -> object: + raise RustUpstreamError(429, "rate limited") + + with pytest.raises(APIError, match="rate limited") as caught: + runtime.invoke( + native_call=fail, + fallback=lambda: pytest.fail("fallback must not run"), + adapt=str, + mode=runtime.FallbackMode.PYTHON, + context=context(), + ) + + assert caught.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_ainvoke_handles_native_success() -> None: + async def native() -> int: + return 3 + + async def fallback() -> str: + pytest.fail("fallback must not run") + + assert ( + await runtime.ainvoke( + native_call=native, + fallback=fallback, + adapt=str, + mode=runtime.FallbackMode.PYTHON, + context=context(), + ) + == "3" + ) + + +def test_required_mode_rejects_unavailable_bridge() -> None: + with pytest.raises(RuntimeError, match="is unavailable"): + runtime.invoke( + native_call=None, + fallback=lambda: pytest.fail("fallback must not run"), + adapt=str, + mode=runtime.FallbackMode.RUST_REQUIRED, + context=context(), + ) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index dcaa6cfd602..c843a66a1c1 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8579,6 +8579,398 @@ class TestConsumedRequestTagsStamp: assert CONSUMED_REQUEST_TAGS_METADATA_KEY not in request_kwargs["metadata"] +class TestClaudeCodeSubagentSessionRouterBinding: + class _RewriteStrategy: + def __init__(self, routed_model: str = "cheap-model") -> None: + self.routed_model = routed_model + + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + return PreRoutingHookResponse( + model=self.routed_model, + messages=messages, + routing_decision={ + "router_model_name": "smart-router", + "router_type": "complexity", + "routed_model": self.routed_model, + "cause": "heuristic_scorer", + }, + ) + + @classmethod + def _router( + cls, + cheap_response: str = "cheap response", + fallbacks: list[dict[str, list[str]]] | None = None, + ) -> "litellm.Router": + from litellm.types.router import TaggedPreRoutingStrategy + + router = litellm.Router( + model_list=[ + { + "model_name": "cheap-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": cheap_response}, + }, + { + "model_name": "expensive-model", + "litellm_params": {"model": "openai/gpt-4o", "mock_response": "expensive response"}, + }, + ], + fallbacks=fallbacks, + num_retries=0, + ) + router.complexity_routers = { + "smart-router": (TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy()),), + "premium-router": (TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy("expensive-model")),), + } + return router + + @staticmethod + def _request_kwargs( + *, + key_hash: str = "key-hash-a", + app: str = "cli", + agent_id: str | None = None, + fallback_depth: int | None = None, + ) -> dict: + headers = { + "X-Claude-Code-Session-Id": "session-1234", + "x-app": app, + **({"x-claude-code-agent-id": agent_id} if agent_id is not None else {}), + } + return { + "metadata": {"user_api_key_hash": key_hash}, + "proxy_server_request": {"headers": headers}, + **({"fallback_depth": fallback_depth} if fallback_depth is not None else {}), + } + + @pytest.mark.asyncio + async def test_subagent_concrete_model_uses_the_main_sessions_router(self): + router = self._router() + + await router.acompletion( + model="smart-router", + messages=[{"role": "user", "content": "main turn"}], + **self._request_kwargs(), + ) + subagent_kwargs = self._request_kwargs(agent_id="agent-1234") + + response = await router.acompletion( + model="expensive-model", + messages=[{"role": "user", "content": "subagent turn"}], + **subagent_kwargs, + ) + + assert response.choices[0].message.content == "cheap response" + assert subagent_kwargs["metadata"]["model_group"] == "smart-router" + assert subagent_kwargs["metadata"]["routing_decision"]["router_model_name"] == "smart-router" + + @pytest.mark.asyncio + async def test_main_thread_side_calls_to_a_plain_model_keep_the_session_router(self): + router = self._router() + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + await router.async_pre_routing_hook(model="expensive-model", request_kwargs=self._request_kwargs()) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + + assert response is not None + assert response.model == "cheap-model" + + @pytest.mark.asyncio + async def test_redis_cleanup_failure_does_not_reject_a_subagent_request(self): + from litellm.caching.caching import RedisCache + + router = self._router() + del router.complexity_routers["smart-router"] + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(return_value="smart-router") + redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) + router._update_redis_cache(cache=redis_cache) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + + assert response is None + redis_cache.async_delete_cache.assert_awaited_once() + + @pytest.mark.asyncio + async def test_redis_read_failure_does_not_reject_a_subagent_request(self): + from litellm.caching.caching import RedisCache + + router = self._router() + request_kwargs = self._request_kwargs(agent_id="agent-1234") + cache_key = router._claude_code_session_router_cache_key(request_kwargs) + assert cache_key is not None + await router._claude_code_session_router_cache.in_memory_cache.async_set_cache( + cache_key, + "smart-router", + ) + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(side_effect=Exception("Redis circuit breaker is open")) + router._update_redis_cache(cache=redis_cache) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=request_kwargs, + ) + + assert response is None + assert "model_group" not in request_kwargs["metadata"] + redis_cache.async_get_cache.assert_awaited_once() + + @pytest.mark.asyncio + async def test_redis_write_failures_do_not_reject_main_or_subagent_requests(self): + from litellm.caching.caching import RedisCache + + router = self._router() + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(return_value="smart-router") + redis_cache.async_set_cache = AsyncMock(side_effect=Exception("redis unavailable")) + router._update_redis_cache(cache=redis_cache) + + main_response = await router.async_pre_routing_hook( + model="smart-router", + request_kwargs=self._request_kwargs(), + ) + subagent_response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + + assert main_response is not None + assert main_response.model == "cheap-model" + assert subagent_response is not None + assert subagent_response.model == "cheap-model" + assert redis_cache.async_set_cache.await_count == 2 + + @pytest.mark.asyncio + async def test_subagents_follow_the_main_threads_latest_router_across_workers(self): + from types import SimpleNamespace + + from litellm.caching.caching import RedisCache + + shared_binding = SimpleNamespace(value=None) + shared_redis = MagicMock(spec=RedisCache) + shared_redis.async_get_cache = AsyncMock(side_effect=lambda key, **_: shared_binding.value) + shared_redis.async_set_cache = AsyncMock( + side_effect=lambda key, value, **_: setattr(shared_binding, "value", value) + ) + main_worker, subagent_worker = self._router(), self._router() + main_worker._update_redis_cache(cache=shared_redis) + subagent_worker._update_redis_cache(cache=shared_redis) + + await main_worker.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + first = await subagent_worker.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + await main_worker.async_pre_routing_hook(model="premium-router", request_kwargs=self._request_kwargs()) + second = await subagent_worker.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + + assert first is not None + assert first.model == "cheap-model" + assert second is not None + assert second.model == "expensive-model" + assert shared_binding.value == "premium-router" + + @pytest.mark.asyncio + async def test_no_pre_routing_strategies_means_no_session_cache_traffic(self): + from litellm.caching.caching import RedisCache + + router = self._router() + router.complexity_routers.clear() + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_get_cache = AsyncMock(return_value=None) + redis_cache.async_set_cache = AsyncMock() + redis_cache.async_delete_cache = AsyncMock() + router._update_redis_cache(cache=redis_cache) + + for request_kwargs in (self._request_kwargs(), self._request_kwargs(agent_id="agent-1234")): + response = await router.async_pre_routing_hook(model="expensive-model", request_kwargs=request_kwargs) + assert response is None + + redis_cache.async_get_cache.assert_not_awaited() + redis_cache.async_set_cache.assert_not_awaited() + redis_cache.async_delete_cache.assert_not_awaited() + + @pytest.mark.asyncio + async def test_session_bindings_do_not_evict_router_rate_limit_state(self): + router = self._router() + assert router._update_usage(deployment_id="deployment-id", parent_otel_span=None) == 1 + + for session_index in range(201): + request_kwargs = self._request_kwargs() + request_kwargs["proxy_server_request"]["headers"]["X-Claude-Code-Session-Id"] = ( + f"session-{session_index:04d}" + ) + await router.async_pre_routing_hook(model="smart-router", request_kwargs=request_kwargs) + + assert router._update_usage(deployment_id="deployment-id", parent_otel_span=None) == 2 + + @pytest.mark.asyncio + async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): + router = self._router() + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(app="cli-bg"), + ) + await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(fallback_depth=1), + ) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234"), + ) + + assert response is not None + assert response.model == "cheap-model" + + @pytest.mark.asyncio + async def test_subagent_fallback_does_not_reapply_the_session_router(self): + router = self._router() + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(agent_id="agent-1234", fallback_depth=1), + ) + + assert response is None + + @pytest.mark.asyncio + async def test_subagent_can_fallback_to_its_original_requested_model(self): + router = self._router( + cheap_response="litellm.RateLimitError", + fallbacks=[{"cheap-model": ["expensive-model"]}], + ) + + await router.acompletion( + model="smart-router", + messages=[{"role": "user", "content": "main turn"}], + **self._request_kwargs(), + ) + subagent_kwargs = self._request_kwargs(agent_id="agent-1234") + + response = await router.acompletion( + model="expensive-model", + messages=[{"role": "user", "content": "subagent turn"}], + **subagent_kwargs, + ) + + assert response.choices[0].message.content == "expensive response" + assert subagent_kwargs["metadata"]["routing_decision"]["routed_model"] == "cheap-model" + + @pytest.mark.asyncio + async def test_subagent_can_use_the_bound_router_name_fallback(self): + router = self._router( + cheap_response="litellm.RateLimitError", + fallbacks=[{"smart-router": ["expensive-model"]}], + ) + + await router.acompletion( + model="smart-router", + messages=[{"role": "user", "content": "main turn"}], + **self._request_kwargs(), + ) + + response = await router.acompletion( + model="expensive-model", + messages=[{"role": "user", "content": "subagent turn"}], + **self._request_kwargs(agent_id="agent-1234"), + ) + + assert response.choices[0].message.content == "expensive response" + + @pytest.mark.asyncio + async def test_anthropic_subagent_four_fallback_hops_use_each_current_model_chain(self): + from litellm.types.router import TaggedPreRoutingStrategy + + failing_groups = ("cheap-model", "fallback-1", "fallback-2", "fallback-3") + router = litellm.Router( + model_list=[ + *( + { + "model_name": group, + "litellm_params": { + "model": "anthropic/claude-3-haiku-20240307", + "mock_response": "litellm.RateLimitError", + }, + } + for group in failing_groups + ), + { + "model_name": "requested-model", + "litellm_params": { + "model": "anthropic/claude-3-haiku-20240307", + "mock_response": "requested response", + }, + }, + { + "model_name": "fallback-4", + "litellm_params": { + "model": "anthropic/claude-3-haiku-20240307", + "mock_response": "fourth fallback response", + }, + }, + ], + fallbacks=[ + {"smart-router": ["fallback-1"]}, + {"fallback-1": ["fallback-2"]}, + {"fallback-2": ["fallback-3"]}, + {"fallback-3": ["fallback-4"]}, + ], + num_retries=0, + max_fallbacks=4, + ) + router.complexity_routers = { + "smart-router": (TaggedPreRoutingStrategy(tags=(), strategy=self._RewriteStrategy()),) + } + main_kwargs = self._request_kwargs() + main_kwargs["litellm_metadata"] = main_kwargs.pop("metadata") + await router.async_pre_routing_hook(model="smart-router", request_kwargs=main_kwargs) + subagent_kwargs = self._request_kwargs(agent_id="agent-1234") + subagent_kwargs["litellm_metadata"] = subagent_kwargs.pop("metadata") + + response = await router.aanthropic_messages( + model="requested-model", + messages=[{"role": "user", "content": "subagent turn"}], + max_tokens=64, + **subagent_kwargs, + ) + + assert response["content"][0]["text"] == "fourth fallback response" + + @pytest.mark.asyncio + async def test_session_router_binding_is_scoped_to_the_authenticated_key(self): + router = self._router() + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(key_hash="key-hash-b", agent_id="agent-1234"), + ) + + assert response is None + + class TestAutoRouterMaxInputCharsWiring: """`auto_router_max_input_chars` on the deployment has to reach the AutoRouter that embeds prompts. diff --git a/tests/vector_store_tests/test_azure_ai_vector_store.py b/tests/vector_store_tests/test_azure_ai_vector_store.py index 58e45f259ab..d1fc8436fc9 100644 --- a/tests/vector_store_tests/test_azure_ai_vector_store.py +++ b/tests/vector_store_tests/test_azure_ai_vector_store.py @@ -1,10 +1,19 @@ -import pytest -import litellm import json import os +from unittest.mock import MagicMock + +import httpx +import pytest +import respx + +import litellm +from litellm.llms.azure_ai.vector_stores.transformation import AzureAIVectorStoreConfig +from litellm.types.utils import EmbeddingResponse +from litellm.vector_stores import ( + asearch as vector_store_asearch, +) from litellm.vector_stores import ( search as vector_store_search, - asearch as vector_store_asearch, ) @@ -30,10 +39,108 @@ async def test_basic_search_vector_store(sync_mode): if sync_mode: response = vector_store_search(query=default_query, **base_request_args) else: - response = await vector_store_asearch( - query=default_query, **base_request_args - ) + response = await vector_store_asearch(query=default_query, **base_request_args) except litellm.InternalServerError: pytest.skip("Skipping test due to litellm.InternalServerError") print("litellm response=", json.dumps(response, indent=4, default=str)) + + +class RecordingEmbeddingExecutor: + def __init__(self, response): + self.response = response + self.calls = [] + + def embed(self, model, query, configuration): + self.calls.append((model, query, dict(configuration))) + return self.response + + async def aembed(self, model, query, configuration): + self.calls.append((model, query, dict(configuration))) + return self.response + + +ALIAS_QUERY_VECTOR = [0.5, -0.25, 0.125] +ALIAS_EMBEDDING_RESPONSE = EmbeddingResponse( + data=[{"embedding": ALIAS_QUERY_VECTOR, "index": 0, "object": "embedding"}] +) +STORE_EMBEDDINGS_URL = "https://embedding.example/v1/embeddings" + + +def _transform_kwargs(executor): + logging_obj = MagicMock() + logging_obj.model_call_details = {} + return { + "vector_store_id": "my-vector-index", + "query": "what is azure search?", + "vector_store_search_optional_params": {"top_k": 2}, + "api_base": "https://azure-kb-search.search.windows.net", + "litellm_logging_obj": logging_obj, + "litellm_params": { + "litellm_embedding_model": "multilingual-e5-large", + "azure_search_vector_field": "embedding", + }, + "embedding_executor": executor, + } + + +@pytest.mark.asyncio +async def test_transform_uses_injected_executor_without_embedding_config(respx_mock: respx.MockRouter): + executor = RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE) + config = AzureAIVectorStoreConfig() + transform_kwargs = _transform_kwargs(executor) + + url, sync_body = config.transform_search_vector_store_request(**transform_kwargs) + _, async_body = await config.atransform_search_vector_store_request(**transform_kwargs) + + assert respx_mock.calls.call_count == 0 + assert executor.calls == [("multilingual-e5-large", "what is azure search?", {})] * 2 + assert ( + url == "https://azure-kb-search.search.windows.net/indexes/my-vector-index/docs/search?api-version=2024-07-01" + ) + assert sync_body == async_body + assert sync_body["vectorQueries"] == [ + {"vector": ALIAS_QUERY_VECTOR, "fields": "embedding", "kind": "vector", "k": 2} + ] + assert sync_body["top"] == 2 + logging_details = transform_kwargs["litellm_logging_obj"].model_call_details + assert logging_details["embedding_model"] == "multilingual-e5-large" + assert logging_details["top_k"] == 2 + + +def test_transform_falls_back_to_sdk_embedding_without_executor( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = respx_mock.post(STORE_EMBEDDINGS_URL).mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": ALIAS_QUERY_VECTOR}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + ) + transform_kwargs = _transform_kwargs(None) + transform_kwargs["litellm_params"] = { + "litellm_embedding_model": "openai/text-embedding-3-small", + "litellm_embedding_config": {"api_base": "https://embedding.example/v1", "api_key": "store-key"}, + } + + _, body = AzureAIVectorStoreConfig().transform_search_vector_store_request(**transform_kwargs) + + embedding_request = embedding_route.calls.last.request + assert embedding_request.headers["authorization"] == "Bearer store-key" + assert json.loads(embedding_request.read())["input"] == ["what is azure search?"] + assert body["vectorQueries"][0]["vector"] == ALIAS_QUERY_VECTOR + assert body["vectorQueries"][0]["fields"] == "contentVector" + + +def test_transform_requires_embedding_model(): + transform_kwargs = _transform_kwargs(RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE)) + transform_kwargs["litellm_params"] = {"litellm_embedding_config": {"api_key": "store-key"}} + + with pytest.raises(ValueError, match="litellm_embedding_model is required"): + AzureAIVectorStoreConfig().transform_search_vector_store_request(**transform_kwargs) diff --git a/tests/vector_store_tests/test_milvus_vector_store.py b/tests/vector_store_tests/test_milvus_vector_store.py index 6627f6006d1..2ba9168b49f 100644 --- a/tests/vector_store_tests/test_milvus_vector_store.py +++ b/tests/vector_store_tests/test_milvus_vector_store.py @@ -3,16 +3,19 @@ Tests for Milvus Vector Store """ import json -import os from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx import litellm +from litellm import Router +from litellm.llms.milvus.vector_stores.transformation import MilvusVectorStoreConfig +from litellm.types.utils import EmbeddingResponse from litellm.vector_stores import asearch as vector_store_asearch from litellm.vector_stores import search as vector_store_search - # Mock response from actual Milvus API MOCK_MILVUS_SEARCH_RESPONSE = { "code": 0, @@ -98,7 +101,7 @@ class TestMilvusVectorStore: mock_response.json.return_value = MOCK_MILVUS_SEARCH_RESPONSE mock_response.text = json.dumps(MOCK_MILVUS_SEARCH_RESPONSE) - with patch("litellm.embedding") as mock_embedding: + with patch("litellm.aembedding", new_callable=AsyncMock) as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE with patch( @@ -147,16 +150,10 @@ class TestMilvusVectorStore: else: # Fallback: check for json kwarg or in args request_data = call_args.kwargs.get("json") - if ( - request_data is None - and len(call_args.args) > 0 - and isinstance(call_args.args[0], dict) - ): + if request_data is None and len(call_args.args) > 0 and isinstance(call_args.args[0], dict): request_data = call_args.args[0] - assert ( - request_data is not None - ), f"Could not extract request data. Call args: {call_args}" + assert request_data is not None, f"Could not extract request data. Call args: {call_args}" print("Request data:", json.dumps(request_data, indent=2, default=str)) # Validate request structure @@ -213,9 +210,7 @@ class TestMilvusVectorStore: with patch("litellm.embedding") as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE - with patch( - "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_post.return_value = mock_response # Make the search request @@ -252,16 +247,10 @@ class TestMilvusVectorStore: else: # Fallback: check for json kwarg or in args request_data = call_args.kwargs.get("json") - if ( - request_data is None - and len(call_args.args) > 0 - and isinstance(call_args.args[0], dict) - ): + if request_data is None and len(call_args.args) > 0 and isinstance(call_args.args[0], dict): request_data = call_args.args[0] - assert ( - request_data is not None - ), f"Could not extract request data. Call args: {call_args}" + assert request_data is not None, f"Could not extract request data. Call args: {call_args}" # Validate request structure assert "collectionName" in request_data @@ -316,11 +305,7 @@ class TestMilvusVectorStore: if request_data_str: return json.loads(request_data_str) request_data = call_args.kwargs.get("json") - if ( - request_data is None - and len(call_args.args) > 0 - and isinstance(call_args.args[0], dict) - ): + if request_data is None and len(call_args.args) > 0 and isinstance(call_args.args[0], dict): request_data = call_args.args[0] return request_data @@ -334,9 +319,7 @@ class TestMilvusVectorStore: with patch("litellm.embedding") as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE - with patch( - "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_post.return_value = mock_response vector_store_search( @@ -375,9 +358,7 @@ class TestMilvusVectorStore: with patch("litellm.embedding") as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE - with patch( - "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_post.return_value = mock_response vector_store_search( @@ -413,9 +394,7 @@ class TestMilvusVectorStore: with patch("litellm.embedding") as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE - with patch( - "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_post.return_value = mock_response vector_store_search( @@ -492,3 +471,247 @@ if __name__ == "__main__": test.test_basic_search_with_mock_sync() print("\nāœ… All mock tests passed!") + + +class RecordingEmbeddingExecutor: + def __init__(self, response): + self.response = response + self.calls = [] + + def embed(self, model, query, configuration): + self.calls.append((model, query, dict(configuration))) + return self.response + + async def aembed(self, model, query, configuration): + self.calls.append((model, query, dict(configuration))) + return self.response + + +ALIAS_QUERY_VECTOR = [0.5, -0.25, 0.125] +ALIAS_EMBEDDING_RESPONSE = EmbeddingResponse( + data=[{"embedding": ALIAS_QUERY_VECTOR, "index": 0, "object": "embedding"}] +) +OPENAI_EMBEDDINGS_URL = "https://api.openai.com/v1/embeddings" +MILVUS_SEARCH_URL = "https://milvus.example/v2/vectordb/entities/search" +ALIAS_SEARCH_KWARGS = { + "query": "what is machine learning?", + "vector_store_id": "book_2", + "custom_llm_provider": "milvus", + "api_base": "https://milvus.example", + "api_key": "mock_milvus_api_key", + "litellm_embedding_model": "multilingual-e5-large", + "milvus_text_field": "book_intro_text", +} + + +def _alias_router(): + return Router( + model_list=[ + { + "model_name": "multilingual-e5-large", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "deployment-key", + }, + } + ] + ) + + +def _mock_embedding_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post(OPENAI_EMBEDDINGS_URL).mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": ALIAS_QUERY_VECTOR}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + ) + + +def _mock_search_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post(MILVUS_SEARCH_URL).mock(return_value=httpx.Response(200, json=MOCK_MILVUS_SEARCH_RESPONSE)) + + +def _assert_alias_resolved(embedding_route: respx.Route, search_route: respx.Route, response): + embedding_request = embedding_route.calls.last.request + assert embedding_request.headers["authorization"] == "Bearer deployment-key" + embedding_body = json.loads(embedding_request.read()) + assert embedding_body["model"] == "text-embedding-3-small" + assert embedding_body["input"] == ["what is machine learning?"] + search_request = search_route.calls.last.request + assert search_request.headers["authorization"] == "Bearer mock_milvus_api_key" + assert json.loads(search_request.read())["data"] == [ALIAS_QUERY_VECTOR] + assert len(response["data"]) == len(MOCK_MILVUS_SEARCH_RESPONSE["data"]) + assert response["data"][0]["content"][0]["text"] == MOCK_MILVUS_SEARCH_RESPONSE["data"][0]["book_intro_text"] + + +def test_router_search_resolves_bare_embedding_alias_sync( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = _alias_router().vector_store_search(**ALIAS_SEARCH_KWARGS) + + _assert_alias_resolved(embedding_route, search_route, response) + + +@pytest.mark.asyncio +async def test_router_search_resolves_bare_embedding_alias_async( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = await _alias_router().avector_store_search(**ALIAS_SEARCH_KWARGS) + + _assert_alias_resolved(embedding_route, search_route, response) + + +def test_sdk_search_with_router_kwarg_resolves_bare_embedding_alias_sync( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = litellm.vector_stores.search(router=_alias_router(), **ALIAS_SEARCH_KWARGS) + + _assert_alias_resolved(embedding_route, search_route, response) + + +@pytest.mark.asyncio +async def test_sdk_search_with_router_kwarg_resolves_bare_embedding_alias_async( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = await litellm.vector_stores.asearch(router=_alias_router(), **ALIAS_SEARCH_KWARGS) + + _assert_alias_resolved(embedding_route, search_route, response) + + +def _team_alias_router(): + return Router( + model_list=[ + { + "model_name": "team-a-embedder", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "deployment-key", + }, + "model_info": {"team_id": "team-a", "team_public_model_name": "multilingual-e5-large"}, + } + ] + ) + + +@pytest.mark.asyncio +async def test_sdk_search_with_router_kwarg_resolves_team_alias_from_request_metadata( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = await litellm.vector_stores.asearch( + router=_team_alias_router(), metadata={"user_api_key_team_id": "team-a"}, **ALIAS_SEARCH_KWARGS + ) + + _assert_alias_resolved(embedding_route, search_route, response) + + +@pytest.mark.asyncio +async def test_sdk_search_with_router_kwarg_rejects_team_alias_without_team_metadata( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + _mock_search_route(respx_mock) + + with pytest.raises(litellm.APIConnectionError): + await litellm.vector_stores.asearch(router=_team_alias_router(), **ALIAS_SEARCH_KWARGS) + + assert embedding_route.call_count == 0 + + +@pytest.mark.asyncio +async def test_transform_uses_injected_executor_without_embedding_config(respx_mock: respx.MockRouter): + executor = RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE) + config = MilvusVectorStoreConfig() + logging_obj = MagicMock() + logging_obj.model_call_details = {} + transform_kwargs = { + "vector_store_id": "book_2", + "query": ["what is", "milvus?"], + "vector_store_search_optional_params": {"limit": 3}, + "api_base": "https://milvus.example", + "litellm_logging_obj": logging_obj, + "litellm_params": {"litellm_embedding_model": "multilingual-e5-large", "milvus_db_name": "docs"}, + "embedding_executor": executor, + } + + url, sync_body = config.transform_search_vector_store_request(**transform_kwargs) + _, async_body = await config.atransform_search_vector_store_request(**transform_kwargs) + + assert respx_mock.calls.call_count == 0 + assert executor.calls == [("multilingual-e5-large", "what is milvus?", {})] * 2 + assert url == MILVUS_SEARCH_URL + assert sync_body == async_body + assert sync_body == { + "collectionName": "book_2", + "data": [ALIAS_QUERY_VECTOR], + "annsField": "book_intro_vector", + "limit": 3, + "dbName": "docs", + } + assert logging_obj.model_call_details["input"] == "what is milvus?" + assert logging_obj.model_call_details["embedding_model"] == "multilingual-e5-large" + + +def test_transform_falls_back_to_sdk_embedding_without_executor_or_config( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + embedding_route = _mock_embedding_route(respx_mock) + logging_obj = MagicMock() + logging_obj.model_call_details = {} + + _, body = MilvusVectorStoreConfig().transform_search_vector_store_request( + vector_store_id="book_2", + query="q", + vector_store_search_optional_params={}, + api_base="https://milvus.example", + litellm_logging_obj=logging_obj, + litellm_params={"litellm_embedding_model": "openai/text-embedding-3-small"}, + ) + + embedding_request = embedding_route.calls.last.request + assert embedding_request.headers["authorization"] == "Bearer env-key" + assert json.loads(embedding_request.read())["input"] == ["q"] + assert body["data"] == [ALIAS_QUERY_VECTOR] + + +def test_transform_requires_embedding_model(): + with pytest.raises(ValueError, match="litellm_embedding_model is required"): + MilvusVectorStoreConfig().transform_search_vector_store_request( + vector_store_id="book_2", + query="q", + vector_store_search_optional_params={}, + api_base="https://milvus.example", + litellm_logging_obj=MagicMock(), + litellm_params={"litellm_embedding_config": {"api_key": "store-key"}}, + embedding_executor=RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE), + ) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 6273fbce595..a2e63f19881 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,12 +1,12 @@ { "LIT001": { - "limit": 22358 + "limit": 22334 }, "LIT002": { - "limit": 26774 + "limit": 26765 }, "LIT003": { - "limit": 269 + "limit": 261 }, "LIT004": { "limit": 40 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1039 + "limit": 1038 }, "LIT007": { "limit": 0 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16494 + "limit": 16482 }, "LIT011": { - "limit": 5535 + "limit": 5520 }, "LIT012": { "limit": 4495 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts index bbdf4697315..c4d7f45b7cc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts @@ -58,6 +58,7 @@ const dedupe = (models: string[]): string[] => Array.from(new Set(models)); const COMPLEXITY_TYPE_LABELS: Record = { llm: "LLM Classifier", heuristic_first: "Heuristic first", + hybrid: "Hybrid", custom: "Custom classifier", }; diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 96c93306611..a29527a20fa 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -35,6 +35,7 @@ import { heuristicScoringRole, usesLlmClassifier, DEFAULT_HEURISTIC_FIRST_MAX_TIER, + DEFAULT_HYBRID_BOUNDARY_MARGIN, HEURISTIC_FIRST_MAX_TIER_KEYS, effectiveClassifierType, } from "./ComplexityRouterConfig"; @@ -43,9 +44,14 @@ const DEFAULT_SCORING_EXPLANATION = "The router scores each request across 7 dimensions: token count, code presence, reasoning markers, technical " + "terms, simple indicators, multi-step patterns, and question complexity. The weighted score determines the tier:"; +const HEURISTIC_V2_EXPLANATION = + "The router estimates success probability for all four tiers with the bundled calibrated model, then selects " + + "the first tier that meets its trained threshold. It runs locally with no classifier API call."; + const CLASSIFIER_TIMEOUT_ID = "classifier-timeout-ms"; const CLASSIFIER_CONTEXT_WINDOW_SIZE_ID = "classifier-context-window-size"; const CLASSIFIER_CONTEXT_BUDGET_CHARS_ID = "classifier-context-budget-chars"; +const HYBRID_BOUNDARY_MARGIN_ID = "hybrid-boundary-margin"; const CUSTOM_PROMPT_WITH_HEURISTIC_FALLBACK = "This router classifies with your own prompt, so the tier comes from whatever rubric it states. The four tier " + @@ -62,6 +68,7 @@ const CUSTOM_PROMPT_WITH_DEFAULT_MODEL_FALLBACK = * at all, so the panel must not keep implying a score is involved on either router. */ const scoringExplanation = (value: ComplexityRouterConfigValue): string => { + if (value.classifier_type === "heuristic_v2") return HEURISTIC_V2_EXPLANATION; const usesCustomPrompt = usesLlmClassifier(value.classifier_type) && Boolean(value.classifier_llm_config?.system_prompt?.trim()); if (!usesCustomPrompt) return DEFAULT_SCORING_EXPLANATION; @@ -179,6 +186,17 @@ const ClassifierTypeRadios: React.FC<{ + + + + + + ); @@ -247,6 +277,8 @@ const ClassificationMethodConfig: React.FC = ({ classifierType === "heuristic_first" ? value.heuristic_first_max_tier ?? DEFAULT_HEURISTIC_FIRST_MAX_TIER : undefined, + hybrid_boundary_margin: + classifierType === "hybrid" ? value.hybrid_boundary_margin ?? DEFAULT_HYBRID_BOUNDARY_MARGIN : undefined, }; onChange(nextValue); }; @@ -255,6 +287,13 @@ const ClassificationMethodConfig: React.FC = ({ onChange({ ...value, heuristic_first_max_tier: tier }); }; + const handleHybridBoundaryMarginChange = (raw: string) => { + setDraft({ id: HYBRID_BOUNDARY_MARGIN_ID, raw }); + const parsed = Number(raw); + if (raw.trim() === "" || !Number.isFinite(parsed)) return; + onChange({ ...value, hybrid_boundary_margin: Math.min(1, Math.max(0, parsed)) }); + }; + const handleClassificationPromptChange = (classificationPrompt: string | undefined) => { onChange({ ...value, classification_prompt: classificationPrompt }); }; @@ -375,6 +414,30 @@ const ClassificationMethodConfig: React.FC = ({ )} + {classifierType === "hybrid" && ( +
+ Boundary margin + handleHybridBoundaryMarginChange(event.target.value)} + onBlur={() => setDraft(null)} + className="w-full" + /> +

+ A score further than this from every tier boundary routes on the scorer's own tier, however expensive + that tier is. A score closer than this, and anything the scorer found no signal for at all, goes to the + classifier to break the tie +

+
+ )} +
How often to classify { expect(onChange).toHaveBeenCalledWith(expectedValue); }); + it("selects heuristic v2 without requiring a classifier model or showing weighted scoring", () => { + const onChange = vi.fn(); + const { rerender } = renderWithProviders( + , + ); + + fireEvent.click(screen.getByText("Advanced: Classification Method")); + fireEvent.click(screen.getByText("Heuristic v2")); + + expect(onChange).toHaveBeenCalledWith( + expect.objectContaining({ + classifier_type: "heuristic_v2", + classifier_llm_config: undefined, + }), + ); + + const heuristicV2Value: ComplexityRouterConfigValue = { ...defaultValue, classifier_type: "heuristic_v2" }; + rerender(); + + expect(screen.queryByText("Classifier Model")).not.toBeInTheDocument(); + expect(screen.queryByText("Advanced scoring")).not.toBeInTheDocument(); + expect(screen.getByText(/estimates success probability for all four tiers/)).toBeInTheDocument(); + expect(screen.queryByText(/Score < 0.15/)).not.toBeInTheDocument(); + }); + it("should show classifier fields and use the configured values when classifier_type is llm", () => { const llmValue: ComplexityRouterConfigValue = { ...defaultValue, diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 6062705aa82..73ab4f0abef 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -128,7 +128,7 @@ export interface ClassifierLLMConfig { system_prompt?: string; } -export type ClassifierType = "heuristic" | "llm" | "heuristic_first"; +export type ClassifierType = "heuristic" | "heuristic_v2" | "llm" | "heuristic_first" | "hybrid"; /** * Whether this router can call classifier_llm_config.model. Mirrors the backend's @@ -136,7 +136,7 @@ export type ClassifierType = "heuristic" | "llm" | "heuristic_first"; * control and payload key, so a new chaining type cannot strip knobs the operator set. */ export const usesLlmClassifier = (classifierType: ClassifierType): boolean => - classifierType === "llm" || classifierType === "heuristic_first"; + classifierType === "llm" || classifierType === "heuristic_first" || classifierType === "hybrid"; export type ClassifierFallback = "heuristic" | "default_model"; @@ -161,7 +161,9 @@ export const heuristicScoringRoleFor = ( classifierType: ClassifierType, classifierFallback: ClassifierFallback | undefined, ): HeuristicScoringRole => { - if (classifierType === "heuristic" || classifierType === "heuristic_first") return "decides"; + if (classifierType === "heuristic_v2") return "never"; + if (classifierType === "heuristic" || classifierType === "heuristic_first" || classifierType === "hybrid") + return "decides"; return (classifierFallback ?? DEFAULT_CLASSIFIER_FALLBACK) === "heuristic" ? "fallback_only" : "never"; }; @@ -188,13 +190,19 @@ const builtInTierInfo = (rowId: string): { label: string; description: string; e return builtIn ? TIER_DESCRIPTIONS[builtIn] : undefined; }; +const tierConfigIntroText = (value: ComplexityRouterConfigValue): string => { + if (value.classifier_type === "heuristic_v2") { + return "The complexity router classifies each request with a calibrated local four-tier model (no API calls). Configure which model(s) handle each tier."; + } + if (heuristicScoringRole(value) === "never") { + return "The complexity router classifies each request with your classifier model and routes it to that tier. Configure which model(s) handle each tier."; + } + return "The complexity router automatically classifies requests by complexity using rule-based scoring (no API calls, <1ms latency). Configure which model(s) handle each tier."; +}; + const TierConfigIntro: React.FC<{ value: ComplexityRouterConfigValue }> = ({ value }) => ( <> - - {heuristicScoringRole(value) === "never" - ? "The complexity router classifies each request with your classifier model and routes it to that tier. Configure which model(s) handle each tier." - : "The complexity router automatically classifies requests by complexity using rule-based scoring (no API calls, <1ms latency). Configure which model(s) handle each tier."} - + {tierConfigIntroText(value)} {restrictedBy(value, "displayNames")?.reason ?? @@ -397,6 +405,8 @@ export interface ComplexityRouterConfigValue { classification_prompt?: string; /** Highest tier the scorer may decide alone under heuristic_first. Required by that type, rejected by the others. */ heuristic_first_max_tier?: string; + /** How near a tier boundary a score may land before hybrid defers to the classifier. Required by that type, rejected by the others. */ + hybrid_boundary_margin?: number; classification_mode?: ClassificationMode; session_affinity?: boolean; modality_routing?: boolean; @@ -509,6 +519,9 @@ export const effectiveTierLabel = (tier: keyof ComplexityTiers, tierLabels: Comp export const DEFAULT_HEURISTIC_FIRST_MAX_TIER = "SIMPLE"; +/** What the Hybrid radio starts at. Required by that type, so the form always has a value to send. */ +export const DEFAULT_HYBRID_BOUNDARY_MARGIN = 0.03; + /** * Tiers the heuristic_first threshold may name. The top tier is excluded because it would short * circuit every request and leave the classifier unreachable, which the backend rejects. diff --git a/ui/litellm-dashboard/src/components/add_model/HeuristicScoringConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/HeuristicScoringConfig.test.tsx index 9dccd767e49..329aef3f9d4 100644 --- a/ui/litellm-dashboard/src/components/add_model/HeuristicScoringConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/HeuristicScoringConfig.test.tsx @@ -165,6 +165,7 @@ describe("ClassificationMethodConfig scorer gating", () => { it.each([ ["heuristic decides the tier", "heuristic" as ClassifierType, undefined, true], + ["heuristic v2 decides without the weighted scorer", "heuristic_v2" as ClassifierType, undefined, false], ["an LLM classifier falls back to the heuristic", "llm" as ClassifierType, "heuristic" as ClassifierFallback, true], [ "an LLM classifier falls back to the default model", diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 03b7432e4ec..81acb7b7e50 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -342,6 +342,7 @@ const AddAutoRouterTab: React.FC = ({ planModeMinTier: complexityRouterConfig.plan_mode_min_tier, classificationPrompt: complexityRouterConfig.classification_prompt, heuristicFirstMaxTier: complexityRouterConfig.heuristic_first_max_tier, + hybridBoundaryMargin: complexityRouterConfig.hybrid_boundary_margin, classificationMode: complexityRouterConfig.classification_mode, tierLabels: complexityRouterConfig.tier_labels, classifierType: complexityRouterConfig.classifier_type, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index af87cc6c8fb..c6370541a45 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -111,6 +111,21 @@ describe("buildComplexityRouterConfig", () => { expect(config.classifier_llm_config).toBeUndefined(); }); + it("emits heuristic_v2 without classifier-only fields", () => { + const trainedParams: BuildComplexityRouterConfigParams = { + ...baseParams, + classifierType: "heuristic_v2", + classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 }, + classifierContextWindowSize: 5, + classifierFallback: "heuristic", + }; + const config = buildComplexityRouterConfig(trainedParams); + expect(config.classifier_type).toBe("heuristic_v2"); + expect(config.classifier_llm_config).toBeUndefined(); + expect(config.classifier_context_window_size).toBeUndefined(); + expect(config.classifier_fallback).toBeUndefined(); + }); + it("includes classifier_context_window_size and classifier_context_budget_chars only when classifier_type is llm", () => { const params: BuildComplexityRouterConfigParams = { ...baseParams, @@ -713,6 +728,10 @@ describe("getClassifierModelError", () => { expect(getClassifierModelError({ classifier_type: "heuristic" })).toBeNull(); }); + it("stays quiet for a heuristic v2 router, which runs locally", () => { + expect(getClassifierModelError({ classifier_type: "heuristic_v2" })).toBeNull(); + }); + it("blocks an LLM classifier with no model, which the router cannot start without", () => { expect(getClassifierModelError({ classifier_type: "llm" })).toBe( "Please select a classifier model, or switch back to Heuristic", @@ -791,13 +810,46 @@ describe("heuristic_first", () => { }); it("omits heuristic_first_max_tier on every other classifier type, which the backend rejects it on", () => { - for (const classifierType of ["heuristic", "llm"] as const) { + for (const classifierType of ["heuristic", "heuristic_v2", "llm"] as const) { const config = buildComplexityRouterConfig({ ...heuristicFirstParams, classifierType }); expect(config.heuristic_first_max_tier).toBeUndefined(); } }); }); +describe("hybrid", () => { + const hybridParams: BuildComplexityRouterConfigParams = { + ...baseParams, + classifierType: "hybrid", + hybridBoundaryMargin: 0.03, + classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 }, + classifierFallback: "default_model", + }; + + it("emits hybrid_boundary_margin, zero included since exactly-on-a-boundary is a real setting", () => { + expect(buildComplexityRouterConfig(hybridParams).hybrid_boundary_margin).toBe(0.03); + expect(buildComplexityRouterConfig({ ...hybridParams, hybridBoundaryMargin: 0 }).hybrid_boundary_margin).toBe(0); + }); + + it("keeps every classifier key the operator set, since hybrid still calls the classifier", () => { + const config = buildComplexityRouterConfig(hybridParams); + expect(config.classifier_type).toBe("hybrid"); + expect(config.classifier_llm_config).toEqual({ model: "gpt-4o-mini", timeout_ms: 3000 }); + expect(config.classifier_fallback).toBe("default_model"); + }); + + it("omits hybrid_boundary_margin on every other classifier type, which the backend rejects it on", () => { + for (const classifierType of ["heuristic", "llm", "heuristic_first"] as const) { + const config = buildComplexityRouterConfig({ + ...hybridParams, + classifierType, + ...(classifierType === "heuristic_first" && { heuristicFirstMaxTier: "SIMPLE" }), + }); + expect(config.hybrid_boundary_margin).toBeUndefined(); + } + }); +}); + describe("classification_mode", () => { it("emits user_turn", () => { const config = buildComplexityRouterConfig({ ...baseParams, classificationMode: "user_turn" }); @@ -905,10 +957,12 @@ describe("buildComplexityRouterConfig with an edited tier set", () => { dimensionWeights: { length: 1 }, reasoningOverrideMinScore: 0.5, heuristicFirstMaxTier: "SIMPLE", + hybridBoundaryMargin: 0.03, customTechnicalKeywords: ["kubernetes"], }; const emittingType = key === "heuristic_first_max_tier" ? "heuristic_first" : "llm"; - expect(buildComplexityRouterConfig({ ...baseParams, ...loaded, classifierType: emittingType })).toHaveProperty(key); + const typeForKey = key === "hybrid_boundary_margin" ? "hybrid" : emittingType; + expect(buildComplexityRouterConfig({ ...baseParams, ...loaded, classifierType: typeForKey })).toHaveProperty(key); expect(build(loaded)).not.toHaveProperty(key); }); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index af34dc92c0f..94c5badf6d1 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -107,6 +107,7 @@ export interface BuildComplexityRouterConfigParams { classifierFallback: ClassifierFallback | undefined; classificationPrompt: string | undefined; heuristicFirstMaxTier: string | undefined; + hybridBoundaryMargin?: number; classificationMode: ClassificationMode | undefined; sessionAffinity: boolean; modalityRouting?: boolean; @@ -163,6 +164,7 @@ export interface ComplexityRouterConfigPayload { classifier_fallback?: ClassifierFallback; classification_prompt?: string; heuristic_first_max_tier?: string; + hybrid_boundary_margin?: number; classification_mode: ClassificationMode; session_affinity: boolean; deployment_affinity: boolean; @@ -352,6 +354,7 @@ const classifierWireFields = ( classifierLlmConfig, classifierFallback, heuristicFirstMaxTier, + hybridBoundaryMargin, classifierContextWindowSize, classifierContextBudgetChars, classifierContextIncludeAssistantTurns, @@ -360,6 +363,7 @@ const classifierWireFields = ( | "classifierLlmConfig" | "classifierFallback" | "heuristicFirstMaxTier" + | "hybridBoundaryMargin" | "classifierContextWindowSize" | "classifierContextBudgetChars" | "classifierContextIncludeAssistantTurns" @@ -371,6 +375,8 @@ const classifierWireFields = ( classifierFallback !== undefined && { classifier_fallback: classifierFallback }), ...(effectiveType === "heuristic_first" && heuristicFirstMaxTier?.trim() && { heuristic_first_max_tier: heuristicFirstMaxTier }), + ...(effectiveType === "hybrid" && + hybridBoundaryMargin !== undefined && { hybrid_boundary_margin: hybridBoundaryMargin }), ...(usesLlmClassifier(effectiveType) && classifierContextWindowSize !== undefined && { classifier_context_window_size: classifierContextWindowSize, @@ -399,6 +405,7 @@ export const buildComplexityRouterConfig = ({ classifierFallback, classificationPrompt, heuristicFirstMaxTier, + hybridBoundaryMargin, classificationMode, sessionAffinity, modalityRouting, @@ -444,6 +451,7 @@ export const buildComplexityRouterConfig = ({ classifierLlmConfig, classifierFallback, heuristicFirstMaxTier, + hybridBoundaryMargin, classifierContextWindowSize, classifierContextBudgetChars, classifierContextIncludeAssistantTurns, diff --git a/ui/litellm-dashboard/src/components/add_model/tier_rows.ts b/ui/litellm-dashboard/src/components/add_model/tier_rows.ts index d2706a82037..3c5b149f4da 100644 --- a/ui/litellm-dashboard/src/components/add_model/tier_rows.ts +++ b/ui/litellm-dashboard/src/components/add_model/tier_rows.ts @@ -122,10 +122,10 @@ export const CUSTOM_TIER_RESTRICTIONS = { reason: "Session pinning escalates along the built-in tier ladder, which your tier set replaces", }, heuristicClassifier: { - omit: ["heuristic_first_max_tier"], + omit: ["heuristic_first_max_tier", "hybrid_boundary_margin"], reason: "The heuristic scorer only produces the built-in tiers, so an edited set needs the LLM classifier. " + - "Heuristic first is out for the same reason: its local scorer decides the cheap traffic", + "Heuristic first and hybrid are out for the same reason: their local scorer decides the traffic it is sure of", }, heuristicScoring: { omit: [ diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index 3e63cbd3b32..fae6744d3ac 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -520,16 +520,21 @@ describe("managed keys survive an untouched open-and-save", () => { }; // tier_definitions, fallback_tier and classification_prompt cannot sit beside heuristic_first, which - // this fixture uses, so no single stored config can hold every managed key. They get their own round - // trip below. - const CUSTOM_TIER_ONLY_KEYS = new Set(["tier_definitions", "fallback_tier", "classification_prompt"]); + // this fixture uses, and hybrid_boundary_margin belongs to the sibling hybrid type, so no single + // stored config can hold every managed key. Each gets its own round trip below. + const KEYS_ANOTHER_CLASSIFIER_TYPE_OWNS = new Set([ + "tier_definitions", + "fallback_tier", + "classification_prompt", + "hybrid_boundary_margin", + ]); it("carries every managed key a built-in router can hold through hydrate then save", () => { const hydrated = hydrateComplexityRouterConfig(STORED_ALL_MANAGED, undefined); const saved = buildUpdatedComplexityRouterConfig(STORED_ALL_MANAGED, hydrated); const dropped = [...MANAGED_COMPLEXITY_ROUTER_KEYS] - .filter((key) => !CUSTOM_TIER_ONLY_KEYS.has(key)) + .filter((key) => !KEYS_ANOTHER_CLASSIFIER_TYPE_OWNS.has(key)) .filter((key) => saved[key] === undefined); expect(dropped).toEqual([]); }); @@ -583,6 +588,18 @@ describe("managed keys survive an untouched open-and-save", () => { ); }); + it("round-trips a hybrid router's margin, which save requires and the backend rejects without", () => { + const storedHybrid: Record = { + ...STORED_ALL_MANAGED, + classifier_type: "hybrid", + hybrid_boundary_margin: 0.05, + }; + delete storedHybrid.heuristic_first_max_tier; + const hydrated = hydrateComplexityRouterConfig(storedHybrid, undefined); + expect(hydrated.hybrid_boundary_margin).toBe(0.05); + expect(buildUpdatedComplexityRouterConfig(storedHybrid, hydrated).hybrid_boundary_margin).toBe(0.05); + }); + it("round-trips the heuristic_first threshold, which save requires and the backend rejects without", () => { const hydrated = hydrateComplexityRouterConfig(STORED_ALL_MANAGED, undefined); expect(hydrated.heuristic_first_max_tier).toBe("SIMPLE"); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index e18582f77a0..7aa5f4beb1f 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -89,6 +89,7 @@ export interface StoredComplexityRouterConfig { plan_mode_min_tier?: unknown; classification_prompt?: unknown; heuristic_first_max_tier?: unknown; + hybrid_boundary_margin?: unknown; tier_labels?: unknown; classifier_type?: ClassifierType; classifier_llm_config?: ClassifierLLMConfig; @@ -167,6 +168,8 @@ export const hydrateComplexityRouterConfig = ( typeof parsedConfig.heuristic_first_max_tier === "string" && parsedConfig.heuristic_first_max_tier.trim() !== "" ? parsedConfig.heuristic_first_max_tier : undefined, + hybrid_boundary_margin: + typeof parsedConfig.hybrid_boundary_margin === "number" ? parsedConfig.hybrid_boundary_margin : undefined, classification_mode: parsedConfig.classification_mode === "user_turn" || parsedConfig.classification_mode === "every_request" ? parsedConfig.classification_mode @@ -214,6 +217,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "classifier_fallback", "classification_prompt", "heuristic_first_max_tier", + "hybrid_boundary_margin", "classification_mode", "session_affinity", "modality_routing", @@ -303,6 +307,7 @@ export const buildUpdatedComplexityRouterConfig = ( planModeMinTier: value.plan_mode_min_tier, classificationPrompt: value.classification_prompt, heuristicFirstMaxTier: value.heuristic_first_max_tier, + hybridBoundaryMargin: value.hybrid_boundary_margin, classificationMode: value.classification_mode, tierLabels: value.tier_labels, classifierType: value.classifier_type, diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx index aa1d45a859e..ffef01010b7 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx @@ -84,7 +84,9 @@ function describeReasoningOverride(tierLabel: string | undefined, floor: number const CONSTANT_CAUSE_LABELS: Record = { heuristic_scorer: "Heuristic scorer", + heuristic_v2: "Heuristic v2", heuristic_first_short_circuit: "Heuristic scorer, classifier skipped", + hybrid_short_circuit: "Heuristic scorer, score clear of every boundary", classifier_plugin: "Custom classifier plugin", semantic_keyword_match: "Semantic keyword match", session_affinity_pin: "Pinned to session", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 522a2de574b..be81450b941 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29322,6 +29322,8 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Rpm */ rpm?: number | null; + /** Rust */ + rust?: boolean | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; /** S3 Encryption Key Id */ @@ -34558,7 +34560,7 @@ export interface components { * @enum {string} */ classifier_fallback: "heuristic" | "default_model"; - /** @description Configuration for the LLM classifier; required when classifier_type is 'llm' or 'heuristic_first' */ + /** @description Configuration for the LLM classifier; required when classifier_type is 'llm', 'heuristic_first' or 'hybrid' */ classifier_llm_config?: components["schemas"]["ClassifierLLMConfig"] | null; /** * Classifier Plugin @@ -34573,11 +34575,11 @@ export interface components { classifier_plugin_timeout_ms: number; /** * Classifier Type - * @description Classification strategy: local regex/keyword scoring, an LLM call, a custom classifier plugin, or 'heuristic_first', which scores locally and only pays for the LLM classifier when the local scorer does not confidently land a cheap tier + * @description Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, an LLM call, a custom classifier plugin, 'heuristic_first', which scores locally and only pays for the LLM classifier when the local scorer does not confidently land a cheap tier, or 'hybrid', which trusts the local scorer everywhere except when its score lands near a tier boundary * @default heuristic * @enum {string} */ - classifier_type: "heuristic" | "llm" | "custom" | "heuristic_first"; + classifier_type: "heuristic" | "heuristic_v2" | "llm" | "custom" | "heuristic_first" | "hybrid"; /** * Code Keywords * @description Keywords indicating code-related content @@ -34638,11 +34640,22 @@ export interface components { * @description The highest tier the local scorer may decide on its own; required when classifier_type is 'heuristic_first' and rejected otherwise. A request whose heuristic tier is at or below this one skips the LLM classifier and routes straight to that heuristic tier, so the classifier call is only paid for on traffic the scorer could not place cheaply. The scorer must also have produced at least one signal: a prompt where no dimension fired scores 0.0 and would otherwise land SIMPLE by default rather than by evidence, which is how a chained router would silently send unclassified traffic to the cheapest model. Names a built-in tier, and may not name the highest one, since that would make the LLM classifier unreachable. */ heuristic_first_max_tier?: string | null; + /** + * Heuristic V2 Artifact + * @description Success-probability artifact used by classifier_type 'heuristic_v2'. The bundled UltraFeedback artifact is selected by default; an inline trained artifact may replace it + * @default ultrafeedback + */ + heuristic_v2_artifact: components["schemas"]["TrainedTierArtifact"] | "ultrafeedback"; /** * Housekeeping Patterns * @description Additional case-sensitive literal sentinels that mark a request as client housekeeping, on top of the built-in conversation-title ones. For clients whose wording the built-ins don't cover, or after a client release changes its strings. */ housekeeping_patterns?: string[] | null; + /** + * Hybrid Boundary Margin + * @description How close to a tier boundary a heuristic score has to land before the LLM classifier breaks the tie; required when classifier_type is 'hybrid' and rejected otherwise. Everything further than this from every active boundary routes on the scorer's own tier with no classifier call, at any tier, which is what separates 'hybrid' from 'heuristic_first' and its cheap-tier ceiling. A prompt where no dimension fired still goes to the classifier, since the scorer has no opinion to be near a boundary with. 0 escalates only scores sitting exactly on a boundary. + */ + hybrid_boundary_margin?: number | null; /** * Keyword Tier Rules * @description Rules that force a specific tier when their keywords match the prompt @@ -34776,6 +34789,12 @@ export interface components { } & { [key: string]: unknown; }; + /** + * RequestType + * @description Fixed v0 taxonomy. User-extensible types come in v1. + * @enum {string} + */ + RequestType: "code_generation" | "code_understanding" | "technical_design" | "analytical_reasoning" | "writing" | "factual_lookup" | "general"; /** ResetSpendRequest */ ResetSpendRequest: { /** Reset To */ @@ -35853,7 +35872,7 @@ export interface components { * Cause * @enum {string} */ - cause?: "heuristic_scorer" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; + cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "hybrid_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; /** Classifier Cost */ classifier_cost?: number; /** Classifier Model */ @@ -36736,6 +36755,33 @@ export interface components { [key: string]: unknown; }; }; + /** TierCohortStatistic */ + TierCohortStatistic: { + /** Cohort */ + cohort: string; + /** Observations */ + observations: number; + /** Successes */ + successes: number; + /** Tier */ + tier: number; + }; + /** TierDataset */ + TierDataset: { + /** License */ + license: string; + /** Name */ + name: string; + /** Rows */ + rows: number; + /** + * Success Definition + * @default quality score meets the dataset success threshold + */ + success_definition: string; + /** Url */ + url: string; + }; /** * TierDefinition * @description An operator-defined tier: the name the LLM classifier must return and its rubric description. @@ -36752,6 +36798,25 @@ export interface components { */ name: string; }; + /** TierDomainStatistic */ + TierDomainStatistic: { + /** Observations */ + observations: number; + request_type: components["schemas"]["RequestType"]; + /** Successes */ + successes: number; + /** Tier */ + tier: number; + }; + /** TierGlobalStatistic */ + TierGlobalStatistic: { + /** Observations */ + observations: number; + /** Successes */ + successes: number; + /** Tier */ + tier: number; + }; /** * TokenCountDetailsResponse * @description Response structure for token count details with modality breakdown. @@ -37021,6 +37086,57 @@ export interface components { } & { [key: string]: unknown; }; + /** TrainedTierArtifact */ + TrainedTierArtifact: { + /** + * Cohort Prior Mass + * @default 20 + */ + cohort_prior_mass: number; + /** + * Cohort Statistics + * @default [] + */ + cohort_statistics: components["schemas"]["TierCohortStatistic"][]; + /** + * Datasets + * @default [] + */ + datasets: components["schemas"]["TierDataset"][]; + /** + * Domain Prior Mass + * @default 200 + */ + domain_prior_mass: number; + /** + * Domain Statistics + * @default [] + */ + domain_statistics: components["schemas"]["TierDomainStatistic"][]; + /** Global Statistics */ + global_statistics: components["schemas"]["TierGlobalStatistic"][]; + /** + * Routing Threshold + * @default 0.75 + */ + routing_threshold: number; + /** + * Schema Version + * @default 1 + * @constant + */ + schema_version: 1; + /** + * Split Method + * @default sha256(prompt): 70% train, 15% validation, 15% test + */ + split_method: string; + /** + * Success Definition + * @default quality score meets the dataset success threshold + */ + success_definition: string; + }; /** TransformRequestBody */ TransformRequestBody: { call_type: components["schemas"]["CallTypes"]; @@ -39196,6 +39312,8 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Rpm */ rpm?: number | null; + /** Rust */ + rust?: boolean | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; /** S3 Encryption Key Id */