mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
fix(vector_stores): stop sharing mutable default state across registry instances
VectorStoreRegistry and VectorStoreIndexRegistry both declared their constructor argument with a mutable `= []` default. That list is created once at def-time and shared by every instance built without an explicit argument, and several methods append into it: add_vector_store_to_registry, update_vector_store_in_registry, load_vector_stores_from_config and upsert_vector_store_index. A brand-new registry therefore inherited a previous instance's vector stores, including their credential references consumed by get_credentials_for_vector_store. The proxy reaches the no-arg path in ProxyConfig.load_config, where VectorStoreRegistry() is constructed and then immediately fed load_vector_stores_from_config. delete_vector_store_from_registry made it worse by rebinding self.vector_stores to a fresh list rather than mutating, so a delete never removed the entry from the shared default and any later no-arg registry resurrected it. Both constructors now take Optional[...] = None and copy into their own list, which also stops a caller-supplied list from being mutated behind the caller's back. Also fixes delete_vector_store_index, which was an unconditional no-op: it compared each stored LiteLLM_ManagedVectorStoreIndex against the index-name string, and a model is never equal to a str, so the comprehension always returned the list unchanged. It now filters on index_name, matching get_vector_store_index_by_name and is_vector_store_index.
This commit is contained in:
parent
5d4c4d0fce
commit
8040927673
2 changed files with 100 additions and 8 deletions
|
|
@ -25,8 +25,8 @@ else:
|
|||
|
||||
|
||||
class VectorStoreIndexRegistry:
|
||||
def __init__(self, vector_store_indexes: List[LiteLLM_ManagedVectorStoreIndex] = []):
|
||||
self.vector_store_indexes: List[LiteLLM_ManagedVectorStoreIndex] = vector_store_indexes
|
||||
def __init__(self, vector_store_indexes: Optional[List[LiteLLM_ManagedVectorStoreIndex]] = None):
|
||||
self.vector_store_indexes: List[LiteLLM_ManagedVectorStoreIndex] = list(vector_store_indexes or ())
|
||||
|
||||
def get_vector_store_indexes(self) -> List[LiteLLM_ManagedVectorStoreIndex]:
|
||||
"""
|
||||
|
|
@ -55,11 +55,11 @@ class VectorStoreIndexRegistry:
|
|||
return
|
||||
self.vector_store_indexes.append(vector_store_index)
|
||||
|
||||
def delete_vector_store_index(self, vector_store_index: str):
|
||||
def delete_vector_store_index(self, index_name: str):
|
||||
"""
|
||||
Deletes a vector store index from the registry
|
||||
"""
|
||||
self.vector_store_indexes = [index for index in self.vector_store_indexes if index != vector_store_index]
|
||||
self.vector_store_indexes = [index for index in self.vector_store_indexes if index.index_name != index_name]
|
||||
|
||||
def is_vector_store_index(self, vector_store_index_name: str) -> bool:
|
||||
"""
|
||||
|
|
@ -94,8 +94,8 @@ class VectorStoreIndexRegistry:
|
|||
|
||||
|
||||
class VectorStoreRegistry:
|
||||
def __init__(self, vector_stores: List[LiteLLM_ManagedVectorStore] = []):
|
||||
self.vector_stores: List[LiteLLM_ManagedVectorStore] = vector_stores
|
||||
def __init__(self, vector_stores: Optional[List[LiteLLM_ManagedVectorStore]] = None):
|
||||
self.vector_stores: List[LiteLLM_ManagedVectorStore] = list(vector_stores or ())
|
||||
self.vector_store_ids_to_vector_store_map: Dict[str, LiteLLM_ManagedVectorStore] = {}
|
||||
|
||||
def _extract_tool_params(self, tool: Dict) -> VectorStoreToolParams:
|
||||
|
|
|
|||
|
|
@ -16,9 +16,16 @@ from datetime import datetime, timezone
|
|||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import litellm
|
||||
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
|
||||
from litellm.types.vector_stores import (
|
||||
IndexCreateLiteLLMParams,
|
||||
LiteLLM_ManagedVectorStore,
|
||||
LiteLLM_ManagedVectorStoreIndex,
|
||||
)
|
||||
from litellm.vector_stores.main import search
|
||||
from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
|
||||
from litellm.vector_stores.vector_store_registry import (
|
||||
VectorStoreIndexRegistry,
|
||||
VectorStoreRegistry,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
|
|
@ -187,3 +194,88 @@ def test_search_uses_registry_credentials():
|
|||
assert getattr(called_params, "aws_region_name") == "us-east-1"
|
||||
finally:
|
||||
litellm.vector_store_registry = original_registry
|
||||
|
||||
|
||||
def _managed_vector_store(vector_store_id: str) -> LiteLLM_ManagedVectorStore:
|
||||
return LiteLLM_ManagedVectorStore(
|
||||
vector_store_id=vector_store_id,
|
||||
custom_llm_provider="openai",
|
||||
vector_store_name=f"store_{vector_store_id}",
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
|
||||
def _managed_vector_store_index(index_name: str) -> LiteLLM_ManagedVectorStoreIndex:
|
||||
return LiteLLM_ManagedVectorStoreIndex(
|
||||
id=f"id_{index_name}",
|
||||
index_name=index_name,
|
||||
litellm_params=IndexCreateLiteLLMParams(
|
||||
vector_store_index=index_name,
|
||||
vector_store_name=f"store_{index_name}",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_registries_do_not_share_default_storage():
|
||||
"""
|
||||
A registry built without arguments must own its own storage.
|
||||
|
||||
Regression test: both registries used a mutable `= []` default, so the list was created once at
|
||||
def-time and shared by every no-arg instance. Mutating one registry leaked into the next one,
|
||||
which on the proxy meant a fresh registry inherited another's vector stores (and their
|
||||
credential references).
|
||||
"""
|
||||
first_registry = VectorStoreRegistry()
|
||||
first_registry.add_vector_store_to_registry(_managed_vector_store("vs_leak"))
|
||||
assert first_registry.vector_stores[0]["vector_store_id"] == "vs_leak"
|
||||
|
||||
assert VectorStoreRegistry().vector_stores == []
|
||||
|
||||
first_index_registry = VectorStoreIndexRegistry()
|
||||
first_index_registry.upsert_vector_store_index(_managed_vector_store_index("idx_leak"))
|
||||
assert first_index_registry.vector_store_indexes[0].index_name == "idx_leak"
|
||||
|
||||
assert VectorStoreIndexRegistry().vector_store_indexes == []
|
||||
|
||||
|
||||
def test_registries_do_not_alias_caller_list():
|
||||
"""
|
||||
The constructor must copy the caller's list rather than alias it.
|
||||
|
||||
Otherwise registry mutations write back into a list the caller still owns; the proxy passes
|
||||
lists freshly read from the DB, and `load_vector_stores_from_config` appends into whatever it
|
||||
was handed.
|
||||
"""
|
||||
caller_vector_stores = [_managed_vector_store("vs_1")]
|
||||
registry = VectorStoreRegistry(caller_vector_stores)
|
||||
registry.add_vector_store_to_registry(_managed_vector_store("vs_2"))
|
||||
|
||||
assert len(registry.vector_stores) == 2
|
||||
assert len(caller_vector_stores) == 1
|
||||
|
||||
caller_indexes = [_managed_vector_store_index("idx_1")]
|
||||
index_registry = VectorStoreIndexRegistry(caller_indexes)
|
||||
index_registry.upsert_vector_store_index(_managed_vector_store_index("idx_2"))
|
||||
|
||||
assert len(index_registry.vector_store_indexes) == 2
|
||||
assert len(caller_indexes) == 1
|
||||
|
||||
|
||||
def test_delete_vector_store_index_removes_by_name():
|
||||
"""
|
||||
Regression test: delete_vector_store_index compared each stored index object against the
|
||||
index-name string, which is never equal, so the delete was an unconditional no-op.
|
||||
"""
|
||||
registry = VectorStoreIndexRegistry(
|
||||
[_managed_vector_store_index("idx_a"), _managed_vector_store_index("idx_b")]
|
||||
)
|
||||
|
||||
registry.delete_vector_store_index("idx_a")
|
||||
|
||||
assert [index.index_name for index in registry.vector_store_indexes] == ["idx_b"]
|
||||
assert registry.get_vector_store_index_by_name("idx_a") is None
|
||||
assert registry.is_vector_store_index("idx_b") is True
|
||||
|
||||
registry.delete_vector_store_index("idx_does_not_exist")
|
||||
assert [index.index_name for index in registry.vector_store_indexes] == ["idx_b"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue