mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
Merge 620f586dbb into dab2deb5ed
This commit is contained in:
commit
c5468e84fb
2 changed files with 54 additions and 6 deletions
|
|
@ -32,8 +32,10 @@ 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] = (
|
||||
vector_store_indexes if vector_store_indexes is not None else []
|
||||
)
|
||||
|
||||
def get_vector_store_indexes(self) -> list[LiteLLM_ManagedVectorStoreIndex]:
|
||||
"""
|
||||
|
|
@ -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] = vector_stores if vector_stores is not None else []
|
||||
self.vector_store_ids_to_vector_store_map: dict[str, LiteLLM_ManagedVectorStore] = {}
|
||||
|
||||
def _extract_tool_params(self, tool: dict) -> VectorStoreToolParams:
|
||||
|
|
|
|||
|
|
@ -8,12 +8,20 @@ from fastapi.testclient import TestClient
|
|||
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, 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)
|
||||
|
|
@ -249,3 +257,41 @@ async def test_config_owned_store_survives_db_liveness_check_while_missing_db_st
|
|||
prisma_client.db.litellm_managedvectorstorestable.find_unique.assert_awaited_once_with(
|
||||
where={"vector_store_id": "vs_from_db"}
|
||||
)
|
||||
|
||||
|
||||
class TestRegistriesDoNotShareDefaultState:
|
||||
"""#38874: default-constructed registries must not share one list."""
|
||||
|
||||
def test_vector_store_registries_do_not_share_a_list(self):
|
||||
first: Final = VectorStoreRegistry()
|
||||
second: Final = VectorStoreRegistry()
|
||||
assert first.vector_stores is not second.vector_stores
|
||||
|
||||
def test_index_registries_do_not_share_a_list(self):
|
||||
first: Final = VectorStoreIndexRegistry()
|
||||
second: Final = VectorStoreIndexRegistry()
|
||||
assert first.vector_store_indexes is not second.vector_store_indexes
|
||||
|
||||
def test_adding_to_one_registry_leaves_the_next_one_empty(self):
|
||||
VectorStoreRegistry().add_vector_store_to_registry(
|
||||
LiteLLM_ManagedVectorStore(vector_store_id="vs-leak", custom_llm_provider="bedrock")
|
||||
)
|
||||
assert VectorStoreRegistry().vector_stores == []
|
||||
|
||||
def test_upserting_into_one_index_registry_leaves_the_next_one_empty(self):
|
||||
VectorStoreIndexRegistry().upsert_vector_store_index(
|
||||
LiteLLM_ManagedVectorStoreIndex(
|
||||
id="idx-leak",
|
||||
index_name="leak",
|
||||
litellm_params=IndexCreateLiteLLMParams(vector_store_index="vs-leak", vector_store_name="leak"),
|
||||
)
|
||||
)
|
||||
assert VectorStoreIndexRegistry().get_vector_store_indexes() == []
|
||||
|
||||
def test_a_supplied_list_is_still_used_as_given(self):
|
||||
supplied: Final = [LiteLLM_ManagedVectorStore(vector_store_id="vs-1", custom_llm_provider="bedrock")]
|
||||
assert VectorStoreRegistry(vector_stores=supplied).vector_stores is supplied
|
||||
|
||||
def test_a_supplied_empty_list_is_still_used_as_given(self):
|
||||
supplied: Final[list[LiteLLM_ManagedVectorStore]] = []
|
||||
assert VectorStoreRegistry(vector_stores=supplied).vector_stores is supplied
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue