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