diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index c112bf2bb22..d02f5878396 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -128,6 +128,9 @@ jobs: - name: check_fastuuid_usage run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py + - name: check_py310_typing_imports + run: uv run --no-sync python ./tests/code_coverage_tests/check_py310_typing_imports.py + - name: check_e2e_no_raw_requests run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py @@ -145,3 +148,33 @@ jobs: - name: documentation_test_api_docs run: uv run --no-sync python ./tests/documentation_tests/test_api_docs.py + + python-310-import-smoke: + runs-on: ubuntu-latest + timeout-minutes: 15 + + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.10" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Install dependencies + run: uv sync --frozen --extra proxy --python 3.10 + + - run: uv run --no-sync python --version + + - name: Import litellm + run: uv run --no-sync python -c "import litellm" + + - name: Check litellm CLI + run: uv run --no-sync litellm --version diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index aa2cb96a9f7..3f96531cf6f 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 4125 + "limit": 4124 }, "reportFunctionMemberAccess": { "limit": 7 @@ -99,7 +99,7 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44364 + "limit": 44362 }, "reportUnknownLambdaType": { "limit": 109 @@ -117,13 +117,13 @@ "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/__init__.py b/litellm/__init__.py index 44f2e7c1f02..41a3789ab0d 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -424,6 +424,10 @@ anthropic_beta_headers_url: str = os.getenv( "LITELLM_ANTHROPIC_BETA_HEADERS_URL", "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/anthropic_beta_headers_config.json", ) +autorouter_presets_url: str = os.getenv( + "LITELLM_AUTOROUTER_PRESETS_URL", + "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/proxy/public_endpoints/autorouter_presets.json", +) suppress_debug_info: bool = False dynamodb_table_name: Optional[str] = None s3_callback_params: Optional[Dict] = None diff --git a/litellm/_logging.py b/litellm/_logging.py index 14cda772234..c73b5175a31 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -17,7 +17,11 @@ from litellm.constants import ( from litellm.litellm_core_utils.env_utils import get_env_int from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.litellm_core_utils.secret_redaction import redact_string, redact_structured_value +from litellm.litellm_core_utils.secret_redaction import ( + redact_internal_details, + redact_string, + redact_structured_value, +) set_verbose = False @@ -89,6 +93,14 @@ def redact_secrets(value: str) -> str: return _redact_string(value) +def redact_internal_details_from_client_message(value: str) -> str: + """Public API: redact_secrets() plus filesystem paths, internal hostnames, and an + embedded traceback, for a string about to leave the process in an HTTP response.""" + if not _ENABLE_SECRET_REDACTION: + return value + return redact_internal_details(value) + + def _substituted_color_message(record: logging.LogRecord) -> str | None: """Render a record's ``color_message`` against its args, or None if absent. diff --git a/litellm/constants.py b/litellm/constants.py index c7b74e176db..ef9329b9dfc 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1449,6 +1449,7 @@ RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl" CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" +SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( "Truncation is a DB storage safeguard. " 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 2d737bc34e7..587da997f94 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -10,9 +10,9 @@ import asyncio import math import uuid from collections.abc import AsyncIterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast -from typing_extensions import ReadOnly +from typing_extensions import Never, ReadOnly import litellm from litellm._logging import verbose_logger diff --git a/litellm/litellm_core_utils/secret_redaction.py b/litellm/litellm_core_utils/secret_redaction.py index e93ab155786..b62226a6a19 100644 --- a/litellm/litellm_core_utils/secret_redaction.py +++ b/litellm/litellm_core_utils/secret_redaction.py @@ -92,6 +92,27 @@ def redact_string(value: str) -> str: return _SECRET_RE.sub(_REDACTED, value) +_UNIX_SYSTEM_PATH: Final = r"/(?:etc|var|opt|usr|home|root|private|Users|tmp|mnt|srv)/[^\s'\"\)\]}>,]+" +_WINDOWS_DRIVE_PATH: Final = r"[A-Za-z]:\\[^\s'\"\)\]}>,]+" +_PRIVATE_OR_LOOPBACK_IPV4: Final = ( + r"\b(?:10(?:\.\d{1,3}){3}|172\.(?:1[6-9]|2\d|3[01])(?:\.\d{1,3}){2}|192\.168(?:\.\d{1,3}){2}|127(?:\.\d{1,3}){3})\b" +) +_INTERNAL_SUFFIX_HOSTNAME: Final = r"\b[A-Za-z0-9-]+(?:\.[A-Za-z0-9-]+)*\.(?:internal|local|corp|lan|intra|private)\b" +_INTERNAL_DETAIL_RE: Final = re.compile( + "|".join((_UNIX_SYSTEM_PATH, _WINDOWS_DRIVE_PATH, _PRIVATE_OR_LOOPBACK_IPV4, _INTERNAL_SUFFIX_HOSTNAME)), + re.IGNORECASE, +) +_TRACEBACK_MARKER: Final = "Traceback (most recent call last):" + + +def redact_internal_details(value: str) -> str: + """Drop an embedded traceback and scrub filesystem paths and internal hostnames, + on top of redact_string(). For client-facing messages only: server logs keep this detail.""" + marker_index: Final = value.find(_TRACEBACK_MARKER) + without_traceback: Final = value[:marker_index].rstrip() if marker_index != -1 else value + return _INTERNAL_DETAIL_RE.sub(_REDACTED, redact_string(without_traceback)) + + def redact_structured_value(key: str | None, value: str) -> str: """Scrub *value* as it appeared under *key* inside a structured record. diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index f1c7451796d..5c30ff4747a 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -17,7 +17,10 @@ from typing import TYPE_CHECKING, Any, Final, Optional from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -313,9 +316,14 @@ class A2AGuardrailHandler(BaseTranslation): return responses_so_far + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + _, valid_parsed = self._parse_streaming_responses(responses_so_far) + combined_text, _ = self._collect_text_from_parsed_chunks(valid_parsed) + return StreamingScanKey(texts=(combined_text,)) + def _parse_streaming_responses( self, - responses_so_far: list[object], + responses_so_far: Sequence[object], ) -> tuple[list[dict[str, object] | None], list[tuple[int, dict[str, object]]]]: """Parse JSON-RPC items, returning aligned parsed list and valid entries.""" parsed: Final[list[dict[str, object] | None]] = [None] * len(responses_so_far) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index c7d12e5cf3a..c23797f72af 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -26,7 +26,10 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im LiteLLMAnthropicMessagesAdapter, is_provider_native_tool_dict, ) -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.llms.base_llm.guardrail_translation.utils import ( anthropic_tool_name, anthropic_tool_names, @@ -36,6 +39,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, scoped_structured_message_indices, + stream_item_fingerprint, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -1176,6 +1180,25 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["model"] = response_model return inputs + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + stream_ended: Final = self._check_streaming_has_ended(responses_so_far) + return StreamingScanKey( + texts=(self.get_streaming_string_so_far(responses_so_far),), + tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (), + stream_ended=stream_ended, + ) + + @classmethod + def _streamed_tool_use_fingerprints(cls, responses_so_far: Sequence[object]) -> tuple[str, ...]: + return tuple( + stream_item_fingerprint(block) + for item in responses_so_far + for event in cls._iter_sse_events(item) + if event.get("type") == "content_block_start" + and isinstance(block := event.get("content_block"), Mapping) + and block.get("type") == "tool_use" + ) + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Parse streaming responses and extract accumulated text content. 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/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 220fcedb0f8..b28daf73bc4 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -35,6 +35,22 @@ class StreamTransformSink: holdback_per_choice: dict[int, int] = field(default_factory=dict) +@dataclass(frozen=True, slots=True) +class StreamingScanKey: + """What a streaming guardrail round would hand to ``apply_guardrail``. Two keys + compare equal when the round would scan the same content again; ``stream_ended`` + stays out of the comparison and only says whether the handler is on its + end-of-stream path, where an empty payload is still scanned today.""" + + texts: tuple[str, ...] + tool_calls: tuple[str, ...] = () + stream_ended: bool = field(default=False, compare=False) + + @property + def has_nothing_to_scan(self) -> bool: + return not self.stream_ended and not any(self.texts) and not self.tool_calls + + class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( @@ -151,6 +167,9 @@ class BaseTranslation(ABC): """ return responses_so_far + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + return None + def build_block_sse_chunks( self, exc: "ModifyResponseException", diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 9b6f9c47105..8dee262001d 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -4,6 +4,8 @@ import json from collections.abc import Callable, Iterator, Sequence from typing import Any, Final, TypeVar +from pydantic import BaseModel + from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage @@ -130,6 +132,16 @@ def stream_item_field(item: object, field: str) -> object | None: return getattr(item, field, None) +def stream_item_fingerprint(item: object) -> str: + plain: Final = item.model_dump() if isinstance(item, BaseModel) else item + return json.dumps(plain, sort_keys=True, default=str) + + +def stream_item_items(item: object, field: str) -> tuple[object, ...]: + value: Final = stream_item_field(item, field) + return tuple(value) if isinstance(value, (list, tuple)) else () + + def blocked_chat_stream_usage(original_response: object) -> tuple[int, int]: """ ``(prompt_tokens, completion_tokens)`` for a synthetic guardrail-blocked 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/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 3cd6ee54069..0f6966b0ae2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -28,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 @@ -68,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, @@ -274,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 @@ -538,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, @@ -9687,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, @@ -9707,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, ) @@ -9730,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, @@ -9744,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, @@ -9804,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, @@ -9820,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, @@ -9840,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, ) @@ -9860,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/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index 8c306faa036..ac934ad0cb5 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -2,6 +2,7 @@ from typing import Final from httpx import Headers +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -16,16 +17,18 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None: """ Session id to send as `x-session-affinity`, or None when the caller gave none. - Deliberately does not fall back to `litellm_trace_id`: that is generated per - request (`str(uuid.uuid4())` when absent), so using it pins every request to a - different Fireworks node and prompt caching never hits. + Deliberately does not fall back to `litellm_trace_id`, and ignores session ids the + proxy generated for a request that had none: both are per request, so using them + pins every request to a different Fireworks node and prompt caching never hits. """ params: Final = litellm_params + metadata: Final = params.get("metadata") + if isinstance(metadata, dict) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY): + return None for key in ("litellm_session_id", "session_id"): value = params.get(key) if value: return str(value) - metadata: Final = params.get("metadata") if isinstance(metadata, dict): value = metadata.get("session_id") if value: 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/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 96a5ed663fc..d41c8557d72 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -26,6 +26,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, + StreamingScanKey, StreamTransformSink, ) from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -39,6 +40,8 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( role_out_of_guardrail_scope, scoped_structured_message_indices, stream_item_field, + stream_item_fingerprint, + stream_item_items, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -503,12 +506,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """Block-only streaming path: run the guardrail so an in-flight BLOCK can terminate the stream. Text rewrites are not propagated to the client here (see ``_process_streaming_transform`` for the incremental_diff path).""" - # check if the stream has ended - has_stream_ended = False - for chunk in responses_so_far: - if chunk.choices and chunk.choices[0].finish_reason is not None: - has_stream_ended = True - break + has_stream_ended: Final = self._first_choice_has_finished(responses_so_far) if has_stream_ended: # convert to model response @@ -706,8 +704,33 @@ class OpenAIChatCompletionsHandler(BaseTranslation): indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback) } + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream)) + stream_ended: Final = self._first_choice_has_finished(responses_so_far) + return StreamingScanKey( + texts=tuple(self._combine_streaming_texts(chunks).values()), + tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (), + stream_ended=stream_ended, + ) + + @staticmethod + def _streamed_tool_call_fingerprints(responses_so_far: Sequence[object]) -> tuple[str, ...]: + return tuple( + stream_item_fingerprint(tool_call) + for chunk in responses_so_far + for choice in _stream_chunk_choices(chunk) + for tool_call in stream_item_items(stream_item_field(choice, "delta"), "tool_calls") + ) + + @staticmethod + def _first_choice_has_finished(responses_so_far: Sequence[object]) -> bool: + first_choices: Final = tuple( + choices[0] for choices in (_stream_chunk_choices(chunk) for chunk in responses_so_far) if choices + ) + return any(stream_item_field(choice, "finish_reason") is not None for choice in first_choices) + def _combine_streaming_texts( - self, responses_so_far: list["ModelResponseStream"] + self, responses_so_far: Sequence["ModelResponseStream"] ) -> dict[tuple[int, int | None], str]: """ Combine all streaming chunks into complete text per choice. 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..a0db7aadb9e 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 @@ -44,11 +44,17 @@ from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( OpenAiResponsesToChatCompletionStreamIterator, ) -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, stream_item_field, + stream_item_fingerprint, + stream_item_items, ) +from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -62,7 +68,6 @@ from litellm.types.llms.openai import ( ContentPartDonePartOutputText, ErrorEvent, ErrorEventError, - OpenAIMcpServerTool, OutputItemAddedEvent, OutputItemDoneEvent, OutputTextDeltaEvent, @@ -157,23 +162,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 +202,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 +215,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 +226,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 +248,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 +275,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, @@ -645,18 +598,55 @@ class OpenAIResponsesHandler(BaseTranslation): ) return responses_so_far - def _check_streaming_has_ended(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> bool: + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: """ Check if the streaming has ended. """ if not responses_so_far: return False - terminal_types: Final = { - ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, - ResponsesAPIStreamEvents.RESPONSE_FAILED.value, - ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, - } - return responses_so_far[-1].get("type") in terminal_types + terminal_types: Final = frozenset( + ( + ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, + ResponsesAPIStreamEvents.RESPONSE_FAILED.value, + ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, + ) + ) + return stream_item_field(responses_so_far[-1], "type") in terminal_types + + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + if not responses_so_far or not hasattr(responses_so_far[-1], "get"): + return None + last_event: Final = responses_so_far[-1] + last_event_type: Final = stream_item_field(last_event, "type") + if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value: + return None + if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value: + return self._completed_response_scan_key(stream_item_field(last_event, "response")) + return StreamingScanKey( + texts=(self.get_streaming_string_so_far(responses_so_far),), + stream_ended=self._check_streaming_has_ended(responses_so_far), + ) + + @staticmethod + def _completed_response_scan_key(response: object) -> StreamingScanKey: + output_items: Final = stream_item_items(response, "output") + message_items: Final = tuple( + item for item in output_items if stream_item_field(item, "type") != "function_call" + ) + return StreamingScanKey( + texts=tuple( + text + for item in message_items + for part in stream_item_items(item, "content") + if isinstance(text := stream_item_field(part, "text"), str) and text + ), + tool_calls=tuple( + stream_item_fingerprint(item) + for item in output_items + if stream_item_field(item, "type") == "function_call" + ), + stream_ended=True, + ) def build_stream_error_items( self, @@ -681,7 +671,7 @@ class OpenAIResponsesHandler(BaseTranslation): ), ) - def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str: + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Get the string so far from the responses so far. @@ -693,12 +683,16 @@ class OpenAIResponsesHandler(BaseTranslation): """ keyed_events: Final = tuple( ( - (event.get("item_id"), event.get("output_index"), event.get("content_index")), - event.get("text"), - event.get("delta"), + ( + stream_item_field(event, "item_id"), + stream_item_field(event, "output_index"), + stream_item_field(event, "content_index"), + ), + stream_item_field(event, "text"), + stream_item_field(event, "delta"), ) for event in responses_so_far - if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str) + if isinstance(stream_item_field(event, "text"), str) or isinstance(stream_item_field(event, "delta"), str) ) def part_text(part_key: tuple[object, object, object]) -> str: 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/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/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py index 84f714db449..6f1e5baa109 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py @@ -121,6 +121,23 @@ class TokenEndpointClient: return Ok(ExchangedToken(access_token=parsed.access_token, expires_in=parsed.expires_in)) +class _KeyGuard: + """The per-key single-flight lock plus the invalidation generation that lock protects. + + Both live on one object so their lifetimes cannot diverge. `get_or_compute` binds the guard to + a local for its whole critical section, which keeps the weak map's entry alive for as long as + that compute could still write; an `invalidate` overlapping the compute therefore reaches the + very same object and its bump is guaranteed to be observed. Conversely a guard nobody holds is + collectible precisely because no write is outstanding for it to fence. + """ + + __slots__ = ("__weakref__", "generation", "lock") + + def __init__(self) -> None: + self.lock = asyncio.Lock() + self.generation = 0 + + class ExchangedTokenCache: """Memoizes the final token string per key, single-flighting concurrent misses on one lock.""" @@ -129,7 +146,7 @@ class ExchangedTokenCache: max_size_in_memory=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE, default_ttl=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, ) - self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() + self._guards: weakref.WeakValueDictionary[str, _KeyGuard] = weakref.WeakValueDictionary() async def get_or_compute( self, @@ -144,28 +161,50 @@ class ExchangedTokenCache: guaranteeing the token it gets back was minted for the *current* inputs: a stored entry whose fingerprint differs reads as a miss and is re-minted over. That keeps eviction addressable without the key having to encode the credential material it protects. + + An `invalidate` landing while `compute` is in flight wins over that compute's write. The + token is still returned to the caller it was minted for, but it is not stored, so the next + resolution re-mints rather than serving a bearer that predates the invalidation for the + rest of its TTL. """ cached = self._get(cache_key, fingerprint) if cached is not None: return Ok(cached) - async with self._lock(cache_key): + guard = self._guard(cache_key) + async with guard.lock: cached = self._get(cache_key, fingerprint) if cached is not None: return Ok(cached) + generation = guard.generation match await compute(): case Ok(token): - self._cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped - cache_key, - (fingerprint, token.access_token), - ttl=_cache_ttl_seconds(token.expires_in), - ) + if guard.generation == generation: + self._store(cache_key, fingerprint, token) return Ok(token.access_token) case Error(err): return Error(err) def invalidate(self, cache_key: str) -> None: - """Evict one cached token so the next `get_or_compute` re-mints (e.g. after an upstream 401).""" + """Evict one cached token so the next `get_or_compute` re-mints (e.g. after an upstream 401). + + Bumping the guard's generation is what makes the eviction stick against a compute already + awaiting the token endpoint: that compute snapshotted the old generation and so skips its + write. No guard means no compute is in flight, since an in-flight one pins its own. + + Stays synchronous: callers invalidate from plain `def`s. + """ self._cache.delete_cache(cache_key) # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + guard = self._guards.get(cache_key) + if guard is None: + return + guard.generation += 1 + + def _store(self, cache_key: str, fingerprint: str, token: ExchangedToken) -> None: + self._cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + cache_key, + (fingerprint, token.access_token), + ttl=_cache_ttl_seconds(token.expires_in), + ) def _get(self, cache_key: str, fingerprint: str) -> str | None: """The stored token, or None when absent or minted for different inputs. @@ -180,12 +219,12 @@ class ExchangedTokenCache: return None return token if stored_fingerprint == fingerprint else None - def _lock(self, cache_key: str) -> asyncio.Lock: - lock = self._locks.get(cache_key) - if lock is None: - lock = asyncio.Lock() - self._locks[cache_key] = lock - return lock + def _guard(self, cache_key: str) -> _KeyGuard: + guard = self._guards.get(cache_key) + if guard is None: + guard = _KeyGuard() + self._guards[cache_key] = guard + return guard def _cache_ttl_seconds(expires_in: int | None) -> int: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 989b08b929a..af3ff6714af 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3855,6 +3855,13 @@ if MCP_AVAILABLE: and server.auth_type == MCPAuth.oauth2_token_exchange and oauth2_headers and len(mcp_servers or []) == 1 + and server.server_id + in frozenset( + allowed.server_id + for allowed in await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip + ) + ) ): await global_mcp_server_manager.preflight_token_exchange( server=server, diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 4f6305d88cf..af02c11ad86 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -5,10 +5,10 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict, assert_never +from typing import TYPE_CHECKING, Any, Final, TypedDict from pydantic import ValidationError -from typing_extensions import ReadOnly, Required +from typing_extensions import ReadOnly, Required, assert_never import litellm from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2da7ceb2d50..849e54c65aa 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2594,6 +2594,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.", ) + missing_session_id: Literal["generate", "reject"] | None = Field( + None, + description="What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.", + ) enable_public_model_hub: bool = Field( default=False, description="Public model hub for users to see what models they have access to, supported openai params, etc.", diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 3e4dc07a521..3b8151d1064 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -13,10 +13,10 @@ import os import uuid from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Annotated, Final, TypedDict, assert_never +from typing import Annotated, Final, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query, Request -from typing_extensions import ReadOnly, Required +from typing_extensions import ReadOnly, Required, assert_never import litellm from litellm._logging import verbose_proxy_logger diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 2fad9f933c1..a96d3fb9c85 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -100,9 +100,10 @@ class CliPollData(TypedDict, total=False): class CliSsoStartData(TypedDict): - login_id: str - poll_secret: str - user_code: str + login_id: ReadOnly[str] + poll_secret: ReadOnly[str] + user_code: ReadOnly[str] + verification_uri_complete: ReadOnly[NotRequired[str]] class CliAuthResult(TypedDict): @@ -860,11 +861,22 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None: poll_secret: Final = cli_sso_flow["poll_secret"] user_code: Final = cli_sso_flow["user_code"] - sso_url = f"{base_url}/sso/key/generate?" + urlencode({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": key_id}) + browser_prefills_code: Final = isinstance(cli_sso_flow.get("verification_uri_complete"), str) + sso_url: Final = f"{base_url}/sso/key/generate?" + urlencode( + ( + ("source", LITELLM_CLI_SOURCE_IDENTIFIER), + ("key", key_id), + *((("user_code", user_code),) if browser_prefills_code else ()), + ) + ) click.echo(f"Opening browser to: {sso_url}") click.echo("Please complete the SSO authentication in your browser...") - click.echo(f"Verification code: {user_code}") + click.echo( + f"Verification code: {user_code} (pre-filled in the browser, check it matches)" + if browser_prefills_code + else f"Verification code: {user_code}" + ) click.echo(f"Session ID: {key_id}") # Open browser diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 05ddef822f1..fc83c1ddeed 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3,7 +3,6 @@ import contextlib import json import logging import math -import traceback from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence from datetime import datetime from functools import lru_cache @@ -18,7 +17,7 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse from starlette.types import Receive, Scope, Send import litellm -from litellm._logging import _redact_string, verbose_proxy_logger +from litellm._logging import redact_internal_details_from_client_message, verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE, @@ -3417,7 +3416,7 @@ class ProxyBaseLLMRequestProcessing: else: _code = status.HTTP_500_INTERNAL_SERVER_ERROR raise ProxyException( - message=getattr(e, "message", error_msg), + message=redact_internal_details_from_client_message(getattr(e, "message", error_msg)), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), openai_code=getattr(e, "code", None), @@ -3629,10 +3628,8 @@ class ProxyBaseLLMRequestProcessing: if isinstance(e, HTTPException): raise e - error_traceback: Final = _redact_string(traceback.format_exc()) - error_msg: Final = f"{e}\n\n{error_traceback}" proxy_exception: Final = ProxyException( - message=getattr(e, "message", error_msg), + message=redact_internal_details_from_client_message(getattr(e, "message", str(e))), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), code=getattr(e, "status_code", 500), diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 1682cf12f4e..47f69732e95 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -7,7 +7,9 @@ from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from enum import Enum from types import MappingProxyType -from typing import Final, Literal, Protocol, TypeVar, assert_never +from typing import Final, Literal, Protocol, TypeVar + +from typing_extensions import assert_never import litellm from litellm._logging import verbose_proxy_logger diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 2b04828f0f2..3d2ed641a30 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -514,12 +514,26 @@ async def update_guardrail( guardrail_name: Final = result.get("guardrail_name", "Unknown") try: - IN_MEMORY_GUARDRAIL_HANDLER.update_in_memory_guardrail( - guardrail_id=guardrail_id, guardrail=cast(Guardrail, result) - ) + IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=cast(Guardrail, result)) verbose_proxy_logger.info( "Immediate sync: Successfully updated guardrail '%s' (ID: %s)", guardrail_name, guardrail_id ) + except (ValueError, TypeError) as update_error: + # The new config is invalid (a raising guardrail __init__): + # reinitialize_guardrail already restored the previous live instance, but + # update_guardrail_in_db above already persisted the rejected config to + # the DB. Roll that back too, so the DB and the live guardrail never + # disagree about what's actually enforcing, and surface the rejection to + # the caller instead of a misleading 200. + await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=existing_guardrail, + prisma_client=prisma_client, + ) + raise HTTPException( + status_code=422, + detail=f"Invalid guardrail configuration, update rejected: {update_error}", + ) from update_error except Exception as update_error: verbose_proxy_logger.warning( "Immediate sync: Failed to update '%s' (ID: %s) in memory: %s", diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py index 5e75b7d4d94..c88e6e97a96 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py @@ -1,8 +1,8 @@ from typing import TYPE_CHECKING, Final -from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations +from litellm.types.guardrails import SupportedGuardrailIntegrations -from .crowdstrike_aidr import CrowdStrikeAIDRHandler +from .crowdstrike_aidr import CrowdStrikeAIDRHandler, streaming_params_from_litellm_params if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams @@ -15,17 +15,16 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" if not guardrail_name: raise ValueError("CrowdStrike AIDR guardrail name is required") + streaming_params: Final = streaming_params_from_litellm_params(litellm_params) _crowdstrike_aidr_callback: Final = CrowdStrikeAIDRHandler( guardrail_name=guardrail_name, api_base=litellm_params.api_base, api_key=litellm_params.api_key, - # Exclude during_call to prevent duplicate input events - event_hook=[ - GuardrailEventHooks.pre_call.value, - GuardrailEventHooks.post_call.value, - ], + event_hook=litellm_params.mode, default_on=litellm_params.default_on, fail_on_error=litellm_params.fail_on_error, + streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only, + streaming_sampling_rate=streaming_params.streaming_sampling_rate, ) litellm.logging_callback_manager.add_litellm_callback(_crowdstrike_aidr_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index c8284fac440..f7f500b1adc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -24,8 +24,11 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionToolParam +from litellm.types.proxy.guardrails.guardrail_hooks.crowdstrike_aidr import ( + CrowdStrikeAIDRGuardrailConfigModelOptionalParams, +) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -153,6 +156,21 @@ def _merge_metadata_bags(request_data: Mapping[str, Any]) -> Mapping[str, Any] | return merged if present else None +def streaming_params_from_litellm_params( + litellm_params: LitellmParams, +) -> CrowdStrikeAIDRGuardrailConfigModelOptionalParams: + extras: Final[Mapping[str, object]] = litellm_params.model_extra or {} + nested: Final = litellm_params.optional_params + optional_params: Final[Mapping[str, object]] = {} if nested is None else nested.model_dump() + return CrowdStrikeAIDRGuardrailConfigModelOptionalParams.model_validate( + { + name: value + for name in CrowdStrikeAIDRGuardrailConfigModelOptionalParams.model_fields + if (value := optional_params.get(name, extras.get(name))) is not None + } + ) + + def _messages_since_last_assistant( messages: Sequence[AllMessageValues], ) -> _FilteredMessages: @@ -241,6 +259,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): api_key: str | None = None, api_base: str | None = None, fail_on_error: bool | None = True, + streaming_end_of_stream_only: bool | None = None, + streaming_sampling_rate: int | None = None, **kwargs, ) -> None: """ @@ -250,10 +270,19 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): guardrail_name (str): The name of the guardrail instance. api_key (str | None): The CrowdStrike AIDR API key. Reads from CS_AIDR_TOKEN env var if None. api_base (str | None): The CrowdStrike AIDR API base URL. Reads from CS_AIDR_BASE_URL env var if None. + streaming_end_of_stream_only (bool | None): Scan streamed output once at end of stream instead of + every streaming_sampling_rate chunks. Defaults to False. + streaming_sampling_rate (int | None): Scan the accumulated streamed output every Nth chunk. Defaults to 5. **kwargs: Additional arguments passed to the CustomGuardrail base class. """ self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.fail_on_error = True if fail_on_error is None else fail_on_error + self._set_streaming_params( + CrowdStrikeAIDRGuardrailConfigModelOptionalParams( + streaming_end_of_stream_only=streaming_end_of_stream_only, + streaming_sampling_rate=streaming_sampling_rate, + ) + ) self.api_key = api_key or os.environ.get("CS_AIDR_TOKEN") if not self.api_key: @@ -274,6 +303,15 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): "Initialized CrowdStrike AIDR Guardrail: name=%s, api_base=%s", guardrail_name, self.api_base ) + def _set_streaming_params(self, streaming_params: CrowdStrikeAIDRGuardrailConfigModelOptionalParams) -> None: + self.streaming_end_of_stream_only: bool = streaming_params.streaming_end_of_stream_only or False + self.streaming_sampling_rate: int = streaming_params.streaming_sampling_rate or 5 + + @override + def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: + super().update_in_memory_litellm_params(litellm_params) + self._set_streaming_params(streaming_params_from_litellm_params(litellm_params)) + async def _call_crowdstrike_aidr_guard( self, payload: dict[str, Any], hook_name: str ) -> _GuardChatCompletionsResult: diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 46b00829b74..c6b8df1b493 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -36,6 +36,7 @@ if TYPE_CHECKING: from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, + StreamingScanKey, ) # Call types that stream JSON-RPC events (A2A); guardrail HTTPException is emitted as in-stream error @@ -54,6 +55,9 @@ class _EndpointTranslation(Protocol): @property def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ... + @property + def get_streaming_scan_key(self) -> "Callable[[Sequence[object]], StreamingScanKey | None]": ... + @property def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... @@ -70,6 +74,12 @@ def _chunk_choices(item: object) -> Sequence[object]: return choices +def _is_redundant_scan(scan_key: "StreamingScanKey | None", last_scan_key: "StreamingScanKey | None") -> bool: + if scan_key is None: + return False + return scan_key == last_scan_key or scan_key.has_nothing_to_scan + + class _StreamTerminated(Exception): """Internal signal that the incremental transform stream has already emitted its terminal chunks (block message or in-stream error) and must stop.""" @@ -1011,6 +1021,7 @@ class UnifiedLLMGuardrails(CustomLogger): # Drives how a block terminates the stream: continue the in-progress # message (True) vs emit a standalone block message (False, buffered). chunks_yielded = False + last_scan_key: StreamingScanKey | None = None # rebind-ok: replaced after every scan round async for item in response: chunk_counter += 1 @@ -1052,6 +1063,19 @@ class UnifiedLLMGuardrails(CustomLogger): # Process chunk based on sampling rate if chunk_counter % sampling_rate == 0: + endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far) + if _is_redundant_scan(scan_key, last_scan_key): + verbose_proxy_logger.debug( + "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round", + chunk_counter, + guardrail_to_apply.guardrail_name, + ) + chunks_yielded = True + responses_yielded.append(item) + yield item + continue + verbose_proxy_logger.debug( "Processing streaming chunk %s (sampling_rate=%s) with guardrail %s", chunk_counter, @@ -1067,8 +1091,6 @@ class UnifiedLLMGuardrails(CustomLogger): # string, permanently losing this chunk's content. original_item = copy.deepcopy(item) - endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() - try: await endpoint_translation.process_output_streaming_response( responses_so_far=responses_so_far, @@ -1110,6 +1132,8 @@ class UnifiedLLMGuardrails(CustomLogger): ): yield error_item return + if scan_key is not None: + last_scan_key = scan_key chunks_yielded = True responses_yielded.append(original_item) yield original_item @@ -1136,6 +1160,18 @@ class UnifiedLLMGuardrails(CustomLogger): # preserve the list, not clone every chunk (deepcopy would double # peak memory for large responses). buffered_items: Final = list(responses_so_far) if buffer_until_moderated else None + end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far) + if _is_redundant_scan(end_scan_key, last_scan_key): + verbose_proxy_logger.debug( + "Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all", + guardrail_to_apply.guardrail_name, + ) + for buffered_item in buffered_items or (): + yield buffered_item + for pending_item in pending_end_of_stream_items: + responses_yielded.append(pending_item) + yield pending_item + return try: await endpoint_translation.process_output_streaming_response( diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index bd35782444b..003576ba555 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -826,11 +826,12 @@ class InMemoryGuardrailHandler: Removes old callback from litellm.callbacks and creates fresh instance. If the new config fails to initialize (e.g. an invalid on_flagged - combination), the previous instance is restored rather than left - deleted: initialize_guardrail's own ValueError/TypeError propagate - uncaught, so a caller reaching this point after already deleting the - old instance would otherwise leave the guardrail providing no - protection at all, not merely "still enforcing the old config." + combination or an invalid regex), the previous instance is restored + rather than left deleted, and the failure is re-raised as ValueError so + every init failure reaches callers as one exception type: a caller + reaching this point after already deleting the old instance would + otherwise leave the guardrail providing no protection at all, not + merely "still enforcing the old config." """ guardrail_id: Final = guardrail.get("guardrail_id") if not guardrail_id: @@ -849,7 +850,7 @@ class InMemoryGuardrailHandler: # that was enforcing must never fail open because an update was bad. try: return self.initialize_guardrail(guardrail=guardrail, config_file_path=config_file_path, source=source) - except Exception: + except Exception as init_error: if previous_guardrail is not None: verbose_proxy_logger.exception( "Reinitializing guardrail %s with updated params failed; restoring the previous configuration", @@ -861,7 +862,7 @@ class InMemoryGuardrailHandler: ) except Exception: # noqa: BLE001 # the original failure must propagate even if the restore breaks verbose_proxy_logger.exception("Restoring previous guardrail %s also failed", guardrail_id) - raise + raise ValueError(f"Guardrail initialization failed: {init_error}") from init_error def sync_guardrail_from_db(self, guardrail: Guardrail, config_file_path: str | None = None) -> Guardrail | None: """ diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 20f83085286..1d440448c2f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -16,6 +16,7 @@ from starlette.datastructures import Headers import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging +from litellm._uuid import uuid from litellm.constants import ( CONSUMED_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, @@ -23,6 +24,7 @@ from litellm.constants import ( OTEL_SERVICE_NAME_METADATA_KEYS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, + SESSION_ID_GENERATED_METADATA_KEY, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -40,6 +42,7 @@ from litellm.proxy._types import ( AddTeamCallback, CommonProxyErrors, LitellmDataForBackendLLMCall, + LiteLLMRoutes, LitellmUserRoles, ProxyErrorTypes, ProxyException, @@ -47,6 +50,8 @@ from litellm.proxy._types import ( TeamCallbackMetadata, UserAPIKeyAuth, ) +from litellm.proxy.auth.auth_utils import get_request_route +from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, get_metadata_variable_name_from_kwargs, @@ -715,6 +720,50 @@ def _get_anthropic_session_id_from_metadata(metadata: object) -> str | None: return session_id +def _is_llm_inference_route(request: Request) -> bool: + route: Final = get_request_route(request) + return RouteChecks.is_llm_api_route(route=route) and not RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value + ) + + +def apply_missing_session_id_policy( + data: dict[str, object], # mutable-ok: stamps session ids in place on the request body the pipeline threads through + _metadata_variable_name: str, + general_settings: Mapping[str, object] | None, + request: Request, +) -> None: + policy: Final = general_settings.get("missing_session_id") if general_settings else None + if policy is None or not _is_llm_inference_route(request): + return + metadata: Final = data.get(_metadata_variable_name) + if not isinstance(metadata, dict): + return + if data.get("litellm_session_id") or metadata.get("session_id"): + return + match policy: + case "generate": + session_id: Final = str(data.get("litellm_trace_id") or metadata.get("trace_id") or uuid.uuid4()) + data["litellm_session_id"] = session_id # rebind-ok: data is an out-param + data.setdefault("litellm_trace_id", session_id) + metadata["session_id"] = session_id + metadata[SESSION_ID_GENERATED_METADATA_KEY] = True + case "reject": + raise ProxyException( + message=( + "Request has no session id. Send an `x-litellm-session-id` header or `metadata.session_id`. " + "Required by `general_settings.missing_session_id: reject`." + ), + type=ProxyErrorTypes.bad_request_error, + param="session_id", + code=400, + ) + case _: + verbose_proxy_logger.warning( + "Ignoring unknown general_settings.missing_session_id=%r; expected 'generate' or 'reject'", policy + ) + + def is_claude_code_user_agent(user_agent: str) -> bool: """Claude Code identifies itself as ``claude-cli/ ...``; the IDE extensions and the Agent SDK run through the same CLI and share that prefix.""" @@ -1818,6 +1867,12 @@ async def add_litellm_data_to_request( data=data, _metadata_variable_name=_metadata_variable_name, ) + apply_missing_session_id_policy( + data=data, + _metadata_variable_name=_metadata_variable_name, + general_settings=general_settings, + request=request, + ) # Expose request headers under the metadata field for guardrails (fixes #17477) if _metadata_variable_name in data and isinstance(data[_metadata_variable_name], dict): diff --git a/litellm/proxy/openai_files_endpoints/batch_file_validation.py b/litellm/proxy/openai_files_endpoints/batch_file_validation.py index 0aee5e8cc54..a41bd36d510 100644 --- a/litellm/proxy/openai_files_endpoints/batch_file_validation.py +++ b/litellm/proxy/openai_files_endpoints/batch_file_validation.py @@ -2,7 +2,9 @@ import json from collections.abc import Iterator from dataclasses import dataclass from itertools import chain -from typing import BinaryIO, Final, NoReturn, assert_never +from typing import BinaryIO, Final, NoReturn + +from typing_extensions import assert_never from litellm.proxy._types import ProxyException diff --git a/litellm/proxy/openai_files_endpoints/general_upload_validation.py b/litellm/proxy/openai_files_endpoints/general_upload_validation.py index 9d450cb5b8d..8c59a520272 100644 --- a/litellm/proxy/openai_files_endpoints/general_upload_validation.py +++ b/litellm/proxy/openai_files_endpoints/general_upload_validation.py @@ -8,7 +8,9 @@ extensions, path-traversal filenames) regardless of purpose. from dataclasses import dataclass from pathlib import Path -from typing import BinaryIO, Final, NoReturn, assert_never +from typing import BinaryIO, Final, NoReturn + +from typing_extensions import assert_never from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.path_utils import safe_filename diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 23932ba7c8c..ed247d52ce2 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -261,6 +261,7 @@ class ProxyInitializationHelpers: "app": "litellm.proxy.proxy_server:app", "host": host, "port": port, + "server_header": False, } if log_config is not None: print(f"Using log_config: {log_config}") diff --git a/ui/litellm-dashboard/src/autorouter_presets.json b/litellm/proxy/public_endpoints/autorouter_presets.json similarity index 100% rename from ui/litellm-dashboard/src/autorouter_presets.json rename to litellm/proxy/public_endpoints/autorouter_presets.json diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 4d58a974bb8..94a59828451 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -1,11 +1,13 @@ +import asyncio import json import os import re -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from importlib.resources import files from typing import TYPE_CHECKING, Final, Protocol from fastapi import APIRouter, HTTPException, Request +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict import litellm @@ -28,6 +30,7 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import ) from litellm.types.proxy.public_endpoints.public_endpoints import ( AgentCreateInfo, + AutoRouterPresetRecord, ComplexityScorerDefaults, ProviderCreateInfo, PublicModelHubInfo, @@ -464,6 +467,86 @@ async def get_litellm_blog_posts(): return BlogPostsResponse(posts=posts) +_AUTOROUTER_PRESETS_ADAPTER: Final = TypeAdapter(dict[str, AutoRouterPresetRecord]) + + +def _load_bundled_autorouter_presets() -> Mapping[str, AutoRouterPresetRecord]: + raw: Final = json.loads( + files("litellm.proxy.public_endpoints").joinpath("autorouter_presets.json").read_text(encoding="utf-8") + ) + return _AUTOROUTER_PRESETS_ADAPTER.validate_python(raw) + + +async def _fetch_remote_autorouter_presets(url: str) -> Mapping[str, AutoRouterPresetRecord]: + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.types.llms.custom_http import httpxSpecialProvider + + client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.UI) + response: Final = await client.get(url, timeout=5.0) + response.raise_for_status() + presets: Final = _AUTOROUTER_PRESETS_ADAPTER.validate_python(response.json()) + if not presets: + raise ValueError("remote auto-router preset catalog is empty") + return presets + + +async def _resolve_autorouter_presets( + url: str, + fetch: Callable[[str], Awaitable[Mapping[str, AutoRouterPresetRecord]]], +) -> Mapping[str, AutoRouterPresetRecord]: + if os.getenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", "").lower() == "true": + return _load_bundled_autorouter_presets() + try: + return await fetch(url) + except Exception as e: + verbose_logger.warning( + "LiteLLM: failed to fetch auto-router presets from %s: %s. Serving the bundled catalog for the life of this process.", + url, + str(e), + ) + return _load_bundled_autorouter_presets() + + +class _AutoRouterPresetsCache: + presets: Mapping[str, AutoRouterPresetRecord] | None = None + lock: asyncio.Lock | None = None + + +async def get_autorouter_presets( + url: str, + fetch: Callable[[str], Awaitable[Mapping[str, AutoRouterPresetRecord]]] = _fetch_remote_autorouter_presets, +) -> Mapping[str, AutoRouterPresetRecord]: + cached: Final = _AutoRouterPresetsCache.presets + if cached is not None: + return cached + if _AutoRouterPresetsCache.lock is None: + _AutoRouterPresetsCache.lock = asyncio.Lock() + async with _AutoRouterPresetsCache.lock: + held: Final = _AutoRouterPresetsCache.presets + if held is not None: + return held + resolved: Final = await _resolve_autorouter_presets(url=url, fetch=fetch) + _AutoRouterPresetsCache.presets = resolved + return resolved + + +@router.get( + "/public/autorouter_presets", + tags=["public", "auto router"], # mutable-ok: FastAPI route tags take a list + response_model=dict[str, AutoRouterPresetRecord], +) +async def get_public_autorouter_presets() -> Mapping[str, AutoRouterPresetRecord]: + """ + Return the auto-router preset catalog the dashboard's template picker renders. + + Resolved once per process, like the model cost map: fetched from ``litellm.autorouter_presets_url`` + (override with ``LITELLM_AUTOROUTER_PRESETS_URL``) on the first request, falling back to the + catalog bundled with the package on any failure. Set ``LITELLM_LOCAL_AUTOROUTER_PRESETS=True`` + to serve the bundled catalog only. A restart picks up a newly published catalog. + """ + return await get_autorouter_presets(url=litellm.autorouter_presets_url) + + @router.get( "/public/endpoints", tags=["public"], 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/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index c48750cea72..bb7dfafb297 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -55,6 +55,10 @@ router: Final = APIRouter() SPEND_LOGS_PAGINATION_COUNT_CAP: Final = 10000 +_SESSION_GROUP_KEY_SQL: Final = "COALESCE(NULLIF(session_id, ''), request_id), api_key" +_MCP_CALL_TYPES_SQL: Final = "('call_mcp_tool', 'list_mcp_tools')" +_AGENT_CALL_TYPE_SQL: Final = "'asend_message'" + _INTERNAL_HEALTH_CHECK_API_KEYS: Final = ( LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, hash_token(token=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME), @@ -144,21 +148,16 @@ class _DailyTagSpendRow(TypedDict): total_spend: float -class _SessionCountAggregate(TypedDict): - session_id: int - - -class _SessionCountRow(TypedDict): - session_id: str - _count: _SessionCountAggregate - - class _SessionSpendRow(TypedDict): session_id: str + api_key: ReadOnly[str] + session_total_count: ReadOnly[int] session_total_spend: float mcp_tool_call_count: int mcp_tool_call_spend: float session_cache_hit_count: ReadOnly[int] + session_llm_count: ReadOnly[int] + session_agent_count: ReadOnly[int] class _SpendSumAggregate(TypedDict, total=False): @@ -242,18 +241,6 @@ async def _count_spend_logs(prisma_client: PrismaClient, where: Mapping[str, obj return await _spend_logs_table(prisma_client).count(where=where) -async def _count_logs_per_session( - prisma_client: PrismaClient, session_ids: Sequence[str | None] -) -> Sequence[_SessionCountRow]: - """Count spend log rows per session for the given session ids.""" - rows: Final = await _spend_logs_table(prisma_client).group_by( - by=["session_id"], - where={"session_id": {"in": session_ids}}, - count={"session_id": True}, - ) - return cast(Sequence[_SessionCountRow], rows) # cast-ok: group_by(count=) shape is fixed by the by/count args - - async def _find_team_row(prisma_client: PrismaClient, team_id: str) -> _SupportsModelDump | None: """Read a single team row as a Prisma model instance.""" return await _team_table(prisma_client).find_unique(where={"team_id": team_id}) @@ -2290,6 +2277,10 @@ async def ui_view_spend_logs( default=False, description="Exclude LiteLLM internal health check requests from results", ), + group_by_session: bool = fastapi.Query( + default=False, + description="Paginate over sessions instead of raw logs: one representative row per session, total counts sessions", + ), ): """ View spend logs with pagination support. @@ -2644,12 +2635,16 @@ async def ui_view_spend_logs( else: _order_expr = order_column + joined_conditions: Final = " AND ".join(sql_conditions) + session_grouping: Final = group_by_session is True + count_group_clause: Final = f"GROUP BY {_SESSION_GROUP_KEY_SQL}" if session_grouping else "" count_query: Final = f""" SELECT COUNT(*) AS total_count FROM ( SELECT 1 FROM "LiteLLM_SpendLogs" - WHERE {" AND ".join(sql_conditions)} + WHERE {joined_conditions} + {count_group_clause} LIMIT ${p} ) AS bounded_matches """ @@ -2660,21 +2655,36 @@ async def ui_view_spend_logs( total_is_capped: Final = raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP total_records: Final = SPEND_LOGS_PAGINATION_COUNT_CAP if total_is_capped else raw_total - sql_query: Final = f""" - SELECT - request_id, call_type, api_key, spend, total_tokens, + select_columns: Final = """request_id, call_type, api_key, spend, total_tokens, prompt_tokens, completion_tokens, "startTime", "endTime", "completionStartTime", model, model_id, model_group, custom_llm_provider, api_base, "user", metadata, cache_hit, cache_key, request_tags, team_id, organization_id, end_user, requester_ip_address, session_id, status, mcp_namespaced_tool_name, agent_id, - COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms + COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms""" + sql_query: Final = ( + f""" + SELECT * FROM ( + SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL}) + {select_columns} + FROM "LiteLLM_SpendLogs" + WHERE {joined_conditions} + ORDER BY {_SESSION_GROUP_KEY_SQL}, call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC + ) AS session_representatives + ORDER BY {_order_expr} {_sql_dir}{_nulls_clause}, request_id + LIMIT ${p} OFFSET ${p + 1} + """ + if session_grouping + else f""" + SELECT + {select_columns} FROM "LiteLLM_SpendLogs" - WHERE {" AND ".join(sql_conditions)} + WHERE {joined_conditions} ORDER BY {_order_expr} {_sql_dir}{_nulls_clause} LIMIT ${p} OFFSET ${p + 1} """ + ) sql_params.extend([page_size, skip]) data: Final = await prisma_client.db.query_raw(sql_query, *sql_params) @@ -4075,11 +4085,12 @@ async def _build_ui_spend_logs_response( Build the paginated response for the UI spend-logs endpoint. When ``enrich_session_counts`` is ``True`` (the default for the v1/UI - endpoint), each row is enriched with ``session_total_count`` so the - frontend knows which sessions are expandable (multi-call sessions). - For every row that carries a ``session_id``, a single ``GROUP BY`` query - fetches the total number of logs in each referenced session. Rows without - a ``session_id`` default to ``1``. + endpoint), each row is enriched with ``session_total_count`` plus spend + and call-type aggregates so the frontend knows which sessions are + expandable (multi-call sessions). One ``GROUP BY (session_id, api_key)`` + query serves every referenced session, keyed per api key so two callers + reusing a session id never see each other's totals. Rows without a + ``session_id`` default to ``1``. When ``enrich_session_counts`` is ``False`` (v2 endpoint), rows are serialised without the extra query. @@ -4101,7 +4112,6 @@ async def _build_ui_spend_logs_response( A dict with ``data`` (enriched rows), ``total``, ``page``, ``page_size``, ``total_pages``, and ``total_is_capped``. """ - count_map: dict[str, int] = {} if enrich_session_counts: session_ids: Final[Sequence[str | None]] = list( { @@ -4110,15 +4120,8 @@ async def _build_ui_spend_logs_response( if (row.get("session_id") if isinstance(row, dict) else getattr(row, "session_id", None)) } ) - if session_ids: - # NOTE: This GROUP BY runs on every v1/UI page load. The IN clause - # is bounded by page_size (typically 25-50 distinct session IDs). - # If performance degrades at scale, consider short-lived caching or - # folding the count into the main query via a window function. - counts: Final = await _count_logs_per_session(prisma_client, session_ids) - count_map = {r["session_id"]: r["_count"]["session_id"] for r in counts if r.get("session_id")} - session_spend_map: dict[str, dict[str, int | float]] = {} + session_spend_map: dict[tuple[str, str], dict[str, int | float]] = {} if enrich_session_counts and session_ids: from prisma.errors import PrismaError @@ -4130,38 +4133,46 @@ async def _build_ui_spend_logs_response( { (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None)) for row in data - if (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None)) + if (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None)) is not None } ) rows: Final[Sequence[_SessionSpendRow]] = await _query_raw( prisma_client, - """ - SELECT session_id, + f""" + SELECT session_id, api_key, + COUNT(*)::int AS session_total_count, COALESCE(SUM(spend), 0)::double precision AS session_total_spend, COUNT(*) FILTER ( - WHERE call_type IN ('call_mcp_tool', 'list_mcp_tools') + WHERE call_type IN {_MCP_CALL_TYPES_SQL} )::int AS mcp_tool_call_count, COALESCE(SUM(spend) FILTER ( - WHERE call_type IN ('call_mcp_tool', 'list_mcp_tools') + WHERE call_type IN {_MCP_CALL_TYPES_SQL} ), 0)::double precision AS mcp_tool_call_spend, - COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count + COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count, + COUNT(*) FILTER ( + WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL} + )::int AS session_llm_count, + COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count FROM "LiteLLM_SpendLogs" WHERE session_id = ANY($1::text[]) AND api_key = ANY($2::text[]) - GROUP BY session_id + GROUP BY session_id, api_key """, session_ids, authorized_api_keys, ) session_spend_map = { - row["session_id"]: { + (row["session_id"], row["api_key"]): { + "session_total_count": int(row.get("session_total_count") or 0), "session_total_spend": float(row.get("session_total_spend") or 0.0), "mcp_tool_call_count": int(row.get("mcp_tool_call_count") or 0), "mcp_tool_call_spend": float(row.get("mcp_tool_call_spend") or 0.0), "session_cache_hit_count": int(row.get("session_cache_hit_count") or 0), + "session_llm_count": int(row.get("session_llm_count") or 0), + "session_agent_count": int(row.get("session_agent_count") or 0), } for row in rows - if row.get("session_id") + if row.get("session_id") and row.get("api_key") is not None } except PrismaError: verbose_proxy_logger.debug( @@ -4174,14 +4185,17 @@ async def _build_ui_spend_logs_response( for row in data: row_dict = dict(row) if isinstance(row, dict) else row.model_dump() sid = row_dict.get("session_id") - row_dict["session_total_count"] = count_map.get(sid, 1) if sid else 1 - session_stats = session_spend_map.get(sid) if sid else None + row_api_key = row_dict.get("api_key") + session_stats = session_spend_map.get((sid, row_api_key)) if sid and row_api_key is not None else None + row_dict["session_total_count"] = int(session_stats["session_total_count"]) if session_stats else 1 if session_stats: row_dict["session_total_spend"] = session_stats["session_total_spend"] if session_stats["mcp_tool_call_count"]: row_dict["mcp_tool_call_count"] = session_stats["mcp_tool_call_count"] row_dict["mcp_tool_call_spend"] = session_stats["mcp_tool_call_spend"] row_dict["session_cache_hit_count"] = session_stats["session_cache_hit_count"] + row_dict["session_llm_count"] = session_stats["session_llm_count"] + row_dict["session_agent_count"] = session_stats["session_agent_count"] enriched.append(row_dict) response_data: list = enriched else: 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/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7871c85220c..f271655f5e3 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -17,6 +17,7 @@ from typing_extensions import TypeIs import litellm from litellm.constants import ( + EMPTY_MAPPING, LITELLM_MAX_STREAMING_DURATION_SECONDS, STREAM_SSE_DONE_STRING, ) @@ -273,6 +274,9 @@ class BaseResponsesAPIStreamingIterator: self._hidden_params["additional_headers"] = process_response_headers( self.response.headers or {} ) # GUARANTEE OPENAI HEADERS IN RESPONSE + self._raw_response_headers: Mapping[str, str] = MappingProxyType( + dict(self.response.headers or {}) # mutable-ok: immediately frozen by MappingProxyType + ) def _check_max_streaming_duration(self) -> None: """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" @@ -446,6 +450,7 @@ class BaseResponsesAPIStreamingIterator: except Exception: # Fallback to original if serialization fails pass + self._restore_provider_response_headers(logging_response) end_time: Final = datetime.now() if is_async: @@ -480,6 +485,41 @@ class BaseResponsesAPIStreamingIterator: ) self._run_post_success_hooks(end_time=end_time) + def _restore_provider_response_headers(self, logging_response: object) -> None: + """Re-apply the provider's response headers to the copy handed to logging callbacks. + + ``model_validate(model_dump())`` above drops pydantic private attributes, so the + ``_hidden_params`` the provider transform set on the nested response are lost. Returns early + when that copy fell back to the original event, so logging-only state never lands on the + object the caller is iterating. + """ + if logging_response is self.completed_response: + return + target: Final[object] = getattr(logging_response, "response", None) + existing_hidden: Final[object] = getattr(target, "_hidden_params", None) + if not isinstance(existing_hidden, Mapping): + return + existing: Final[Mapping[str, object]] = existing_hidden + source_hidden: Final[object] = getattr( + getattr(self.completed_response, "response", None), "_hidden_params", None + ) + source: Final[Mapping[str, object]] = source_hidden if isinstance(source_hidden, Mapping) else EMPTY_MAPPING + processed: Final[object] = source.get("additional_headers") or self._hidden_params.get("additional_headers") + raw: Final[object] = source.get("headers") or self._raw_response_headers + headers: Final[Mapping[str, object]] = processed if isinstance(processed, Mapping) else EMPTY_MAPPING + raw_headers: Final[Mapping[str, object]] = raw if isinstance(raw, Mapping) else EMPTY_MAPPING + # rebuild by value and let existing keys win: sharing the source dicts would alias what the proxy + # splats into the client's HTTP headers, and copying non-header keys would carry response_cost + setattr( # noqa: B010 # target is typed object here, so a plain attribute store does not type check + target, + "_hidden_params", + { # mutable-ok: the cost calculator writes optional_params into _hidden_params + "additional_headers": {**headers}, # mutable-ok: fresh copy, logging callbacks may mutate it + "headers": {**raw_headers}, # mutable-ok: fresh copy, logging callbacks may mutate it + **existing, + }, + ) + def _handle_logging_completed_response(self): """Base implementation - should be overridden by subclasses""" 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 8e0cad39561..ad8b67d5e8f 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -274,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/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a123a75dd81..430efe339a2 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -26,7 +26,11 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast from pydantic import BaseModel, create_model from litellm._logging import verbose_router_logger -from litellm.constants import EMPTY_MAPPING, RETURN_RAW_MODEL_NAME_METADATA_KEY +from litellm.constants import ( + EMPTY_MAPPING, + RETURN_RAW_MODEL_NAME_METADATA_KEY, + SESSION_ID_GENERATED_METADATA_KEY, +) 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.internal_call_metadata import forwarded_internal_call_metadata @@ -799,6 +803,7 @@ class ClassificationOutcome(NamedTuple): "reasoning_override", "llm_classifier", "heuristic_first_short_circuit", + "hybrid_short_circuit", "housekeeping", "classifier_plugin", "classifier_fallback", @@ -1241,6 +1246,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. @@ -1367,6 +1381,8 @@ class ComplexityRouter(CustomLogger): 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) @@ -1418,6 +1434,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, @@ -2677,7 +2716,7 @@ class ComplexityRouter(CustomLogger): """Resolve a client-supplied session_id.""" for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): session_id = metadata.get("session_id") - if session_id is not None: + if session_id is not None and not metadata.get(SESSION_ID_GENERATED_METADATA_KEY): return str(session_id) return None diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index fdf2a3a0b39..0ae0db63fad 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -43,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, ...]] = ( @@ -627,12 +627,13 @@ class ComplexityRouterConfig(BaseModel): ) # Classifier strategy - classifier_type: Literal["heuristic", "heuristic_v2", "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, the bundled trained four-tier heuristic, " - "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" + "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( @@ -644,7 +645,10 @@ class ComplexityRouterConfig(BaseModel): ) 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, @@ -659,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=( @@ -1135,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: @@ -1257,7 +1291,7 @@ 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_v2", "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 four built-in tiers, as does heuristic_v2" 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/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index b1e9dbdefa8..39d3e25aacb 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -21,7 +21,7 @@ from typing_extensions import TypedDict from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import AllMessageValues @@ -265,7 +265,7 @@ class DeploymentAffinityCheck(CustomLogger): @staticmethod def _get_session_id_from_metadata_dict(metadata: dict) -> str | None: session_id: Final = metadata.get("session_id") - if session_id is None: + if session_id is None or metadata.get(SESSION_ID_GENERATED_METADATA_KEY): return None return str(session_id) diff --git a/litellm/types/llms/gemini_audio_transcription.py b/litellm/types/llms/gemini_audio_transcription.py index cb12e0f45b8..f7e74ba4bf8 100644 --- a/litellm/types/llms/gemini_audio_transcription.py +++ b/litellm/types/llms/gemini_audio_transcription.py @@ -1,7 +1,7 @@ -from typing import Literal, Required +from typing import Literal from pydantic import BaseModel, ConfigDict -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict class GeminiTranscriptionAudioInput(TypedDict): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py index f47c38af3e3..6beca030a3a 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py @@ -4,7 +4,18 @@ from .base import GuardrailConfigModel class CrowdStrikeAIDRGuardrailConfigModelOptionalParams(BaseModel): - pass + streaming_end_of_stream_only: bool | None = Field( + default=None, + description="If False (default when unset), post_call scans the accumulated streamed response every " + "streaming_sampling_rate chunks and an in-flight block stops the stream. If True, the guard runs once " + "over the assembled response at end of stream, so flagged content may already have reached the client.", + ) + streaming_sampling_rate: int | None = Field( + default=None, + ge=1, + description="When streaming_end_of_stream_only is False, scan the accumulated streamed response every Nth " + "chunk. Defaults to 5 when unset.", + ) class CrowdStrikeAIDRGuardrailConfigModel(GuardrailConfigModel[CrowdStrikeAIDRGuardrailConfigModelOptionalParams]): diff --git a/litellm/types/proxy/public_endpoints/public_endpoints.py b/litellm/types/proxy/public_endpoints/public_endpoints.py index c7f80a61e0f..fa73926305b 100644 --- a/litellm/types/proxy/public_endpoints/public_endpoints.py +++ b/litellm/types/proxy/public_endpoints/public_endpoints.py @@ -1,7 +1,7 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Any, Literal -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict class PublicModelHubInfo(BaseModel): @@ -73,6 +73,44 @@ class SupportedEndpointsResponse(BaseModel): endpoints: list[SupportedEndpoint] +class AutoRouterPresetTiers(BaseModel): + """Exactly the four built-in tiers the dashboard's preset prefill can apply. + + extra="forbid" on purpose: a tier name this dashboard cannot apply would grey out or crash the + picker, so such a catalog is rejected wholesale and the bundled one serves instead. + """ + + model_config = ConfigDict(extra="forbid") + + SIMPLE: Sequence[str] + MEDIUM: Sequence[str] + COMPLEX: Sequence[str] + REASONING: Sequence[str] + + +class AutoRouterPresetConfig(BaseModel): + """The complexity_router_config a preset prefills. + + Only tiers is validated, because every dashboard consumer dereferences it; everything else + passes through verbatim with unknown fields kept (extra="allow"), so a catalog published after + this proxy shipped still serves its new fields intact. + """ + + model_config = ConfigDict(extra="allow") + + tiers: AutoRouterPresetTiers + + +class AutoRouterPresetRecord(BaseModel): + """One auto-router preset as served to the dashboard's template picker.""" + + model_config = ConfigDict(extra="allow") + + label: str + description: str + complexity_router_config: AutoRouterPresetConfig + + class ComplexityScorerDefaults(BaseModel): """The complexity router's shipped heuristic scorer defaults. diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 09c01873a9d..ee6f09e05dc 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2852,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/ruff-strict-budget.json b/ruff-strict-budget.json index a390c61dcf2..4fcf650a8bc 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -177,7 +177,7 @@ "limit": 8 }, "RUF019": { - "limit": 31 + "limit": 27 }, "RUF046": { "limit": 4 diff --git a/tests/code_coverage_tests/check_py310_typing_imports.py b/tests/code_coverage_tests/check_py310_typing_imports.py new file mode 100644 index 00000000000..0cd4d089890 --- /dev/null +++ b/tests/code_coverage_tests/check_py310_typing_imports.py @@ -0,0 +1,150 @@ +import ast +import os +import sys +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +PY311_PLUS_TYPING_NAMES: Final[frozenset[str]] = frozenset( + { + "NotRequired", + "Required", + "Self", + "LiteralString", + "Never", + "assert_never", + "assert_type", + "reveal_type", + "TypeVarTuple", + "Unpack", + "dataclass_transform", + "override", + "TypeAliasType", + "get_original_bases", + "ReadOnly", + "TypeIs", + "NoDefault", + "get_protocol_members", + "is_protocol", + "evaluate_forward_ref", + "TypeForm", + } +) + + +@dataclass(frozen=True, slots=True) +class TypingImportViolation: + file: str + line: int + name: str + + +def _walk_with_ancestors( + node: ast.AST, ancestors: tuple[tuple[ast.AST, str], ...] = () +) -> Iterator[tuple[ast.AST, tuple[tuple[ast.AST, str], ...]]]: + yield node, ancestors + for field_name, field_value in ast.iter_fields(node): + if isinstance(field_value, ast.AST): + yield from _walk_with_ancestors(field_value, (*ancestors, (node, field_name))) + elif isinstance(field_value, list): + for child in field_value: + if isinstance(child, ast.AST): + yield from _walk_with_ancestors(child, (*ancestors, (node, field_name))) + + +def _is_sys_version_info(node: ast.AST) -> bool: + return ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "sys" + and node.attr == "version_info" + ) + + +def _is_version_guarded(ancestors: tuple[tuple[ast.AST, str], ...]) -> bool: + nearest_if: Final[tuple[ast.If, str] | None] = next( + ( + (ancestor, field_name) + for ancestor, field_name in reversed(ancestors) + if isinstance(ancestor, ast.If) + ), + None, + ) + if nearest_if is None: + return False + enclosing_if, branch = nearest_if + test: Final[ast.expr] = enclosing_if.test + if not isinstance(test, ast.Compare) or len(test.ops) != 1 or not _is_sys_version_info(test.left): + return False + operator: Final[ast.cmpop] = test.ops[0] + return (isinstance(operator, (ast.Gt, ast.GtE)) and branch == "body") or ( + isinstance(operator, (ast.Lt, ast.LtE)) and branch == "orelse" + ) + + +def scan_file(file_path: str | os.PathLike[str]) -> tuple[TypingImportViolation, ...]: + path: Final[Path] = Path(file_path) + tree: Final[ast.Module] = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + return tuple( + violation + for node, ancestors in _walk_with_ancestors(tree) + if not _is_version_guarded(ancestors) + for violation in _violations_for_node(node, path) + ) + + +def _violations_for_node( + node: ast.AST, path: Path +) -> tuple[TypingImportViolation, ...]: + if isinstance(node, ast.ImportFrom) and node.module == "typing": + return tuple( + TypingImportViolation(file=str(path), line=node.lineno, name=alias.name) + for alias in node.names + if alias.name in PY311_PLUS_TYPING_NAMES + ) + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "typing" + and node.attr in PY311_PLUS_TYPING_NAMES + ): + return (TypingImportViolation(file=str(path), line=node.lineno, name=node.attr),) + return () + + +def scan_directory(base_dir: str | os.PathLike[str] = ".") -> tuple[TypingImportViolation, ...]: + base_path: Final[Path] = Path(base_dir) + return tuple( + violation + for directory in ( + base_path / "litellm", + base_path / "enterprise", + base_path / "litellm-proxy-extras" / "litellm_proxy_extras", + ) + if directory.exists() + for path in directory.rglob("*.py") + for violation in scan_file(path) + ) + + +def main() -> None: + violations: Final[tuple[TypingImportViolation, ...]] = scan_directory() + if violations: + message: Final[str] = "\n".join( + ( + "Python 3.10-incompatible typing imports found:", + *( + f"{violation.file}:{violation.line}: {violation.name} is unavailable in Python 3.10; " + "import it from typing_extensions instead because litellm supports Python 3.10" + for violation in violations + ), + ) + ) + sys.stdout.write(f"{message}\n") + raise RuntimeError("Import Python 3.10-incompatible typing names from typing_extensions instead") + sys.stdout.write("No Python 3.10-incompatible typing imports found.\n") + + +if __name__ == "__main__": + main() 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/logs/logsPagination.spec.ts b/tests/e2e/ui/tests/logs/logsPagination.spec.ts new file mode 100644 index 00000000000..416e1c3171f --- /dev/null +++ b/tests/e2e/ui/tests/logs/logsPagination.spec.ts @@ -0,0 +1,150 @@ +import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; +import { Page } from "../../fixtures/pages"; +import { CHAT_MODEL_A, createVirtualKey, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic"; + +/** + * Session-grouped pagination (#38060): a page of N rows must render exactly N session rows, a + * session must never straddle pages, and two callers reusing one session id stay separate rows. + * All traffic is generated per run behind a unique key alias or session id, so concurrent specs + * cannot decide the outcome. + */ + +const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + +/** Every tab stays mounted, so the DOM holds four tables at once; scope to the visible one. */ +const requestLogsRows = (page: PlaywrightPage): Locator => + page.locator("table").filter({ visible: true }).first().locator("tbody tr"); + +const visibleTestId = (page: PlaywrightPage, id: string): Locator => page.getByTestId(id).filter({ visible: true }); + +async function openLogs(page: PlaywrightPage): Promise { + await navigateToPage(page, Page.Logs); + await dismissFeedbackPopup(page); + await expect(visibleTestId(page, "datatable-search")).toBeVisible({ timeout: 20_000 }); +} + +async function openFilterDrawer(page: PlaywrightPage): Promise { + await visibleTestId(page, "datatable-filters-trigger").click(); + const drawer = page.getByRole("dialog", { name: "Filters" }); + await expect(drawer).toBeVisible({ timeout: 10_000 }); + return drawer; +} + +async function applyKeyAliasFilter(page: PlaywrightPage, drawer: Locator, alias: string): Promise { + await drawer.getByRole("combobox", { name: "Search a key alias" }).click(); + await page.keyboard.type(alias); + await page.getByRole("option", { name: alias, exact: true }).first().click(); + await drawer.getByRole("button", { name: "Apply Filters" }).click(); + await expect(drawer).not.toBeVisible({ timeout: 10_000 }); +} + +async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): Promise { + await visibleTestId(page, "pagination-page-size").click(); + await page.getByRole("option", { name: size, exact: true }).click(); +} + +test.describe("Logs page session-grouped pagination", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("a 25-row page renders exactly 25 session rows and no session straddles pages", async ({ page, request }) => { + const suffix = uniqueSuffix(); + const alias = `e2e-logs-pgn-${suffix}`; + const mine = await createVirtualKey(request, { key_alias: alias }); + + const soloIds: string[] = []; + for (let i = 0; i < 26; i++) { + soloIds.push( + await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-pgn-solo-${i}-${suffix}`, + apiKey: mine.key, + }), + ); + } + const sessionA = `sess-pgn-a-${suffix}`; + const sessionB = `sess-pgn-b-${suffix}`; + let lastSessionCallId = ""; + for (let i = 0; i < 7; i++) { + lastSessionCallId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-pgn-a-${i}-${suffix}`, + apiKey: mine.key, + traceId: sessionA, + }); + } + for (let i = 0; i < 3; i++) { + lastSessionCallId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-pgn-b-${i}-${suffix}`, + apiKey: mine.key, + traceId: sessionB, + }); + } + await waitForSpendLog(request, lastSessionCallId); + await waitForSpendLog(request, soloIds[soloIds.length - 1]); + + // 36 calls in 28 session groups: 26 solos plus sessions of 7 and 3. + await openLogs(page); + const drawer = await openFilterDrawer(page); + await applyKeyAliasFilter(page, drawer, alias); + await setRowsPerPage(page, "25"); + + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 28", { timeout: 30_000 }); + await expect(requestLogsRows(page)).toHaveCount(25); + // The sessions are the newest groups, so their single representative rows sit on page 1. + await expect(requestLogsRows(page).filter({ hasText: sessionA })).toHaveCount(1); + await expect(requestLogsRows(page).filter({ hasText: sessionA })).toContainText("7"); + await expect(requestLogsRows(page).filter({ hasText: sessionB })).toHaveCount(1); + + await visibleTestId(page, "pagination-next").click(); + + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 26-28 of 28", { timeout: 30_000 }); + await expect(requestLogsRows(page)).toHaveCount(3); + await expect(requestLogsRows(page).filter({ hasText: sessionA })).toHaveCount(0); + await expect(requestLogsRows(page).filter({ hasText: sessionB })).toHaveCount(0); + }); + + test("two keys reusing one session id stay separate rows", async ({ page, request }) => { + const suffix = uniqueSuffix(); + const mine = await createVirtualKey(request, { key_alias: `e2e-logs-pgn-mine-${suffix}` }); + const theirs = await createVirtualKey(request, { key_alias: `e2e-logs-pgn-theirs-${suffix}` }); + const sharedSession = `sess-pgn-shared-${suffix}`; + + let lastId = ""; + for (let i = 0; i < 2; i++) { + lastId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-pgn-shared-mine-${i}-${suffix}`, + apiKey: mine.key, + traceId: sharedSession, + }); + } + lastId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-pgn-shared-theirs-${suffix}`, + apiKey: theirs.key, + traceId: sharedSession, + }); + await waitForSpendLog(request, lastId); + + await openLogs(page); + const drawer = await openFilterDrawer(page); + await drawer.getByPlaceholder("Enter session ID…").fill(sharedSession); + await drawer.getByRole("button", { name: "Apply Filters" }).click(); + await expect(drawer).not.toBeVisible({ timeout: 10_000 }); + + // One row per caller: reusing a session id must not merge two keys' activity into one row. + await expect(requestLogsRows(page).filter({ hasText: sharedSession })).toHaveCount(2, { timeout: 30_000 }); + + // And each row carries ITS key's totals: two calls badge the first key's row, + // while the other key's single call renders as a plain LLM row. + const mineRow = requestLogsRows(page).filter({ hasText: sharedSession }).filter({ hasText: mine.token }); + const theirsRow = requestLogsRows(page).filter({ hasText: sharedSession }).filter({ hasText: theirs.token }); + await expect(mineRow).toHaveCount(1); + await expect(theirsRow).toHaveCount(1); + await expect(mineRow.getByText("2", { exact: true })).toBeVisible(); + await expect(theirsRow.getByText("LLM", { exact: true })).toBeVisible(); + }); +}); 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/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py index 1a7fb1f3e41..abae26e02cd 100644 --- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py +++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py @@ -1,5 +1,5 @@ import httpx -from openai import OpenAI, BadRequestError, APIStatusError +from openai import OpenAI, BadRequestError, NotFoundError, APIStatusError import pytest @@ -105,10 +105,9 @@ def test_streaming_response(): assert len(collected_chunks) > 0 -def test_bad_request_error(): +def test_model_not_found_error(): client = get_test_client() - with pytest.raises(BadRequestError): - # Trigger error with invalid model name + with pytest.raises(NotFoundError): client.responses.create(model="non-existent-model", input="This should fail") 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/rust-python-harness/AGENTS.md b/tests/rust-python-harness/AGENTS.md new file mode 100644 index 00000000000..e9b17027ddc --- /dev/null +++ b/tests/rust-python-harness/AGENTS.md @@ -0,0 +1,44 @@ +# Expected Structure + +```text +tests/rust-python-harness/ +├── __main__.py +│ +├── strategies/ +│ ├── e2e_parity/ +│ │ ├── runner.py +│ │ ├── sdk/ +│ │ │ ├── ocr/ +│ │ │ ├── messages/ +│ │ │ ├── chat_completions/ +│ │ │ └── responses/ +│ │ └── gateway/ +│ │ +│ ├── trace_parity/ +│ │ ├── runner.py +│ │ ├── sdk/ +│ │ └── gateway/ +│ │ +│ └── unit_tests/ +│ ├── runner.py +│ ├── mapping_validator.py +│ ├── python_runner.py +│ └── rust_runner.py +│ +└── shared/ + ├── parity/ + ├── tracing/ + └── reporting/ +``` + +- Run locally only; no CI integration +- `__main__.py` selects strategies and combines their reports; each strategy also runs independently +- `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses +- `trace_parity/` compares mapped operations, call counts, and required execution ordering +- E2E and trace runners share orchestration across `sdk/` and `gateway/`; surface-specific execution lives in those folders +- `unit_tests/runner.py` combines mapping validation, Python test runs, and native Rust test runs +- `mapping_validator.py` matches Python/Rust tests by agreed names or annotations and reports missing or ambiguous counterparts +- `python_runner.py` runs existing Python tests with Rust disabled and enabled in separate processes, verifies backend selection, and compares results +- `rust_runner.py` runs Cargo tests; native Rust unit tests stay beside their implementation +- `shared/` contains reusable parity, tracing, and reporting machinery +- Keep fixtures with their owning API and existing Python tests in their current locations diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 15b7ae9d07a..1778eca25ef 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1013,6 +1013,48 @@ def test_an_unmapped_exception_with_no_model_or_provider_is_a_connection_error(q assert "boom" in raised.value.message +def _raise_and_map( + model: str | None, original_exception: Exception, custom_llm_provider: str | None +) -> None: + """Calls exception_type() from inside the except block, as litellm/main.py does, + so traceback.format_exc() has a real stack.""" + try: + raise original_exception + except type(original_exception) as caught: + exception_type( + model=model, + original_exception=caught, + custom_llm_provider=custom_llm_provider, + ) + + +def test_an_unmapped_exception_message_keeps_traceback_for_sdk_callers(quiet_exception_mapping): + """Direct SDK callers debug unmapped provider exceptions with this traceback; + only the proxy's response boundary strips it.""" + with pytest.raises(litellm.APIConnectionError) as raised: + _raise_and_map( + model="MiniMax-M2.5", + original_exception=RuntimeError("socket hung up"), + custom_llm_provider="minimax", + ) + + assert "Traceback (most recent call last)" in raised.value.message + assert "test_exception_mapping_utils.py" in raised.value.message + + +def test_an_unmapped_exception_with_no_model_or_provider_message_keeps_traceback( + quiet_exception_mapping, +): + with pytest.raises(litellm.APIConnectionError) as raised: + _raise_and_map( + model=None, + original_exception=ValueError("boom"), + custom_llm_provider=None, + ) + + assert "Traceback (most recent call last)" in raised.value.message + + CONTEXT_WINDOW_MESSAGE = "This model's maximum context length is 4096 tokens." CONTENT_POLICY_MESSAGE = ( '{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}' diff --git a/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py b/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py new file mode 100644 index 00000000000..dd7e8fadcd8 --- /dev/null +++ b/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py @@ -0,0 +1,35 @@ +"""Tests for litellm/llms/a2a/chat/guardrail_translation/handler.py.""" + +import json + +from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey + + +def _text_event(text: str) -> str: + return json.dumps( + { + "jsonrpc": "2.0", + "id": "req-1", + "result": {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": text}]}, + } + ) + + +def _status_event() -> str: + return json.dumps({"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "status-update", "status": {}}}) + + +class TestA2AGuardrailHandlerStreamingScanKey: + def test_key_joins_the_text_of_every_message_event(self): + key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hello "), _text_event("world")]) + assert key == StreamingScanKey(texts=("hello world",)) + + def test_events_without_text_leave_the_key_unchanged(self): + handler = A2AGuardrailHandler() + events = [_text_event("hello")] + assert handler.get_streaming_scan_key(events + [_status_event()]) == handler.get_streaming_scan_key(events) + + def test_unparseable_items_are_ignored(self): + key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hi"), "not json", b"bytes"]) + assert key.texts == ("hi",) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 0fe7730e91e..3044a321aa6 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -13,6 +13,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.llms.anthropic.chat.guardrail_translation.handler import ( AnthropicMessagesHandler, ) @@ -1991,3 +1992,56 @@ class TestStructuredWriteBackKeepsToolResults: } later_blocks = [b for m in messages[tool_use_index + 1 :] for b in self._blocks(m)] assert {"type": "text", "text": "Now fetch the page."} in later_blocks + + +class TestAnthropicMessagesHandlerStreamingScanKey: + """get_streaming_scan_key mirrors what process_output_streaming_response would scan""" + + @staticmethod + def _sse(event_type, data): + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() + + def _text_delta(self, text): + return self._sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ) + + def test_key_is_empty_before_any_text_arrives(self): + head = self._sse("message_start", {"type": "message_start", "message": {"stop_reason": None}}) + key = AnthropicMessagesHandler().get_streaming_scan_key([head]) + assert key == StreamingScanKey(texts=("",)) + + def test_key_accumulates_text_deltas(self): + key = AnthropicMessagesHandler().get_streaming_scan_key([self._text_delta("hello "), self._text_delta("world")]) + assert key.texts == ("hello world",) + assert key.stream_ended is False + + def _stop(self, stop_reason): + return self._sse( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": stop_reason, "stop_sequence": None}, "usage": {}}, + ) + + def test_stop_without_tool_use_scans_the_same_payload(self): + handler = AnthropicMessagesHandler() + open_key = handler.get_streaming_scan_key([self._text_delta("hi")]) + ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), self._stop("end_turn")]) + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_tool_use_blocks_enter_the_key_once_the_stream_has_ended(self): + handler = AnthropicMessagesHandler() + tool_use = self._sse( + "content_block_start", + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}}, + }, + ) + open_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use]) + ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use, self._stop("tool_use")]) + assert open_key == StreamingScanKey(texts=("hi",)) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key 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/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/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index d7cc89868af..ec8725db5f7 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -8,6 +8,7 @@ import litellm from litellm import get_model_info, supports_reasoning, supports_vision from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import ( ChatCompletionMessageToolCall, @@ -235,6 +236,21 @@ def test_get_fireworks_session_id_prefers_litellm_session_id_over_trace_id(): ) +def test_get_fireworks_session_id_ignores_proxy_generated_session_id(): + """general_settings.missing_session_id: generate stamps a fresh id per request; sending it + as x-session-affinity would pin every request to a different node.""" + assert ( + get_fireworks_session_id( + { + "litellm_session_id": "generated-1", + "litellm_trace_id": "generated-1", + "metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}, + } + ) + is None + ) + + def test_handle_message_content_with_tool_calls(): config = FireworksAIConfig() message = Message( 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/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 7dd6065063a..cebab2512d0 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -12,6 +12,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.llms.openai.chat.guardrail_translation.handler import ( OpenAIChatCompletionsHandler, ) @@ -1643,3 +1644,74 @@ class TestCheckStreamingHasEnded: ) ] assert handler._check_streaming_has_ended(chunks) is True + + +class TestStreamingScanKey: + """get_streaming_scan_key identifies what a sampled round would scan so the + unified hook can skip rounds that would re-scan already-cleared text""" + + @staticmethod + def _chunk(content, finish_reason=None, index=0): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + return ModelResponseStream( + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)] + ) + + def test_key_carries_accumulated_text_and_open_stream(self): + handler = OpenAIChatCompletionsHandler() + key = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")]) + assert key == StreamingScanKey(texts=("hello",)) + + def test_chunks_without_text_leave_the_key_unchanged(self): + handler = OpenAIChatCompletionsHandler() + before = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")]) + after = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo"), self._chunk(None)]) + assert after == before + + def test_finish_chunk_without_tool_calls_scans_the_same_payload(self): + handler = OpenAIChatCompletionsHandler() + open_key = handler.get_streaming_scan_key([self._chunk("hi")]) + ended_key = handler.get_streaming_scan_key([self._chunk("hi"), self._chunk(None, finish_reason="stop")]) + assert open_key.stream_ended is False + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_tool_calls_only_enter_the_key_once_the_stream_has_ended(self): + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + handler = OpenAIChatCompletionsHandler() + tool_call = ChatCompletionDeltaToolCall( + id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}') + ) + tool_chunk = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason=None)] + ) + open_key = handler.get_streaming_scan_key([self._chunk("hi"), tool_chunk]) + ended_key = handler.get_streaming_scan_key( + [self._chunk("hi"), tool_chunk, self._chunk(None, finish_reason="stop")] + ) + assert open_key == StreamingScanKey(texts=("hi",)) + assert ended_key.texts == ("hi",) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key + + def test_text_after_the_first_choice_finishes_still_changes_the_key(self): + handler = OpenAIChatCompletionsHandler() + first_done = [self._chunk("a", index=0), self._chunk("b", finish_reason="stop", index=0)] + key_at_first_finish = handler.get_streaming_scan_key(first_done) + key_after_more_text = handler.get_streaming_scan_key(first_done + [self._chunk("y", index=1)]) + assert key_at_first_finish.stream_ended is True + assert key_after_more_text.stream_ended is True + assert key_after_more_text != key_at_first_finish + + def test_non_stream_items_are_ignored(self): + handler = OpenAIChatCompletionsHandler() + key = handler.get_streaming_scan_key([self._chunk("hi"), b"data: [DONE]"]) + assert key.texts == ("hi",) 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..295121167d6 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""" @@ -1537,3 +1731,93 @@ class TestBuildBlockSseChunks: dones = [payload for payload in payloads if payload["type"] == "response.output_item.done"] assert len(dones) == 1 assert dones[0]["item"]["content"][0]["text"] == "Blocked by policy." + + +class TestOpenAIResponsesHandlerStreamingScanKey: + """get_streaming_scan_key mirrors what process_output_streaming_response would scan""" + + @staticmethod + def _delta(sequence_number, text): + return { + "type": "response.output_text.delta", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + } + + def test_no_events_yields_no_key(self): + assert OpenAIResponsesHandler().get_streaming_scan_key([]) is None + + def test_key_accumulates_deltas_while_the_stream_is_open(self): + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey + + key = OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hel"), self._delta(1, "lo")]) + assert key == StreamingScanKey(texts=("hello",)) + + def test_typed_delta_events_accumulate_like_dicts(self): + from litellm.types.llms.openai import OutputTextDeltaEvent + + events = [ + OutputTextDeltaEvent( + type="response.output_text.delta", + item_id="msg_1", + output_index=0, + content_index=0, + delta=text, + sequence_number=i, + ) + for i, text in enumerate(("hel", "lo")) + ] + key = OpenAIResponsesHandler().get_streaming_scan_key(events) + assert key.texts == ("hello",) + assert key.stream_ended is False + + def test_events_without_text_leave_the_key_unchanged(self): + handler = OpenAIResponsesHandler() + events = [self._delta(0, "hi")] + quiet = events + [{"type": "response.in_progress", "sequence_number": 1}] + assert handler.get_streaming_scan_key(quiet) == handler.get_streaming_scan_key(events) + + @staticmethod + def _completed(sequence_number, output): + return {"type": "response.completed", "sequence_number": sequence_number, "response": {"output": output}} + + def test_completed_event_keys_on_the_final_output_text(self): + handler = OpenAIResponsesHandler() + message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]} + open_key = handler.get_streaming_scan_key([self._delta(0, "hi")]) + ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message])]) + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_completed_event_with_a_function_call_changes_the_key(self): + handler = OpenAIResponsesHandler() + message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]} + function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"} + open_key = handler.get_streaming_scan_key([self._delta(0, "hi")]) + ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message, function_call])]) + assert ended_key.texts == ("hi",) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key + + def test_completed_event_reads_every_output_text_part(self): + from litellm.types.responses.main import GenericResponseOutputItem, OutputText + + item = GenericResponseOutputItem( + type="message", + id="msg_1", + status="completed", + role="assistant", + content=[ + OutputText(type="output_text", text="one", annotations=[]), + OutputText(type="output_text", text="two", annotations=[]), + ], + ) + key = OpenAIResponsesHandler().get_streaming_scan_key([self._completed(0, [item])]) + assert key.texts == ("one", "two") + + def test_output_item_done_round_is_never_deduped(self): + done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}} + assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None 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/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/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py index 5f277db2f72..db3a1a386a3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py @@ -7,6 +7,7 @@ cache's hit/single-flight behavior. Each assertion fails under a real mutation o """ import asyncio +import gc import json from unittest.mock import AsyncMock, MagicMock, patch @@ -359,6 +360,110 @@ async def test_cache_invalidate_only_evicts_the_named_key(): assert calls == 2 +@pytest.mark.asyncio +async def test_cache_invalidate_mid_compute_is_not_overwritten_by_that_compute(): + """A bearer minted before an invalidation must never be served after it. + + The compute is suspended at the token endpoint when the invalidation lands, so its write is + the one that would resurrect the evicted bearer for the rest of its TTL. The caller it was + minted for still gets it; the *cache* is what the invalidation is about. + """ + cache = ExchangedTokenCache() + mint_started, release_mint = asyncio.Event(), asyncio.Event() + + async def slow_mint(): + mint_started.set() + await release_mint.wait() + return _ok_token("bearer-minted-before-invalidation") + + async def re_mint(): + return _ok_token("bearer-minted-after-invalidation") + + in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp")) + await mint_started.wait() + + assert not in_flight.done() + cache.invalidate("slot") + release_mint.set() + + raced = await in_flight + assert isinstance(raced, Ok) and raced.ok == "bearer-minted-before-invalidation" + + after = await cache.get_or_compute("slot", re_mint, fingerprint="fp") + assert isinstance(after, Ok) and after.ok == "bearer-minted-after-invalidation" + + +@pytest.mark.asyncio +async def test_cache_invalidate_mid_compute_survives_garbage_collection(): + """The record of an invalidation must outlive a collection cycle taken mid-compute. + + Per-key state is held weakly so idle keys do not accumulate. If the state a compute checks + before writing were collectible while that compute is suspended, the check would read as + "nothing was invalidated" and the stale write would land; the running compute has to pin it. + """ + cache = ExchangedTokenCache() + mint_started, release_mint = asyncio.Event(), asyncio.Event() + + async def slow_mint(): + mint_started.set() + await release_mint.wait() + return _ok_token("bearer-minted-before-invalidation") + + async def re_mint(): + return _ok_token("bearer-minted-after-invalidation") + + in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp")) + await mint_started.wait() + + assert not in_flight.done() + cache.invalidate("slot") + gc.collect() + release_mint.set() + await in_flight + + after = await cache.get_or_compute("slot", re_mint, fingerprint="fp") + assert isinstance(after, Ok) and after.ok == "bearer-minted-after-invalidation" + + +@pytest.mark.asyncio +async def test_cache_stores_a_compute_that_started_after_the_invalidation(): + """Only the mint that predates the invalidation loses its write. + + A caller queued behind the single-flight lock computes after the eviction, so its token is + fresh and must be cached; otherwise the fix would trade one stale bearer for re-minting on + every subsequent resolution. + """ + cache = ExchangedTokenCache() + mint_started, release_mint = asyncio.Event(), asyncio.Event() + + async def slow_mint(): + mint_started.set() + await release_mint.wait() + return _ok_token("bearer-minted-before-invalidation") + + async def re_mint(): + return _ok_token("bearer-minted-after-invalidation") + + async def must_not_run(): + pytest.fail("the mint that followed the invalidation should have been cached") + + in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp")) + await mint_started.wait() + queued = asyncio.create_task(cache.get_or_compute("slot", re_mint, fingerprint="fp")) + await asyncio.sleep(0) + + assert not queued.done() + cache.invalidate("slot") + release_mint.set() + + raced, fresh = await asyncio.gather(in_flight, queued) + assert isinstance(raced, Ok) and raced.ok == "bearer-minted-before-invalidation" + assert isinstance(fresh, Ok) and fresh.ok == "bearer-minted-after-invalidation" + + served = await cache.get_or_compute("slot", must_not_run, fingerprint="fp") + assert isinstance(served, Ok) and served.ok == "bearer-minted-after-invalidation" + + @pytest.mark.asyncio async def test_cache_does_not_store_a_failed_compute(): cache = ExchangedTokenCache() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 82f74cda835..3f6d8f8837c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -8198,6 +8198,79 @@ class TestPreemptive401ModeAware: await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False) +def _make_obo_server(alias: str) -> MCPServer: + return MCPServer( + server_id=f"id-{alias}", + name=alias, + alias=alias, + server_name=alias, + url=f"https://{alias}.test/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.test/token", + client_id="cid", + client_secret="csecret", + mcp_info={"server_name": alias}, + ) + + +class TestOboPreflightScopedToAllowedServers: + """The connect-time OBO exchange is an outbound IdP call whose result is cached, so it must + only run for a server the caller's key resolves to through the allowed set, not for any + server the requested path happens to name.""" + + SUBJECT_HEADERS = {"Authorization": "Bearer upstream-subject-token"} + + async def _run(self, requested: MCPServer, allowed: list[MCPServer], user_api_key_auth: UserAPIKeyAuth | None): + from litellm.proxy._experimental.mcp_server import server as server_module + + allowed_lookup = AsyncMock(return_value=allowed) + preflight = AsyncMock() + with ( + patch.object( # test-quality-ok: route handler reads the module-level manager, no injection seam + server_module.global_mcp_server_manager, "get_mcp_server_by_name", return_value=requested + ), + patch.object( # test-quality-ok: the exchanger is the observable; a real one would call an IdP + server_module.global_mcp_server_manager, "preflight_token_exchange", preflight + ), + patch.object( # test-quality-ok: allowed-set resolution needs the DB; the test controls its answer + server_module, "_get_allowed_mcp_servers", allowed_lookup + ), + ): + await server_module._raise_preemptive_401_for_unauthenticated_servers( + scope={"type": "http", "method": "POST", "path": f"/mcp/{requested.alias}", "headers": []}, + mcp_servers=[requested.alias], + oauth2_headers=self.SUBJECT_HEADERS, + mcp_server_auth_headers=None, + user_api_key_auth=user_api_key_auth, + client_ip="10.0.0.7", + ) + return allowed_lookup, preflight + + @pytest.mark.asyncio + async def test_unentitled_key_never_reaches_the_exchanger(self): + requested = _make_obo_server("obo_tools") + key = UserAPIKeyAuth(api_key="sk-plain-only") + + allowed_lookup, preflight = await self._run( + requested, allowed=[_make_obo_server("plain_tools")], user_api_key_auth=key + ) + + preflight.assert_not_awaited() + allowed_lookup.assert_awaited_once_with( + user_api_key_auth=key, mcp_servers=[requested.alias], client_ip="10.0.0.7" + ) + + @pytest.mark.asyncio + async def test_entitled_key_still_exchanges_at_connect(self): + requested = _make_obo_server("obo_tools") + key = UserAPIKeyAuth(api_key="sk-obo") + + _, preflight = await self._run(requested, allowed=[requested], user_api_key_auth=key) + + preflight.assert_awaited_once_with(server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key) + + @pytest.mark.asyncio async def test_post_mcp_call_guardrails_return_the_rewritten_result(): """The result a post_mcp_call guardrail rewrote must be what the caller sends back.""" diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 1d0a99b8e0a..821323e722c 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -59,6 +59,7 @@ def _mock_cli_sso_start_response( login_id: str = "cli-session-uuid-456", poll_secret: str = "poll-secret", user_code: str = "ABCD-EFGH", + **extra_fields: object, ) -> Mock: mock_response = Mock() mock_response.status_code = 200 @@ -66,6 +67,7 @@ def _mock_cli_sso_start_response( "login_id": login_id, "poll_secret": poll_secret, "user_code": user_code, + **extra_fields, } mock_response.raise_for_status = Mock() return mock_response @@ -333,7 +335,9 @@ class TestLoginCommand: call_args = mock_browser.call_args[0][0] assert "https://test.example.com/sso/key/generate" in call_args assert "cli-test-uuid-123" in call_args + assert "user_code" not in call_args assert "Verification code: ABCD-EFGH" in result.output + assert "pre-filled in the browser" not in result.output mock_post.assert_called_once() mock_get.assert_called() assert mock_get.call_args.kwargs["headers"] == {"x-litellm-cli-poll-secret": "poll-secret"} @@ -347,6 +351,72 @@ class TestLoginCommand: # Verify commands were shown mock_show_commands.assert_called_once() + def test_login_prefills_the_code_in_the_browser_when_the_proxy_advertises_it( + self, isolated_home, secret_vault_factory + ) -> None: + vault = secret_vault_factory() + poll_response = Mock() + poll_response.status_code = 200 + poll_response.json.return_value = { + "status": "ready", + "key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt", + "user_id": "test-user-123", + "team_id": "team-1", + "teams": ["team-1"], + } + start_response = _mock_cli_sso_start_response( + login_id="cli-test-uuid-123", + verification_uri_complete=( + "https://internal-hostname.example.com/sso/key/generate" + "?source=litellm-cli&key=cli-test-uuid-123&user_code=ABCD-EFGH" + ), + ) + + with ( + patch("webbrowser.open") as mock_browser, + patch("requests.post", return_value=start_response), + patch("requests.get", return_value=poll_response), + ): + result = self.runner.invoke(login, obj={"base_url": "https://test.example.com", "secret_vault": vault}) + + assert result.exit_code == 0, result.output + assert json.loads(vault.blob)["key"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt" + assert json.loads((isolated_home / ".litellm" / "token.json").read_text())["user_id"] == "test-user-123" + opened_url = mock_browser.call_args[0][0] + assert opened_url.startswith("https://test.example.com/sso/key/generate?") + assert "internal-hostname" not in opened_url + assert "key=cli-test-uuid-123" in opened_url + assert "user_code=ABCD-EFGH" in opened_url + assert "Verification code: ABCD-EFGH (pre-filled in the browser, check it matches)" in result.output + + def test_login_keeps_the_code_out_of_the_url_when_the_proxy_sends_a_non_url_verification_uri( + self, secret_vault_factory + ) -> None: + poll_response = Mock() + poll_response.status_code = 200 + poll_response.json.return_value = { + "status": "ready", + "key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt", + "user_id": "test-user-123", + "team_id": "team-1", + "teams": ["team-1"], + } + for advertised in (None, True): + start_response = _mock_cli_sso_start_response(verification_uri_complete=advertised) + + with ( + patch("webbrowser.open") as mock_browser, + patch("requests.post", return_value=start_response), + patch("requests.get", return_value=poll_response), + ): + result = self.runner.invoke( + login, obj={"base_url": "https://test.example.com", "secret_vault": secret_vault_factory()} + ) + + assert result.exit_code == 0, result.output + assert "user_code" not in mock_browser.call_args[0][0] + assert "pre-filled in the browser" not in result.output + def test_login_timeout(self): """Test login timeout scenario""" mock_context = Mock() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index 476d443d8d8..cb6772977ec 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -310,8 +310,8 @@ def _make_stream_chunk(content: str, finish_reason=None): @pytest.mark.asyncio async def test_openai_moderation_streaming_default_uses_sampled_cadence(): """Default config samples every 5th streamed chunk and runs a final aggregate - pass after the stream ends. 10 chunks → sampled at chunks 5 and 10 → 2 in-stream - calls, plus 1 final = 3 total. + pass after the stream ends. 10 chunks are sampled at 5 and 10; the end-of-stream + round is skipped because chunk 10 already scanned the full text, for 2 total calls """ import litellm @@ -370,8 +370,9 @@ async def test_openai_moderation_streaming_default_uses_sampled_cadence(): ): pass - assert patched_make_request.await_count == 3, ( - f"Expected 3 moderation calls (2 sampled at chunks 5 / 10 + 1 final), " + assert patched_make_request.await_count == 2, ( + f"Expected 2 moderation calls (2 sampled at chunks 5 / 10; " + f"the end-of-stream round is skipped because chunk 10 already scanned the full text), " f"got {patched_make_request.await_count}" ) @@ -448,7 +449,8 @@ async def test_openai_moderation_streaming_end_of_stream_only_opt_in_calls_moder @pytest.mark.asyncio async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disabled(): """With streaming_end_of_stream_only=False and streaming_sampling_rate=2, - moderation runs every 2nd chunk during the stream, plus once more at end. + moderation runs every 2nd chunk during the stream. The terminal chunk scan covers + the final aggregate, for 3 total calls """ import litellm @@ -509,9 +511,8 @@ async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disab ): pass - # 6 chunks, sampling_rate=2 → in-stream calls at chunks 2, 4, 6 (3 calls), - # plus the final aggregate pass after the stream ends (1 call) = 4 total. - assert patched_make_request.await_count == 4, ( - f"Expected 4 moderation calls (3 sampled + 1 final aggregate), " + assert patched_make_request.await_count == 3, ( + f"Expected 3 moderation calls (3 sampled; the end-of-stream round is skipped " + f"because chunk 6 already scanned the full text), " f"got {patched_make_request.await_count}" ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index ec7854b9a35..1b50ea53db2 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -3,7 +3,9 @@ from unittest.mock import patch import httpx import pytest from fastapi import HTTPException +from pydantic import ValidationError +import litellm from litellm.exceptions import Timeout from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr import initialize_guardrail @@ -12,8 +14,8 @@ from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr.crowdstrike_aidr CrowdStrikeAIDRHandler, ) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 -from litellm.types.guardrails import Guardrail, LitellmParams -from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse +from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams +from litellm.types.utils import Delta, GenericGuardrailAPIInputs, ModelResponse, ModelResponseStream @pytest.fixture @@ -1578,3 +1580,139 @@ async def test_unparseable_transformed_response_fails_closed_under_fail_open() - assert exc_info.value.status_code == 500 assert "failing closed" in exc_info.value.detail["error"] + + +def _initialize_from_config(**litellm_params_kwargs: object) -> CrowdStrikeAIDRHandler: + litellm_params = LitellmParams( + guardrail="crowdstrike_aidr", + api_key="pts_crowdstrike_tokenid", + api_base="https://api.crowdstrike.com/aidr/aiguard", + default_on=True, + **litellm_params_kwargs, + ) + guardrail = Guardrail(guardrail_name="crowdstrike-aidr-guard", litellm_params=litellm_params) + return initialize_guardrail(litellm_params=litellm_params, guardrail=guardrail) + + +@pytest.mark.parametrize( + ("mode", "runs_pre_call", "runs_post_call"), + [("post_call", False, True), ("pre_call", True, False), (["pre_call", "post_call"], True, True)], +) +def test_initialize_guardrail_honors_configured_mode( + mode: str | list[str], runs_pre_call: bool, runs_post_call: bool +) -> None: + handler = _initialize_from_config(mode=mode) + + assert handler.should_run_guardrail({}, GuardrailEventHooks.pre_call) is runs_pre_call + assert handler.should_run_guardrail({}, GuardrailEventHooks.post_call) is runs_post_call + + +def test_initialize_guardrail_rejects_unsupported_mode_instead_of_running_other_hooks() -> None: + with pytest.raises(ValueError, match="during_call is not in the supported event hooks"): + _initialize_from_config(mode="during_call") + + +def test_initialize_guardrail_defaults_streaming_params() -> None: + handler = _initialize_from_config(mode="post_call") + + assert handler.streaming_end_of_stream_only is False + assert handler.streaming_sampling_rate == 5 + + +@pytest.mark.parametrize( + "configured", + [ + {"streaming_end_of_stream_only": True, "streaming_sampling_rate": 50}, + {"optional_params": {"streaming_end_of_stream_only": True, "streaming_sampling_rate": 50}}, + ], +) +def test_initialize_guardrail_forwards_streaming_params(configured: dict[str, object]) -> None: + handler = _initialize_from_config(mode="post_call", **configured) + + assert handler.streaming_end_of_stream_only is True + assert handler.streaming_sampling_rate == 50 + + +def test_initialize_guardrail_rejects_non_positive_sampling_rate() -> None: + with pytest.raises(ValidationError): + _initialize_from_config(mode="post_call", streaming_sampling_rate=0) + + +def test_update_in_memory_litellm_params_reapplies_streaming_params() -> None: + handler = _initialize_from_config(mode="post_call") + + handler.update_in_memory_litellm_params( + LitellmParams( + guardrail="crowdstrike_aidr", + mode="post_call", + streaming_end_of_stream_only=True, + streaming_sampling_rate=7, + ) + ) + + assert handler.streaming_end_of_stream_only is True + assert handler.streaming_sampling_rate == 7 + + +def _stream_chunk(content: str, finish_reason: str | None) -> ModelResponseStream: + return ModelResponseStream( + model="gpt-4", + choices=[ + litellm.StreamingChoices( + index=0, delta=Delta(role="assistant", content=content), finish_reason=finish_reason + ) + ], + ) + + +async def _guard_calls_for_stream(handler: CrowdStrikeAIDRHandler, chunk_texts: list[str]) -> int: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails + + async def stream(): + for i, content in enumerate(chunk_texts): + yield _stream_chunk(content, "stop" if i == len(chunk_texts) - 1 else None) + + calls = 0 + + def _allow(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return httpx.Response( + status_code=200, json={"result": {"blocked": False, "transformed": False}}, request=request + ) + + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": handler, + "metadata": {"guardrails": ["crowdstrike-aidr-guard"]}, + } + async with httpx.AsyncClient(transport=httpx.MockTransport(_allow)) as client: + await handler.async_handler.close() + handler.async_handler.client = client + async for _ in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/chat/completions"), + response=stream(), + request_data=request_data, + ): + pass + return calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("configured", "expected_calls"), + [ + ({}, 3), + ({"streaming_sampling_rate": 2}, 6), + ({"streaming_end_of_stream_only": True}, 1), + ({"streaming_end_of_stream_only": True, "streaming_sampling_rate": 2}, 1), + ], +) +async def test_streaming_params_from_config_control_output_scan_cadence( + configured: dict[str, object], expected_calls: int +) -> None: + """10 chunks: default samples at 5 and 10 plus the final pass, rate 2 samples 5 times plus final, end-of-stream scans once.""" + handler = _initialize_from_config(mode="post_call", **configured) + + assert await _guard_calls_for_stream(handler, list("ABCDEFGHIJ")) == expected_calls diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 523ec1a37b4..83cc9ae8bb9 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -1517,7 +1517,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: @pytest.mark.asyncio async def test_streaming_default_uses_sampled_cadence(self): - """Default samples every 5th chunk + final pass: 10 chunks → calls at 5, 10, and final = 3.""" + """Default samples every 5th chunk. For 10 chunks, sampled scans at 5 and 10 + cover the full text, so the end-of-stream round is skipped and there are 2 calls + """ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -1566,8 +1568,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: ): pass - assert mock_post.await_count == 3, ( - f"Expected 3 guardrail calls (2 sampled at chunks 5 / 10 + 1 final), " + assert mock_post.await_count == 2, ( + f"Expected 2 guardrail calls (2 sampled at chunks 5 / 10; " + f"the end-of-stream round is skipped because chunk 10 already scanned the full text), " f"got {mock_post.await_count}" ) for call in mock_post.await_args_list: @@ -1631,7 +1634,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: @pytest.mark.asyncio async def test_streaming_sampling_rate_override(self): - """sampling_rate=2 on 6 chunks → in-stream at 2,4,6 plus final = 4 calls.""" + """sampling_rate=2 on 6 chunks. Scans at 2, 4, and 6 cover the full text, so + the end-of-stream round is skipped and there are 3 calls + """ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -1680,8 +1685,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: ): pass - assert mock_post.await_count == 4, ( - f"Expected 4 guardrail calls (3 sampled + 1 final aggregate), " + assert mock_post.await_count == 3, ( + f"Expected 3 guardrail calls (3 sampled; the end-of-stream round is skipped " + f"because chunk 6 already scanned the full text), " f"got {mock_post.await_count}" ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 8cad1c634a9..a28a2a71613 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1971,3 +1971,271 @@ class TestStreamingGuardrailInformationBucket: assert recorded[0]["guardrail_name"] == "audit-recorder" assert recorded[0]["guardrail_status"] == "success" assert request_data["metadata"]["user_api_key_user_id"] == "user-1" + + +class _ScanCountingGuardrail(CustomGuardrail): + """Pass-through guardrail that records every response-side scan payload.""" + + def __init__(self, *, sampling_rate=5, end_of_stream_only=False, buffer_until_moderated=False): + super().__init__(guardrail_name="scan-counter") + self.streaming_sampling_rate = sampling_rate + self.streaming_end_of_stream_only = end_of_stream_only + self.streaming_buffer_until_moderated = buffer_until_moderated + self.guardrail_config = {} + self.scans: tuple[dict[str, object], ...] = () + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.scans = ( + *self.scans, + { + "texts": list(inputs.get("texts") or []), + "tool_calls": list(inputs.get("tool_calls") or []), + "model": inputs.get("model"), + }, + ) + return inputs + + +def _responses_delta(sequence_number, text): + return { + "type": "response.output_text.delta", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + } + + +def _responses_tail(sequence_number, text): + return [ + { + "type": "response.output_text.done", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "text": text, + }, + { + "type": "response.completed", + "sequence_number": sequence_number + 1, + "response": { + "model": "gpt-5.6", + "output": [{"type": "message", "content": [{"type": "output_text", "text": text}]}], + }, + }, + ] + + +class TestStreamingScanDedup: + """A sampled round whose scan payload matches the previous round (or carries + no text yet) is skipped, so a stream is never re-scanned for output the + guardrail already cleared. Regression for LIT-6692.""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self, monkeypatch): + monkeypatch.setattr( + unified_module, + "endpoint_guardrail_translation_mappings", + load_guardrail_translation_mappings(), + ) + + @pytest.mark.asyncio + async def test_chat_terminal_chunk_on_sampled_index_is_scanned_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 3 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] + + @pytest.mark.asyncio + async def test_chat_round_with_unchanged_text_is_skipped(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [ + _stream_chunk("a"), + _stream_chunk("b"), + _stream_chunk("c"), + _stream_chunk(None), + _stream_chunk(None), + _stream_chunk(None), + _stream_chunk("d", finish_reason="stop"), + ] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 7 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abcd"]] + + @pytest.mark.asyncio + async def test_chat_finish_chunk_right_after_a_sampled_round_is_not_rescanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), _stream_chunk(None, finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 4 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] + + @pytest.mark.asyncio + async def test_chat_finish_chunk_carrying_tool_calls_is_still_scanned(self): + from litellm.types.utils import ChatCompletionDeltaToolCall, Function + + guardrail = _ScanCountingGuardrail(sampling_rate=3) + tool_call = ChatCompletionDeltaToolCall( + id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}') + ) + finish = ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason="tool_calls") + ] + ) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), finish] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 4 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abc"]] + assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"] + + @pytest.mark.asyncio + async def test_chat_second_choice_finishing_later_still_gets_the_end_scan(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [ + _stream_chunk("a", index=0), + _stream_chunk("x", index=1), + _stream_chunk("b", finish_reason="stop", index=0), + _stream_chunk("y", index=1), + _stream_chunk("z", finish_reason="stop", index=1), + ] + + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(guardrail.scans) == 2 + assert any("yz" in text for text in guardrail.scans[-1]["texts"]) + + @pytest.mark.asyncio + async def test_responses_completed_event_on_sampled_index_is_scanned_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(8)] + full_text = "".join(f"t{i}" for i in range(8)) + chunks = deltas + _responses_tail(8, full_text) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 10 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], [full_text]] + assert guardrail.scans[-1]["model"] == "gpt-5.6" + + @pytest.mark.asyncio + async def test_responses_completed_right_after_a_sampled_round_is_not_rescanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + chunks = deltas + _responses_tail(5, "t0t1t2t3t4") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 7 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"]] + + @pytest.mark.asyncio + async def test_responses_completed_carrying_a_function_call_is_still_scanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + completed = { + "type": "response.completed", + "sequence_number": 5, + "response": { + "model": "gpt-5.6", + "output": [ + {"type": "message", "content": [{"type": "output_text", "text": "t0t1t2t3t4"}]}, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + ], + }, + } + chunks = deltas + [completed] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 6 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], ["t0t1t2t3t4"]] + assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"] + + @pytest.mark.asyncio + async def test_responses_round_with_unchanged_text_is_skipped(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + quiet = [{"type": "response.in_progress", "sequence_number": i} for i in range(5, 10)] + chunks = deltas + quiet + _responses_tail(10, "t0t1t2t3t4") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 12 + assert guardrail.scans == ({"texts": ["t0t1t2t3t4"], "tool_calls": [], "model": None},) + + @pytest.mark.asyncio + async def test_responses_tool_call_done_event_is_still_scanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2) + tool_call_done = { + "type": "response.output_item.done", + "sequence_number": 1, + "output_index": 1, + "item": { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + } + chunks = [_responses_delta(0, "hi"), tool_call_done] + _responses_tail(2, "hi") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 4 + assert len(guardrail.scans) == 2 + assert [call["function"]["name"] for call in guardrail.scans[0]["tool_calls"]] == ["get_weather"] + assert guardrail.scans[1]["texts"] == ["hi"] + + @pytest.mark.asyncio + async def test_anthropic_skips_empty_round_and_terminal_duplicate(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2) + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages") + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]] + + @pytest.mark.asyncio + async def test_end_of_stream_only_still_scans_exactly_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2, end_of_stream_only=True) + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages") + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]] + + @pytest.mark.asyncio + async def test_buffer_until_moderated_still_scans_exactly_once_and_releases_every_chunk(self): + guardrail = _ScanCountingGuardrail(sampling_rate=1, buffer_until_moderated=True) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index a222e22f6d0..23439bf7b23 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -104,7 +104,7 @@ def mock_in_memory_handler(mocker): mock_handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL mock_handler.get_source.return_value = "config" mock_handler.initialize_guardrail = mocker.Mock() - mock_handler.update_in_memory_guardrail = mocker.Mock() + mock_handler.sync_guardrail_from_db = mocker.Mock() mock_handler.delete_in_memory_guardrail = mocker.Mock() mock_handler.reconcile_db_guardrails = mocker.Mock(return_value=[]) return mock_handler @@ -1036,13 +1036,15 @@ async def test_create_guardrail_endpoint( "scenario,expected_result,expected_exception", [ ("success_with_sync", "test-db-guardrail", None), - ("success_sync_fails", "test-db-guardrail", None), + ("success_sync_fails_unexpected_error", "test-db-guardrail", None), + ("sync_fails_invalid_config", None, HTTPException), ("database_failure", None, HTTPException), ("no_prisma_client", None, HTTPException), ], ids=[ "success_with_immediate_sync", - "success_but_sync_fails", + "success_but_sync_fails_with_unexpected_error", + "sync_rejects_invalid_config", "database_error", "missing_prisma_client", ], @@ -1062,6 +1064,7 @@ async def test_update_guardrail_endpoint( mock_logger = None if scenario == "success_with_sync": mock_prisma_client = mocker.Mock() + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock() mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", @@ -1072,10 +1075,13 @@ async def test_update_guardrail_endpoint( mock_in_memory_handler, ) - elif scenario == "success_sync_fails": + elif scenario == "success_sync_fails_unexpected_error": + # A non-ValueError/TypeError failure is not a config-rejection signal, + # so it keeps the pre-existing swallow-and-warn behavior rather than + # rolling back the DB write. mock_prisma_client = mocker.Mock() - mock_in_memory_handler.update_in_memory_guardrail.side_effect = Exception( - "Sync failed" + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock( + side_effect=Exception("Sync failed") ) mock_logger = mocker.patch( "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" @@ -1091,6 +1097,25 @@ async def test_update_guardrail_endpoint( mock_in_memory_handler, ) + elif scenario == "sync_fails_invalid_config": + # Regression for the PUT half of the fix: a TypeError from the sync (the + # in-place update_in_memory_guardrail raised exactly this on every PUT) + # must roll back the DB write and surface a 422, not persist the + # rejected config with a 200. + mock_prisma_client = mocker.Mock() + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock( + side_effect=TypeError("vars() argument must have __dict__ attribute") + ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: reused pattern + mocker.patch( # test-quality-ok: reused pattern + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + mocker.patch( # test-quality-ok: reused pattern + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + elif scenario == "database_failure": mock_prisma_client = mocker.Mock() mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception( @@ -1119,6 +1144,16 @@ async def test_update_guardrail_endpoint( assert "Database error" in str(exc_info.value.detail) elif scenario == "no_prisma_client": assert "Prisma client not initialized" in str(exc_info.value.detail) + elif scenario == "sync_fails_invalid_config": + assert exc_info.value.status_code == 422 + assert "update rejected" in str(exc_info.value.detail) + # Rolled back: update_guardrail_in_db is called once for the + # rejected write and once more to restore the previous config. + assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2 + assert ( + mock_guardrail_registry.update_guardrail_in_db.call_args.kwargs["guardrail"] + == MOCK_DB_GUARDRAIL + ) else: result = await update_guardrail( @@ -1134,11 +1169,11 @@ async def test_update_guardrail_endpoint( prisma_client=mocker.ANY, ) - mock_in_memory_handler.update_in_memory_guardrail.assert_called_once_with( - guardrail_id="test-guardrail-id", guardrail=mocker.ANY + mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with( + guardrail=mocker.ANY ) - if scenario == "success_sync_fails": + if scenario == "success_sync_fails_unexpected_error": assert mock_logger is not None mock_logger.warning.assert_called_once() assert "Failed to update" in str(mock_logger.warning.call_args) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 24742e1bac2..56661b5b843 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -913,3 +913,96 @@ def test_reinitialize_guardrail_restores_previous_on_failure(): assert restored.guardrail_name == "restore-me" finally: registry_module.guardrail_initializer_registry.pop("restore_test", None) + + +def test_reinitialize_guardrail_raises_value_error_for_non_value_error_init_failures(): + """Regression for the LIT-6479 fix's 422 path: a constructor failure that is not + already a ValueError/TypeError (re.error from an invalid regex has neither in its + MRO) must still surface as ValueError, so the PUT/PATCH endpoints' rollback+422 + catch is exhaustive instead of warn-and-200 persisting a broken config.""" + import re + + from litellm.proxy.guardrails import guardrail_registry as registry_module + + def _initializer(litellm_params, guardrail): + if litellm_params.api_key == "bad-regex": + re.compile("([") + return CustomGuardrail( + guardrail_name=guardrail["guardrail_name"], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + registry_module.guardrail_initializer_registry["regex_test"] = _initializer + try: + handler = InMemoryGuardrailHandler() + created = handler.initialize_guardrail( + guardrail={ + "guardrail_name": "regex-me", + "litellm_params": {"guardrail": "regex_test", "mode": "pre_call", "api_key": "ok"}, + }, + ) + guardrail_id = created["guardrail_id"] + + with pytest.raises(ValueError, match="Guardrail initialization failed") as excinfo: + handler.reinitialize_guardrail( + guardrail={ + "guardrail_id": guardrail_id, + "guardrail_name": "regex-me", + "litellm_params": {"guardrail": "regex_test", "mode": "pre_call", "api_key": "bad-regex"}, + }, + ) + + assert isinstance(excinfo.value.__cause__, re.error) + assert guardrail_id in handler.IN_MEMORY_GUARDRAILS + restored = handler.guardrail_id_to_custom_guardrail[guardrail_id] + assert restored is not None and restored.guardrail_name == "regex-me" + finally: + registry_module.guardrail_initializer_registry.pop("regex_test", None) + + +def test_sync_guardrail_from_db_applies_db_dict_params_to_live_instance(): + """ + Regression for PUT /guardrails/{id}: the DB row arrives with litellm_params as + a plain jsonb dict, and the in-place update_in_memory_guardrail cast it to + LitellmParams without constructing one, so vars() raised and the running proxy + kept enforcing the stale config forever. The PUT endpoint now routes through + sync_guardrail_from_db, which must rebuild the live instance from the dict: + new blocked words compiled in, old ones gone, and the event hook re-derived + from mode (the base-class setattr path wrote self.mode while dispatch reads + self.event_hook, so only a full re-init applies a mode change). + """ + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + + handler = InMemoryGuardrailHandler() + gid = "66666666-6666-6666-6666-666666666666" + + def db_guardrail(word: str, mode: str) -> Guardrail: + return Guardrail( + guardrail_id=gid, + guardrail_name="cf-put-sync", + litellm_params={ + "guardrail": "litellm_content_filter", + "mode": mode, + "default_on": True, + "blocked_words": [{"keyword": word, "action": "BLOCK"}], + }, + ) + + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + handler.sync_guardrail_from_db(db_guardrail("foobarblock", "pre_call")) + handler.sync_guardrail_from_db(db_guardrail("quxnewblock", "during_call")) + + instance = handler.guardrail_id_to_custom_guardrail[gid] + assert isinstance(instance, ContentFilterGuardrail) + assert instance._check_blocked_words("hello QUXNEWBLOCK") is not None + assert instance._check_blocked_words("hello FOOBARBLOCK") is None + assert instance.event_hook == GuardrailEventHooks.during_call + assert instance.should_run_guardrail(data={}, event_type=GuardrailEventHooks.during_call) is True + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 8006f64ba41..551e27a18f5 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -1077,3 +1077,243 @@ def test_public_mcp_hub_does_not_expose_upstream_url(): assert all("url" not in item for item in data) assert secret_url not in response.text app.dependency_overrides.clear() + + + +@pytest.fixture +def reset_autorouter_presets_cache(): + from litellm.proxy.public_endpoints.public_endpoints import _AutoRouterPresetsCache + + _AutoRouterPresetsCache.presets = None + _AutoRouterPresetsCache.lock = None + yield + _AutoRouterPresetsCache.presets = None + _AutoRouterPresetsCache.lock = None + + +def test_get_autorouter_presets_local_mode_serves_bundled_catalog( + monkeypatch, reset_autorouter_presets_cache +): + monkeypatch.setenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", "True") + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + response = client.get("/public/autorouter_presets") + + assert response.status_code == 200 + payload = response.json() + assert "anthropic_family" in payload + for preset in payload.values(): + assert isinstance(preset["label"], str) + assert isinstance(preset["description"], str) + assert "tiers" in preset["complexity_router_config"] + + +@pytest.mark.asyncio +async def test_get_autorouter_presets_fetches_once_per_process( + monkeypatch, reset_autorouter_presets_cache +): + from litellm.proxy.public_endpoints.public_endpoints import ( + _AUTOROUTER_PRESETS_ADAPTER, + get_autorouter_presets, + ) + + monkeypatch.delenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", raising=False) + remote = _AUTOROUTER_PRESETS_ADAPTER.validate_python( + { + "remote_only": { + "label": "Remote Only", + "description": "from the remote catalog", + "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}}, + } + } + ) + calls = [] + + async def fake_fetch(url): + calls.append(url) + return remote + + first = await get_autorouter_presets(url="https://example.test/presets.json", fetch=fake_fetch) + second = await get_autorouter_presets(url="https://example.test/presets.json", fetch=fake_fetch) + + assert first == remote + assert second == remote + assert calls == ["https://example.test/presets.json"] + + +@pytest.mark.asyncio +async def test_get_autorouter_presets_single_flight_on_concurrent_cold_start( + monkeypatch, reset_autorouter_presets_cache +): + import asyncio + + from litellm.proxy.public_endpoints.public_endpoints import ( + _AUTOROUTER_PRESETS_ADAPTER, + get_autorouter_presets, + ) + + monkeypatch.delenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", raising=False) + remote = _AUTOROUTER_PRESETS_ADAPTER.validate_python( + { + "remote_only": { + "label": "Remote Only", + "description": "from the remote catalog", + "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}}, + } + } + ) + calls = [] + + async def slow_fetch(url): + calls.append(url) + await asyncio.sleep(0.05) + return remote + + results = await asyncio.gather( + get_autorouter_presets(url="https://example.test/presets.json", fetch=slow_fetch), + get_autorouter_presets(url="https://example.test/presets.json", fetch=slow_fetch), + get_autorouter_presets(url="https://example.test/presets.json", fetch=slow_fetch), + ) + + assert all(result == remote for result in results) + assert len(calls) == 1 + + +@pytest.mark.asyncio +async def test_get_autorouter_presets_caches_bundled_fallback_on_remote_failure( + monkeypatch, reset_autorouter_presets_cache +): + from litellm.proxy.public_endpoints.public_endpoints import get_autorouter_presets + + monkeypatch.delenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", raising=False) + calls = [] + + async def broken_fetch(url): + calls.append(url) + raise ValueError("remote catalog unavailable") + + first = await get_autorouter_presets(url="https://example.test/presets.json", fetch=broken_fetch) + second = await get_autorouter_presets(url="https://example.test/presets.json", fetch=broken_fetch) + + assert "anthropic_family" in first + assert second == first + assert len(calls) == 1 + + +@pytest.mark.asyncio +async def test_autorouter_presets_adapter_rejects_wrong_shapes(): + from pydantic import ValidationError + + from litellm.proxy.public_endpoints.public_endpoints import _AUTOROUTER_PRESETS_ADAPTER + + with pytest.raises(ValidationError): + _AUTOROUTER_PRESETS_ADAPTER.validate_python({"bad": {"label": "no description or config"}}) + with pytest.raises(ValidationError): + _AUTOROUTER_PRESETS_ADAPTER.validate_python(["not", "a", "mapping"]) + with pytest.raises(ValidationError): + _AUTOROUTER_PRESETS_ADAPTER.validate_python( + {"no_tiers": {"label": "L", "description": "D", "complexity_router_config": {}}} + ) + with pytest.raises(ValidationError): + _AUTOROUTER_PRESETS_ADAPTER.validate_python( + { + "missing_builtin_tier": { + "label": "L", + "description": "D", + "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"]}}, + } + } + ) + with pytest.raises(ValidationError): + _AUTOROUTER_PRESETS_ADAPTER.validate_python( + { + "unknown_tier_name": { + "label": "L", + "description": "D", + "complexity_router_config": { + "tiers": { + "SIMPLE": ["m1"], + "MEDIUM": ["m2"], + "COMPLEX": ["m3"], + "REASONING": ["m4"], + "ULTRA": ["m5"], + } + }, + } + } + ) + with pytest.raises(ValidationError): + _AUTOROUTER_PRESETS_ADAPTER.validate_python( + { + "bad_tiers": { + "label": "L", + "description": "D", + "complexity_router_config": {"tiers": "not-a-mapping"}, + } + } + ) + + +def test_get_autorouter_presets_passes_unknown_catalog_fields_through( + monkeypatch, reset_autorouter_presets_cache +): + from litellm.proxy.public_endpoints.public_endpoints import ( + _AUTOROUTER_PRESETS_ADAPTER, + _AutoRouterPresetsCache, + ) + + monkeypatch.delenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", raising=False) + _AutoRouterPresetsCache.presets = _AUTOROUTER_PRESETS_ADAPTER.validate_python( + { + "future_preset": { + "label": "Future", + "description": "carries fields this proxy version does not know", + "complexity_router_config": { + "tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}, + "future_config_knob": 3, + }, + "icon": "sparkles", + } + } + ) + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + response = client.get("/public/autorouter_presets") + + assert response.status_code == 200 + served = response.json()["future_preset"] + assert served["icon"] == "sparkles" + assert served["complexity_router_config"]["future_config_knob"] == 3 + assert served["complexity_router_config"]["tiers"]["SIMPLE"] == ["m1"] + + +@pytest.mark.asyncio +async def test_fetch_remote_autorouter_presets_parses_and_rejects_empty(monkeypatch): + import litellm.llms.custom_httpx.http_handler as http_handler_module + from litellm.proxy.public_endpoints.public_endpoints import _fetch_remote_autorouter_presets + + catalog = { + "remote_only": { + "label": "Remote Only", + "description": "from the remote catalog", + "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}}, + } + } + response = MagicMock() + response.raise_for_status = MagicMock() + response.json = MagicMock(return_value=catalog) + client = MagicMock() + client.get = AsyncMock(return_value=response) + monkeypatch.setattr(http_handler_module, "get_async_httpx_client", lambda llm_provider: client) + + presets = await _fetch_remote_autorouter_presets("https://example.test/presets.json") + assert presets["remote_only"].label == "Remote Only" + response.raise_for_status.assert_called_once() + + response.json = MagicMock(return_value={}) + with pytest.raises(ValueError, match="empty"): + await _fetch_remote_autorouter_presets("https://example.test/presets.json") diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index a0dcbf802ef..30b086bab61 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3959,18 +3959,18 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts(): ] mock_prisma = MagicMock() - mock_prisma.db.litellm_spendlogs.group_by = AsyncMock( - return_value=[ - {"session_id": session_id, "_count": {"session_id": 2}}, - ] - ) + mock_prisma.db.litellm_spendlogs.group_by = AsyncMock() mock_prisma.db.query_raw = AsyncMock( return_value=[ { "session_id": session_id, + "api_key": api_key, + "session_total_count": 2, "session_total_spend": 15.0, "mcp_tool_call_count": 1, "mcp_tool_call_spend": 10.0, + "session_llm_count": 1, + "session_agent_count": 0, } ] ) @@ -3995,6 +3995,8 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts(): assert rows[0]["mcp_tool_call_spend"] == 10.0 assert rows[1]["mcp_tool_call_count"] == 1 assert rows[1]["mcp_tool_call_spend"] == 10.0 + assert rows[0]["session_llm_count"] == 1 + assert rows[0]["session_agent_count"] == 0 # Every row in the session carries the full session spend, not just its own assert rows[0]["session_total_spend"] == 15.0 @@ -4003,13 +4005,126 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts(): # Row without a session_id defaults to 1 assert rows[2]["session_total_count"] == 1 - # group_by should have been called with the session_id - mock_prisma.db.litellm_spendlogs.group_by.assert_called_once_with( - by=["session_id"], - where={"session_id": {"in": [session_id]}}, - count={"session_id": True}, + # The count is folded into the single aggregate query; no separate group_by call. + mock_prisma.db.litellm_spendlogs.group_by.assert_not_called() + + +@pytest.mark.asyncio +async def test_build_ui_spend_logs_response_key_split_session_gets_per_key_aggregates(): + """ + Two keys reusing one session id are separate rows under grouped pagination, + and each row must carry ITS key's totals, never the combined session's: + the aggregate query and its lookup are keyed by (session_id, api_key). + """ + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + _build_ui_spend_logs_response, ) + session_id = "sess-shared" + dict_rows = [ + {"request_id": "req-a", "session_id": session_id, "call_type": "completion", "api_key": "key-a"}, + {"request_id": "req-b", "session_id": session_id, "call_type": "completion", "api_key": "key-b"}, + ] + + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + { + "session_id": session_id, + "api_key": "key-a", + "session_total_count": 2, + "session_total_spend": 0.2, + "mcp_tool_call_count": 0, + "mcp_tool_call_spend": 0.0, + "session_cache_hit_count": 1, + "session_llm_count": 2, + "session_agent_count": 0, + }, + { + "session_id": session_id, + "api_key": "key-b", + "session_total_count": 1, + "session_total_spend": 0.7, + "mcp_tool_call_count": 0, + "mcp_tool_call_spend": 0.0, + "session_cache_hit_count": 0, + "session_llm_count": 1, + "session_agent_count": 0, + }, + ] + ) + + result = await _build_ui_spend_logs_response( + prisma_client=mock_prisma, + data=dict_rows, + total_records=2, + page=1, + page_size=50, + total_pages=1, + enrich_session_counts=True, + ) + + rows = result["data"] + assert [(r["session_total_count"], r["session_total_spend"]) for r in rows] == [(2, 0.2), (1, 0.7)] + assert [r["session_cache_hit_count"] for r in rows] == [1, 0] + assert [r["session_llm_count"] for r in rows] == [2, 1] + + aggregate_sql = mock_prisma.db.query_raw.mock_calls[0][1][0] + assert "GROUP BY session_id, api_key" in aggregate_sql + + +@pytest.mark.asyncio +async def test_build_ui_spend_logs_response_empty_api_key_keeps_session_aggregates(): + """ + The spend-log schema defaults api_key to an empty string, which is a real + group value and not a missing one: a multi-call session logged under an + empty key must keep its count and spend instead of degrading to a plain + single-call row. + """ + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + _build_ui_spend_logs_response, + ) + + session_id = "sess-keyless" + dict_rows = [ + {"request_id": "req-1", "session_id": session_id, "call_type": "completion", "api_key": ""}, + ] + + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + { + "session_id": session_id, + "api_key": "", + "session_total_count": 3, + "session_total_spend": 0.09, + "mcp_tool_call_count": 0, + "mcp_tool_call_spend": 0.0, + "session_cache_hit_count": 0, + "session_llm_count": 3, + "session_agent_count": 0, + } + ] + ) + + result = await _build_ui_spend_logs_response( + prisma_client=mock_prisma, + data=dict_rows, + total_records=1, + page=1, + page_size=50, + total_pages=1, + enrich_session_counts=True, + ) + + row = result["data"][0] + assert row["session_total_count"] == 3 + assert row["session_total_spend"] == 0.09 + + # The empty key must reach the aggregate's authorized-keys filter too. + _, call_args, _ = mock_prisma.db.query_raw.mock_calls[0] + assert call_args[2] == [""] + @pytest.mark.asyncio async def test_build_ui_spend_logs_response_sums_multi_round_session_spend(): @@ -4033,14 +4148,13 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_spend(): ] mock_prisma = MagicMock() - mock_prisma.db.litellm_spendlogs.group_by = AsyncMock( - return_value=[{"session_id": session_id, "_count": {"session_id": 3}}] - ) # The raw aggregate query returns the full session spend (0.01 + 0.02 + 0.03). mock_prisma.db.query_raw = AsyncMock( return_value=[ { "session_id": session_id, + "api_key": api_key, + "session_total_count": 3, "session_total_spend": 0.06, "mcp_tool_call_count": 0, "mcp_tool_call_spend": 0.0, @@ -4089,13 +4203,12 @@ async def test_build_ui_spend_logs_response_session_cache_hit_count(): ] mock_prisma = MagicMock() - mock_prisma.db.litellm_spendlogs.group_by = AsyncMock( - return_value=[{"session_id": session_id, "_count": {"session_id": 2}}] - ) mock_prisma.db.query_raw = AsyncMock( return_value=[ { "session_id": session_id, + "api_key": api_key, + "session_total_count": 2, "session_total_spend": 0.05, "mcp_tool_call_count": 0, "mcp_tool_call_spend": 0.0, diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py index ef68d9ce178..27e633099f0 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py @@ -274,6 +274,9 @@ async def test_spend_logs_ui_uses_bounded_count_not_full_scan(monkeypatch): "the page query must not carry a window count that forces a full-window " f"scan. SQL was:\n{page_sql}" ) + assert "GROUP BY" not in count_sql and "DISTINCT ON" not in page_sql, ( + "without group_by_session the endpoint must keep raw per-call pagination" + ) assert response["total"] == 137 assert response["total_is_capped"] is False @@ -499,3 +502,106 @@ async def test_global_spend_report_team_group_forwards_team_id(monkeypatch): params = mock_prisma.db.query_raw.call_args[0][1:] assert "team_x" in params, "team_id must be forwarded into the DB query params" assert "sl.team_id = $3" in sql, f"team query must filter on team_id. SQL was:\n{sql}" + + +@pytest.mark.asyncio +async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): + """ + With group_by_session=true, /spend/logs/ui must page and count SESSIONS, + not raw calls: the page query returns one representative row per session + (DISTINCT ON the session group key, preferring non-MCP calls, newest + first) and the bounded count counts groups. Otherwise the UI collapses a + server page of N calls into fewer visible rows while the footer still + claims N (issue #38060). + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + SPEND_LOGS_PAGINATION_COUNT_CAP, + ui_view_spend_logs, + ) + + page_rows = [ + {"request_id": "req-1", "metadata": "{}", "session_id": None}, + {"request_id": "req-2", "metadata": "{}", "session_id": None}, + ] + mock_prisma = _make_ui_spend_logs_mock(count_total=12, page_rows=page_rows) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + response = await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=1, + page_size=50, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + group_by_session=True, + ) + + group_key = "COALESCE(NULLIF(session_id, ''), request_id), api_key" + + count_call = mock_prisma.db.query_raw.call_args_list[0] + count_sql = count_call[0][0] + assert f"GROUP BY {group_key}" in count_sql, f"grouped total must count sessions. SQL was:\n{count_sql}" + assert "COUNT(*) OVER ()" not in count_sql + assert "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded" + assert count_call[0][-1] == SPEND_LOGS_PAGINATION_COUNT_CAP + 1 + + page_sql = mock_prisma.db.query_raw.call_args_list[1][0][0] + assert f"DISTINCT ON ({group_key})" in page_sql, f"page must return one row per session. SQL was:\n{page_sql}" + assert f"ORDER BY {group_key}, call_type IN ('call_mcp_tool', 'list_mcp_tools'), \"startTime\" DESC" in page_sql, ( + "the session representative must prefer the newest non-MCP call" + ) + assert "COUNT(*) OVER ()" not in page_sql + + assert response["total"] == 12 + assert response["total_is_capped"] is False + assert response["total_pages"] == 1 + + +@pytest.mark.asyncio +async def test_spend_logs_ui_request_id_lookup_with_grouping_returns_exact_row(monkeypatch): + """ + A request_id lookup with group_by_session=true must still resolve the + exact requested row: the filter runs before grouping, so the row is its + own group's representative and deep links keep working. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ui_view_spend_logs + + target_row = {"request_id": "req-deep-link", "metadata": "{}", "session_id": None} + mock_prisma = _make_ui_spend_logs_mock(count_total=1, page_rows=[target_row]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + response = await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id="req-deep-link", + start_date=None, + end_date=None, + page=1, + page_size=1, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + group_by_session=True, + ) + + page_call = mock_prisma.db.query_raw.call_args_list[1] + assert "request_id = $" in page_call[0][0], "the request_id equality filter must survive grouping" + assert "req-deep-link" in page_call[0] + assert [row["request_id"] for row in response["data"]] == ["req-deep-link"] + assert response["total"] == 1 diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index df14224af5c..6d6aad22ca3 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3013,7 +3013,7 @@ class TestHandleLLMApiExceptionDictDetail: assert "NotFoundError" in proxy_exc.message async def test_exception_with_status_code_propagates(self): - """Exception with a statically-set status_code should propagate it.""" + """Exception with a statically-set status_code should propagate it and its message.""" from litellm.llms.vertex_ai.common_utils import VertexAIError exc = VertexAIError( @@ -3022,12 +3022,30 @@ class TestHandleLLMApiExceptionDictDetail: ) proxy_exc = await self._invoke(exc) assert proxy_exc.code == "429" + assert proxy_exc.message == "Rate limit exceeded" async def test_exception_without_status_code_defaults_to_500(self): - """Exception with no status_code attribute defaults to 500.""" + """Exception with no status_code attribute defaults to 500; a message with nothing + to redact still reaches the client, since routes raise plain exceptions as validation text.""" exc = ValueError("Something broke") proxy_exc = await self._invoke(exc) assert proxy_exc.code == "500" + assert proxy_exc.message == "Something broke" + + async def test_unclassified_exception_redacts_internal_details_from_client_message(self): + """Regression for LIT-6747: an unclassified exception's credential, path, and host + must not reach the client.""" + exc = RuntimeError( + "Failed to connect to postgresql://litellm_internal:S3cr3tPGPass@10.20.30.40:5432/litellm_prod " + "(config file /etc/litellm/secrets/db.yaml)" + ) + proxy_exc = await self._invoke(exc) + assert proxy_exc.code == "500" + assert "S3cr3tPGPass" not in proxy_exc.message + assert "litellm_internal" not in proxy_exc.message + assert "10.20.30.40" not in proxy_exc.message + assert "/etc/litellm/secrets/db.yaml" not in proxy_exc.message + assert "REDACTED" in proxy_exc.message async def test_already_normalized_proxy_exception_is_honored(self): """A ProxyException raised mid-request (e.g. a guardrail block) is already @@ -3244,6 +3262,42 @@ class TestStreamCloseOnDisconnect: assert upstream.aclosed + async def test_async_streaming_data_generator_redacts_internal_details_on_error( + self, + ): + """Regression for LIT-6747: a mid-stream exception must not hand its raw text or a + traceback to serialize_error.""" + + class FailingUpstream: + def __aiter__(self): + return self + + async def __anext__(self): + raise RuntimeError( + "Failed to connect to postgresql://litellm_internal:S3cr3tPGPass@10.20.30.40:5432/litellm_prod " + "(config file /etc/litellm/secrets/db.yaml)" + ) + + ProxyLogging._callback_capabilities_cache.clear() + captured: list = [] + gen = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=FailingUpstream(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + request_data={"model": "mock-model"}, + proxy_logging_obj=ProxyLogging(user_api_key_cache=MagicMock()), + serialize_chunk=lambda c: "data: x\n\n", + serialize_error=lambda e: captured.append(e) or "data: error\n\n", + ) + + await gen.__anext__() + + assert len(captured) == 1 + message = captured[0].message + assert "S3cr3tPGPass" not in message + assert "10.20.30.40" not in message + assert "/etc/litellm/secrets/db.yaml" not in message + assert "Traceback (most recent call last)" not in message + @staticmethod def _request_that_disconnects() -> Request: async def receive(): diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index ee0e2014951..8366e5546a9 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -41,7 +41,9 @@ from litellm.litellm_core_utils.get_provider_specific_headers import ( from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( TRUSTED_CALLBACK_VARS_FIELD, ) +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import CredentialItem @@ -7719,3 +7721,177 @@ def test_stamped_model_access_groups_survive_the_litellm_metadata_merge(): } assert get_litellm_metadata_from_kwargs(kwargs)[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"] + + +def _request_for(path: str) -> MagicMock: + request = MagicMock(spec=Request) + request.scope = {"path": path} + request.url = MagicMock() + request.url.path = path + request.url.__str__.return_value = f"http://localhost{path}" + request.method = "POST" + request.query_params = {} + request.headers = {"Content-Type": "application/json"} + request.client = MagicMock() + request.client.host = "127.0.0.1" + return request + + +def _spend_log_session_id(data: dict[str, object]) -> str: + """Resolve session_id the way LiteLLM_SpendLogs does: standard_logging_payload.trace_id.""" + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_session_id_for_spend_log + + metadata = data["metadata"] + assert isinstance(metadata, dict) + litellm_params = get_litellm_params( + litellm_session_id=str(data["litellm_session_id"]) if "litellm_session_id" in data else None, + litellm_trace_id=str(data["litellm_trace_id"]) if "litellm_trace_id" in data else None, + metadata=metadata, + ) + trace_id = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=SimpleNamespace(litellm_trace_id="per-call-random-trace-id"), + litellm_params=litellm_params, + ) + return _get_session_id_for_spend_log(kwargs={}, standard_logging_payload={"trace_id": trace_id}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_correlation_in_logs", [False, True]) +async def test_missing_session_id_generate_makes_spend_log_and_callback_session_ids_agree( + monkeypatch: pytest.MonkeyPatch, request_correlation_in_logs: bool +): + """Without a session header, SpendLogs.session_id and the metadata.session_id that Langfuse logs + must be the same generated id, so cross-referencing the two by session_id works. The id is marked + as generated so affinity consumers (Fireworks x-session-affinity, router session pins) skip it.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", request_correlation_in_logs) + data = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]} + + updated = await add_litellm_data_to_request( + data=data, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "generate"}, + ) + + callback_session_id = updated["metadata"]["session_id"] + assert isinstance(callback_session_id, str) and len(callback_session_id) == 36 + assert _spend_log_session_id(updated) == callback_session_id + assert updated["metadata"][SESSION_ID_GENERATED_METADATA_KEY] is True + assert get_fireworks_session_id( + {"litellm_session_id": updated["litellm_session_id"], "metadata": updated["metadata"]} + ) is None + + +@pytest.mark.asyncio +async def test_missing_session_id_unset_keeps_legacy_divergence(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + ) + + assert "session_id" not in updated["metadata"] + assert "litellm_session_id" not in updated + assert _spend_log_session_id(updated) == "per-call-random-trace-id" + + +@pytest.mark.asyncio +async def test_missing_session_id_generate_reuses_traceparent_trace_id(): + """A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it.""" + request = _request_for("/v1/chat/completions") + request.headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "generate"}, + ) + + assert updated["metadata"]["session_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert _spend_log_session_id(updated) == "4bf92f3577b34da6a3ce929d0e0e4736" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", ["generate", "reject"]) +async def test_missing_session_id_policy_keeps_client_supplied_session_id(policy: str): + request = _request_for("/v1/chat/completions") + request.headers = {"x-litellm-session-id": "client-session-1"} + + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": policy}, + ) + + assert updated["litellm_session_id"] == "client-session-1" + assert updated["metadata"]["session_id"] == "client-session-1" + assert _spend_log_session_id(updated) == "client-session-1" + assert SESSION_ID_GENERATED_METADATA_KEY not in updated["metadata"] + assert ( + get_fireworks_session_id({"litellm_session_id": "client-session-1", "metadata": updated["metadata"]}) + == "client-session-1" + ) + + +@pytest.mark.asyncio +async def test_missing_session_id_reject_accepts_body_metadata_session_id(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": [], "metadata": {"session_id": "body-session-1"}}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert updated["metadata"]["session_id"] == "body-session-1" + + +@pytest.mark.asyncio +async def test_missing_session_id_reject_returns_400_without_session_id(): + with pytest.raises(ProxyException) as exc_info: + await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert exc_info.value.code == "400" + assert exc_info.value.param == "session_id" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/mcp/", "/mcp/tools", "/key/health"]) +async def test_missing_session_id_policy_skips_non_inference_routes(path: str): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o"}, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert "session_id" not in updated["metadata"] + + +@pytest.mark.asyncio +async def test_missing_session_id_unknown_value_is_ignored(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "typo"}, + ) + + assert "session_id" not in updated["metadata"] diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 3e70dee23b7..7b3528f3a68 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -95,6 +95,7 @@ class TestProxyInitializationHelpers: assert args["app"] == "litellm.proxy.proxy_server:app" assert args["host"] == "localhost" assert args["port"] == 8000 + assert args["server_header"] is False # Test with log_config args = ProxyInitializationHelpers._get_default_unvicorn_init_args( 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_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 9edcaaef034..c226c0b4d09 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -6,7 +6,7 @@ completion_start_time = end_time.""" import json from datetime import datetime from typing import Optional -from unittest.mock import Mock +from unittest.mock import Mock, patch import httpx import pytest @@ -378,3 +378,162 @@ def test_stamp_responses_usage_cost_survives_calculator_failure(): _stamp_responses_usage_cost(response, logging_obj) assert getattr(response.usage, "cost", None) is None + + +def _capture_dispatch(logged: list): + """Record the object handed to the success handlers. + + ``Mock(spec=LiteLLMLoggingObj).dispatch_success_handlers`` is an AsyncMock whose side effect + only runs when the coroutine is awaited, so capture with a plain function instead. + """ + + async def _noop() -> None: + return None + + def _dispatch(result, **kwargs): + logged.append(result) + return _noop() + + return _dispatch + + +def _headers_config(*, transform_hidden_params: Optional[dict] = None) -> Mock: + """Config whose completed event carries a real ResponsesAPIResponse, so the logging copy + performs a genuine model_dump/model_validate round trip.""" + mock_config = Mock(spec=BaseResponsesAPIConfig) + + def _transform(model, parsed_chunk, logging_obj): + evt_type = parsed_chunk.get("type") + if evt_type != "response.completed": + stub = Mock() + stub.type = evt_type + return stub + response = ResponsesAPIResponse( + id="resp_headers", + created_at=1, + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + if transform_hidden_params is not None: + response._hidden_params.update(transform_hidden_params) + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=response, + ) + + mock_config.transform_streaming_response.side_effect = _transform + return mock_config + + +def _make_header_iterator( + *, + headers: dict, + config: Mock, + logging_obj: LiteLLMLoggingObj, +) -> ResponsesAPIStreamingIterator: + async def aiter_bytes(): + yield _sse_event({"type": "response.completed"}) + + mock_response = Mock() + mock_response.headers = headers + mock_response.aiter_bytes = aiter_bytes + + return ResponsesAPIStreamingIterator( + response=mock_response, + model="gpt-4o-mini", + responses_api_provider_config=config, + logging_obj=logging_obj, + litellm_metadata={}, + custom_llm_provider="azure", + ) + + +@pytest.mark.asyncio +async def test_streaming_logging_response_carries_provider_response_headers(): + """LIT-6055: the provider headers the iterator captured must reach the logged response, so + custom loggers can read Azure's apim-request-id from the callback payload.""" + logging_obj = _logging_obj_stub() + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator( + headers={"apim-request-id": "azure-correlation-1", "x-ms-region": "East US 2"}, + config=_headers_config(), + logging_obj=logging_obj, + ) + async for _ in iterator: + pass + + assert len(logged) == 1 + hidden_params = logged[0].response._hidden_params + assert hidden_params["additional_headers"]["llm_provider-apim-request-id"] == "azure-correlation-1" + assert hidden_params["additional_headers"]["llm_provider-x-ms-region"] == "East US 2" + assert hidden_params["headers"]["apim-request-id"] == "azure-correlation-1" + # the proxy builds the client's response headers from the iterator's own dict, so the logged + # response must hold copies rather than alias it + assert hidden_params["additional_headers"] is not iterator._hidden_params["additional_headers"] + assert hidden_params["headers"] is not iterator._raw_response_headers + + +@pytest.mark.asyncio +async def test_streaming_logging_copy_preserves_transform_hidden_params(): + """LIT-6055: model_validate(model_dump()) drops pydantic private attributes, so headers a + provider transform already set on the response (fake_stream) must be re-applied.""" + logging_obj = _logging_obj_stub() + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator( + headers={}, + config=_headers_config( + transform_hidden_params={ + "additional_headers": {"llm_provider-apim-request-id": "from-transform"}, + "headers": {"apim-request-id": "from-transform"}, + "response_cost": 0.5, + } + ), + logging_obj=logging_obj, + ) + async for _ in iterator: + pass + + assert len(logged) == 1 + hidden_params = logged[0].response._hidden_params + assert hidden_params["additional_headers"]["llm_provider-apim-request-id"] == "from-transform" + assert hidden_params["headers"]["apim-request-id"] == "from-transform" + assert iterator.completed_response is not logged[0] + # only the header keys travel: response_cost would short-circuit the cost calculator + assert "response_cost" not in hidden_params + + +@pytest.mark.asyncio +async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched(): + """LIT-6055: when the logging copy falls back to the original event, the header restore must + not stamp logging-only state onto the object the caller is iterating.""" + logging_obj = _logging_obj_stub() + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator( + headers={"apim-request-id": "azure-correlation-1"}, + config=_headers_config(), + logging_obj=logging_obj, + ) + async for _ in iterator: + pass + + assert len(logged) == 1 + iterator._completed_response_logged = False + logged.clear() + with patch.object(type(iterator.completed_response), "model_dump", side_effect=ValueError("cannot serialize")): + iterator._log_completed_response(is_async=True) + + assert logged == [iterator.completed_response] + assert iterator.completed_response.response._hidden_params == {} diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 371631d6297..d7d02544efb 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -16,7 +16,7 @@ import litellm from litellm import Router from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.router_strategy.complexity_router.complexity_router import ( _CLASSIFICATION_CURRENT_MESSAGE_ONLY, _CLASSIFICATION_WITH_CONVERSATION, @@ -4274,6 +4274,26 @@ class TestSessionAffinity: assert first.model == "o1-preview" assert second.model == "gpt-4o-mini" + @pytest.mark.asyncio + async def test_proxy_generated_session_id_never_pins(self, mock_router_instance, session_affinity_config): + """A session id the proxy generated for a request that had none is per request, so + it must not create a pin even with session_affinity enabled.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + request_kwargs = {"metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}} + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.model == "o1-preview" + assert second.model == "gpt-4o-mini" + @pytest.mark.asyncio async def test_can_be_enabled_to_pin_every_later_turn(self, mock_router_instance, session_affinity_config): """Regression: session_affinity=True is the opt-in, so a shared session_id reuses the @@ -10089,6 +10109,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.""" diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py index a3772a276fa..cf48888600e 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py @@ -7,7 +7,7 @@ import json import litellm from litellm.caching.dual_cache import DualCache -from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) @@ -180,6 +180,47 @@ async def test_async_session_id_affinity_priority_over_user_key(): assert filtered[0]["model_info"]["id"] == "deployment-2" +@pytest.mark.asyncio +async def test_proxy_generated_session_id_does_not_pin_a_deployment(): + """A session id the proxy generated for a request that had none is per request, so a + pin stored under it must be ignored and none must be written.""" + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=123, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + ) + healthy_deployments = [ + {"model_name": "model_group", "litellm_params": {"model": "model_1"}, "model_info": {"id": "deployment-1"}}, + {"model_name": "model_group", "litellm_params": {"model": "model_2"}, "model_info": {"id": "deployment-2"}}, + ] + await cache.async_set_cache( + DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1"), + {"model_id": "deployment-2"}, + ) + request_kwargs = { + "metadata": {"user_api_key_hash": "user1", "session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True} + } + + filtered = await callback.async_filter_deployments( + model="model_group", healthy_deployments=healthy_deployments, messages=[], request_kwargs=request_kwargs + ) + await callback.async_pre_call_deployment_hook( + kwargs={ + "metadata": {**request_kwargs["metadata"], "deployment_model_name": "model_group"}, + "model_info": {"id": "deployment-1"}, + }, + call_type=None, + ) + + assert len(filtered) == 2 + assert await cache.async_get_cache( + DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1") + ) == {"model_id": "deployment-2"} + + MOCK_RESPONSES_API_RESPONSE = { "id": "resp_mock-resp-456", "object": "response", 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/test_check_py310_typing_imports.py b/tests/test_litellm/test_check_py310_typing_imports.py new file mode 100644 index 00000000000..de370326091 --- /dev/null +++ b/tests/test_litellm/test_check_py310_typing_imports.py @@ -0,0 +1,86 @@ +import sys +from pathlib import Path +from typing import Final + +_CODE_COVERAGE_DIR: Final[Path] = Path(__file__).resolve().parents[1] / "code_coverage_tests" +sys.path.insert(0, str(_CODE_COVERAGE_DIR)) # test-quality-ok: required to import checker from its source directory +import check_py310_typing_imports as checker # noqa: E402 # load checker from its source directory + + +def _scan(tmp_path: Path, source: str) -> tuple[object, ...]: + file_path = tmp_path / "fixture.py" + file_path.write_text(source, encoding="utf-8") + return checker.scan_file(file_path) + + +def test_typing_import_flags_python_311_name(tmp_path: Path) -> None: + violations = _scan(tmp_path, "from typing import NotRequired, TypedDict\n") + assert tuple(violation.name for violation in violations) == ("NotRequired",) + + +def test_typing_extensions_import_passes(tmp_path: Path) -> None: + assert _scan(tmp_path, "from typing_extensions import NotRequired\n") == () + + +def test_typing_attribute_flags_python_311_name(tmp_path: Path) -> None: + violations = _scan(tmp_path, "import typing\nx: typing.Self\n") + assert tuple(violation.name for violation in violations) == ("Self",) + + +def test_version_guarded_typing_import_passes(tmp_path: Path) -> None: + source = ( + "import sys\n" + "if sys.version_info >= (3, 11):\n" + " from typing import NotRequired\n" + "else:\n" + " from typing_extensions import NotRequired\n" + ) + assert _scan(tmp_path, source) == () + + +def test_python_310_branch_flags_typing_import(tmp_path: Path) -> None: + source = ( + "import sys\n" + "if sys.version_info >= (3, 11):\n" + " from typing_extensions import NotRequired\n" + "else:\n" + " from typing import NotRequired\n" + ) + violations = _scan(tmp_path, source) + assert tuple(violation.name for violation in violations) == ("NotRequired",) + + +def test_python_310_branch_is_exempt_for_less_than_guard(tmp_path: Path) -> None: + source = ( + "import sys\n" + "if sys.version_info < (3, 11):\n" + " from typing_extensions import NotRequired\n" + "else:\n" + " from typing import NotRequired\n" + ) + assert _scan(tmp_path, source) == () + + +def test_nearest_if_controls_version_guard(tmp_path: Path) -> None: + source = ( + "if sys.version_info >= (3, 11):\n" + " from typing import Self\n" + " x = 1\n" + "if True:\n" + " from typing import Self\n" + ) + violations = _scan(tmp_path, source) + assert tuple((violation.name, violation.line) for violation in violations) == (("Self", 5),) + + +def test_scan_directory_includes_proxy_extras(tmp_path: Path) -> None: + file_path = tmp_path / "litellm-proxy-extras" / "litellm_proxy_extras" / "m.py" + file_path.parent.mkdir(parents=True) + file_path.write_text("from typing import NotRequired\n", encoding="utf-8") + + violations = checker.scan_directory(tmp_path) + assert tuple((violation.name, violation.file) for violation in violations) == (("NotRequired", str(file_path)),) + + +def test_python_310_typing_name_passes(tmp_path: Path) -> None: + assert _scan(tmp_path, "from typing import Optional\n") == () diff --git a/tests/test_litellm/test_redact_string_in_error_paths.py b/tests/test_litellm/test_redact_string_in_error_paths.py index 6404db91acf..07d1ec5f523 100644 --- a/tests/test_litellm/test_redact_string_in_error_paths.py +++ b/tests/test_litellm/test_redact_string_in_error_paths.py @@ -172,8 +172,6 @@ class TestLLMHTTPHandlerRealtimeRedaction: class TestProxyStreamingDataGeneratorRedaction: - """Test _redact_string on traceback.format_exc() — the pattern at common_request_processing.py:1733.""" - def test_redact_traceback_format_exc(self): try: raise RuntimeError( 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/test_litellm/test_secret_redaction.py b/tests/test_litellm/test_secret_redaction.py index 303efcab9c7..9fa748edec1 100644 --- a/tests/test_litellm/test_secret_redaction.py +++ b/tests/test_litellm/test_secret_redaction.py @@ -2,6 +2,7 @@ import logging import logging.config import sys import time +import traceback from collections.abc import Callable from io import StringIO from typing import Final @@ -13,11 +14,12 @@ from litellm._logging import ( JsonFormatter, _redact_string, _secret_filter, + redact_internal_details_from_client_message, verbose_logger, verbose_proxy_logger, verbose_router_logger, ) -from litellm.litellm_core_utils.secret_redaction import redact_string +from litellm.litellm_core_utils.secret_redaction import redact_internal_details, redact_string SECRET = "sk-proj-abc123def456ghi789jklmnopqrst" @@ -657,3 +659,65 @@ def test_json_formatter_redacts_non_string_extra_values(extra): assert output.strip(), "no record captured" assert SECRET not in output, f"non-string extra leaked a secret: {output}" assert "REDACTED" in output + + +@pytest.mark.parametrize( + "text,leaked", + ( + ("config file /etc/litellm/secrets/db.yaml", "/etc/litellm/secrets/db.yaml"), + ("home dir /Users/admin/.litellm/master_key.txt", "/Users/admin/.litellm/master_key.txt"), + ("cache at /var/cache/litellm/tokens.db", "/var/cache/litellm/tokens.db"), + ("path C:\\Users\\admin\\secrets.env", "C:\\Users\\admin\\secrets.env"), + ("connecting to host 10.20.30.40", "10.20.30.40"), + ("connecting to host 192.168.1.5", "192.168.1.5"), + ("connecting to host 172.16.0.9", "172.16.0.9"), + ("connecting to host 127.0.0.1", "127.0.0.1"), + ("connecting to db-primary.internal", "db-primary.internal"), + ("connecting to redis.corp", "redis.corp"), + ), +) +def test_redact_internal_details_catches_paths_and_hostnames(text, leaked): + result = redact_internal_details(text) + assert leaked not in result, f"{leaked!r} was not redacted" + assert "REDACTED" in result + + +def test_redact_internal_details_leaves_public_hostnames_and_routes_alone(): + """litellm's own error messages rely on routes like /v1/models staying legible.""" + safe_strings = ( + "call https://api.openai.com/v1/chat/completions", + "/chat/completions: Invalid model name passed in model=gpt-9", + "Call `/v1/models` to view available models for your key", + "reducto:// file IDs are not accepted through the proxy OCR API", + ) + for text in safe_strings: + assert redact_internal_details(text) == text + + +def test_redact_internal_details_layers_on_top_of_credential_redaction(): + text = "postgresql://litellm_internal:S3cr3tPGPass@10.20.30.40:5432/litellm_prod" + result = redact_internal_details(text) + assert "S3cr3tPGPass" not in result + assert "10.20.30.40" not in result + + +def test_redact_internal_details_drops_embedded_traceback(): + """Regression for LIT-6747: the traceback exception_type() embeds for SDK callers + must never reach an HTTP client.""" + try: + raise RuntimeError("socket hung up") + except RuntimeError: + raw_tb = traceback.format_exc() + message = f"litellm.APIConnectionError: MinimaxException - socket hung up\n{raw_tb}" + + result = redact_internal_details(message) + + assert result == "litellm.APIConnectionError: MinimaxException - socket hung up" + assert "Traceback (most recent call last)" not in result + assert __file__.split("/")[-1] not in result + + +def test_redact_internal_details_from_client_message_respects_disable_flag(): + with patch("litellm._logging._ENABLE_SECRET_REDACTION", False): # test-quality-ok: the opt-out flag is the SUT + text = "config file /etc/litellm/secrets/db.yaml" + assert redact_internal_details_from_client_message(text) == text 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 839f7be2d50..f3c4c7760c6 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22358 + "limit": 22334 }, "LIT002": { - "limit": 26772 + "limit": 26763 }, "LIT003": { "limit": 261 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1039 + "limit": 1038 }, "LIT007": { "limit": 0 @@ -27,12 +27,12 @@ "limit": 0 }, "LIT010": { - "limit": 16486 + "limit": 16480 }, "LIT011": { - "limit": 5521 + "limit": 5520 }, "LIT012": { - "limit": 4495 + "limit": 4489 } } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets.ts new file mode 100644 index 00000000000..82f9cc2da99 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets.ts @@ -0,0 +1,16 @@ +import { AutoRouterPreset, hydratePresets } from "@/lib/autorouter_presets"; +import { getAutoRouterPresets } from "@/components/networking"; +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +const presetKeys = createQueryKeys("autoRouterPresets"); + +export const useAutoRouterPresets = () => { + const options = { + queryKey: presetKeys.list({}), + queryFn: async () => hydratePresets(await getAutoRouterPresets()), + staleTime: 24 * 60 * 60 * 1000, + gcTime: 24 * 60 * 60 * 1000, + }; + return useQuery(options); +}; 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 c2ed4b1f55a..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"; @@ -50,6 +51,7 @@ const HEURISTIC_V2_EXPLANATION = 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 " + @@ -213,6 +215,18 @@ const ClassifierTypeRadios: React.FC<{ + + + ); @@ -263,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); }; @@ -271,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 }); }; @@ -391,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 - classifierType === "llm" || classifierType === "heuristic_first"; + classifierType === "llm" || classifierType === "heuristic_first" || classifierType === "hybrid"; export type ClassifierFallback = "heuristic" | "default_model"; @@ -162,7 +162,8 @@ export const heuristicScoringRoleFor = ( classifierFallback: ClassifierFallback | undefined, ): HeuristicScoringRole => { if (classifierType === "heuristic_v2") return "never"; - if (classifierType === "heuristic" || classifierType === "heuristic_first") return "decides"; + if (classifierType === "heuristic" || classifierType === "heuristic_first" || classifierType === "hybrid") + return "decides"; return (classifierFallback ?? DEFAULT_CLASSIFIER_FALLBACK) === "heuristic" ? "fallback_only" : "never"; }; @@ -404,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; @@ -516,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/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index ba380662403..7769042a0c7 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -9,11 +9,19 @@ import { getSubmitBlockedReason } from "./add_auto_router_tab"; import { buildModelAvailability } from "@/lib/autorouter_presets"; import { testAutoRouterRouting } from "../networking"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; -import { getAllPresets, getPresetByKey, getRequiredModelsInPreset } from "@/lib/autorouter_presets"; +import { AutoRouterPreset, getRequiredModelsInPreset } from "@/lib/autorouter_presets"; +import { BUNDLED_PRESETS, LOADED_PRESETS_QUERY, useAutoRouterPresets } from "../../../tests/mocks/autoRouterPresets"; vi.mock( "@/app/(dashboard)/hooks/autoRouter/useComplexityScorerDefaults", async () => await import("../../../tests/mocks/complexityScorerDefaults"), ); +vi.mock( + "@/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets", + async () => await import("../../../tests/mocks/autoRouterPresets"), +); + +const getAllPresets = (): AutoRouterPreset[] => BUNDLED_PRESETS; +const getPresetByKey = (key: string): AutoRouterPreset | undefined => BUNDLED_PRESETS.find((p) => p.key === key); const ANTHROPIC_PRESET = getPresetByKey("anthropic_family")!; const ANTHROPIC_TIERS = ANTHROPIC_PRESET.complexity_router_config.tiers; @@ -1142,3 +1150,52 @@ describe("getSubmitBlockedReason", () => { ); }); }); + +describe("preset catalog fetch states", () => { + afterEach(() => vi.mocked(useAutoRouterPresets).mockReturnValue(LOADED_PRESETS_QUERY)); + + it("keeps showing cached presets without the error banner when only a refetch fails", () => { + vi.mocked(useAutoRouterPresets).mockReturnValue({ + ...LOADED_PRESETS_QUERY, + isError: true, + } as never); + renderWithProviders(); + + expect(screen.queryByText(/Could not load templates/)).not.toBeInTheDocument(); + + openTemplateDropdown(); + expect(screen.queryAllByRole("option").length).toBeGreaterThan(1); + }); + + it("shows a loading hint while the catalog fetch is pending", () => { + vi.mocked(useAutoRouterPresets).mockReturnValue({ + ...LOADED_PRESETS_QUERY, + data: undefined, + isPending: true, + } as never); + renderWithProviders(); + + expect(screen.getByText("Loading templates...")).toBeInTheDocument(); + }); + + it("degrades to Custom Configuration with a retry hint that refetches the catalog", async () => { + const refetch = vi.fn(); + vi.mocked(useAutoRouterPresets).mockReturnValue({ + ...LOADED_PRESETS_QUERY, + data: undefined, + isError: true, + refetch, + } as never); + renderWithProviders(); + + expect(await screen.findByText(/Could not load templates/)).toBeInTheDocument(); + + openTemplateDropdown(); + const options = screen.queryAllByRole("option"); + expect(options).toHaveLength(1); + expect(options[0]).toHaveTextContent("Custom Configuration"); + + fireEvent.click(screen.getByRole("button", { name: "Retry" })); + expect(refetch).toHaveBeenCalled(); + }); +}); 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..ab5ea4cbfc9 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 @@ -50,8 +50,6 @@ import AutoRouterConnectionTest from "./auto_router_connection_test"; import AutoRouterRoutingTest from "./AutoRouterRoutingTest"; import { toast } from "@/lib/toast"; import { - getAllPresets, - getPresetByKey, getMissingModelsInPreset, getReferencedModelsError, buildEmptyPrefill, @@ -62,6 +60,7 @@ import { PresetPrefill, AutoRouterPreset, } from "@/lib/autorouter_presets"; +import { useAutoRouterPresets } from "@/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; interface AddAutoRouterTabProps { @@ -102,9 +101,7 @@ const presetDisabledHint = (availability: PresetAvailability): string | null => // caller-specific missing-model reason gets the alarming red treatment. const isPresetHintAlarming = (availability: PresetAvailability): boolean => availability.kind === "missing_models"; -// getAllPresets() already returns a stable, module-level array (see autorouter_presets.ts), so -// this is resolved once at import time rather than re-called from inside the component every render. -const presets = getAllPresets(); +const NO_PRESETS: AutoRouterPreset[] = []; // A one-line summary of what's configured, shown when the detailed section is collapsed so a // caller can see the shape of the config without opening it. @@ -229,6 +226,14 @@ const AddAutoRouterTab: React.FC = ({ }); const modelsLoading = groupsLoading || deploymentsLoading; const modelInfo = React.useMemo(() => data ?? [], [data]); + const { + data: presetsData, + isPending: presetsPending, + isError: presetsError, + refetch: refetchPresets, + } = useAutoRouterPresets(); + const presets = presetsData ?? NO_PRESETS; + const presetsUnavailable = presetsError && presetsData === undefined; // react-query keeps the last successful list around when a later refetch fails, so isError alone // can't tell "never loaded" apart from "loaded, then a background refetch errored" - only the // former leaves us with nothing trustworthy to verify a preset's models against. @@ -277,7 +282,7 @@ const AddAutoRouterTab: React.FC = ({ presets .map((preset) => ({ preset, availability: presetAvailability(preset) })) .sort((a, b) => Number(b.availability.kind === "available") - Number(a.availability.kind === "available")), - [presetAvailability], + [presets, presetAvailability], ); const templateItems = React.useMemo( @@ -307,7 +312,7 @@ const AddAutoRouterTab: React.FC = ({ return; } - const preset = getPresetByKey(presetKey); + const preset = presets.find((p) => p.key === presetKey); // Refuse to apply a preset whose models are not verified available. The dropdown disables // these options, so this is a guard against a stale click resolving after the list changed. if (!preset) return; @@ -342,6 +347,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, @@ -537,6 +543,15 @@ const AddAutoRouterTab: React.FC = ({
)} + {presetsPending &&
Loading templates...
} + {presetsUnavailable && ( +
+ Could not load templates, so only Custom Configuration is shown.{" "} + +
+ )} {requiresTeamScope && ( 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 2d8f55b2b64..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 @@ -817,6 +817,39 @@ describe("heuristic_first", () => { }); }); +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" }); @@ -924,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/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index db7baf5e171..eaffa2b4802 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -90,6 +90,7 @@ import type { } from "@/app/(dashboard)/caching/_components/coordination_redis_settings/types"; import { MCP_TOOLS_PREVIEW_FORBIDDEN_MESSAGE } from "./mcp_tools/constants"; import type { ComplexityRouterConfigPayload } from "./add_model/build_complexity_router_config"; +import type { AutoRouterPresetsResponse } from "@/lib/autorouter_presets"; import type { VectorStoreIndex } from "@/app/(dashboard)/vector-stores/_components/IndexesTab"; import type { RoutingDecision } from "./view_logs/LogDetailsDrawer/RoutingDecisionCard"; import { @@ -410,6 +411,15 @@ export const getComplexityScorerDefaults = async (): Promise => { + /** + * Fetch the auto-router preset catalog from the proxy's public endpoint. The template picker + * renders from this rather than from a copy in the dashboard, so a catalog update propagates + * without a dashboard release. + */ + return await apiClient.get(`/public/autorouter_presets`); +}; + export const getAgentCreateMetadata = async (): Promise => { /** * Fetch agent type metadata from the proxy's public endpoint. @@ -2042,6 +2052,7 @@ interface UiSpendLogsParams { min_spend?: number; max_spend?: number; exclude_internal_health_checks?: boolean; + group_by_session?: boolean; } interface UiSpendLogsCallOptions { 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 66e80dcc655..ffef01010b7 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx @@ -86,6 +86,7 @@ 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/components/view_logs/RequestLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx index 22e3f635b50..c368a43dd1e 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx @@ -138,33 +138,39 @@ describe("RequestLogsPanel", () => { respondWith([]); }); - describe("multi-call session collapsing", () => { - const sessionRows = [ - logEntry({ request_id: "req-mcp", call_type: "call_mcp_tool", session_id: "sess-1", session_total_count: 3 }), - logEntry({ request_id: "req-llm", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), - logEntry({ request_id: "req-llm-2", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), - ]; - - it("collapses a multi-call session to a single representative row", async () => { - respondWith(sessionRows); + describe("server-grouped session pagination (#38060)", () => { + it("requests session-grouped pages of 10 rows by default", async () => { renderPanel(); - await waitFor(() => expect(row("req-mcp") ?? row("req-llm") ?? row("req-llm-2")).not.toBeNull()); - - const rendered = ["req-mcp", "req-llm", "req-llm-2"].filter((id) => row(id) !== null); - expect(rendered).toHaveLength(1); + await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled()); + expect(lastCall()?.params?.group_by_session).toBe(true); + expect(lastCall()?.page_size).toBe(10); }); - it("prefers an LLM call over an MCP call as the session's representative", async () => { - respondWith(sessionRows); + it("renders every row the server returns without client-side collapsing", async () => { + respondWith([ + logEntry({ request_id: "req-a", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), + logEntry({ request_id: "req-b", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), + logEntry({ request_id: "req-c", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), + ]); renderPanel(); - await waitFor(() => expect(row("req-llm")).not.toBeNull()); - expect(row("req-mcp")).toBeNull(); + await waitFor(() => expect(row("req-a")).not.toBeNull()); + expect(row("req-b")).not.toBeNull(); + expect(row("req-c")).not.toBeNull(); }); - it("shows the session's call count and composition on the representative row", async () => { - respondWith(sessionRows); + it("shows the session's call count on the server-picked representative row", async () => { + respondWith([ + logEntry({ + request_id: "req-llm", + call_type: "acompletion", + session_id: "sess-1", + session_total_count: 3, + session_llm_count: 2, + mcp_tool_call_count: 1, + }), + ]); renderPanel(); await waitFor(() => expect(row("req-llm")).not.toBeNull()); @@ -296,6 +302,7 @@ describe("RequestLogsPanel", () => { if (!byIdCall) throw new Error("expected a by-id uiSpendLogsCall"); expect(byIdCall.page).toBe(1); expect(byIdCall.page_size).toBe(1); + expect(byIdCall.params?.group_by_session).toBeUndefined(); }); it("closing the drawer removes ?log_id= from the URL and closes the drawer", async () => { diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx index 52ea78abf5e..2ebabf64ab0 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx @@ -10,7 +10,7 @@ import type { KeyResponse } from "../key_team_helpers/key_list"; import { keyInfoV1Call, uiSpendLogsCall } from "../networking"; import KeyInfoView from "../templates/key_info_view"; import type { LogEntry } from "./columns"; -import { AGENT_CALL_TYPES, MCP_CALL_TYPES } from "./constants"; +import { LOGS_PAGE_SIZE_OPTIONS } from "./constants"; import { DEFAULT_LOGS_SORTING, formatLogsWindow, @@ -24,7 +24,7 @@ import { LogDetailsDrawer } from "./LogDetailsDrawer"; import { LiveTailBanner, LogsTableToolbar } from "./LogsTableToolbar"; import { RequestLogsTable } from "./RequestLogsTable"; -const PAGE_SIZE = 50; +const PAGE_SIZE = LOGS_PAGE_SIZE_OPTIONS[0]; const DEFAULT_INTERVAL = { value: 24, unit: "hours" }; interface RequestLogsPanelProps { @@ -35,12 +35,6 @@ interface RequestLogsPanelProps { isActive: boolean; } -interface SessionComposition { - llm: number; - agent: number; - mcp: number; -} - export default function RequestLogsPanel({ accessToken, token, userRole, userID, isActive }: RequestLogsPanelProps) { const [pagination, setPagination] = useState({ pageIndex: 0, pageSize: PAGE_SIZE }); const [sorting, setSorting] = useState(DEFAULT_LOGS_SORTING); @@ -157,49 +151,7 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, const isDrawerOpen = displayLog !== null || displaySessionId !== null; - const rows = useMemo(() => { - const searchedLogs = filteredLogs.data; - - const sessionCompositionById = searchedLogs.reduce>((acc, log) => { - if (!log.session_id) return acc; - if (!acc[log.session_id]) { - acc[log.session_id] = { llm: 0, agent: 0, mcp: 0 }; - } - if (MCP_CALL_TYPES.includes(log.call_type)) { - acc[log.session_id].mcp += 1; - } else if (AGENT_CALL_TYPES.includes(log.call_type)) { - acc[log.session_id].agent += 1; - } else { - acc[log.session_id].llm += 1; - } - return acc; - }, {}); - - const sessionRepresentativeMap = new Map(); - for (const log of searchedLogs) { - if (!log.session_id || (log.session_total_count || 1) <= 1) continue; - const isMcp = MCP_CALL_TYPES.includes(log.call_type); - const existing = sessionRepresentativeMap.get(log.session_id); - if (!existing || (existing.isMcp && !isMcp)) { - sessionRepresentativeMap.set(log.session_id, { requestId: log.request_id, isMcp }); - } - } - - return searchedLogs - .map((log) => { - const sessionComposition = log.session_id ? sessionCompositionById[log.session_id] : undefined; - return { - ...log, - session_llm_count: sessionComposition?.llm ?? undefined, - session_mcp_count: sessionComposition?.mcp ?? undefined, - session_agent_count: sessionComposition?.agent ?? undefined, - }; - }) - .filter((log) => { - if (!log.session_id || (log.session_total_count || 1) <= 1) return true; - return sessionRepresentativeMap.get(log.session_id)?.requestId === log.request_id; - }); - }, [filteredLogs.data]); + const rows: LogEntry[] = filteredLogs.data; const searchTerm = useMemo(() => { const entry = columnFilters.find((filter) => filter.id === LOG_FILTER_IDS.REQUEST_ID); @@ -258,13 +210,12 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, ); const handleSessionClick = useCallback( - (sessionId: string) => { - if (!sessionId) return; - const log = rows.find((candidate) => candidate.session_id === sessionId) ?? null; + (log: LogEntry) => { + if (!log.session_id) return; setSelectedLog(log); - openSession(sessionId, log?.request_id ?? null); + openSession(log.session_id, log.request_id); }, - [rows, openSession], + [openSession], ); const handleSelectLog = useCallback( diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx index 7146cf33847..4159b3b699b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx @@ -8,6 +8,7 @@ import { DataTable, DataTableFilterDrawer, DataTableToolbar } from "@/components import type { Team } from "../key_team_helpers/key_list"; import type { LogEntry } from "./columns"; +import { LOGS_PAGE_SIZE_OPTIONS } from "./constants"; import { LOG_FILTER_LABELS, type LogsWindow } from "./log_filter_logic"; import { RequestLogsFilters } from "./RequestLogsFilters"; import { getRequestLogsTableColumns } from "./RequestLogsTableColumns"; @@ -28,7 +29,7 @@ interface RequestLogsTableProps { onRefresh: () => void; onRowClick: (log: LogEntry) => void; onKeyHashClick: (keyHash: string) => void; - onSessionClick: (sessionId: string) => void; + onSessionClick: (log: LogEntry) => void; teams: Team[]; logsWindow: LogsWindow; toolbarChildren?: ReactNode; @@ -91,6 +92,7 @@ export function RequestLogsTable({ paginationMode="server" pagination={pagination} onPaginationChange={onPaginationChange} + pageSizeOptions={LOGS_PAGE_SIZE_OPTIONS} rowCount={rowCount} filterMode="server" columnFilters={columnFilters} diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx index 3d0ea03c5d5..ce59c62f1c4 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx @@ -85,13 +85,19 @@ describe("row action cells", () => { expect(deps.onKeyHashClick).toHaveBeenCalledWith("sk-hash-9"); }); - it("reports the session id from the session cell", async () => { + it("reports the clicked row from the session cell, so two rows sharing a session id stay distinguishable", async () => { const user = userEvent.setup(); const deps = { onKeyHashClick: vi.fn(), onSessionClick: vi.fn() }; - renderRows([logEntry({ request_id: "req-sess", session_id: "sess-42" })], deps); + renderRows( + [ + logEntry({ request_id: "req-key-a", session_id: "sess-42", api_key: "key-a" }), + logEntry({ request_id: "req-key-b", session_id: "sess-42", api_key: "key-b" }), + ], + deps, + ); - await user.click(screen.getByText("sess-42")); - expect(deps.onSessionClick).toHaveBeenCalledWith("sess-42"); + await user.click(screen.getAllByText("sess-42")[1]); + expect(deps.onSessionClick).toHaveBeenCalledWith(expect.objectContaining({ request_id: "req-key-b" })); }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx index cf776515bd4..b9058d02a6b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx @@ -13,7 +13,7 @@ import { AgentBadge, AgentIcon, LlmBadge, McpBadge, SparkleIcon, WrenchIcon } fr export interface RequestLogsTableColumnsDeps { onKeyHashClick: (keyHash: string) => void; - onSessionClick: (sessionId: string) => void; + onSessionClick: (log: LogEntry) => void; } const readMetaString = (metadata: Record | undefined, key: string): string | undefined => { @@ -61,7 +61,7 @@ export const getRequestLogsTableColumns = ({ const isAgent = AGENT_CALL_TYPES.includes(log.call_type); const sessionLlmCount = log.session_llm_count ?? (isMcp || isAgent ? 0 : sessionCount); const sessionAgentCount = log.session_agent_count ?? (isAgent ? sessionCount : 0); - const sessionMcpCount = log.session_mcp_count ?? (isMcp ? sessionCount : 0); + const sessionMcpCount = log.mcp_tool_call_count ?? (isMcp ? sessionCount : 0); if (isMcp) return ; if (isAgent && sessionCount <= 1) return ; @@ -113,7 +113,7 @@ export const getRequestLogsTableColumns = ({ header: "Session ID", size: 120, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => onSessionClick(row.original)} />, }, { id: "request_id", diff --git a/ui/litellm-dashboard/src/components/view_logs/columns.tsx b/ui/litellm-dashboard/src/components/view_logs/columns.tsx index eef957922d7..2f3a3681352 100644 --- a/ui/litellm-dashboard/src/components/view_logs/columns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/columns.tsx @@ -46,6 +46,5 @@ export type LogEntry = { mcp_tool_call_count?: number; mcp_tool_call_spend?: number; session_llm_count?: number; - session_mcp_count?: number; session_agent_count?: number; }; diff --git a/ui/litellm-dashboard/src/components/view_logs/constants.ts b/ui/litellm-dashboard/src/components/view_logs/constants.ts index 5b0b1d0fee3..1c17f398e35 100644 --- a/ui/litellm-dashboard/src/components/view_logs/constants.ts +++ b/ui/litellm-dashboard/src/components/view_logs/constants.ts @@ -12,6 +12,9 @@ export const ERROR_CODE_OPTIONS: { label: string; value: string }[] = [ { label: "529 - Overloaded", value: "529" }, ]; +/** Page sizes the logs tables offer; the first entry is the default. */ +export const LOGS_PAGE_SIZE_OPTIONS = [10, 25, 50, 100]; + /** Call types that represent MCP tool invocations (shared across columns, index, drawer). */ export const MCP_CALL_TYPES = ["call_mcp_tool", "list_mcp_tools"]; diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx index 3b8d96596de..acd5be06593 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx @@ -181,6 +181,7 @@ export function useLogFilterLogic({ sort_by: sortBy, sort_order: sortOrder, exclude_internal_health_checks: excludeInternalHealthChecks, + group_by_session: true, }, }); }, diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index b5802ffa2dc..e3132067f4c 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -1,7 +1,9 @@ import { describe, it, expect } from "vitest"; +import bundledPresets from "../../../../litellm/proxy/public_endpoints/autorouter_presets.json"; import { - getAllPresets, - getPresetByKey, + hydratePresets, + AutoRouterPreset, + AutoRouterPresetsResponse, getRequiredModelsInPreset, getMissingModelsInPreset, getRequiredModels, @@ -18,8 +20,13 @@ import { DEFAULT_ESCALATION_KEYWORDS } from "@/components/add_model/EscalationKe const groupsOnly = (models: Iterable) => buildModelAvailability(models, []); +// Hydrated from the real bundled catalog so a catalog edit flows into these expectations. +const PRESETS = hydratePresets(bundledPresets as AutoRouterPresetsResponse); +const getAllPresets = (): AutoRouterPreset[] => PRESETS; +const getPresetByKey = (key: string): AutoRouterPreset | undefined => PRESETS.find((p) => p.key === key); + describe("autorouter_presets", () => { - it("loads exactly the bundled presets", () => { + it("hydrates exactly the bundled presets", () => { const presets = getAllPresets(); expect(presets.map((p) => p.label).sort()).toEqual(["Anthropic Family", "Gemini Family", "Lite", "OpenAI Family"]); // Every preset carries all four fields the UI relies on; a JSON typo dropping one fails here. diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index 721bd6f2b2a..30f8ed99e74 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -20,7 +20,6 @@ import { } from "@/components/add_model/complexity_router_tiers"; import { DEFAULT_ESCALATION_KEYWORDS } from "@/components/add_model/EscalationKeywords"; import { DEFAULT_MATCH_THRESHOLD } from "@/components/add_model/SemanticKeywordMatching"; -import presetsRaw from "@/autorouter_presets.json"; // `key` is the stable JSON object key (e.g. "anthropic_family"); `label` is display text and // never an identity. @@ -31,16 +30,10 @@ export interface AutoRouterPreset { complexity_router_config: ComplexityRouterConfigPayload; } -// The bundled JSON is a developer-authored, build-time asset, so it is trusted at the import -// boundary rather than re-validated at runtime (resolveJsonModule widens its string literals, -// hence this one cast). autorouter_presets.test.ts pins the parsed shape, so a JSON typo fails CI. -const RAW = presetsRaw as Record>; +export type AutoRouterPresetsResponse = Record>; -const PRESETS: AutoRouterPreset[] = Object.entries(RAW).map(([key, preset]) => ({ key, ...preset })); - -export const getAllPresets = (): AutoRouterPreset[] => PRESETS; - -export const getPresetByKey = (key: string): AutoRouterPreset | undefined => PRESETS.find((p) => p.key === key); +export const hydratePresets = (raw: AutoRouterPresetsResponse): AutoRouterPreset[] => + Object.entries(raw).map(([key, preset]) => ({ key, ...preset })); // Generalized over ComplexityRouterConfigPayload so the same accessors check either a preset's own // bundled config or a caller's actually-built config - the two need to agree, since a preset only diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 79e44de4f1f..491c4fb6a44 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -12129,6 +12129,31 @@ export interface paths { patch?: never; trace?: never; }; + "/public/autorouter_presets": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Public Autorouter Presets + * @description Return the auto-router preset catalog the dashboard's template picker renders. + * + * Resolved once per process, like the model cost map: fetched from ``litellm.autorouter_presets_url`` + * (override with ``LITELLM_AUTOROUTER_PRESETS_URL``) on the first request, falling back to the + * catalog bundled with the package on any failure. Set ``LITELLM_LOCAL_AUTOROUTER_PRESETS=True`` + * to serve the bundled catalog only. A restart picks up a newly published catalog. + */ + get: operations["get_public_autorouter_presets_public_autorouter_presets_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/public/complexity_router/scorer_defaults": { parameters: { query?: never; @@ -23391,6 +23416,49 @@ export interface components { /** Tier Definitions */ tier_definitions: components["schemas"]["TierDefinition"][]; }; + /** + * AutoRouterPresetConfig + * @description The complexity_router_config a preset prefills. + * + * Only tiers is validated, because every dashboard consumer dereferences it; everything else + * passes through verbatim with unknown fields kept (extra="allow"), so a catalog published after + * this proxy shipped still serves its new fields intact. + */ + AutoRouterPresetConfig: { + tiers: components["schemas"]["AutoRouterPresetTiers"]; + } & { + [key: string]: unknown; + }; + /** + * AutoRouterPresetRecord + * @description One auto-router preset as served to the dashboard's template picker. + */ + AutoRouterPresetRecord: { + complexity_router_config: components["schemas"]["AutoRouterPresetConfig"]; + /** Description */ + description: string; + /** Label */ + label: string; + } & { + [key: string]: unknown; + }; + /** + * AutoRouterPresetTiers + * @description Exactly the four built-in tiers the dashboard's preset prefill can apply. + * + * extra="forbid" on purpose: a tier name this dashboard cannot apply would grey out or crash the + * picker, so such a catalog is rejected wholesale and the bundled one serves instead. + */ + AutoRouterPresetTiers: { + /** Complex */ + COMPLEX: string[]; + /** Medium */ + MEDIUM: string[]; + /** Reasoning */ + REASONING: string[]; + /** Simple */ + SIMPLE: string[]; + }; /** * AutoRouterRoutingTestRequest * @description A single request to classify against a complexity-router config that need not be saved yet. @@ -25704,6 +25772,11 @@ export interface components { * @description Number of trusted reverse proxies/load balancers in front of the gateway that append to X-Forwarded-For. When set (and mcp_trusted_proxy_ranges validates the direct peer), the client IP for MCP access control is read this many entries from the right of the chain instead of the spoofable leftmost value, defeating append-style X-Forwarded-For forgery. */ mcp_xff_num_trusted_hops?: number | null; + /** + * Missing Session Id + * @description What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id. + */ + missing_session_id?: ("generate" | "reject") | null; /** * Model List Healthy Only * @description When true, `/models`, `/v1/models/{id}` and `/model/info` hide models whose backing deployments are all unhealthy, for every caller, without needing `healthy_only=true` per request. Requires `background_health_checks: true`, and keeps deployment health state cached without turning on `enable_health_check_routing`, so routing is unaffected. With no health state nothing is hidden. Hiding is presentation-only, a hidden model can still be called. @@ -34560,7 +34633,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 @@ -34575,11 +34648,11 @@ export interface components { classifier_plugin_timeout_ms: number; /** * Classifier Type - * @description Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, 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" | "heuristic_v2" | "llm" | "custom" | "heuristic_first"; + classifier_type: "heuristic" | "heuristic_v2" | "llm" | "custom" | "heuristic_first" | "hybrid"; /** * Code Keywords * @description Keywords indicating code-related content @@ -34651,6 +34724,11 @@ export interface components { * @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 @@ -35867,7 +35945,7 @@ export interface components { * Cause * @enum {string} */ - cause?: "heuristic_scorer" | "heuristic_v2" | "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 */ @@ -54674,6 +54752,28 @@ export interface operations { }; }; }; + get_public_autorouter_presets_public_autorouter_presets_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: components["schemas"]["AutoRouterPresetRecord"]; + }; + }; + }; + }; + }; get_complexity_scorer_defaults_public_complexity_router_scorer_defaults_get: { parameters: { query?: never; @@ -56750,6 +56850,8 @@ export interface operations { sort_order?: string | null; /** @description Exclude LiteLLM internal health check requests from results */ exclude_internal_health_checks?: boolean; + /** @description Paginate over sessions instead of raw logs: one representative row per session, total counts sessions */ + group_by_session?: boolean; }; header?: never; path?: never; @@ -56862,6 +56964,8 @@ export interface operations { sort_order?: string | null; /** @description Exclude LiteLLM internal health check requests from results */ exclude_internal_health_checks?: boolean; + /** @description Paginate over sessions instead of raw logs: one representative row per session, total counts sessions */ + group_by_session?: boolean; }; header?: never; path?: never; diff --git a/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts b/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts new file mode 100644 index 00000000000..f73faaa70fe --- /dev/null +++ b/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts @@ -0,0 +1,16 @@ +import { vi } from "vitest"; +import bundledPresets from "../../../../litellm/proxy/public_endpoints/autorouter_presets.json"; +import { hydratePresets, type AutoRouterPresetsResponse } from "@/lib/autorouter_presets"; + +// Derived from the real bundled catalog so a preset edit there flows into test expectations +// instead of redding on a stale copy. Exported as vi.fn so a test can override the query state. +export const BUNDLED_PRESETS = hydratePresets(bundledPresets as AutoRouterPresetsResponse); + +export const LOADED_PRESETS_QUERY = { + data: BUNDLED_PRESETS, + isPending: false, + isError: false, + refetch: vi.fn(), +}; + +export const useAutoRouterPresets = vi.fn(() => LOADED_PRESETS_QUERY);