mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(vector_stores): surface config-declared vector stores in
/vector_store/list (#25947) Vector stores declared under `vector_store_registry:` in proxy_config.yaml are loaded into `litellm.vector_store_registry` in memory only — they are never written to the `litellm_managedvectorstorestable` DB table. But `list_vector_stores` treats the DB as the source of truth and runs a reconcile step that deletes any in-memory entry missing from the DB, on the assumption that a disappeared DB row means the entry was deleted at runtime. That assumption doesn't hold for config-declared entries: they're *supposed* to live only in memory, so the reconcile step silently pruned them on every /vector_store/list call and the UI at `/ui/?page=vector-stores` always appeared empty, as reported in #25947. Fix: - Mark entries that came from `load_vector_stores_from_config` with `from_litellm_config=True` (new optional field on `LiteLLM_ManagedVectorStore`). - Teach `list_vector_stores` to skip the "missing from DB => delete" branch for those entries; they are included in the response and left untouched in the in-memory registry. - Stale *runtime* entries (API-created rows that the DB no longer knows about) are still pruned — that path is unchanged. Test coverage: - `test_load_vector_stores_from_config_marks_from_litellm_config` - `test_add_vector_store_to_registry_does_not_mark_from_litellm_config` - `test_list_vector_stores_returns_config_loaded_even_when_db_empty` (direct repro of #25947) - `test_list_vector_stores_still_prunes_stale_runtime_entries` (regression guard — the old behaviour is correct for non-config entries and must not regress) - `test_list_vector_stores_mixed_config_and_db` (config and DB entries coexist without masking each other) `pytest tests/test_litellm/vector_stores/test_vector_store_registry.py tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py` — 53 passed (5 new + 48 pre-existing). Fixes #25947
This commit is contained in:
sakenuGOD
parent
b8f7d61400
commit
b2c5b3267c
5 changed files with 252 additions and 0 deletions
|
|
@ -537,6 +537,16 @@ async def list_vector_stores(
|
|||
if not vector_store_id:
|
||||
continue
|
||||
|
||||
# Vector stores declared in proxy_config.yaml live only in
|
||||
# memory — they are never persisted to the DB table, so the
|
||||
# "in memory but not in DB => deleted" heuristic does not
|
||||
# apply to them. Surface them in the response and leave the
|
||||
# in-memory registry alone.
|
||||
if vector_store.get("from_litellm_config") is True:
|
||||
if vector_store_id not in vector_store_map:
|
||||
vector_store_map[vector_store_id] = vector_store
|
||||
continue
|
||||
|
||||
# If vector store is in memory but NOT in database, it was deleted
|
||||
if vector_store_id not in db_vector_store_ids:
|
||||
verbose_proxy_logger.info(
|
||||
|
|
|
|||
|
|
@ -43,6 +43,14 @@ class LiteLLM_ManagedVectorStore(TypedDict, total=False):
|
|||
team_id: Optional[str]
|
||||
user_id: Optional[str]
|
||||
|
||||
# provenance
|
||||
# True when this entry came from proxy_config.yaml's vector_store_registry
|
||||
# rather than from the database or a runtime API call. Consumers that
|
||||
# reconcile in-memory state against the DB should treat config-loaded
|
||||
# entries as valid even when absent from the DB table, otherwise they get
|
||||
# silently pruned on every list call and never appear to the user.
|
||||
from_litellm_config: Optional[bool]
|
||||
|
||||
|
||||
class LiteLLM_ManagedVectorStoreListResponse(TypedDict, total=False):
|
||||
"""Response format for listing vector stores"""
|
||||
|
|
|
|||
|
|
@ -470,6 +470,7 @@ class VectorStoreRegistry:
|
|||
),
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
from_litellm_config=True,
|
||||
)
|
||||
self.vector_stores.append(litellm_managed_vector_store)
|
||||
|
||||
|
|
|
|||
|
|
@ -1837,3 +1837,187 @@ async def test_create_vector_store_in_db_raises_when_no_db():
|
|||
|
||||
assert exc_info.value.status_code == 500
|
||||
assert "database not connected" in exc_info.value.detail.lower()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Regression coverage for issue #25947 — vector stores declared only in
|
||||
# proxy_config.yaml must survive the list endpoint's DB-vs-memory reconcile
|
||||
# step (they are never written to the DB, so the old
|
||||
# "in memory but not in DB => delete" rule silently erased them every time).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
from litellm.proxy.vector_store_endpoints.management_endpoints import (
|
||||
list_vector_stores,
|
||||
)
|
||||
from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
|
||||
|
||||
|
||||
class _TestVectorStoreListConfigLoaded:
|
||||
# Prefixed with underscore so pytest does not collect the helper — tests
|
||||
# live as module-level functions below.
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _isolated_registry():
|
||||
"""Swap litellm.vector_store_registry for the duration of the test."""
|
||||
original = litellm.vector_store_registry
|
||||
litellm.vector_store_registry = VectorStoreRegistry(vector_stores=[])
|
||||
try:
|
||||
yield litellm.vector_store_registry
|
||||
finally:
|
||||
litellm.vector_store_registry = original
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _permissive_user():
|
||||
"""A user key that passes every feature-access / team check."""
|
||||
return UserAPIKeyAuth(
|
||||
api_key="sk-test",
|
||||
user_id="u",
|
||||
team_id=None,
|
||||
metadata={},
|
||||
user_role="proxy_admin",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_vector_stores_returns_config_loaded_even_when_db_empty(
|
||||
_isolated_registry, _permissive_user
|
||||
):
|
||||
"""
|
||||
Reproduces #25947: a vector store declared in proxy_config.yaml sits
|
||||
only in the in-memory registry (never in the DB table), yet
|
||||
/vector_store/list must still surface it to the UI.
|
||||
"""
|
||||
_isolated_registry.load_vector_stores_from_config(
|
||||
[
|
||||
{
|
||||
"vector_store_name": "pgvector-config",
|
||||
"litellm_params": {
|
||||
"custom_llm_provider": "pg_vector",
|
||||
"vector_store_id": "vs-from-config",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
assert _isolated_registry.vector_stores[0]["from_litellm_config"] is True
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.vector_store_endpoints.management_endpoints"
|
||||
".check_feature_access_for_user",
|
||||
new=AsyncMock(return_value=None),
|
||||
), patch(
|
||||
"litellm.proxy.vector_store_endpoints.management_endpoints"
|
||||
".VectorStoreRegistry._get_vector_stores_from_db",
|
||||
new=AsyncMock(return_value=[]),
|
||||
), patch(
|
||||
"litellm.proxy.vector_store_endpoints.management_endpoints"
|
||||
"._check_vector_store_access",
|
||||
return_value=True,
|
||||
):
|
||||
response = await list_vector_stores(user_api_key_dict=_permissive_user)
|
||||
|
||||
ids_returned = [vs.get("vector_store_id") for vs in response["data"]]
|
||||
assert "vs-from-config" in ids_returned, (
|
||||
"config-loaded vector store was missing from the listing response"
|
||||
)
|
||||
|
||||
# And crucially it must still be in the in-memory registry —
|
||||
# before the fix the reconcile step would have evicted it.
|
||||
in_memory_ids = [
|
||||
vs.get("vector_store_id") for vs in _isolated_registry.vector_stores
|
||||
]
|
||||
assert "vs-from-config" in in_memory_ids
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_vector_stores_still_prunes_stale_runtime_entries(
|
||||
_isolated_registry, _permissive_user
|
||||
):
|
||||
"""
|
||||
Guardrail for the fix: a store that was NOT config-loaded and is no
|
||||
longer in the DB must still be evicted from the in-memory cache —
|
||||
the old behaviour was correct for that case and must not regress.
|
||||
"""
|
||||
_isolated_registry.add_vector_store_to_registry(
|
||||
LiteLLM_ManagedVectorStore(
|
||||
vector_store_id="vs-stale-runtime",
|
||||
custom_llm_provider="openai",
|
||||
vector_store_name="stale",
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.vector_store_endpoints.management_endpoints"
|
||||
".check_feature_access_for_user",
|
||||
new=AsyncMock(return_value=None),
|
||||
), patch(
|
||||
"litellm.proxy.vector_store_endpoints.management_endpoints"
|
||||
".VectorStoreRegistry._get_vector_stores_from_db",
|
||||
new=AsyncMock(return_value=[]),
|
||||
), patch(
|
||||
"litellm.proxy.vector_store_endpoints.management_endpoints"
|
||||
"._check_vector_store_access",
|
||||
return_value=True,
|
||||
):
|
||||
response = await list_vector_stores(user_api_key_dict=_permissive_user)
|
||||
|
||||
ids_returned = [vs.get("vector_store_id") for vs in response["data"]]
|
||||
assert "vs-stale-runtime" not in ids_returned
|
||||
|
||||
in_memory_ids = [
|
||||
vs.get("vector_store_id") for vs in _isolated_registry.vector_stores
|
||||
]
|
||||
assert "vs-stale-runtime" not in in_memory_ids, (
|
||||
"stale runtime entries should still be pruned from memory"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_vector_stores_mixed_config_and_db(
|
||||
_isolated_registry, _permissive_user
|
||||
):
|
||||
"""
|
||||
Both DB-resident and config-declared stores should appear in the
|
||||
response simultaneously; the config entry must not mask or be
|
||||
masked by the DB one.
|
||||
"""
|
||||
_isolated_registry.load_vector_stores_from_config(
|
||||
[
|
||||
{
|
||||
"vector_store_name": "config-store",
|
||||
"litellm_params": {
|
||||
"custom_llm_provider": "pg_vector",
|
||||
"vector_store_id": "vs-config",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
db_only = LiteLLM_ManagedVectorStore(
|
||||
vector_store_id="vs-db",
|
||||
custom_llm_provider="openai",
|
||||
vector_store_name="db-store",
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.vector_store_endpoints.management_endpoints"
|
||||
".check_feature_access_for_user",
|
||||
new=AsyncMock(return_value=None),
|
||||
), patch(
|
||||
"litellm.proxy.vector_store_endpoints.management_endpoints"
|
||||
".VectorStoreRegistry._get_vector_stores_from_db",
|
||||
new=AsyncMock(return_value=[db_only]),
|
||||
), patch(
|
||||
"litellm.proxy.vector_store_endpoints.management_endpoints"
|
||||
"._check_vector_store_access",
|
||||
return_value=True,
|
||||
):
|
||||
response = await list_vector_stores(user_api_key_dict=_permissive_user)
|
||||
|
||||
ids_returned = {vs.get("vector_store_id") for vs in response["data"]}
|
||||
assert ids_returned == {"vs-config", "vs-db"}
|
||||
|
|
|
|||
|
|
@ -180,3 +180,52 @@ def test_search_uses_registry_credentials():
|
|||
assert getattr(called_params, "aws_region_name") == "us-east-1"
|
||||
finally:
|
||||
litellm.vector_store_registry = original_registry
|
||||
|
||||
|
||||
def test_load_vector_stores_from_config_marks_from_litellm_config():
|
||||
"""
|
||||
load_vector_stores_from_config must tag each loaded entry with
|
||||
from_litellm_config=True so downstream code can distinguish
|
||||
config-declared stores from DB/runtime stores (issue #25947).
|
||||
"""
|
||||
# Pass an explicit empty list — VectorStoreRegistry's __init__ has a
|
||||
# mutable default argument ([]) that is shared across instances within
|
||||
# the same process.
|
||||
registry = VectorStoreRegistry(vector_stores=[])
|
||||
registry.load_vector_stores_from_config(
|
||||
[
|
||||
{
|
||||
"vector_store_name": "demo",
|
||||
"litellm_params": {
|
||||
"custom_llm_provider": "pg_vector",
|
||||
"vector_store_id": "vs-demo",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert len(registry.vector_stores) == 1
|
||||
loaded = registry.vector_stores[0]
|
||||
assert loaded["vector_store_id"] == "vs-demo"
|
||||
assert loaded.get("from_litellm_config") is True
|
||||
|
||||
|
||||
def test_add_vector_store_to_registry_does_not_mark_from_litellm_config():
|
||||
"""
|
||||
Programmatic inserts (DB reload, API create) must NOT be tagged
|
||||
as from_litellm_config. Only entries declared in proxy_config.yaml
|
||||
should carry the flag.
|
||||
"""
|
||||
registry = VectorStoreRegistry(vector_stores=[])
|
||||
registry.add_vector_store_to_registry(
|
||||
LiteLLM_ManagedVectorStore(
|
||||
vector_store_id="vs-runtime",
|
||||
custom_llm_provider="openai",
|
||||
vector_store_name="runtime-store",
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
)
|
||||
|
||||
assert len(registry.vector_stores) == 1
|
||||
assert registry.vector_stores[0].get("from_litellm_config") is not True
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue