diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 22d27bc3266..99baa257624 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -32,8 +32,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: list[LiteLLM_ManagedVectorStoreIndex] | None = None): + self.vector_store_indexes: list[LiteLLM_ManagedVectorStoreIndex] = list(vector_store_indexes or ()) def get_vector_store_indexes(self) -> list[LiteLLM_ManagedVectorStoreIndex]: """ @@ -66,7 +66,9 @@ class VectorStoreIndexRegistry: """ 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 != vector_store_index + ] def is_vector_store_index(self, vector_store_index_name: str) -> bool: """ @@ -101,8 +103,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: list[LiteLLM_ManagedVectorStore] | None = 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: @@ -384,16 +386,15 @@ class VectorStoreRegistry: return vector_stores_to_run def _get_vector_store_ids_from_tool_calls( - self, tools: list[dict] | None = None, vector_store_ids: list[str] = [] + self, tools: list[dict] | None = None, vector_store_ids: list[str] | None = None ) -> list[str]: """ Returns the vector store ids from the tool calls """ - if tools: - for tool in tools: - if "vector_store_ids" in tool: - vector_store_ids.extend(tool["vector_store_ids"]) - return vector_store_ids + return [ + *(vector_store_ids or ()), + *(vs_id for tool in (tools or ()) if "vector_store_ids" in tool for vs_id in tool["vector_store_ids"]), + ] def load_vector_stores_from_config(self, vector_stores_config: list[dict]): """ diff --git a/tests/test_litellm/vector_stores/test_vector_store_registry.py b/tests/test_litellm/vector_stores/test_vector_store_registry.py index f19c3706845..76da2f2b85d 100644 --- a/tests/test_litellm/vector_stores/test_vector_store_registry.py +++ b/tests/test_litellm/vector_stores/test_vector_store_registry.py @@ -11,9 +11,16 @@ from datetime import datetime, timezone from unittest.mock import MagicMock 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) @@ -182,3 +189,108 @@ 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"] + + +def test_get_vector_store_ids_from_tool_calls_does_not_share_or_mutate(): + """ + Regression test: _get_vector_store_ids_from_tool_calls took `vector_store_ids: list[str] = []` + and extended it in place, so the def-time default accumulated ids across every call that + omitted the argument, and a caller-supplied list was mutated behind the caller's back. + """ + registry = VectorStoreRegistry() + tools = [{"vector_store_ids": ["vs_a"]}] + + first = registry._get_vector_store_ids_from_tool_calls(tools=tools) + second = registry._get_vector_store_ids_from_tool_calls(tools=tools) + + assert first == ["vs_a"] + assert second == ["vs_a"] + + caller_ids = ["vs_seed"] + result = registry._get_vector_store_ids_from_tool_calls(tools=tools, vector_store_ids=caller_ids) + + assert result == ["vs_seed", "vs_a"] + assert caller_ids == ["vs_seed"]