From 692f3b513ca9ac2fafb225c24d58f7ee5152eae6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:43:01 -0700 Subject: [PATCH 01/31] fix(proxy-extras): recover the v2 migration resolver from concurrent migrate deploy deadlocks Two instances racing prisma migrate deploy on one database deadlock on CREATE INDEX CONCURRENTLY: the victim gets P3018 with 40P01 and the survivor then sees the failed ledger row as P3009. Both were treated as unrecoverable, so neither instance came up. Roll the deadlocked migration's ledger row back and retry the deploy on P3018, consult the failed row's logs in _prisma_migrations to do the same on P3009, and retry a deadlock reported without a Prisma error code. Genuinely broken migrations still fail fast. --- .../litellm_proxy_extras/utils.py | 91 ++++++++++- .../tests/test_setup_database_fail_fast.py | 141 ++++++++++++++++++ 2 files changed, 228 insertions(+), 4 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b8032dd0d28..fb948afd200 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -40,6 +40,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 @@ -262,6 +264,48 @@ 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) -> str: + """Logs recorded on the migration's failed _prisma_migrations row. + + P3009 stderr does not carry the original failure, so this is the only + way to tell a migration that lost a deadlock race against a concurrent + migrate deploy from one whose SQL is genuinely broken. Returns "" when + psycopg is missing, the DB is unreachable, or no failed row exists. + """ + database_url = os.getenv("DATABASE_URL") + if not database_url: + return "" + + try: + import psycopg + except ImportError: + return "" + + cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url) + try: + with psycopg.connect( + cleaned_url, connect_timeout=10, autocommit=True + ) as conn: + row = conn.execute( + "SELECT logs FROM _prisma_migrations " + "WHERE migration_name = %s AND finished_at IS NULL " + "AND rolled_back_at IS NULL", + (migration_name,), + ).fetchone() + except (psycopg.OperationalError, psycopg.DatabaseError): + return "" + return (row[0] or "") if row else "" + @staticmethod def _resolve_specific_migration(migration_name: str): """Mark a specific migration as applied""" @@ -658,7 +702,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. @@ -764,6 +809,22 @@ class ProxyExtrasDBManager: f"Detail: {resolve_err}" ) from resolve_err continue + if migration_match and _MIGRATION_DEADLOCK_MARKER in ( + ProxyExtrasDBManager._failed_migration_logs( + migration_match.group(1) + ) + ): + logger.info( + "Migration %s lost a deadlock race 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}" @@ -809,11 +870,33 @@ 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 + raise RuntimeError( "Database migration failed and cannot be auto-recovered. " f"Manual intervention required.\n\nPrisma error:\n{stderr}" @@ -821,9 +904,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." + "exhausted by timeouts, deadlock retries, or repeated " + "idempotent-recovery continues). Check database connectivity, " + "load, and _prisma_migrations ledger state." ) finally: os.chdir(original_dir) 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..c4347a91dce 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,144 @@ 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_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 From 5192b2162c987f260c9c33700343a73ea4676749 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 14:04:12 -0700 Subject: [PATCH 02/31] fix(proxy-extras): schema-qualify the _prisma_migrations logs lookup for non-public Prisma schemas --- litellm-proxy-extras/litellm_proxy_extras/utils.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index fb948afd200..97b6b1c667c 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -292,14 +292,22 @@ class ProxyExtrasDBManager: return "" 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( - "SELECT logs FROM _prisma_migrations " - "WHERE migration_name = %s AND finished_at IS NULL " - "AND rolled_back_at IS NULL", + 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): From bfa5eac76b18ae9e3965d5f942b2fd0382e6b796 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 10:54:39 -0700 Subject: [PATCH 03/31] fix(vector-store): resolve embedding aliases for search --- .../proxy/vector_store_endpoints/endpoints.py | 13 ++-- .../management_endpoints.py | 48 +++++++++++--- .../test_vector_store_endpoints.py | 66 ++++++++++++------- 3 files changed, 89 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index a59d7a277cc..e0b6cf8817a 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -70,17 +70,17 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( # time, instead of at row-creation time. The resolved # ``api_key`` / ``api_base`` / ``api_version`` lives only in # this per-request ``data`` dict and is never persisted. - # Legacy rows that already carry a resolved (cleartext) - # ``litellm_embedding_config`` skip the lookup and pass through - # unchanged so the embed call keeps working. + # Legacy rows that carry a resolved config are refreshed when the + # embedding model is an alias so the provider-qualified model is used. embedding_model: Final = litellm_params.get("litellm_embedding_model") - if embedding_model and not litellm_params.get("litellm_embedding_config"): + if embedding_model: from litellm.proxy.proxy_server import prisma_client - resolved_config: Final = await _resolve_embedding_config( + embedding_resolution: Final = await _resolve_embedding_config( embedding_model=embedding_model, prisma_client=prisma_client ) - if resolved_config: + if embedding_resolution: + resolved_model, resolved_config = embedding_resolution # Build a fresh dict via spread instead of mutating # ``litellm_params`` in place — the registry hands back # a reference to its cached object, so an in-place @@ -88,6 +88,7 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( # in-memory cache for the lifetime of the process. litellm_params = { **litellm_params, + "litellm_embedding_model": resolved_model, "litellm_embedding_config": resolved_config, } data.update(litellm_params) diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 183a03cc13c..1930ec4aaba 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -10,7 +10,7 @@ All /vector_store management endpoints import copy import json -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, TypeAlias from fastapi import APIRouter, Depends, HTTPException @@ -49,6 +49,7 @@ from litellm.types.vector_stores import ( from litellm.vector_stores.vector_store_registry import VectorStoreRegistry router: Final = APIRouter() +EmbeddingResolution: TypeAlias = tuple[str, dict[str, object]] def _vector_store_table(prisma_client: "PrismaClient") -> "TableActions[_VectorStoreRow]": @@ -155,7 +156,19 @@ async def _fetch_and_authorize_vector_store( return typed -def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, object] | None: +def _provider_qualified_embedding_model( + fallback: str, + model: object, + custom_llm_provider: object, +) -> str: + if not isinstance(model, str) or not model: + return fallback + if "/" in model or not isinstance(custom_llm_provider, str) or not custom_llm_provider: + return model + return f"{custom_llm_provider}/{model}" + + +def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> EmbeddingResolution | None: """ Resolve embedding config from router's config-defined models. @@ -168,7 +181,7 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d llm_router: The LiteLLM router instance Returns: - Dictionary with api_key, api_base, and api_version if model found, None otherwise + Provider-qualified model and its connection config if found, otherwise None """ if not embedding_model or llm_router is None: return None @@ -218,12 +231,21 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d if project_id: embedding_config["project_id"] = project_id + resolved_model: Final = _provider_qualified_embedding_model( + fallback=embedding_model, + model=getattr(litellm_params, "model", None), + custom_llm_provider=getattr(litellm_params, "custom_llm_provider", None), + ) + # 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 + return ( + resolved_model, + embedding_config, + ) except Exception as e: verbose_proxy_logger.debug("Error resolving embedding config from router for model %s: %s", model_name, e) continue @@ -233,7 +255,7 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d async def _resolve_embedding_config_from_db( embedding_model: str, prisma_client: "PrismaClient" -) -> dict[str, object] | None: +) -> EmbeddingResolution | None: """ Resolve embedding config from database model configuration. @@ -246,7 +268,7 @@ async def _resolve_embedding_config_from_db( prisma_client: The Prisma client instance Returns: - Dictionary with api_key, api_base, and api_version if model found, None otherwise + Provider-qualified model and its connection config if found, otherwise None """ if not embedding_model: return None @@ -315,7 +337,15 @@ async def _resolve_embedding_config_from_db( model_name, list(embedding_config.keys()), ) - return embedding_config + resolved_model: Final = _provider_qualified_embedding_model( + fallback=embedding_model, + model=decrypted_params.get("model"), + custom_llm_provider=decrypted_params.get("custom_llm_provider"), + ) + return ( + resolved_model, + embedding_config, + ) except Exception as e: verbose_proxy_logger.debug("Error resolving embedding config for model %s: %s", model_name, e) continue @@ -325,7 +355,7 @@ async def _resolve_embedding_config_from_db( async def _resolve_embedding_config( embedding_model: str, prisma_client: "PrismaClient | None", llm_router: "Router | None" = None -) -> dict[str, object] | None: +) -> EmbeddingResolution | None: """ Resolve embedding config from either router (config-defined) or database models. @@ -343,7 +373,7 @@ async def _resolve_embedding_config( 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 + Provider-qualified model and its connection config if found, otherwise None """ if not embedding_model: return None 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 eae6f90863a..1484adb258f 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 @@ -514,7 +514,7 @@ async def test_update_request_data_resolves_embedding_config_at_use_time(): "vector_store_id": "test_store", "custom_llm_provider": "azure_ai", "litellm_params": { - "litellm_embedding_model": "azure/text-embedding-3-large", + "litellm_embedding_model": "multilingual-e5-large", # Note: no litellm_embedding_config persisted }, } @@ -534,24 +534,22 @@ async def test_update_request_data_resolves_embedding_config_at_use_time(): patch.object(litellm, "vector_store_registry", mock_registry), patch( "litellm.proxy.vector_store_endpoints.endpoints._resolve_embedding_config", - new=AsyncMock(return_value=resolved), + new=AsyncMock(return_value=("azure/multilingual-e5-large", resolved)), ), ): result = await _update_request_data_with_litellm_managed_vector_store_registry( data={}, vector_store_id="test_store" ) - assert result["litellm_embedding_model"] == "azure/text-embedding-3-large" + assert result["litellm_embedding_model"] == "azure/multilingual-e5-large" assert result["litellm_embedding_config"] == resolved @pytest.mark.asyncio -async def test_update_request_data_passes_through_legacy_embedding_config(): +async def test_update_request_data_preserves_legacy_embedding_config_when_model_not_resolved(): """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.""" + ``litellm_params``. Preserve it when the model cannot be resolved.""" legacy_config = { "api_key": "legacy-cleartext-key", "api_base": "https://legacy-azure.example", @@ -571,7 +569,7 @@ async def test_update_request_data_passes_through_legacy_embedding_config(): mock_vector_store ) - resolve_mock = AsyncMock() + resolve_mock = AsyncMock(return_value=None) with ( patch.object(litellm, "vector_store_registry", mock_registry), @@ -585,7 +583,7 @@ async def test_update_request_data_passes_through_legacy_embedding_config(): ) assert result["litellm_embedding_config"] == legacy_config - resolve_mock.assert_not_awaited() + resolve_mock.assert_awaited_once() class TestCheckVectorStorePermission: @@ -2010,6 +2008,7 @@ async def test_resolve_embedding_config_from_db(): # Mock database model with litellm_params mock_db_model = MagicMock() mock_db_model.litellm_params = { + "model": "openai/text-embedding-3-small", "api_key": "test-api-key", "api_base": "https://api.openai.com", "api_version": "2024-01-01", @@ -2028,9 +2027,11 @@ async def test_resolve_embedding_config_from_db(): ) 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" + resolved_model, resolved_config = result + assert resolved_model == "openai/text-embedding-3-small" + assert resolved_config["api_key"] == "test-api-key" + assert resolved_config["api_base"] == "https://api.openai.com" + assert resolved_config["api_version"] == "2024-01-01" mock_prisma_client.db.litellm_proxymodeltable.find_first.assert_called_once_with( where={"model_name": "text-embedding-ada-002"} ) @@ -2164,6 +2165,8 @@ def test_resolve_embedding_config_from_router(): 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_litellm_params.model = "text-embedding-3-small" + mock_litellm_params.custom_llm_provider = "openai" mock_deployment = MagicMock(spec=Deployment) mock_deployment.litellm_params = mock_litellm_params @@ -2176,9 +2179,11 @@ def test_resolve_embedding_config_from_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" + resolved_model, resolved_config = result + assert resolved_model == "openai/text-embedding-3-small" + assert resolved_config["api_key"] == "config-api-key" + assert resolved_config["api_base"] == "https://config-api-base.com" + assert resolved_config["api_version"] == "2024-02-01" mock_router.get_deployment_by_model_group_name.assert_called_once_with( model_group_name="text-embedding-ada-002" @@ -2197,6 +2202,8 @@ def test_resolve_embedding_config_from_router_with_provider_prefix(): 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_litellm_params.model = "text-embedding-3-large" + mock_litellm_params.custom_llm_provider = "azure" mock_deployment = MagicMock(spec=Deployment) mock_deployment.litellm_params = mock_litellm_params @@ -2209,9 +2216,11 @@ def test_resolve_embedding_config_from_router_with_provider_prefix(): ) 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" + resolved_model, resolved_config = result + assert resolved_model == "azure/text-embedding-3-large" + assert resolved_config["api_key"] == "azure-api-key" + assert resolved_config["api_base"] == "https://azure-endpoint.openai.azure.com" + assert resolved_config["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 @@ -2239,6 +2248,8 @@ def test_resolve_embedding_config_from_router_handles_os_environ(): 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_litellm_params.model = "text-embedding-3-small" + mock_litellm_params.custom_llm_provider = "openai" mock_deployment = MagicMock(spec=Deployment) mock_deployment.litellm_params = mock_litellm_params @@ -2254,9 +2265,11 @@ def test_resolve_embedding_config_from_router_handles_os_environ(): ) 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 + resolved_model, resolved_config = result + assert resolved_model == "openai/text-embedding-3-small" + assert resolved_config["api_key"] == "resolved-from-env" + assert resolved_config["api_base"] == "https://direct-url.com" + assert "api_version" not in resolved_config mock_get_secret.assert_called_once_with("os.environ/OPENAI_API_KEY") @@ -2274,6 +2287,8 @@ async def test_resolve_embedding_config_tries_router_then_db(): 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_litellm_params.model = "text-embedding-3-small" + mock_litellm_params.custom_llm_provider = "openai" mock_deployment = MagicMock(spec=Deployment) mock_deployment.litellm_params = mock_litellm_params @@ -2290,7 +2305,9 @@ async def test_resolve_embedding_config_tries_router_then_db(): ) assert result is not None - assert result["api_key"] == "router-api-key" + resolved_model, resolved_config = result + assert resolved_model == "openai/text-embedding-3-small" + assert resolved_config["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() @@ -2345,6 +2362,7 @@ async def test_resolve_embedding_config_falls_back_to_db(): # DB has the model mock_db_model = MagicMock() mock_db_model.litellm_params = { + "model": "openai/text-embedding-3-small", "api_key": "db-api-key", "api_base": "https://db-api-base.com", } @@ -2363,7 +2381,9 @@ async def test_resolve_embedding_config_falls_back_to_db(): ) assert result is not None - assert result["api_key"] == "db-api-key" + resolved_model, resolved_config = result + assert resolved_model == "openai/text-embedding-3-small" + assert resolved_config["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() From 5635811726ed05811abe5a242645dafd48eca9a0 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 12:15:03 -0700 Subject: [PATCH 04/31] fix(vector-store): route embeddings through router --- .../base_llm/vector_store/transformation.py | 61 +- litellm/llms/custom_httpx/llm_http_handler.py | 6 + .../valkey/vector_stores/transformation.py | 39 +- .../proxy/vector_store_endpoints/endpoints.py | 32 +- .../management_endpoints.py | 289 +-------- litellm/router.py | 56 +- litellm/vector_stores/main.py | 29 +- .../test_router_embedding_integration.py | 94 ++- .../test_valkey_transformation.py | 34 +- .../test_vector_store_endpoints.py | 576 +++++------------- uv.lock | 22 +- 11 files changed, 469 insertions(+), 769 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 02a51a8bace..772e4f849a0 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -1,10 +1,14 @@ +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 typing import TYPE_CHECKING, Any, NoReturn, Protocol, runtime_checkable import httpx from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import EmbeddingResponse from litellm.types.vector_stores import ( VECTOR_STORE_OPENAI_PARAMS, BaseVectorStoreAuthCredentials, @@ -17,6 +21,7 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.router import Router from ..chat.transformation import BaseLLMException as _BaseLLMException @@ -27,6 +32,58 @@ 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 + ) + + +@dataclass(frozen=True, slots=True) +class RouterVectorStoreEmbeddingExecutor: + router: Router + metadata: Mapping[str, object] + + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + if configuration: + return LiteLLMVectorStoreEmbeddingExecutor().embed(model, query, configuration) + 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 + metadata=dict(self.metadata), # mutable-ok: Router metadata requires a concrete dict + ) + + async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + if configuration: + return await LiteLLMVectorStoreEmbeddingExecutor().aembed(model, query, configuration) + 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 + metadata=dict(self.metadata), # mutable-ok: Router metadata requires a concrete dict + ) + + class BaseVectorStoreConfig: def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]: return [] @@ -172,6 +229,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 @@ -184,6 +242,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 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 834f7d564a2..118656b81a2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -70,6 +70,7 @@ from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeech from litellm.llms.base_llm.vector_store.transformation import ( BaseDirectVectorStoreConfig, BaseVectorStoreConfig, + VectorStoreEmbeddingExecutor, ) from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, @@ -9683,6 +9684,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, @@ -9702,6 +9704,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, ) @@ -9797,6 +9800,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, @@ -9812,6 +9816,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, @@ -9831,6 +9836,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, ) 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/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index e0b6cf8817a..3fc67181d5b 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -14,9 +14,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, @@ -65,32 +62,9 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( data["litellm_credential_name"] = vector_store_to_run.get("litellm_credential_name") if "litellm_params" in vector_store_to_run: - litellm_params = vector_store_to_run.get("litellm_params", {}) or {} - # Resolve ``litellm_embedding_config`` here, at request-handling - # time, instead of at row-creation time. The resolved - # ``api_key`` / ``api_base`` / ``api_version`` lives only in - # this per-request ``data`` dict and is never persisted. - # Legacy rows that carry a resolved config are refreshed when the - # embedding model is an alias so the provider-qualified model is used. - embedding_model: Final = litellm_params.get("litellm_embedding_model") - if embedding_model: - from litellm.proxy.proxy_server import prisma_client - - embedding_resolution: Final = await _resolve_embedding_config( - embedding_model=embedding_model, prisma_client=prisma_client - ) - if embedding_resolution: - resolved_model, resolved_config = embedding_resolution - # Build a fresh dict via spread instead of mutating - # ``litellm_params`` in place — the registry hands back - # a reference to its cached object, so an in-place - # update would persist the resolved cleartext into the - # in-memory cache for the lifetime of the process. - litellm_params = { - **litellm_params, - "litellm_embedding_model": resolved_model, - "litellm_embedding_config": resolved_config, - } + litellm_params: Final = ( + vector_store_to_run.get("litellm_params", {}) or {} + ) # mutable-ok: request execution merges persisted params into a mutable body data.update(litellm_params) return data diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 1930ec4aaba..8ca45f736ae 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -10,7 +10,7 @@ All /vector_store management endpoints import copy import json -from typing import TYPE_CHECKING, Any, Final, TypeAlias +from typing import TYPE_CHECKING, Any, Final from fastapi import APIRouter, Depends, HTTPException @@ -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, @@ -49,7 +43,6 @@ from litellm.types.vector_stores import ( from litellm.vector_stores.vector_store_registry import VectorStoreRegistry router: Final = APIRouter() -EmbeddingResolution: TypeAlias = tuple[str, dict[str, object]] def _vector_store_table(prisma_client: "PrismaClient") -> "TableActions[_VectorStoreRow]": @@ -65,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: """ @@ -156,264 +127,6 @@ async def _fetch_and_authorize_vector_store( return typed -def _provider_qualified_embedding_model( - fallback: str, - model: object, - custom_llm_provider: object, -) -> str: - if not isinstance(model, str) or not model: - return fallback - if "/" in model or not isinstance(custom_llm_provider, str) or not custom_llm_provider: - return model - return f"{custom_llm_provider}/{model}" - - -def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> EmbeddingResolution | 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: - Provider-qualified model and its connection config if found, otherwise None - """ - 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 - - resolved_model: Final = _provider_qualified_embedding_model( - fallback=embedding_model, - model=getattr(litellm_params, "model", None), - custom_llm_provider=getattr(litellm_params, "custom_llm_provider", None), - ) - - # 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 ( - resolved_model, - 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" -) -> EmbeddingResolution | 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: - Provider-qualified model and its connection config if found, otherwise None - """ - 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()), - ) - resolved_model: Final = _provider_qualified_embedding_model( - fallback=embedding_model, - model=decrypted_params.get("model"), - custom_llm_provider=decrypted_params.get("custom_llm_provider"), - ) - return ( - resolved_model, - 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 -) -> EmbeddingResolution | 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: - Provider-qualified model and its connection config if found, otherwise None - """ - 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 ######################################################## diff --git a/litellm/router.py b/litellm/router.py index 471a1116f44..9e0e267f21c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -84,6 +84,9 @@ 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, +) 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 @@ -6319,6 +6322,34 @@ class Router: client: object | None = None, **kwargs, ): + if call_type == "vector_store_search": + metadata: Final = self._vector_store_request_metadata(kwargs) + provider_kwargs: Final = ( + { + "custom_llm_provider": custom_llm_provider + } # mutable-ok: provider kwargs are expanded into the request + if custom_llm_provider is not None + else MappingProxyType({}) + ) + search_kwargs: Final = { # mutable-ok: the routed request requires dynamic keyword arguments + **kwargs, + **provider_kwargs, + "_direct_vector_store_embedding_executor": RouterVectorStoreEmbeddingExecutor( + router=self, + metadata=metadata, + ), + } + model: Final = search_kwargs.get("model") + if isinstance(model, str) and model: + routed_kwargs: Final = { # mutable-ok: model must be removed before expanding routed kwargs + key: value for key, value in search_kwargs.items() if key != "model" + } + return self._generic_api_call_with_fallbacks( + model=model, + original_function=original_function, + **routed_kwargs, + ) + return original_function(**search_kwargs) return self._generic_api_call_with_fallbacks(original_function=original_function, **kwargs) return sync_wrapper @@ -6512,10 +6543,21 @@ 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, - **kwargs, + **vector_store_kwargs, ) elif call_type in ("afile_delete", "afile_content"): return await self._ageneric_api_call_with_fallbacks( @@ -6551,6 +6593,18 @@ class Router: return async_wrapper + @staticmethod + 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 cast( # cast-ok: isinstance validates the runtime dict boundary + "dict[str, object]", litellm_metadata + ) + metadata: Final = kwargs.get("metadata") + if isinstance(metadata, dict): + return cast("dict[str, object]", metadata) # cast-ok: isinstance validates the runtime dict boundary + return MappingProxyType({}) + async def _init_vector_store_api_endpoints( self, original_function: Callable, diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index 9b0ff71730a..89c3319ca5a 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -15,6 +15,10 @@ 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 ( + LiteLLMVectorStoreEmbeddingExecutor, + VectorStoreEmbeddingExecutor, +) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( @@ -35,6 +39,14 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# +def _direct_vector_store_embedding_executor(value: object) -> VectorStoreEmbeddingExecutor: + if value is None: + return LiteLLMVectorStoreEmbeddingExecutor() + if isinstance(value, VectorStoreEmbeddingExecutor): + return value + raise TypeError("Invalid direct vector store embedding executor") + + def mock_vector_store_search_response( mock_results: list[VectorStoreSearchResult] | None = None, ): @@ -285,7 +297,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) + ) + 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() @@ -308,6 +325,7 @@ async def asearch( extra_body=extra_body, timeout=timeout, custom_llm_provider=custom_llm_provider, + _direct_vector_store_embedding_executor=embedding_executor, **kwargs, ) @@ -363,12 +381,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) + ) + 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: @@ -445,6 +467,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/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py index 75dacbaf08e..5c01587a6fe 100644 --- a/tests/router_unit_tests/test_router_embedding_integration.py +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -5,17 +5,107 @@ 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 +from unittest.mock import AsyncMock, MagicMock, patch import pytest - from litellm import Router +from litellm.llms.base_llm.vector_store.transformation import ( + LiteLLMVectorStoreEmbeddingExecutor, + RouterVectorStoreEmbeddingExecutor, +) +from litellm.types.utils import EmbeddingResponse 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): + response = EmbeddingResponse(data=[{"embedding": [0.1], "index": 0, "object": "embedding"}]) + sdk_executor = LiteLLMVectorStoreEmbeddingExecutor() + + with ( + patch("litellm.embedding", return_value=response) as embedding, + patch("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 + 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("litellm.embedding", return_value=response) as explicit_embedding: + assert router_executor.embed("openai/model", "query", {"api_key": "store-key"}) is response + explicit_embedding.assert_called_once_with(model="openai/model", input=["query"], api_key="store-key") + mock_router.embedding.assert_called_once() + + with patch("litellm.aembedding", new=AsyncMock(return_value=response)) as explicit_aembedding: + assert await router_executor.aembed("openai/model", "query", {"api_key": "store-key"}) is response + explicit_aembedding.assert_awaited_once_with(model="openai/model", input=["query"], api_key="store-key") + def test_embedding_with_deployment_specific_headers(self): """ Test that deployment-specific headers are propagated. 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/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 1484adb258f..ad411e874ca 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,98 @@ 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()) - 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() + wrapped_create = router.factory_function(create_original, call_type="vector_store_create") + with patch.object( # test-quality-ok: fallback dispatch is the boundary this wrapper delegates to + router, "_generic_api_call_with_fallbacks", return_value="created" + ) as fallback: + assert wrapped_create(name="store") == "created" + fallback.assert_called_once_with(original_function=create_original, 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 + 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_called_once_with(model="openai/model", input=["query"], api_key="store-key") + explicit_aembedding.assert_awaited_once_with(model="openai/model", input=["query"], api_key="store-key") @pytest.mark.asyncio @@ -82,10 +162,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 +177,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,89 +615,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": "multilingual-e5-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=("azure/multilingual-e5-large", 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/multilingual-e5-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_preserves_legacy_embedding_config_when_model_not_resolved(): - """A vector store row created by an older proxy version may already - carry a fully-resolved ``litellm_embedding_config`` in its persisted - ``litellm_params``. Preserve it when the model cannot be resolved.""" - 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(return_value=None) - - 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_awaited_once() - class TestCheckVectorStorePermission: """Test suite for check_vector_store_permission function.""" @@ -2001,60 +2055,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 = { - "model": "openai/text-embedding-3-small", - "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 - resolved_model, resolved_config = result - assert resolved_model == "openai/text-embedding-3-small" - assert resolved_config["api_key"] == "test-api-key" - assert resolved_config["api_base"] == "https://api.openai.com" - assert resolved_config["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 @@ -2071,14 +2072,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 @@ -2089,10 +2082,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 = {} @@ -2113,280 +2102,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_litellm_params.model = "text-embedding-3-small" - mock_litellm_params.custom_llm_provider = "openai" - - 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 - resolved_model, resolved_config = result - assert resolved_model == "openai/text-embedding-3-small" - assert resolved_config["api_key"] == "config-api-key" - assert resolved_config["api_base"] == "https://config-api-base.com" - assert resolved_config["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_litellm_params.model = "text-embedding-3-large" - mock_litellm_params.custom_llm_provider = "azure" - - 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 - resolved_model, resolved_config = result - assert resolved_model == "azure/text-embedding-3-large" - assert resolved_config["api_key"] == "azure-api-key" - assert resolved_config["api_base"] == "https://azure-endpoint.openai.azure.com" - assert resolved_config["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_litellm_params.model = "text-embedding-3-small" - mock_litellm_params.custom_llm_provider = "openai" - - 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 - resolved_model, resolved_config = result - assert resolved_model == "openai/text-embedding-3-small" - assert resolved_config["api_key"] == "resolved-from-env" - assert resolved_config["api_base"] == "https://direct-url.com" - assert "api_version" not in resolved_config - - 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_litellm_params.model = "text-embedding-3-small" - mock_litellm_params.custom_llm_provider = "openai" - - 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 - resolved_model, resolved_config = result - assert resolved_model == "openai/text-embedding-3-small" - assert resolved_config["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 = { - "model": "openai/text-embedding-3-small", - "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 - resolved_model, resolved_config = result - assert resolved_model == "openai/text-embedding-3-small" - assert resolved_config["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 @@ -2445,9 +2175,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/uv.lock b/uv.lock index 27be919eea1..8d886044083 100644 --- a/uv.lock +++ b/uv.lock @@ -9441,19 +9441,19 @@ wheels = [ [[package]] name = "tornado" -version = "6.5.7" +version = "6.5.8" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/64/24/95ec527ad67b76d59299e5465b3935d05e4294b7e0290a3924b7487df30b/tornado-6.5.7.tar.gz", hash = "sha256:66c513a76cda70d53907bc27cf1447557699c2e95aa48ba27a442ff61c3ddfc2", size = 519252, upload-time = "2026-06-08T17:34:51.232Z" } +sdist = { url = "https://files.pythonhosted.org/packages/10/d3/343e5bb989d6515b1646cf3d40135d73f3d5e45339bded401b56cdac24dd/tornado-6.5.8.tar.gz", hash = "sha256:9452e1b208a8bd771e2cb1f2ff564985b9b214bdebbe622793e1799e0a6bd23f", size = 520493, upload-time = "2026-08-07T02:12:42.971Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/02/dc/c7043cab6fed8ae159fc1923ce829ada35c4dbd797d408a43858ffaf9639/tornado-6.5.7-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:148b2eb15c2c765a50796172c1e499649b35f30d2e3c3d3e15913cfa56bfb163", size = 448543, upload-time = "2026-06-08T17:34:38.052Z" }, - { url = "https://files.pythonhosted.org/packages/92/4f/090b1431e5a43df696feceffc268c5383cc079ecb5f08ce58f917109aafe/tornado-6.5.7-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9da38de27f1da3b78a966f0dae12b5a1ea9afe72ca805d84ff06508272ddf100", size = 446707, upload-time = "2026-06-08T17:34:39.594Z" }, - { url = "https://files.pythonhosted.org/packages/37/d8/ef374952fd5da67d4463122c2b8e5a96536ec10b4b339254c6dcde81d01c/tornado-6.5.7-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8d759e71906ee783f8867b93bf26a265743da4c1e2f4a018464c1ba019862972", size = 449774, upload-time = "2026-06-08T17:34:41.204Z" }, - { url = "https://files.pythonhosted.org/packages/35/37/d434c73f4c6e014b745b9b37085f34f40c022f007efff3d7fe65991899f3/tornado-6.5.7-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a46347a18f23fb92b396beebe0fb78f61dda0cc302445202c16203d8a18848b", size = 450745, upload-time = "2026-06-08T17:34:42.531Z" }, - { url = "https://files.pythonhosted.org/packages/b6/2b/56b9aff361d7f1ab728a805ec7d7ea835f8807afa9f5cc690ea0e630efb9/tornado-6.5.7-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:7778b30bef919231265e91c69963ce0f49a1e9c07ac900bbe75b19ce2575ba92", size = 450578, upload-time = "2026-06-08T17:34:43.787Z" }, - { url = "https://files.pythonhosted.org/packages/02/30/a7444fb23aa76860a14198fab96ac79f1866b0a6e19e26c4381b0938e50f/tornado-6.5.7-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e726f0c75da7726eec023aa62751ff8878bd2737e34fbdd33b1ae5897d2200f5", size = 449985, upload-time = "2026-06-08T17:34:45.326Z" }, - { url = "https://files.pythonhosted.org/packages/5c/42/5f0e56c01e8d9d36f4e23f367b85ae6cae0c1ecddd5e6977d8388ad27488/tornado-6.5.7-cp39-abi3-win32.whl", hash = "sha256:f8de3bf12d3efdd0cbe7c8887868198f8a91415e3f29fcf258d9b8eb7b1d9ae4", size = 451047, upload-time = "2026-06-08T17:34:46.784Z" }, - { url = "https://files.pythonhosted.org/packages/c9/a4/b393076ffb21b469eec5b328a0534cf03a3b90bfc6b1f09507cdd075d938/tornado-6.5.7-cp39-abi3-win_amd64.whl", hash = "sha256:de942f843533a039ef9fa3d9c88c7cd8a7c94553fb5ad0154270989b3d99a2c4", size = 451485, upload-time = "2026-06-08T17:34:48.248Z" }, - { url = "https://files.pythonhosted.org/packages/71/2e/7b1c769803121b809112cf9a00681c472eae1d80e32d7ec0e0bd61d0d0e1/tornado-6.5.7-cp39-abi3-win_arm64.whl", hash = "sha256:ff934fce95643af5f11efdae618eaa73d469dc588641e5c8d19295a0c65c4796", size = 450506, upload-time = "2026-06-08T17:34:49.702Z" }, + { url = "https://files.pythonhosted.org/packages/f2/d5/007086fd8df5489338e204f65adce33fd4f21a4999dbb2b9cff2f897b5f4/tornado-6.5.8-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:cc6aa787d7cfab7c3d35189dc7a56fbd2399a569624c730c6b55b3d6531d0403", size = 449487, upload-time = "2026-08-07T02:12:28.682Z" }, + { url = "https://files.pythonhosted.org/packages/70/c8/5a24a99495903f594f6a199dd7beead1cbc0a13e2cb9102727bcaaf2a997/tornado-6.5.8-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9715b5eb79735b2bcd454ce216a9275b7c0470e64ea1bf5742f78b2f72b26eeb", size = 447649, upload-time = "2026-08-07T02:12:30.306Z" }, + { url = "https://files.pythonhosted.org/packages/6e/de/f2e733f386b85962d1b1dc82cd63d169b5b4580062b35397eac9244a41fe/tornado-6.5.8-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:547d63f450d570c14fe0e8db2cfb14c9bbd1c2503b4a6612586267955aa47b58", size = 450707, upload-time = "2026-08-07T02:12:31.95Z" }, + { url = "https://files.pythonhosted.org/packages/0b/94/20efeee9a01c141e9ac47c397f81679dfda24b32768fc4fff24e76d36c2c/tornado-6.5.8-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7e2360a0ffbe145eca8af0b19cb7203d79b1a98dd4cccdd6b368f6f49c2e3808", size = 451677, upload-time = "2026-08-07T02:12:33.512Z" }, + { url = "https://files.pythonhosted.org/packages/42/ec/a96ccb8ccf0de2b7bc2c5fa1608a4803735018242e90c4882365a9fd418f/tornado-6.5.8-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:5d242290bdf7ab3151bc1065fdd75c0dcc21cbc7b49f22a4c56329c2d6566d22", size = 451510, upload-time = "2026-08-07T02:12:35.346Z" }, + { url = "https://files.pythonhosted.org/packages/29/b5/93185859245ad3f00e62175f29607346788b696369347f0146e0421286bb/tornado-6.5.8-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7b94ff0e128fe0542f3bd331fb44d06260fc4ac16881545159f34ef08aad4195", size = 450917, upload-time = "2026-08-07T02:12:36.963Z" }, + { url = "https://files.pythonhosted.org/packages/97/cf/fe33cf062834487d34d1559746a4a12521033c22645b6d74d4bca702e018/tornado-6.5.8-cp39-abi3-win32.whl", hash = "sha256:67832909c4779c64942380cb5f044a5c6163d00831472d80e25e115de9917836", size = 451952, upload-time = "2026-08-07T02:12:38.512Z" }, + { url = "https://files.pythonhosted.org/packages/cb/e1/468ad54333e92ccb62627e62cb88e5fc14a2171daa67ed47b1b8542d5b86/tornado-6.5.8-cp39-abi3-win_amd64.whl", hash = "sha256:11881db6b7c168494be2c2d12e65931451bdf7ee718535418ae1d8855dd5a0ee", size = 452391, upload-time = "2026-08-07T02:12:39.971Z" }, + { url = "https://files.pythonhosted.org/packages/ad/3e/cd5e4f06e34cde33b8ef66cf36aa2b5ad46354cc1af7d2136bbe365fee1d/tornado-6.5.8-cp39-abi3-win_arm64.whl", hash = "sha256:68a7468c7e289f8514d7d664101753903217eff1bb6822c6b5994a0b5f5bcb26", size = 451411, upload-time = "2026-08-07T02:12:41.469Z" }, ] [[package]] From 6805d01709f9401f36bba163dc7f9c3192643c51 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 14:05:12 -0700 Subject: [PATCH 05/31] fix(vector-store): preserve aliases with embedding config --- .../base_llm/vector_store/transformation.py | 21 +++++--- .../test_router_embedding_integration.py | 53 ++++++++++++++++--- .../test_vector_store_endpoints.py | 17 +++++- 3 files changed, 75 insertions(+), 16 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 772e4f849a0..e9c925448a8 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -3,7 +3,7 @@ from __future__ import annotations from abc import abstractmethod from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, NoReturn, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, runtime_checkable import httpx @@ -65,22 +65,29 @@ class RouterVectorStoreEmbeddingExecutor: router: Router metadata: Mapping[str, object] + def _embedding_kwargs(self, configuration: Mapping[str, object]) -> dict[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 embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: - if configuration: - return LiteLLMVectorStoreEmbeddingExecutor().embed(model, query, configuration) 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 - metadata=dict(self.metadata), # mutable-ok: Router metadata requires a concrete dict + **self._embedding_kwargs(configuration), # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic ) async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: - if configuration: - return await LiteLLMVectorStoreEmbeddingExecutor().aembed(model, query, configuration) 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 - metadata=dict(self.metadata), # mutable-ok: Router metadata requires a concrete dict + **self._embedding_kwargs(configuration), # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic ) diff --git a/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py index 5c01587a6fe..d5e0c750d88 100644 --- a/tests/router_unit_tests/test_router_embedding_integration.py +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -97,14 +97,53 @@ class TestRouterEmbeddingIntegration: metadata={"user_api_key_team_id": "team-a"}, ) - with patch("litellm.embedding", return_value=response) as explicit_embedding: - assert router_executor.embed("openai/model", "query", {"api_key": "store-key"}) is response - explicit_embedding.assert_called_once_with(model="openai/model", input=["query"], api_key="store-key") - mock_router.embedding.assert_called_once() + alias_router = Router( + model_list=[ + { + "model_name": "team-alias", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "deployment-key", + }, + } + ] + ) + alias_executor = RouterVectorStoreEmbeddingExecutor( + router=alias_router, + metadata={"user_api_key_team_id": "team-a"}, + ) + 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", + } - with patch("litellm.aembedding", new=AsyncMock(return_value=response)) as explicit_aembedding: - assert await router_executor.aembed("openai/model", "query", {"api_key": "store-key"}) is response - explicit_aembedding.assert_awaited_once_with(model="openai/model", input=["query"], api_key="store-key") + with ( + patch("litellm.embedding", return_value=response) as explicit_embedding, + patch("litellm.aembedding", new=AsyncMock(return_value=response)) as explicit_aembedding, + ): + assert alias_executor.embed("team-alias", "sync query", explicit_config) is response + assert await alias_executor.aembed("team-alias", "async query", explicit_config) is response + + sync_kwargs = explicit_embedding.call_args.kwargs + assert sync_kwargs["model"] == "openai/text-embedding-3-small" + assert sync_kwargs["input"] == ["sync query"] + assert sync_kwargs["api_base"] == "https://embedding.example/v1" + assert sync_kwargs["api_key"] == "store-key" + assert sync_kwargs["metadata"]["configured"] is True + assert sync_kwargs["metadata"]["user_api_key_team_id"] == "team-a" + + async_kwargs = explicit_aembedding.await_args.kwargs + assert async_kwargs["model"] == "openai/text-embedding-3-small" + assert async_kwargs["input"] == ["async query"] + assert async_kwargs["api_base"] == "https://embedding.example/v1" + assert async_kwargs["api_key"] == "store-key" + assert async_kwargs["metadata"]["configured"] is True + assert async_kwargs["metadata"]["user_api_key_team_id"] == "team-a" def test_embedding_with_deployment_specific_headers(self): """ 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 ad411e874ca..903d5cb55f3 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 @@ -117,6 +117,7 @@ async def test_vector_store_embedding_executors_preserve_explicit_configuration( 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"}, @@ -140,8 +141,20 @@ async def test_vector_store_embedding_executors_preserve_explicit_configuration( 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_called_once_with(model="openai/model", input=["query"], api_key="store-key") - explicit_aembedding.assert_awaited_once_with(model="openai/model", input=["query"], api_key="store-key") + 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 From 5799a32cdda6647d2f16460d79b8610dbae49d34 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 15:37:50 -0700 Subject: [PATCH 06/31] fix(vector-store): route pre-call searches through router --- .../vector_store_pre_call_hook.py | 22 ++++++++-- .../test_bedrock_knowledgebase_hook.py | 42 +++++++++++++++++++ 2 files changed, 61 insertions(+), 3 deletions(-) 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..aaf5cb080dc 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 @@ -80,10 +80,15 @@ 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, + prisma_client as _prisma_client, + ) prisma_client = _prisma_client + llm_router = _llm_router except ImportError: pass @@ -114,12 +119,23 @@ 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: Final = ( + litellm_logging_obj.model_call_details.get("litellm_params", {}) + if litellm_logging_obj is not None + else {} + ) + request_metadata: Final = ( + request_litellm_params.get("metadata", {}) if isinstance(request_litellm_params, dict) else {} + ) + search_function: Final = ( + llm_router.avector_store_search if llm_router is not None else 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/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 3f9f2bacdd3..06083b77e84 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, From 0cc0c47f8b2a4196beb4e766bcc65fe84743e9ed Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 15:52:57 -0700 Subject: [PATCH 07/31] style(vector-store): satisfy import lint --- .../vector_store_pre_call_hook.py | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) 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 aaf5cb080dc..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 @@ -84,6 +85,8 @@ class VectorStorePreCallHook(CustomLogger): try: from litellm.proxy.proxy_server import ( llm_router as _llm_router, + ) + from litellm.proxy.proxy_server import ( prisma_client as _prisma_client, ) @@ -119,17 +122,20 @@ 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 {} - request_litellm_params: Final = ( - litellm_logging_obj.model_call_details.get("litellm_params", {}) - if litellm_logging_obj is not None - else {} - ) - request_metadata: Final = ( + 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 {} ) - search_function: Final = ( - llm_router.avector_store_search if llm_router is not None else litellm.vector_stores.asearch - ) + 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, From 1cd99a036e0538bb61b280e17639545c75374d81 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 17:29:58 -0700 Subject: [PATCH 08/31] fix(router): route Claude Code subagents through session router --- litellm/router.py | 80 +++++++++++++++++- tests/test_litellm/test_router.py | 129 ++++++++++++++++++++++++++++++ 2 files changed, 208 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index 6e4405ebfef..b49269f2457 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -353,6 +353,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 def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) -> bool: @@ -12546,6 +12548,77 @@ 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 _resolve_claude_code_session_router( + self, + model: str, + registered_model_name: str, + request_kwargs: Mapping[str, object], + ) -> str: + 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 + + agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") + if agent_id is not None: + bound_model: Final = await self.cache.async_get_cache(key=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.cache.async_delete_cache(key=cache_key) + return registered_model_name + await self.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 request_kwargs.get("fallback_depth") not in (None, 0): + return registered_model_name + if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: + await self.cache.async_delete_cache(key=cache_key) + return registered_model_name + await self.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, @@ -12565,7 +12638,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. diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 84f6344be35..d22a1cc04ca 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8316,6 +8316,135 @@ class TestConsumedRequestTagsStamp: assert CONSUMED_REQUEST_TAGS_METADATA_KEY not in request_kwargs["metadata"] +class TestClaudeCodeSubagentSessionRouterBinding: + class _RewriteStrategy: + 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="cheap-model", + messages=messages, + routing_decision={ + "router_model_name": "smart-router", + "router_type": "complexity", + "routed_model": "cheap-model", + "cause": "heuristic_scorer", + }, + ) + + @classmethod + def _router(cls) -> "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"}, + }, + ] + ) + router.complexity_routers = { + "smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy())] + } + 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_direct_model_clears_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 None + + @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_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. From 82d046c42821ca3b24036dd57e362252564c0596 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 17:43:48 -0700 Subject: [PATCH 09/31] fix(router): make Claude session cleanup best effort --- litellm/router.py | 14 ++++++++++++-- tests/test_litellm/test_router.py | 19 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b49269f2457..2bb725b0880 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12578,6 +12578,16 @@ class Router: 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.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 _resolve_claude_code_session_router( self, model: str, @@ -12595,7 +12605,7 @@ class Router: 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.cache.async_delete_cache(key=cache_key) + await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, @@ -12610,7 +12620,7 @@ class Router: if request_kwargs.get("fallback_depth") not in (None, 0): return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: - await self.cache.async_delete_cache(key=cache_key) + await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index d22a1cc04ca..77b73a2d12a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8409,6 +8409,25 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None + @pytest.mark.asyncio + async def test_redis_cleanup_failure_does_not_reject_a_direct_model_request(self): + from litellm.caching.caching import RedisCache + + router = self._router() + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) + + await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) + router._update_redis_cache(cache=redis_cache) + + response = await router.async_pre_routing_hook( + model="expensive-model", + request_kwargs=self._request_kwargs(), + ) + + assert response is None + redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): router = self._router() From e3a61c82da9f8dbe42fdfcc4907af7b3a2901392 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 17:46:41 -0700 Subject: [PATCH 10/31] test(router): register indirect session routing coverage --- tests/code_coverage_tests/router_code_coverage.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index a5e00799519..60b56b7fac6 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -82,6 +82,10 @@ 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 ] From 2d4301589c1e741489f68e6eef3c2d112da91e2a Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Wed, 2 Sep 2026 01:10:31 +0000 Subject: [PATCH 11/31] fix(router): keep serving when Claude Code session router cleanup fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 12 ++++-------- tests/test_litellm/test_router.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+), 8 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 2bb725b0880..d6d9f20085f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12578,15 +12578,11 @@ class Router: 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: + async def _clear_claude_code_session_router(self, cache_key: str) -> None: try: await self.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, - ) + verbose_router_logger.debug("Claude Code session router cleanup skipped for %s: %s", cache_key, e) async def _resolve_claude_code_session_router( self, @@ -12605,7 +12601,7 @@ class Router: 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) + await self._clear_claude_code_session_router(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, @@ -12620,7 +12616,7 @@ class Router: if request_kwargs.get("fallback_depth") not in (None, 0): return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: - await self._delete_claude_code_session_router_binding(cache_key) + await self._clear_claude_code_session_router(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 77b73a2d12a..b66dbf6aa7d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8428,6 +8428,24 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() + @pytest.mark.asyncio + async def test_main_direct_model_still_served_when_cache_delete_fails(self): + router = self._router() + await router.acompletion( + model="smart-router", messages=[{"role": "user", "content": "main turn"}], **self._request_kwargs() + ) + + async def failing_delete(key: str) -> None: + raise Exception("Redis circuit breaker is open — skipping async_delete_cache") + + router.cache.async_delete_cache = failing_delete + + response = await router.acompletion( + model="expensive-model", messages=[{"role": "user", "content": "direct turn"}], **self._request_kwargs() + ) + + assert response.choices[0].message.content == "expensive response" + @pytest.mark.asyncio async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): router = self._router() From 46502f58042a41619648be7acc67108079f737e2 Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Wed, 2 Sep 2026 01:11:23 +0000 Subject: [PATCH 12/31] Revert "fix(router): keep serving when Claude Code session router cleanup fails" This reverts commit 2d4301589c1e741489f68e6eef3c2d112da91e2a. --- litellm/router.py | 12 ++++++++---- tests/test_litellm/test_router.py | 18 ------------------ 2 files changed, 8 insertions(+), 22 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index d6d9f20085f..2bb725b0880 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12578,11 +12578,15 @@ class Router: return None return f"claude_code_session_router:v1:{caller_scope}:{session_id}" - async def _clear_claude_code_session_router(self, cache_key: str) -> None: + async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: try: await self.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.debug("Claude Code session router cleanup skipped for %s: %s", cache_key, e) + verbose_router_logger.warning( + "Failed to delete Claude Code session router binding; " + "the binding may remain until its TTL expires: %s", + e, + ) async def _resolve_claude_code_session_router( self, @@ -12601,7 +12605,7 @@ class Router: 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._clear_claude_code_session_router(cache_key) + await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, @@ -12616,7 +12620,7 @@ class Router: if request_kwargs.get("fallback_depth") not in (None, 0): return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: - await self._clear_claude_code_session_router(cache_key) + await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name await self.cache.async_set_cache( key=cache_key, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b66dbf6aa7d..77b73a2d12a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8428,24 +8428,6 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() - @pytest.mark.asyncio - async def test_main_direct_model_still_served_when_cache_delete_fails(self): - router = self._router() - await router.acompletion( - model="smart-router", messages=[{"role": "user", "content": "main turn"}], **self._request_kwargs() - ) - - async def failing_delete(key: str) -> None: - raise Exception("Redis circuit breaker is open — skipping async_delete_cache") - - router.cache.async_delete_cache = failing_delete - - response = await router.acompletion( - model="expensive-model", messages=[{"role": "user", "content": "direct turn"}], **self._request_kwargs() - ) - - assert response.choices[0].message.content == "expensive response" - @pytest.mark.asyncio async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): router = self._router() From 8a0967443d84ecc02c10caf6ca55385b907b11b2 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 20:38:13 -0700 Subject: [PATCH 13/31] fix(router): isolate Claude session binding cache --- litellm/router.py | 19 +++++++++++++------ tests/test_litellm/test_router.py | 14 ++++++++++++++ 2 files changed, 27 insertions(+), 6 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 2bb725b0880..fb8625bc6ba 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -782,6 +782,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) @@ -1102,8 +1106,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 @@ -12580,7 +12584,7 @@ class Router: async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: try: - await self.cache.async_delete_cache(key=cache_key) + 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; " @@ -12600,14 +12604,14 @@ class Router: agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") if agent_id is not None: - bound_model: Final = await self.cache.async_get_cache(key=cache_key) + bound_model: Final = await self._claude_code_session_router_cache.async_get_cache(key=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.cache.async_set_cache( + await self._claude_code_session_router_cache.async_set_cache( key=cache_key, value=bound_model, ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, @@ -12622,7 +12626,7 @@ class Router: if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name - await self.cache.async_set_cache( + await self._claude_code_session_router_cache.async_set_cache( key=cache_key, value=model, ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, @@ -13450,6 +13454,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/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 77b73a2d12a..a528266d738 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8428,6 +8428,20 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() + @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() From 6adc14b4b12250ddc927682c8153469976e87888 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 21:13:45 -0700 Subject: [PATCH 14/31] fix(router): preserve Claude subagent fallbacks --- litellm/router.py | 4 ++-- tests/test_litellm/test_router.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index fb8625bc6ba..6f1f1bc700b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12601,6 +12601,8 @@ class Router: 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: @@ -12621,8 +12623,6 @@ class Router: if self._request_header(request_kwargs, "x-app") != "cli": return registered_model_name - if request_kwargs.get("fallback_depth") not in (None, 0): - return registered_model_name if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: await self._delete_claude_code_session_router_binding(cache_key) return registered_model_name diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index a528266d738..b413bb18f04 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8464,6 +8464,19 @@ class TestClaudeCodeSubagentSessionRouterBinding: 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_session_router_binding_is_scoped_to_the_authenticated_key(self): router = self._router() From 7b919f89a85ea7671fd0c3ac2fb27d31fb74201b Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 2 Sep 2026 10:00:47 -0700 Subject: [PATCH 15/31] fix(router): track routed model in fallback attempts --- .../router_utils/fallback_event_handlers.py | 5 +-- .../test_fallback_event_handlers.py | 21 +++++++++++ tests/test_litellm/test_router.py | 35 +++++++++++++++++-- 3 files changed, 56 insertions(+), 5 deletions(-) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 3d37ca216a7..0167721f9fe 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -470,10 +470,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/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 894b2d9e74f..9e51a60364b 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", diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b413bb18f04..d255135bf64 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8335,20 +8335,26 @@ class TestClaudeCodeSubagentSessionRouterBinding: ) @classmethod - def _router(cls) -> "litellm.Router": + 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"}, + "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())] @@ -8477,6 +8483,29 @@ class TestClaudeCodeSubagentSessionRouterBinding: 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_session_router_binding_is_scoped_to_the_authenticated_key(self): router = self._router() From d7ee215c57af44219d5a19f5042343ce98e9de2a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:32:04 -0700 Subject: [PATCH 16/31] fix(responses): keep namespace tools intact when a guardrail returns them unchanged Any pre_call guardrail on /v1/responses flattened Codex namespace tools into ns__member functions and wrote the flattened list back to the request, so the model called mcp__server__tool with no namespace and Codex rejected the call as unsupported. The handler now keeps the client's original tools, hands the guardrail a deep copy of the flattened ones, and rebuilds data["tools"] by matching the guardrail's output to the originals by type and name. Unchanged tools go back as the original objects, a dropped or edited namespace member changes only that member, and tools the guardrail injects are still appended. Fixes #39183 --- basedpyright-code-budget.json | 12 +- .../guardrail_translation/handler.py | 114 +++------- .../guardrail_translation/tool_merge.py | 177 +++++++++++++++ .../transformation.py | 161 ++++++++------ ruff-strict-budget.json | 2 +- ...test_openai_responses_guardrail_handler.py | 206 +++++++++++++++++- ...t_openai_responses_guardrail_tool_merge.py | 144 ++++++++++++ .../test_litellm_completion_responses.py | 13 ++ type-discipline-budget.json | 10 +- 9 files changed, 669 insertions(+), 170 deletions(-) create mode 100644 litellm/llms/openai/responses/guardrail_translation/tool_merge.py create mode 100644 tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index da788bf1ce3..e7a069de29a 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 14076 }, "reportArgumentType": { - "limit": 2216 + "limit": 2215 }, "reportAssignmentType": { "limit": 319 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 4128 + "limit": 4127 }, "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/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 1530c154e93..5a5970fb867 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -28,6 +28,7 @@ Output: response.output is List[GenericResponseOutputItem] where each has: - text: str """ +import copy import time import uuid from collections.abc import Mapping, Sequence @@ -36,7 +37,6 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall -from openai.types.responses.tool_param import FunctionToolParam from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -49,6 +49,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, stream_item_field, ) +from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -62,7 +63,6 @@ from litellm.types.llms.openai import ( ContentPartDonePartOutputText, ErrorEvent, ErrorEventError, - OpenAIMcpServerTool, OutputItemAddedEvent, OutputItemDoneEvent, OutputTextDeltaEvent, @@ -157,23 +157,31 @@ class OpenAIResponsesHandler(BaseTranslation): Handles both string input and list of message objects. """ input_data: Final[str | ResponseInputParam | None] = data.get("input") - tools_to_check: Final[list[ChatCompletionToolParam]] = [] if input_data is None: return data structured_messages: Final = self.get_structured_messages(data) + raw_tools: Final = data.get("tools") + original_tools: Final[tuple[Mapping[str, object], ...]] = ( + tuple(raw_tools) if isinstance(raw_tools, list) else () + ) + flattened_tool_groups: Final = tuple( + form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools) + ) + flattened_tools: Final = tuple( + cast(ChatCompletionToolParam, tool) # cast-ok: mcp tools ride along in the guardrail's tool list + for group in flattened_tool_groups + for tool in group + ) + tools_to_check: Final[list[ChatCompletionToolParam]] = list( # mutable-ok: guardrail inputs want a list + copy.deepcopy(flattened_tools) + ) # Handle simple string input if isinstance(input_data, str): inputs = GenericGuardrailAPIInputs(texts=[input_data]) - original_tools: list[dict[str, object]] = [] - - # Extract and transform tools if present - if "tools" in data and data["tools"]: - original_tools = list(data["tools"]) - self._extract_and_transform_tools(data["tools"], tools_to_check) - if tools_to_check: - inputs["tools"] = tools_to_check + if tools_to_check: + inputs["tools"] = tools_to_check if structured_messages: inputs["structured_messages"] = structured_messages # Include model information if available @@ -189,7 +197,9 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts = guardrailed_inputs.get("texts", []) data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data - self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools")) + self._apply_guardrailed_tools_to_data( + data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") + ) verbose_proxy_logger.debug("OpenAI Responses API: Processed string input") return data @@ -200,7 +210,6 @@ class OpenAIResponsesHandler(BaseTranslation): texts_to_check: Final[list[str]] = [] images_to_check: Final[list[str]] = [] task_mappings: Final[list[tuple[int, int | None]]] = [] - original_tools_list: Final[list[dict[str, object]]] = list(data.get("tools") or []) # Step 1: Extract all text content, images, and tools for msg_idx, message in enumerate(input_data): @@ -212,10 +221,6 @@ class OpenAIResponsesHandler(BaseTranslation): task_mappings=task_mappings, ) - # Extract and transform tools if present - if "tools" in data and data["tools"]: - self._extract_and_transform_tools(data["tools"], tools_to_check) - # Step 2: Apply guardrail to all texts in batch if texts_to_check: inputs = GenericGuardrailAPIInputs(texts=texts_to_check) @@ -238,9 +243,7 @@ class OpenAIResponsesHandler(BaseTranslation): guardrailed_texts = guardrailed_inputs.get("texts", []) self._apply_guardrailed_tools_to_data( - data, - original_tools_list, - guardrailed_inputs.get("tools"), + data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") ) # Step 3: Map guardrail responses back to original input structure @@ -267,73 +270,18 @@ class OpenAIResponsesHandler(BaseTranslation): names.append(str(tool["server_label"])) return names - def _extract_and_transform_tools( - self, - tools: list[FunctionToolParam | OpenAIMcpServerTool], - tools_to_check: list[ChatCompletionToolParam], - ) -> None: - """ - Extract and transform tools from Responses API format to Chat Completion format. - - Uses the LiteLLM transformation function to convert Responses API tools - to Chat Completion tools that can be passed to guardrails. - """ - if tools is not None and isinstance(tools, list): - # Transform Responses API tools to Chat Completion tools - ( - transformed_tools, - _, - ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools) - tools_to_check.extend(cast(list[ChatCompletionToolParam], transformed_tools)) - - def _remap_tools_to_responses_api_format(self, guardrailed_tools: list[Any]) -> list[dict[str, object]]: - """ - Remap guardrail-returned tools (Chat Completion format) back to - Responses API request tool format. - """ - return LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools( - guardrailed_tools - ) - - def _merge_tools_after_guardrail( - self, - original_tools: list[dict[str, object]], - remapped: list[dict[str, object]], - ) -> list[dict[str, object]]: - """ - Merge remapped guardrailed tools with original tools that were not sent - to the guardrail (e.g. web_search, web_search_preview), preserving order. - Tools a guardrail appended (``remapped`` longer than ``original_tools``) - have no original slot and are kept so an injected tool is not dropped. - """ - if not original_tools: - return remapped - result: Final[list[dict[str, object]]] = [] - j = 0 - for tool in original_tools: - if isinstance(tool, dict) and tool.get("type") in ( - "web_search", - "web_search_preview", - ): - result.append(tool) - else: - if j < len(remapped): - result.append(remapped[j]) - j += 1 - # Keep guardrail-appended tools that matched no original slot above. - result.extend(remapped[j:]) - return result - def _apply_guardrailed_tools_to_data( self, data: dict, - original_tools: list[dict[str, object]], - guardrailed_tools: list[ChatCompletionToolParam] | None, + original_tools: Sequence[Mapping[str, object]], + flattened_tool_groups: Sequence[Sequence[Mapping[str, object]]], + guardrailed_tools: Sequence[ChatCompletionToolParam] | None, ) -> None: - """Remap guardrailed tools to Responses API format and merge with original, then set data['tools'].""" - if guardrailed_tools is not None: - remapped: Final = self._remap_tools_to_responses_api_format(guardrailed_tools) - data["tools"] = self._merge_tools_after_guardrail(original_tools, remapped) + if guardrailed_tools is None: + return + data["tools"] = list( # mutable-ok: downstream wants a list # rebind-ok: in-place request rewrite + merge_guardrailed_tools(original_tools, flattened_tool_groups, guardrailed_tools) + ) def _extract_input_text_and_images( self, diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py new file mode 100644 index 00000000000..3ae951d3f61 --- /dev/null +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -0,0 +1,177 @@ +from collections.abc import Iterable, Mapping, Sequence +from itertools import accumulate, chain +from types import MappingProxyType +from typing import Final, TypeAlias + +from pydantic import 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: + try: + return _TOOL_ADAPTER.validate_python(value) + 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) + return tuple((key, keys[:position].count(key)) 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,) + if not guardrailed_group: + return () + 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) + 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/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 5f3e88bb12f..9f91d7527cb 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,82 @@ 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"): + # 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, + ) + 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 +1932,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 +1952,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/ruff-strict-budget.json b/ruff-strict-budget.json index 9b1cc977a64..935fed18a79 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -177,7 +177,7 @@ "limit": 8 }, "RUF019": { - "limit": 31 + "limit": 29 }, "RUF046": { "limit": 4 diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 315b6948bd8..d071ef78c2d 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -5,6 +5,8 @@ Tests the handler's ability to process input/output for the Responses API with guardrail transformations. """ +import copy +from collections.abc import Callable from typing import Any, List, Literal, Optional, Tuple from unittest.mock import AsyncMock, MagicMock @@ -19,6 +21,10 @@ from litellm.llms import get_guardrail_translation_mapping from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) +from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.responses.main import GenericResponseOutputItem, OutputText from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs @@ -1287,14 +1293,14 @@ class TestOpenAIResponsesHandlerToolInjection: """A tool a guardrail injects must survive the write-back to Responses format.""" def test_merge_keeps_guardrail_appended_tool(self): - """_merge_tools_after_guardrail must not drop the extra appended tool.""" - handler = OpenAIResponsesHandler() + """merge_guardrailed_tools must not drop the extra appended tool.""" original = [{"type": "function", "name": "a"}] - remapped = [ - {"type": "function", "name": "a"}, - {"type": "function", "name": "b"}, + groups = [form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original)] + guardrailed = [ + *groups[0], + {"type": "function", "function": {"name": "b", "description": "", "parameters": {"type": "object"}}}, ] - merged = handler._merge_tools_after_guardrail(original, remapped) + merged = merge_guardrailed_tools(original, groups, guardrailed) assert [t["name"] for t in merged] == ["a", "b"] @pytest.mark.asyncio @@ -1323,6 +1329,194 @@ class TestOpenAIResponsesHandlerToolInjection: assert "injected_tool" in names +class ToolEditingGuardrail(CustomGuardrail): + """Guardrail that rewrites the flattened chat tools it was handed through ``edit``""" + + def __init__(self, edit: Callable[[list[dict]], list[dict]], **kwargs): + super().__init__(**kwargs) + self.edit = edit + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Any | None = None, + ) -> GenericGuardrailAPIInputs: + inputs["tools"] = self.edit(list(inputs.get("tools") or [])) + return inputs + + +def _codex_request(input_value): + """A Responses API request shaped like what the Codex CLI sends when an MCP server is configured""" + return { + "model": "gpt-5.3-codex", + "input": input_value, + "tools": [ + { + "type": "function", + "name": "get_weather", + "description": "Weather lookup", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + "strict": False, + }, + { + "type": "namespace", + "name": "mcp__confluence", + "description": "Confluence tools", + "tools": [ + { + "type": "function", + "name": "confluence_get_page", + "description": "Get a page", + "parameters": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}, + "strict": False, + }, + { + "type": "function", + "name": "confluence_search", + "description": "Search pages", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, + "strict": False, + }, + ], + }, + { + "type": "custom", + "name": "apply_patch", + "description": "Apply a patch", + "format": {"type": "grammar", "syntax": "lark", "definition": 'start: "x"'}, + }, + {"type": "web_search"}, + ], + } + + +def _tool_named(tools, name): + return next(tool for tool in tools if tool.get("name") == name) + + +class TestOpenAIResponsesHandlerNamespaceTools: + """Codex sends MCP tools as ``namespace`` tools; a guardrail must never flatten them (GH #39183)""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "input_value", + ["hi", [{"role": "user", "content": "hi", "type": "message"}]], + ids=["string_input", "list_input"], + ) + async def test_pass_through_guardrail_leaves_tools_untouched(self, input_value): + data = _codex_request(input_value) + expected_tools = copy.deepcopy(data["tools"]) + + result = await OpenAIResponsesHandler().process_input_messages( + data, MockPassThroughGuardrail(guardrail_name="test") + ) + + assert result["tools"] == expected_tools + + @pytest.mark.asyncio + async def test_appending_guardrail_keeps_namespace_and_adds_tool(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + + result = await OpenAIResponsesHandler().process_input_messages( + data, ToolAppendingGuardrail(guardrail_name="test") + ) + + assert result["tools"][:-1] == expected_tools + assert result["tools"][-1]["type"] == "function" + assert result["tools"][-1]["name"] == "injected_tool" + + @pytest.mark.asyncio + async def test_dropping_one_member_prunes_only_that_member(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + guardrail = ToolEditingGuardrail( + edit=lambda tools: [t for t in tools if t["function"]["name"] != "mcp__confluence__confluence_search"], + guardrail_name="test", + ) + + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + namespace = _tool_named(result["tools"], "mcp__confluence") + assert [member["name"] for member in namespace["tools"]] == ["confluence_get_page"] + assert namespace["tools"][0] == expected_tools[1]["tools"][0] + assert [t for t in result["tools"] if t is not namespace] == [expected_tools[0], *expected_tools[2:]] + + @pytest.mark.asyncio + async def test_editing_a_member_lands_on_that_member_without_the_namespace_prefix(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + + def redact_search(tools): + for tool in tools: + if tool["function"]["name"] == "mcp__confluence__confluence_search": + tool["function"]["description"] = "Confluence tools\n\nREDACTED" + return tools + + result = await OpenAIResponsesHandler().process_input_messages( + data, ToolEditingGuardrail(edit=redact_search, guardrail_name="test") + ) + + namespace = _tool_named(result["tools"], "mcp__confluence") + assert namespace["tools"][0] == expected_tools[1]["tools"][0] + assert namespace["tools"][1] == {**expected_tools[1]["tools"][1], "description": "REDACTED"} + assert {k: v for k, v in namespace.items() if k != "tools"} == { + k: v for k, v in expected_tools[1].items() if k != "tools" + } + + @pytest.mark.asyncio + async def test_dropping_every_member_drops_the_namespace(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + guardrail = ToolEditingGuardrail( + edit=lambda tools: [t for t in tools if not t["function"]["name"].startswith("mcp__confluence__")], + guardrail_name="test", + ) + + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert result["tools"] == [expected_tools[0], *expected_tools[2:]] + + @pytest.mark.asyncio + async def test_edited_top_level_function_is_rewritten_in_place(self): + data = _codex_request("hi") + expected_tools = copy.deepcopy(data["tools"]) + + def rename_weather(tools): + for tool in tools: + if tool["function"]["name"] == "get_weather": + tool["function"]["description"] = "Weather lookup (guarded)" + return tools + + result = await OpenAIResponsesHandler().process_input_messages( + data, ToolEditingGuardrail(edit=rename_weather, guardrail_name="test") + ) + + assert result["tools"][0] == {**expected_tools[0], "description": "Weather lookup (guarded)"} + assert result["tools"][1:] == expected_tools[1:] + + +class TestOpenAIResponsesHandlerMalformedTools: + @pytest.mark.asyncio + async def test_request_tools_that_are_not_a_list_never_reach_the_guardrail(self): + handler = OpenAIResponsesHandler() + seen: list[list[dict]] = [] + + def record(tools): + seen.append(tools) + return tools + + guardrail = ToolEditingGuardrail(edit=record, guardrail_name="test") + data = {"input": "hi", "tools": {"type": "function", "name": "get_weather"}} + + result = await handler.process_input_messages(data, guardrail) + + assert seen == [[]] + assert result["input"] == "hi" + + class TestBuildBlockSseChunks: """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE events""" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py new file mode 100644 index 00000000000..b80dd0b36aa --- /dev/null +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py @@ -0,0 +1,144 @@ +""" +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, +) + + +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_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_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_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/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/type-discipline-budget.json b/type-discipline-budget.json index 52cb9628252..970a44cd4fa 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22364 + "limit": 22340 }, "LIT002": { - "limit": 26777 + "limit": 26770 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1039 + "limit": 1038 }, "LIT007": { "limit": 0 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16507 + "limit": 16503 }, "LIT011": { - "limit": 5535 + "limit": 5534 }, "LIT012": { "limit": 4495 From dc12e4c2b4ad31b1eda1544ecd2e424aabc78151 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:57:01 -0700 Subject: [PATCH 17/31] fix(responses): match guardrail tools by ordinal in one pass Sort the chat-tool keys once and number duplicates with groupby instead of rescanning every preceding key per position, so the guardrail merge stays O(n log n) on client-supplied tool lists. Drop the comment that restated the unsupported-tool warning in the Responses-to-chat transformation. --- .../guardrail_translation/tool_merge.py | 8 ++++++-- .../transformation.py | 4 ---- ...st_openai_responses_guardrail_tool_merge.py | 18 ++++++++++++++++++ 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py index 3ae951d3f61..9fbf9f31a0f 100644 --- a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -1,5 +1,5 @@ from collections.abc import Iterable, Mapping, Sequence -from itertools import accumulate, chain +from itertools import accumulate, chain, groupby from types import MappingProxyType from typing import Final, TypeAlias @@ -47,7 +47,11 @@ def _chat_tool_key(tool: Tool) -> str: def _indexed_keys(tools: Sequence[Tool]) -> tuple[IndexedKey, ...]: keys: Final = tuple(_chat_tool_key(tool) for tool in tools) - return tuple((key, keys[:position].count(key)) for position, key in enumerate(keys)) + 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, ...]: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 9f91d7527cb..b2d1a69e0d8 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1903,10 +1903,6 @@ class LiteLLMCompletionResponsesConfig: 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"): - # 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.", 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 index b80dd0b36aa..4075c209606 100644 --- 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 @@ -56,6 +56,24 @@ def test_duplicate_function_names_are_matched_by_ordinal(): 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) From f4eca10f1d7f832f6300b588434bcd99003f4bb7 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 2 Sep 2026 20:07:57 +0000 Subject: [PATCH 18/31] ci: retrigger checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> From 0f6d983c7057faf13716639869d828d309bcba5b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:08:14 -0700 Subject: [PATCH 19/31] fix(router): skip Claude Code session binding without pre-routing strategies --- litellm/router.py | 2 ++ tests/test_litellm/test_router.py | 20 ++++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index fc4f2d227f5..b1038ca6002 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12636,6 +12636,8 @@ class Router: 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 diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 05dcb664ed5..ef25502a4f9 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8499,6 +8499,26 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() + @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 = {} + 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() From 4705dc6325fc8b2803e59ad19c85f9d07230a9a8 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 2 Sep 2026 20:34:54 +0000 Subject: [PATCH 20/31] fix: retry P3009 when the deadlocked ledger row was already rolled back Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_proxy_extras/utils.py | 46 +++++++---------- .../tests/test_setup_database_fail_fast.py | 51 +++++++++++++++++++ 2 files changed, 70 insertions(+), 27 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 1183757c827..9cc33e6ca56 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -277,22 +277,16 @@ class ProxyExtrasDBManager: pass @staticmethod - def _failed_migration_logs(migration_name: str) -> str: - """Logs recorded on the migration's failed _prisma_migrations row. - - P3009 stderr does not carry the original failure, so this is the only - way to tell a migration that lost a deadlock race against a concurrent - migrate deploy from one whose SQL is genuinely broken. Returns "" when - psycopg is missing, the DB is unreachable, or no failed row exists. - """ + 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 "" + return None try: import psycopg except ImportError: - return "" + return None cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url) ledger_table = psycopg.sql.SQL("{}.{}").format( @@ -314,7 +308,7 @@ class ProxyExtrasDBManager: (migration_name,), ).fetchone() except (psycopg.OperationalError, psycopg.DatabaseError): - return "" + return None return (row[0] or "") if row else "" @staticmethod @@ -825,22 +819,20 @@ class ProxyExtrasDBManager: f"Detail: {resolve_err}" ) from resolve_err continue - if migration_match and _MIGRATION_DEADLOCK_MARKER in ( - ProxyExtrasDBManager._failed_migration_logs( - migration_match.group(1) - ) - ): - logger.info( - "Migration %s lost a deadlock race 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 + 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}" 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 c4347a91dce..70e0f216338 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -350,6 +350,57 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_ 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) From f6eff1bde0f2464ce06c560b55feb7dab4192ab7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:55:16 -0700 Subject: [PATCH 21/31] fix(router): keep Claude Code session bindings across side calls and workers --- litellm/router.py | 9 ++- .../router_code_coverage.py | 1 + tests/test_litellm/test_router.py | 58 +++++++++++++++---- 3 files changed, 56 insertions(+), 12 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b1038ca6002..f4d912eb5be 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12630,6 +12630,12 @@ class Router: e, ) + async def _get_claude_code_session_router_binding(self, cache_key: str) -> object: + session_cache: Final = self._claude_code_session_router_cache + 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) + async def _resolve_claude_code_session_router( self, model: str, @@ -12646,7 +12652,7 @@ class Router: agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") if agent_id is not None: - bound_model: Final = await self._claude_code_session_router_cache.async_get_cache(key=cache_key) + 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 @@ -12664,7 +12670,6 @@ class Router: 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: - 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, diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 60b56b7fac6..0af29f069c6 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -86,6 +86,7 @@ ignored_function_names = [ "_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/test_litellm/test_router.py b/tests/test_litellm/test_router.py index ef25502a4f9..1e5767aa645 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8383,18 +8383,21 @@ class TestConsumedRequestTagsStamp: 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="cheap-model", + model=self.routed_model, messages=messages, routing_decision={ "router_model_name": "smart-router", "router_type": "complexity", - "routed_model": "cheap-model", + "routed_model": self.routed_model, "cause": "heuristic_scorer", }, ) @@ -8422,7 +8425,8 @@ class TestClaudeCodeSubagentSessionRouterBinding: num_retries=0, ) router.complexity_routers = { - "smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy())] + "smart-router": (TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy()),), + "premium-router": (TaggedPreRoutingStrategy(tags=(), strategy=cls._RewriteStrategy("expensive-model")),), } return router @@ -8467,7 +8471,7 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert subagent_kwargs["metadata"]["routing_decision"]["router_model_name"] == "smart-router" @pytest.mark.asyncio - async def test_main_direct_model_clears_the_session_router(self): + 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()) @@ -8478,33 +8482,67 @@ class TestClaudeCodeSubagentSessionRouterBinding: request_kwargs=self._request_kwargs(agent_id="agent-1234"), ) - assert response is None + assert response is not None + assert response.model == "cheap-model" @pytest.mark.asyncio - async def test_redis_cleanup_failure_does_not_reject_a_direct_model_request(self): + 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")) - - await router.async_pre_routing_hook(model="smart-router", request_kwargs=self._request_kwargs()) router._update_redis_cache(cache=redis_cache) response = await router.async_pre_routing_hook( model="expensive-model", - request_kwargs=self._request_kwargs(), + 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_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 = {} + router.complexity_routers.clear() redis_cache = MagicMock(spec=RedisCache) redis_cache.async_get_cache = AsyncMock(return_value=None) redis_cache.async_set_cache = AsyncMock() From 86c8b93bf748ca5a0d5a8e59d366a45f04c7ba34 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 14:14:41 -0700 Subject: [PATCH 22/31] fix(vector-store): embed Milvus and Azure AI Search queries through the request executor Milvus REST and Azure AI Search still embedded the query through the SDK, so a bare Router alias as litellm_embedding_model kept failing after the executor landed for Valkey. Both now share BaseQueryEmbeddingVectorStoreConfig, which embeds through the injected executor, drops the empty litellm_embedding_config requirement, and awaits aembedding on the async path. The Router executor falls back to the SDK for models the Router does not serve, so inline provider configs such as azure/text-embedding-3-large with their own credentials keep working through the proxy. Tests fake OpenAI and Milvus at the HTTP boundary with respx instead of patching litellm.embedding. --- .../azure_ai/vector_stores/transformation.py | 130 +++++----- .../base_llm/vector_store/transformation.py | 119 +++++++++- litellm/llms/custom_httpx/llm_http_handler.py | 46 ++-- .../milvus/vector_stores/transformation.py | 129 +++++----- .../test_router_embedding_integration.py | 183 ++++++++------ .../test_azure_ai_vector_store.py | 119 +++++++++- .../test_milvus_vector_store.py | 223 +++++++++++++++--- uv.lock | 22 +- 8 files changed, 700 insertions(+), 271 deletions(-) diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index 5e61d0a1dd9..044b8f5243c 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -1,10 +1,13 @@ +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, @@ -25,7 +28,7 @@ else: LiteLLMLoggingObj = Any -class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): +class AzureAIVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAzureLLM): """ Configuration for Azure AI Search Vector Store @@ -109,82 +112,71 @@ 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, - ) -> 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, + 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) + 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, + 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) + return self._search_request( + vector_store_id, + query_text, + query_vector, + vector_store_search_optional_params, + api_base, + litellm_logging_obj, + litellm_params, + ) - # Get embedding model from litellm_params (required) - embedding_model: Final = litellm_params.get("litellm_embedding_model") - if not embedding_model: - raise ValueError( - "embedding_model is required in litellm_params for Azure AI Search. " - "Example: litellm_params['embedding_model'] = 'azure/text-embedding-3-large'" - ) - - embedding_config: Final = litellm_params.get("litellm_embedding_config", {}) - if not embedding_config: - raise ValueError( - "embedding_config is required in litellm_params for Azure AI Search. " - "Example: litellm_params['embedding_config'] = {'api_base': 'https://krris-mh44uf7y-eastus2.cognitiveservices.azure.com/', 'api_key': 'os.environ/AZURE_API_KEY', 'api_version': '2025-09-01'}" - ) - - # Get vector field name (defaults to contentVector) + @staticmethod + def _search_request( + vector_store_id: str, + query_text: str, + query_vector: Sequence[float], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, object]]: vector_field: Final = litellm_params.get("azure_search_vector_field", "contentVector") - - # Get top_k (number of results to return) top_k: Final = vector_store_search_optional_params.get("top_k", 10) - - # Generate embedding for the query using litellm.embeddings - try: - embedding_response: Final = litellm.embedding( - model=embedding_model, - input=[query], - **embedding_config, - ) - query_vector: Final = embedding_response.data[0]["embedding"] - except Exception as e: - raise Exception(f"Failed to generate embedding for query: {e}") - - # Azure AI Search endpoint for search - index_name: Final = vector_store_id # vector_store_id is the index name - url: Final = f"{api_base}/indexes/{index_name}/docs/search?api-version=2024-07-01" - - # Build the request body for Azure AI Search with vector search - request_body: Final = { - "search": "*", # Get all documents (filtered by vector similarity) - "vectorQueries": [ - { - "vector": query_vector, - "fields": vector_field, - "kind": "vector", - "k": top_k, # Number of nearest neighbors to return - } - ], - "select": "id,content", # Fields to return (customize based on schema) + litellm_logging_obj.model_call_details["input"] = query_text + litellm_logging_obj.model_call_details["embedding_model"] = litellm_params.get("litellm_embedding_model") + litellm_logging_obj.model_call_details["top_k"] = top_k + return f"{api_base}/indexes/{vector_store_id}/docs/search?api-version=2024-07-01", { + "search": "*", + "vectorQueries": [{"vector": query_vector, "fields": vector_field, "kind": "vector", "k": top_k}], + "select": "id,content", "top": top_k, } - ######################################################### - # Update logging object with details of the request - ######################################################### - litellm_logging_obj.model_call_details["input"] = query - litellm_logging_obj.model_call_details["embedding_model"] = embedding_model - litellm_logging_obj.model_call_details["top_k"] = top_k - - return url, request_body - def transform_search_vector_store_response( self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj ) -> VectorStoreSearchResponse: diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index e9c925448a8..95863266bf7 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -3,9 +3,11 @@ from __future__ import annotations from abc import abstractmethod from collections.abc import Mapping, Sequence 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 @@ -65,7 +67,7 @@ class RouterVectorStoreEmbeddingExecutor: router: Router metadata: Mapping[str, object] - def _embedding_kwargs(self, configuration: Mapping[str, object]) -> dict[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 {}), @@ -76,18 +78,32 @@ class RouterVectorStoreEmbeddingExecutor: "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 embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + embedding_kwargs: Final = self._embedding_kwargs(configuration) + if not self._router_serves(model): + 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 - **self._embedding_kwargs(configuration), # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic + **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 not self._router_serves(model): + 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 - **self._embedding_kwargs(configuration), # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic + **embedding_kwargs, # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic ) @@ -221,6 +237,103 @@ 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, + 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, + 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, + 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 + + def embed_query( + self, + query_text: str, + litellm_params: Mapping[str, object], + embedding_executor: VectorStoreEmbeddingExecutor | None, + ) -> Sequence[float]: + model: Final = self.query_embedding_model(litellm_params) + configuration: Final = self.query_embedding_configuration(litellm_params) + executor: Final = ( + embedding_executor if embedding_executor is not None else LiteLLMVectorStoreEmbeddingExecutor() + ) + 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, + ) -> Sequence[float]: + model: Final = self.query_embedding_model(litellm_params) + configuration: Final = self.query_embedding_configuration(litellm_params) + executor: Final = ( + embedding_executor if embedding_executor is not None else LiteLLMVectorStoreEmbeddingExecutor() + ) + 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 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 118656b81a2..c0ef7456680 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -69,6 +69,7 @@ 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, ) @@ -9728,8 +9729,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - # Check if provider has async transform method - if hasattr(vector_store_provider_config, "atransform_search_vector_store_request"): + if isinstance(vector_store_provider_config, BaseQueryEmbeddingVectorStoreConfig): ( url, request_body, @@ -9741,12 +9741,13 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, + 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, @@ -9857,18 +9858,33 @@ 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, - ) + 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, + 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, + ) 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/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 34f0cd854c4..b0291c692d5 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -1,9 +1,12 @@ +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 ( @@ -36,7 +39,7 @@ MILVUS_OPTIONAL_PARAMS: Final = { } -class MilvusVectorStoreConfig(BaseVectorStoreConfig): +class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): """ Configuration for Milvus Vector Store @@ -117,77 +120,77 @@ 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, - ) -> 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, + 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) + 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, + 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) + 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/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py index d5e0c750d88..96c6ce3708e 100644 --- a/tests/router_unit_tests/test_router_embedding_integration.py +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -5,16 +5,57 @@ These tests simulate real-world scenarios where headers and configuration need to be properly propagated through the router to the LLM API. """ +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, ) -from litellm.types.utils import EmbeddingResponse + +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: @@ -70,48 +111,23 @@ class TestRouterEmbeddingIntegration: ) @pytest.mark.asyncio - async def test_vector_store_embedding_executors_cover_sdk_and_router_paths(self): - response = EmbeddingResponse(data=[{"embedding": [0.1], "index": 0, "object": "embedding"}]) + 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() - with ( - patch("litellm.embedding", return_value=response) as embedding, - patch("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 + 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"}) - 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") + 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"]) - mock_router = MagicMock() - mock_router.embedding.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"}, - ) - - alias_router = Router( - model_list=[ - { - "model_name": "team-alias", - "litellm_params": { - "model": "openai/text-embedding-3-small", - "api_key": "deployment-key", - }, - } - ] - ) - alias_executor = RouterVectorStoreEmbeddingExecutor( - router=alias_router, - metadata={"user_api_key_team_id": "team-a"}, - ) explicit_config = { "api_base": "https://embedding.example/v1", "api_key": "store-key", @@ -121,29 +137,66 @@ class TestRouterEmbeddingIntegration: }, "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"}, + ) - with ( - patch("litellm.embedding", return_value=response) as explicit_embedding, - patch("litellm.aembedding", new=AsyncMock(return_value=response)) as explicit_aembedding, - ): - assert alias_executor.embed("team-alias", "sync query", explicit_config) is response - assert await alias_executor.aembed("team-alias", "async query", explicit_config) is response + 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) - sync_kwargs = explicit_embedding.call_args.kwargs - assert sync_kwargs["model"] == "openai/text-embedding-3-small" - assert sync_kwargs["input"] == ["sync query"] - assert sync_kwargs["api_base"] == "https://embedding.example/v1" - assert sync_kwargs["api_key"] == "store-key" - assert sync_kwargs["metadata"]["configured"] is True - assert sync_kwargs["metadata"]["user_api_key_team_id"] == "team-a" + 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"]) - async_kwargs = explicit_aembedding.await_args.kwargs - assert async_kwargs["model"] == "openai/text-embedding-3-small" - assert async_kwargs["input"] == ["async query"] - assert async_kwargs["api_base"] == "https://embedding.example/v1" - assert async_kwargs["api_key"] == "store-key" - assert async_kwargs["metadata"]["configured"] is True - assert async_kwargs["metadata"]["user_api_key_team_id"] == "team-a" + @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"]) + + 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): """ @@ -251,9 +304,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: @@ -369,9 +420,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}"]) @@ -456,9 +505,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/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..ea3c1883e46 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,175 @@ 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) + + +@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/uv.lock b/uv.lock index 8d886044083..27be919eea1 100644 --- a/uv.lock +++ b/uv.lock @@ -9441,19 +9441,19 @@ wheels = [ [[package]] name = "tornado" -version = "6.5.8" +version = "6.5.7" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/10/d3/343e5bb989d6515b1646cf3d40135d73f3d5e45339bded401b56cdac24dd/tornado-6.5.8.tar.gz", hash = "sha256:9452e1b208a8bd771e2cb1f2ff564985b9b214bdebbe622793e1799e0a6bd23f", size = 520493, upload-time = "2026-08-07T02:12:42.971Z" } +sdist = { url = "https://files.pythonhosted.org/packages/64/24/95ec527ad67b76d59299e5465b3935d05e4294b7e0290a3924b7487df30b/tornado-6.5.7.tar.gz", hash = "sha256:66c513a76cda70d53907bc27cf1447557699c2e95aa48ba27a442ff61c3ddfc2", size = 519252, upload-time = "2026-06-08T17:34:51.232Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f2/d5/007086fd8df5489338e204f65adce33fd4f21a4999dbb2b9cff2f897b5f4/tornado-6.5.8-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:cc6aa787d7cfab7c3d35189dc7a56fbd2399a569624c730c6b55b3d6531d0403", size = 449487, upload-time = "2026-08-07T02:12:28.682Z" }, - { url = "https://files.pythonhosted.org/packages/70/c8/5a24a99495903f594f6a199dd7beead1cbc0a13e2cb9102727bcaaf2a997/tornado-6.5.8-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9715b5eb79735b2bcd454ce216a9275b7c0470e64ea1bf5742f78b2f72b26eeb", size = 447649, upload-time = "2026-08-07T02:12:30.306Z" }, - { url = "https://files.pythonhosted.org/packages/6e/de/f2e733f386b85962d1b1dc82cd63d169b5b4580062b35397eac9244a41fe/tornado-6.5.8-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:547d63f450d570c14fe0e8db2cfb14c9bbd1c2503b4a6612586267955aa47b58", size = 450707, upload-time = "2026-08-07T02:12:31.95Z" }, - { url = "https://files.pythonhosted.org/packages/0b/94/20efeee9a01c141e9ac47c397f81679dfda24b32768fc4fff24e76d36c2c/tornado-6.5.8-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7e2360a0ffbe145eca8af0b19cb7203d79b1a98dd4cccdd6b368f6f49c2e3808", size = 451677, upload-time = "2026-08-07T02:12:33.512Z" }, - { url = "https://files.pythonhosted.org/packages/42/ec/a96ccb8ccf0de2b7bc2c5fa1608a4803735018242e90c4882365a9fd418f/tornado-6.5.8-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:5d242290bdf7ab3151bc1065fdd75c0dcc21cbc7b49f22a4c56329c2d6566d22", size = 451510, upload-time = "2026-08-07T02:12:35.346Z" }, - { url = "https://files.pythonhosted.org/packages/29/b5/93185859245ad3f00e62175f29607346788b696369347f0146e0421286bb/tornado-6.5.8-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7b94ff0e128fe0542f3bd331fb44d06260fc4ac16881545159f34ef08aad4195", size = 450917, upload-time = "2026-08-07T02:12:36.963Z" }, - { url = "https://files.pythonhosted.org/packages/97/cf/fe33cf062834487d34d1559746a4a12521033c22645b6d74d4bca702e018/tornado-6.5.8-cp39-abi3-win32.whl", hash = "sha256:67832909c4779c64942380cb5f044a5c6163d00831472d80e25e115de9917836", size = 451952, upload-time = "2026-08-07T02:12:38.512Z" }, - { url = "https://files.pythonhosted.org/packages/cb/e1/468ad54333e92ccb62627e62cb88e5fc14a2171daa67ed47b1b8542d5b86/tornado-6.5.8-cp39-abi3-win_amd64.whl", hash = "sha256:11881db6b7c168494be2c2d12e65931451bdf7ee718535418ae1d8855dd5a0ee", size = 452391, upload-time = "2026-08-07T02:12:39.971Z" }, - { url = "https://files.pythonhosted.org/packages/ad/3e/cd5e4f06e34cde33b8ef66cf36aa2b5ad46354cc1af7d2136bbe365fee1d/tornado-6.5.8-cp39-abi3-win_arm64.whl", hash = "sha256:68a7468c7e289f8514d7d664101753903217eff1bb6822c6b5994a0b5f5bcb26", size = 451411, upload-time = "2026-08-07T02:12:41.469Z" }, + { url = "https://files.pythonhosted.org/packages/02/dc/c7043cab6fed8ae159fc1923ce829ada35c4dbd797d408a43858ffaf9639/tornado-6.5.7-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:148b2eb15c2c765a50796172c1e499649b35f30d2e3c3d3e15913cfa56bfb163", size = 448543, upload-time = "2026-06-08T17:34:38.052Z" }, + { url = "https://files.pythonhosted.org/packages/92/4f/090b1431e5a43df696feceffc268c5383cc079ecb5f08ce58f917109aafe/tornado-6.5.7-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9da38de27f1da3b78a966f0dae12b5a1ea9afe72ca805d84ff06508272ddf100", size = 446707, upload-time = "2026-06-08T17:34:39.594Z" }, + { url = "https://files.pythonhosted.org/packages/37/d8/ef374952fd5da67d4463122c2b8e5a96536ec10b4b339254c6dcde81d01c/tornado-6.5.7-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8d759e71906ee783f8867b93bf26a265743da4c1e2f4a018464c1ba019862972", size = 449774, upload-time = "2026-06-08T17:34:41.204Z" }, + { url = "https://files.pythonhosted.org/packages/35/37/d434c73f4c6e014b745b9b37085f34f40c022f007efff3d7fe65991899f3/tornado-6.5.7-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a46347a18f23fb92b396beebe0fb78f61dda0cc302445202c16203d8a18848b", size = 450745, upload-time = "2026-06-08T17:34:42.531Z" }, + { url = "https://files.pythonhosted.org/packages/b6/2b/56b9aff361d7f1ab728a805ec7d7ea835f8807afa9f5cc690ea0e630efb9/tornado-6.5.7-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:7778b30bef919231265e91c69963ce0f49a1e9c07ac900bbe75b19ce2575ba92", size = 450578, upload-time = "2026-06-08T17:34:43.787Z" }, + { url = "https://files.pythonhosted.org/packages/02/30/a7444fb23aa76860a14198fab96ac79f1866b0a6e19e26c4381b0938e50f/tornado-6.5.7-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e726f0c75da7726eec023aa62751ff8878bd2737e34fbdd33b1ae5897d2200f5", size = 449985, upload-time = "2026-06-08T17:34:45.326Z" }, + { url = "https://files.pythonhosted.org/packages/5c/42/5f0e56c01e8d9d36f4e23f367b85ae6cae0c1ecddd5e6977d8388ad27488/tornado-6.5.7-cp39-abi3-win32.whl", hash = "sha256:f8de3bf12d3efdd0cbe7c8887868198f8a91415e3f29fcf258d9b8eb7b1d9ae4", size = 451047, upload-time = "2026-06-08T17:34:46.784Z" }, + { url = "https://files.pythonhosted.org/packages/c9/a4/b393076ffb21b469eec5b328a0534cf03a3b90bfc6b1f09507cdd075d938/tornado-6.5.7-cp39-abi3-win_amd64.whl", hash = "sha256:de942f843533a039ef9fa3d9c88c7cd8a7c94553fb5ad0154270989b3d99a2c4", size = 451485, upload-time = "2026-06-08T17:34:48.248Z" }, + { url = "https://files.pythonhosted.org/packages/71/2e/7b1c769803121b809112cf9a00681c472eae1d80e32d7ec0e0bd61d0d0e1/tornado-6.5.7-cp39-abi3-win_arm64.whl", hash = "sha256:ff934fce95643af5f11efdae618eaa73d469dc588641e5c8d19295a0c65c4796", size = 450506, upload-time = "2026-06-08T17:34:49.702Z" }, ] [[package]] From 5561b8438c34054b3fd1475e011ed5cfd6f6ec97 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 14:18:16 -0700 Subject: [PATCH 23/31] fix(responses): keep a namespace's non-function members when every function member is dropped --- .../responses/guardrail_translation/tool_merge.py | 4 ++-- .../test_openai_responses_guardrail_tool_merge.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py index 9fbf9f31a0f..9a34946838a 100644 --- a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -147,11 +147,11 @@ def _merged_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,) - if not guardrailed_group: - return () 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) ) 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 index 4075c209606..cbcddb98235 100644 --- 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 @@ -132,6 +132,19 @@ def test_namespace_keeps_a_non_function_member_when_a_function_member_is_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) From 7cbc74399d683b1113a4236768ac8414e7e9f07d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:24:43 -0700 Subject: [PATCH 24/31] fix(responses): accept pydantic tool objects returned by guardrails --- .../guardrail_translation/tool_merge.py | 5 ++-- ...t_openai_responses_guardrail_tool_merge.py | 23 +++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py index 9a34946838a..b596adfad6f 100644 --- a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -3,7 +3,7 @@ from itertools import accumulate, chain, groupby from types import MappingProxyType from typing import Final, TypeAlias -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.responses.litellm_completion_transformation.transformation import ( @@ -19,8 +19,9 @@ _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(value) + return _TOOL_ADAPTER.validate_python(candidate) except ValidationError: return 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 index cbcddb98235..9c236d81f51 100644 --- 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 @@ -9,6 +9,7 @@ from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) +from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GuardrailToolParam def _groups(tools): @@ -166,6 +167,28 @@ def test_guardrail_output_is_read_once(): 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) From 979c3cb1aa0dc579c38d10f44839b93ecc0d9823 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:29:10 -0700 Subject: [PATCH 25/31] fix(proxy-extras): retry when the advisory lock wait times out behind a concurrent migrate deploy --- .../litellm_proxy_extras/utils.py | 9 ++++++ .../tests/test_setup_database_fail_fast.py | 28 +++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 9cc33e6ca56..d22484bc0e8 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -905,6 +905,15 @@ class ProxyExtrasDBManager: 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}" 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 70e0f216338..406f07eb792 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -432,3 +432,31 @@ def test_v2_bare_deadlock_stderr_retries(monkeypatch, tmp_path): 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) From ed0d3f94422e32c498781fe63722219989987285 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 2 Sep 2026 15:21:37 -0700 Subject: [PATCH 26/31] fix(router): fail open on session binding reads --- litellm/router.py | 13 ++++++-- tests/test_litellm/test_router.py | 50 +++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index f4d912eb5be..2d8b4a43759 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12632,9 +12632,16 @@ class Router: async def _get_claude_code_session_router_binding(self, cache_key: str) -> object: session_cache: Final = self._claude_code_session_router_cache - 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) + 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, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 1e5767aa645..16c4d997ad7 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8504,6 +8504,56 @@ class TestClaudeCodeSubagentSessionRouterBinding: 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 1306a4505a03bdced3ccdf5ac6b26a719867bbbb Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 2 Sep 2026 16:01:15 -0700 Subject: [PATCH 27/31] fix(router): preserve bound router fallbacks for subagents --- .../router_utils/fallback_event_handlers.py | 12 +++++----- .../test_fallback_event_handlers.py | 22 +++++++++++++++++++ tests/test_litellm/test_router.py | 21 ++++++++++++++++++ 3 files changed, 50 insertions(+), 5 deletions(-) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 0167721f9fe..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)) 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 9e51a60364b..a4965c49f07 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -1220,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_router.py b/tests/test_litellm/test_router.py index 16c4d997ad7..0a433d0b0a1 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8679,6 +8679,27 @@ class TestClaudeCodeSubagentSessionRouterBinding: 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_session_router_binding_is_scoped_to_the_authenticated_key(self): router = self._router() From ddc25823740e8ed36a8d5ccddb93c4dee9ff95bc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:15:20 -0700 Subject: [PATCH 28/31] fix(vector-store): build the Router executor from the router kwarg and gate the SDK fallback on explicit config --- .../base_llm/vector_store/transformation.py | 7 +++-- litellm/vector_stores/main.py | 16 +++++------- .../test_router_embedding_integration.py | 19 ++++++++++++++ .../test_vector_store_endpoints.py | 2 +- .../test_milvus_vector_store.py | 26 +++++++++++++++++++ 5 files changed, 58 insertions(+), 12 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index e7c549b7358..3547bf81448 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -86,9 +86,12 @@ class RouterVectorStoreEmbeddingExecutor: ) 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 not self._router_serves(model): + 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, @@ -98,7 +101,7 @@ class RouterVectorStoreEmbeddingExecutor: async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: embedding_kwargs: Final = self._embedding_kwargs(configuration) - if not self._router_serves(model): + 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, diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index ade4d19815b..d530a1c9d52 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -16,7 +16,7 @@ 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 ( - LiteLLMVectorStoreEmbeddingExecutor, + BaseQueryEmbeddingVectorStoreConfig, VectorStoreEmbeddingExecutor, ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler @@ -42,12 +42,10 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# -def _direct_vector_store_embedding_executor(value: object) -> VectorStoreEmbeddingExecutor: - if value is None: - return LiteLLMVectorStoreEmbeddingExecutor() - if isinstance(value, VectorStoreEmbeddingExecutor): - return value - raise TypeError("Invalid direct vector store embedding executor") +def _direct_vector_store_embedding_executor(value: object, router: "Router | None") -> 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) def mock_vector_store_search_response( @@ -302,7 +300,7 @@ async def asearch( Async: Search a vector store for relevant chunks based on a query and file attributes filter. """ embedding_executor: Final = _direct_vector_store_embedding_executor( - kwargs.pop("_direct_vector_store_embedding_executor", None) + kwargs.pop("_direct_vector_store_embedding_executor", None), router ) local_vars: Final = { # mutable-ok: exception logging requires a sanitized mutable snapshot key: value for key, value in locals().items() if key != "embedding_executor" @@ -388,7 +386,7 @@ def search( VectorStoreSearchResponse containing the search results. """ embedding_executor: Final = _direct_vector_store_embedding_executor( - kwargs.pop("_direct_vector_store_embedding_executor", None) + kwargs.pop("_direct_vector_store_embedding_executor", None), router ) local_vars: Final = { # mutable-ok: exception logging requires a sanitized mutable snapshot key: value for key, value in locals().items() if key != "embedding_executor" diff --git a/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py index 96c6ce3708e..2cc9914c9b3 100644 --- a/tests/router_unit_tests/test_router_embedding_integration.py +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -186,6 +186,25 @@ class TestRouterEmbeddingIntegration: 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 ): 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 854ce351f2d..a956aef51d6 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 @@ -52,7 +52,7 @@ def _serialize_litellm_params(litellm_params): 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()) + _direct_vector_store_embedding_executor(object(), None) def test_router_vector_store_search_injects_executor_and_request_metadata(): diff --git a/tests/vector_store_tests/test_milvus_vector_store.py b/tests/vector_store_tests/test_milvus_vector_store.py index ea3c1883e46..3b3bd444b20 100644 --- a/tests/vector_store_tests/test_milvus_vector_store.py +++ b/tests/vector_store_tests/test_milvus_vector_store.py @@ -575,6 +575,32 @@ async def test_router_search_resolves_bare_embedding_alias_async( _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) + + @pytest.mark.asyncio async def test_transform_uses_injected_executor_without_embedding_config(respx_mock: respx.MockRouter): executor = RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE) From f5721ff6f063cd6cfed4285a4bee80e6cf1fc81b Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 2 Sep 2026 16:17:35 -0700 Subject: [PATCH 29/31] test(router): cover Claude session multi-hop fallbacks --- tests/test_litellm/test_router.py | 59 +++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 0a433d0b0a1..db73f0da3c2 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8700,6 +8700,65 @@ class TestClaudeCodeSubagentSessionRouterBinding: 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() From 1df402b80d2d3f4fed8d504a5a73946162c133b6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:53:05 -0700 Subject: [PATCH 30/31] fix(vector-store): carry request metadata into the Router executor built from the router kwarg --- .../base_llm/vector_store/transformation.py | 16 ++++++- litellm/router.py | 11 +---- litellm/vector_stores/main.py | 13 ++++-- .../test_vector_store_endpoints.py | 2 +- .../test_milvus_vector_store.py | 46 +++++++++++++++++++ 5 files changed, 73 insertions(+), 15 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 3547bf81448..9624a721870 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -62,6 +62,19 @@ class LiteLLMVectorStoreEmbeddingExecutor: ) +_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 @@ -312,11 +325,12 @@ class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): 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=MappingProxyType({})) + return RouterVectorStoreEmbeddingExecutor(router=router, metadata=request_metadata) return LiteLLMVectorStoreEmbeddingExecutor() def embed_query( diff --git a/litellm/router.py b/litellm/router.py index 991f1bc2828..8b510e748f5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -87,6 +87,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( ) 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 @@ -6697,15 +6698,7 @@ class Router: @staticmethod 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 cast( # cast-ok: isinstance validates the runtime dict boundary - "dict[str, object]", litellm_metadata - ) - metadata: Final = kwargs.get("metadata") - if isinstance(metadata, dict): - return cast("dict[str, object]", metadata) # cast-ok: isinstance validates the runtime dict boundary - return MappingProxyType({}) + return vector_store_request_metadata(kwargs) async def _init_vector_store_api_endpoints( self, diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index d530a1c9d52..636bdd4b52e 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging 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 @@ -42,10 +43,14 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# -def _direct_vector_store_embedding_executor(value: object, router: "Router | None") -> VectorStoreEmbeddingExecutor: +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) + return BaseQueryEmbeddingVectorStoreConfig.query_embedding_executor( + value, router, vector_store_request_metadata(request_kwargs) + ) def mock_vector_store_search_response( @@ -300,7 +305,7 @@ async def asearch( Async: Search a vector store for relevant chunks based on a query and file attributes filter. """ embedding_executor: Final = _direct_vector_store_embedding_executor( - kwargs.pop("_direct_vector_store_embedding_executor", None), router + 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" @@ -386,7 +391,7 @@ def search( VectorStoreSearchResponse containing the search results. """ embedding_executor: Final = _direct_vector_store_embedding_executor( - kwargs.pop("_direct_vector_store_embedding_executor", None), router + 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" 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 a956aef51d6..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 @@ -52,7 +52,7 @@ def _serialize_litellm_params(litellm_params): 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) + _direct_vector_store_embedding_executor(object(), None, {}) def test_router_vector_store_search_injects_executor_and_request_metadata(): diff --git a/tests/vector_store_tests/test_milvus_vector_store.py b/tests/vector_store_tests/test_milvus_vector_store.py index 3b3bd444b20..2ba9168b49f 100644 --- a/tests/vector_store_tests/test_milvus_vector_store.py +++ b/tests/vector_store_tests/test_milvus_vector_store.py @@ -601,6 +601,52 @@ async def test_sdk_search_with_router_kwarg_resolves_bare_embedding_alias_async( _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) From b0fe71010b89425501587a3e7e66ede9a4e6e775 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:56:35 -0700 Subject: [PATCH 31/31] fix(ollama_chat): stamp finish_reason tool_calls when tool calls streamed before the done chunk (#39010) * fix(ollama_chat): stamp finish_reason tool_calls when tool calls streamed before the done chunk Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ollama_chat): trim finish_reason override comment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mubashir1osmani --- litellm/llms/ollama/chat/transformation.py | 7 +- .../test_streaming_iterator_stop_reason.py | 79 +++++++++++++++++++ .../ollama/test_ollama_chat_transformation.py | 40 ++++++++++ 3 files changed, 124 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_stop_reason.py 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/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/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."""