mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(vector-stores): preserve config registry entries
This commit is contained in:
parent
16b78df82e
commit
57b08e8f7a
5 changed files with 96 additions and 10 deletions
|
|
@ -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": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
4
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
4
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue