Merge ccfb0beada into 49affa7c01

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]):
         """

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"]