mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge 7028599e32 into 176b2e5eb8
This commit is contained in:
commit
bfc7c5e202
2 changed files with 46 additions and 3 deletions
|
|
@ -104,6 +104,7 @@ 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] = {}
|
||||
self.config_vector_store_ids: frozenset[str] = frozenset()
|
||||
|
||||
def _extract_tool_params(self, tool: dict) -> VectorStoreToolParams:
|
||||
"""
|
||||
|
|
@ -338,8 +339,13 @@ class VectorStoreRegistry:
|
|||
break
|
||||
|
||||
# Verify vector store still exists in database (if we have DB access)
|
||||
# This ensures deleted vector stores are removed from cache
|
||||
if vector_store is not None and prisma_client is not None:
|
||||
# This ensures deleted vector stores are removed from cache.
|
||||
# Config-defined vector stores are never in the database, so skip the check for them
|
||||
if (
|
||||
vector_store is not None
|
||||
and prisma_client is not None
|
||||
and vector_store_id not in self.config_vector_store_ids
|
||||
):
|
||||
try:
|
||||
# Check if it still exists in database
|
||||
db_vector_store = await ManagedVectorStoresRepository(prisma_client).table.find_unique(
|
||||
|
|
@ -427,6 +433,7 @@ class VectorStoreRegistry:
|
|||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
self.vector_stores.append(litellm_managed_vector_store)
|
||||
self.config_vector_store_ids = self.config_vector_store_ids | {vector_store_id}
|
||||
|
||||
verbose_logger.debug(
|
||||
"all loaded vector stores = %s",
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from fastapi.testclient import TestClient
|
|||
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import litellm
|
||||
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
|
||||
|
|
@ -182,3 +182,39 @@ def test_search_uses_registry_credentials():
|
|||
assert getattr(called_params, "aws_region_name") == "us-east-1"
|
||||
finally:
|
||||
litellm.vector_store_registry = original_registry
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_vector_store_survives_db_check():
|
||||
"""Config-defined vector stores are not in the DB and must not be evicted when a DB is connected."""
|
||||
registry = VectorStoreRegistry()
|
||||
registry.load_vector_stores_from_config(
|
||||
[
|
||||
{
|
||||
"vector_store_name": "config-kb",
|
||||
"litellm_params": {
|
||||
"vector_store_id": "CONFIGKB",
|
||||
"custom_llm_provider": "bedrock",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
db_store = LiteLLM_ManagedVectorStore(
|
||||
vector_store_id="DBONLY",
|
||||
custom_llm_provider="bedrock",
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
registry.add_vector_store_to_registry(db_store)
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=None)
|
||||
result = await registry.pop_vector_stores_to_run_with_db_fallback(
|
||||
non_default_params={"vector_store_ids": ["CONFIGKB", "DBONLY"]},
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
assert [vs.get("vector_store_id") for vs in result] == ["CONFIGKB"]
|
||||
assert any(vs.get("vector_store_id") == "CONFIGKB" for vs in registry.vector_stores)
|
||||
assert not any(vs.get("vector_store_id") == "DBONLY" for vs in registry.vector_stores)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue