fix(vector-stores): preserve config registry entries

This commit is contained in:
Yujong Lee 2026-09-01 14:36:07 -07:00
parent 16b78df82e
commit 57b08e8f7a
5 changed files with 96 additions and 10 deletions

View file

@ -38535,7 +38535,7 @@
"paths": {
"/v1/vector_store/list": {
"get": {
"description": "List all available vector stores with optional filtering and pagination.\nCombines both in-memory vector stores and those stored in the database.\nDatabase is the source of truth - deleted stores are removed from memory, updated stores sync to memory.\n\nParameters:\n- page: int - Page number for pagination (default: 1)\n- page_size: int - Number of items per page (default: 100)",
"description": "List all available vector stores with optional filtering and pagination.\nCombines both in-memory vector stores and those stored in the database.\nConfig entries remain authoritative; database-backed entries sync from the database.\n\nParameters:\n- page: int - Page number for pagination (default: 1)\n- page_size: int - Number of items per page (default: 100)",
"operationId": "list_vector_stores_v1_vector_store_list_get",
"parameters": [
{
@ -38686,7 +38686,7 @@
},
"/vector_store/list": {
"get": {
"description": "List all available vector stores with optional filtering and pagination.\nCombines both in-memory vector stores and those stored in the database.\nDatabase is the source of truth - deleted stores are removed from memory, updated stores sync to memory.\n\nParameters:\n- page: int - Page number for pagination (default: 1)\n- page_size: int - Number of items per page (default: 100)",
"description": "List all available vector stores with optional filtering and pagination.\nCombines both in-memory vector stores and those stored in the database.\nConfig entries remain authoritative; database-backed entries sync from the database.\n\nParameters:\n- page: int - Page number for pagination (default: 1)\n- page_size: int - Number of items per page (default: 100)",
"operationId": "list_vector_stores_vector_store_list_get",
"parameters": [
{

View file

@ -369,7 +369,7 @@ async def list_vector_stores(
"""
List all available vector stores with optional filtering and pagination.
Combines both in-memory vector stores and those stored in the database.
Database is the source of truth - deleted stores are removed from memory, updated stores sync to memory.
Config entries remain authoritative; database-backed entries sync from the database.
Parameters:
- page: int - Page number for pagination (default: 1)
@ -396,6 +396,7 @@ async def list_vector_stores(
# Process in-memory vector stores
if litellm.vector_store_registry is not None:
in_memory_vector_stores: Final = copy.deepcopy(litellm.vector_store_registry.vector_stores)
config_vector_store_ids: Final = litellm.vector_store_registry.config_vector_store_ids
vector_stores_to_delete_from_memory: Final[list[str]] = []
@ -404,8 +405,9 @@ async def list_vector_stores(
if not vector_store_id:
continue
# If vector store is in memory but NOT in database, it was deleted
if vector_store_id not in db_vector_store_ids:
if vector_store_id in config_vector_store_ids:
vector_store_map[vector_store_id] = vector_store
elif vector_store_id not in db_vector_store_ids:
verbose_proxy_logger.info(
"Vector store %s exists in memory but not in database - marking for deletion from cache",
vector_store_id,
@ -424,7 +426,7 @@ async def list_vector_stores(
# 2. Update in-memory registry with database versions (for updates)
for vector_store in vector_stores_from_db:
vector_store_id = vector_store.get("vector_store_id", None)
if vector_store_id:
if vector_store_id and vector_store_id not in config_vector_store_ids:
litellm.vector_store_registry.update_vector_store_in_registry(
vector_store_id=vector_store_id, updated_data=vector_store
)

View file

@ -105,6 +105,7 @@ class VectorStoreRegistry:
def __init__(self, vector_stores: list[LiteLLM_ManagedVectorStore] = []):
self.vector_stores: list[LiteLLM_ManagedVectorStore] = vector_stores
self.vector_store_ids_to_vector_store_map: dict[str, LiteLLM_ManagedVectorStore] = {}
self.config_vector_store_ids: frozenset[str] = frozenset()
def _extract_tool_params(self, tool: dict) -> VectorStoreToolParams:
"""
@ -342,7 +343,11 @@ class VectorStoreRegistry:
# Verify vector store still exists in database (if we have DB access)
# This ensures deleted vector stores are removed from cache
if vector_store is not None and prisma_client is not None:
if (
vector_store is not None
and prisma_client is not None
and vector_store_id not in self.config_vector_store_ids
):
try:
# Check if it still exists in database
db_vector_store = await ManagedVectorStoresRepository(prisma_client).table.find_unique(
@ -409,7 +414,7 @@ class VectorStoreRegistry:
vector_store_litellm_params: dict[str, Any] = litellm_vector_store_config.get("litellm_params") or {}
vector_store_id = vector_store_litellm_params.get("vector_store_id")
if vector_store_id is None:
if not isinstance(vector_store_id, str) or not vector_store_id:
raise ValueError(
f"vector_store_id is required for initializing vector store, got vector_store_id={vector_store_id}"
)
@ -433,6 +438,7 @@ class VectorStoreRegistry:
updated_at=datetime.now(timezone.utc),
)
self.vector_stores.append(litellm_managed_vector_store)
self.config_vector_store_ids = self.config_vector_store_ids.union((vector_store_id,))
verbose_logger.debug(
"all loaded vector stores = %s",

View file

@ -21,6 +21,7 @@ from litellm.proxy.vector_store_endpoints.endpoints import (
from litellm.proxy.vector_store_endpoints.management_endpoints import (
_check_vector_store_access,
create_vector_store_in_db,
list_vector_stores,
new_vector_store,
)
from litellm.proxy.vector_store_endpoints.utils import (
@ -799,6 +800,83 @@ async def test_config_loaded_milvus_grpc_connection_is_trusted():
assert result[MILVUS_ADMIN_CONFIGURED_CONNECTION] is True
@pytest.mark.asyncio
async def test_list_vector_stores_preserves_config_source_and_evicts_stale_database_cache():
registry = VectorStoreRegistry()
registry.load_vector_stores_from_config(
[
{
"vector_store_name": "configured",
"litellm_params": {
"vector_store_id": "configured",
"custom_llm_provider": "milvus",
"milvus_transport": "grpc",
"api_base": "https://configured-milvus:19530",
"litellm_embedding_model": "team-embedding-alias",
},
}
]
)
registry.add_vector_store_to_registry(
LiteLLM_ManagedVectorStore(
vector_store_id="stale-database-cache",
custom_llm_provider="openai",
)
)
with (
patch.object( # test-quality-ok: installs the isolated registry used by the endpoint
litellm, "vector_store_registry", registry
),
patch( # test-quality-ok: models a proxy without persisted vector stores
"litellm.proxy.proxy_server.prisma_client", None
),
patch( # test-quality-ok: isolates registry synchronization from feature entitlement
"litellm.proxy.vector_store_endpoints.management_endpoints.check_feature_access_for_user",
new=AsyncMock(),
),
patch( # test-quality-ok: isolates registry synchronization from per-user access policy
"litellm.proxy.vector_store_endpoints.management_endpoints._check_vector_store_access",
new=AsyncMock(return_value=True),
),
):
response = await list_vector_stores(
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert registry.config_vector_store_ids == frozenset(("configured",))
assert [store["vector_store_id"] for store in response["data"]] == ["configured"]
assert registry.get_litellm_managed_vector_store_from_registry("configured") is not None
assert registry.get_litellm_managed_vector_store_from_registry("stale-database-cache") is None
@pytest.mark.asyncio
async def test_db_fallback_does_not_evict_config_source():
registry = VectorStoreRegistry()
registry.load_vector_stores_from_config(
[
{
"vector_store_name": "configured",
"litellm_params": {
"vector_store_id": "configured",
"custom_llm_provider": "openai",
},
}
]
)
prisma_client = MagicMock()
prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=None)
result = await registry.pop_vector_stores_to_run_with_db_fallback(
non_default_params={"vector_store_ids": ["configured"]},
prisma_client=prisma_client,
)
assert [store["vector_store_id"] for store in result] == ["configured"]
assert registry.get_litellm_managed_vector_store_from_registry("configured") is not None
prisma_client.db.litellm_managedvectorstorestable.find_unique.assert_not_awaited()
def test_admin_persistence_strips_forged_marker_and_adds_server_marker():
params = prepare_milvus_connection_for_persistence(
custom_llm_provider="milvus/probe",

View file

@ -20285,7 +20285,7 @@ export interface paths {
* List Vector Stores
* @description List all available vector stores with optional filtering and pagination.
* Combines both in-memory vector stores and those stored in the database.
* Database is the source of truth - deleted stores are removed from memory, updated stores sync to memory.
* Config entries remain authoritative; database-backed entries sync from the database.
*
* Parameters:
* - page: int - Page number for pagination (default: 1)
@ -21586,7 +21586,7 @@ export interface paths {
* List Vector Stores
* @description List all available vector stores with optional filtering and pagination.
* Combines both in-memory vector stores and those stored in the database.
* Database is the source of truth - deleted stores are removed from memory, updated stores sync to memory.
* Config entries remain authoritative; database-backed entries sync from the database.
*
* Parameters:
* - page: int - Page number for pagination (default: 1)