mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix(vector_store): /list preserves config-sourced entries
Track config-loaded vector store IDs on VectorStoreRegistry and exempt them from the "in memory but not in DB → deleted" reconciliation in GET /vector_store/list. Previously, any store declared in config.yaml / kustomization was invisible in the UI AND silently evicted from the running registry on every list call, breaking downstream routing until the next config reload. Fixes #38258
This commit is contained in:
parent
947dbbf029
commit
fef88c8665
3 changed files with 76 additions and 0 deletions
|
|
@ -623,6 +623,10 @@ 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)
|
||||
# Vector stores loaded from config.yaml have no DB row; they must be
|
||||
# returned alongside DB-sourced entries and MUST NOT be pruned by the
|
||||
# "in memory but not in DB → deleted" reconciliation below.
|
||||
config_loaded_ids: Final = litellm.vector_store_registry.config_loaded_vector_store_ids
|
||||
|
||||
vector_stores_to_delete_from_memory: Final[list[str]] = []
|
||||
|
||||
|
|
@ -631,6 +635,13 @@ async def list_vector_stores(
|
|||
if not vector_store_id:
|
||||
continue
|
||||
|
||||
# Config-loaded stores: surface them in the response but never
|
||||
# treat them as "deleted from DB" — they never had a DB row.
|
||||
if vector_store_id in config_loaded_ids:
|
||||
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(
|
||||
|
|
|
|||
|
|
@ -97,6 +97,10 @@ 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] = {}
|
||||
# IDs of vector stores that were loaded from config.yaml / kustomization.
|
||||
# These have no DB row, so reconciliation code paths that treat "in memory
|
||||
# but not in DB" as "was deleted" must exempt them.
|
||||
self.config_loaded_vector_store_ids: set[str] = set()
|
||||
|
||||
def _extract_tool_params(self, tool: dict) -> VectorStoreToolParams:
|
||||
"""
|
||||
|
|
@ -418,6 +422,7 @@ class VectorStoreRegistry:
|
|||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
self.vector_stores.append(litellm_managed_vector_store)
|
||||
self.config_loaded_vector_store_ids.add(vector_store_id)
|
||||
|
||||
verbose_logger.debug(
|
||||
"all loaded vector stores = %s",
|
||||
|
|
@ -463,6 +468,9 @@ class VectorStoreRegistry:
|
|||
for vector_store in self.vector_stores
|
||||
if vector_store.get("vector_store_id") != vector_store_id
|
||||
]
|
||||
# If this was a config-loaded id, drop the marker so it can be re-loaded
|
||||
# cleanly on next config reload without being incorrectly protected.
|
||||
self.config_loaded_vector_store_ids.discard(vector_store_id)
|
||||
|
||||
def update_vector_store_in_registry(self, vector_store_id: str, updated_data: LiteLLM_ManagedVectorStore):
|
||||
"""Update or add a vector store in the registry"""
|
||||
|
|
|
|||
57
tests/vector_store_tests/test_config_loaded_vector_stores.py
Normal file
57
tests/vector_store_tests/test_config_loaded_vector_stores.py
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
"""
|
||||
Tests for the config-sourced vector store fix.
|
||||
|
||||
Regression: `/vector_store/list` used to prune config.yaml-loaded stores from
|
||||
the in-memory registry because they have no DB row. The fix tracks
|
||||
`config_loaded_vector_store_ids` on the registry and exempts them from the
|
||||
"in memory but not in DB" reconciliation branch.
|
||||
"""
|
||||
from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
|
||||
|
||||
|
||||
def _config_entry(vs_id: str) -> dict:
|
||||
return {
|
||||
"vector_store_name": f"name-{vs_id}",
|
||||
"litellm_params": {
|
||||
"vector_store_id": vs_id,
|
||||
"custom_llm_provider": "bedrock",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_load_from_config_tracks_id():
|
||||
registry = VectorStoreRegistry()
|
||||
assert registry.config_loaded_vector_store_ids == set()
|
||||
|
||||
registry.load_vector_stores_from_config([_config_entry("vs_a"), _config_entry("vs_b")])
|
||||
|
||||
assert {vs["vector_store_id"] for vs in registry.vector_stores} == {"vs_a", "vs_b"}
|
||||
assert registry.config_loaded_vector_store_ids == {"vs_a", "vs_b"}
|
||||
|
||||
|
||||
def test_delete_drops_config_marker():
|
||||
"""After a config-loaded store is explicitly deleted, the marker must go
|
||||
with it so a later reload can re-create the entry without confusion."""
|
||||
registry = VectorStoreRegistry()
|
||||
registry.load_vector_stores_from_config([_config_entry("vs_a")])
|
||||
assert "vs_a" in registry.config_loaded_vector_store_ids
|
||||
|
||||
registry.delete_vector_store_from_registry("vs_a")
|
||||
|
||||
assert registry.vector_stores == []
|
||||
assert "vs_a" not in registry.config_loaded_vector_store_ids
|
||||
|
||||
|
||||
def test_delete_of_db_only_store_does_not_touch_marker_set():
|
||||
registry = VectorStoreRegistry()
|
||||
registry.load_vector_stores_from_config([_config_entry("vs_config")])
|
||||
# Simulate a DB-sourced entry being added via /vector_store/new
|
||||
registry.add_vector_store_to_registry(
|
||||
{"vector_store_id": "vs_from_db", "custom_llm_provider": "openai"}
|
||||
)
|
||||
|
||||
registry.delete_vector_store_from_registry("vs_from_db")
|
||||
|
||||
assert [vs["vector_store_id"] for vs in registry.vector_stores] == ["vs_config"]
|
||||
# Config marker for the untouched entry is intact.
|
||||
assert registry.config_loaded_vector_store_ids == {"vs_config"}
|
||||
Loading…
Add table
Reference in a new issue