mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(vector-stores): reject configured ID collisions
This commit is contained in:
parent
57b08e8f7a
commit
1a0f4aa4b2
3 changed files with 91 additions and 0 deletions
|
|
@ -137,6 +137,16 @@ def _litellm_params_dict(
|
|||
return {} # mutable-ok: absent persisted parameters normalize to an empty mutable mapping
|
||||
|
||||
|
||||
def _reject_config_vector_store_id(vector_store_id: str) -> None:
|
||||
registry: Final = litellm.vector_store_registry
|
||||
if registry is None or vector_store_id not in registry.config_vector_store_ids:
|
||||
return
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Vector store ID {vector_store_id} is defined in proxy configuration and cannot be managed through the API",
|
||||
)
|
||||
|
||||
|
||||
async def _fetch_and_authorize_vector_store(
|
||||
vector_store_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -210,6 +220,8 @@ async def create_vector_store_in_db(
|
|||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail="Database not connected")
|
||||
|
||||
_reject_config_vector_store_id(vector_store_id)
|
||||
|
||||
# Check if vector store already exists
|
||||
existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique(
|
||||
where={"vector_store_id": vector_store_id}
|
||||
|
|
@ -629,6 +641,7 @@ async def update_vector_store(
|
|||
try:
|
||||
update_data: Final = data.model_dump(exclude_unset=True)
|
||||
vector_store_id: Final[str] = update_data.pop("vector_store_id")
|
||||
_reject_config_vector_store_id(vector_store_id)
|
||||
|
||||
# Per-store access control: anyone authenticated who passes the
|
||||
# premium-feature gate could otherwise update *any* vector store —
|
||||
|
|
|
|||
|
|
@ -479,6 +479,8 @@ class VectorStoreRegistry:
|
|||
"""
|
||||
Delete a vector store from the registry
|
||||
"""
|
||||
if vector_store_id in self.config_vector_store_ids:
|
||||
return
|
||||
self.vector_stores = [
|
||||
vector_store
|
||||
for vector_store in self.vector_stores
|
||||
|
|
@ -487,6 +489,8 @@ class VectorStoreRegistry:
|
|||
|
||||
def update_vector_store_in_registry(self, vector_store_id: str, updated_data: LiteLLM_ManagedVectorStore):
|
||||
"""Update or add a vector store in the registry"""
|
||||
if vector_store_id in self.config_vector_store_ids:
|
||||
return
|
||||
for i, vector_store in enumerate(self.vector_stores):
|
||||
if vector_store.get("vector_store_id") == vector_store_id:
|
||||
self.vector_stores[i] = updated_data
|
||||
|
|
|
|||
|
|
@ -877,6 +877,80 @@ async def test_db_fallback_does_not_evict_config_source():
|
|||
prisma_client.db.litellm_managedvectorstorestable.find_unique.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_vector_store_id_cannot_be_created_in_database():
|
||||
registry = VectorStoreRegistry()
|
||||
registry.config_vector_store_ids = frozenset(("configured",))
|
||||
prisma_client = MagicMock()
|
||||
|
||||
with (
|
||||
patch.object(litellm, "vector_store_registry", registry),
|
||||
pytest.raises(HTTPException, match="defined in proxy configuration") as exc_info,
|
||||
):
|
||||
await create_vector_store_in_db(
|
||||
vector_store_id="configured",
|
||||
custom_llm_provider="milvus",
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
prisma_client.db.litellm_managedvectorstorestable.find_unique.assert_not_called()
|
||||
prisma_client.db.litellm_managedvectorstorestable.create.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_vector_store_id_cannot_be_updated_in_database():
|
||||
from litellm.proxy.vector_store_endpoints.management_endpoints import update_vector_store
|
||||
from litellm.types.vector_stores import VectorStoreUpdateRequest
|
||||
|
||||
registry = VectorStoreRegistry()
|
||||
registry.config_vector_store_ids = frozenset(("configured",))
|
||||
prisma_client = MagicMock()
|
||||
|
||||
with (
|
||||
patch.object(litellm, "vector_store_registry", registry),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
|
||||
patch(
|
||||
"litellm.proxy.vector_store_endpoints.management_endpoints.check_feature_access_for_user",
|
||||
new=AsyncMock(),
|
||||
),
|
||||
pytest.raises(HTTPException, match="defined in proxy configuration") as exc_info,
|
||||
):
|
||||
await update_vector_store(
|
||||
data=VectorStoreUpdateRequest(
|
||||
vector_store_id="configured",
|
||||
vector_store_description="replacement",
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
prisma_client.db.litellm_managedvectorstorestable.find_unique.assert_not_called()
|
||||
prisma_client.db.litellm_managedvectorstorestable.update.assert_not_called()
|
||||
|
||||
|
||||
def test_config_vector_store_cannot_be_replaced_or_deleted_from_registry():
|
||||
configured = LiteLLM_ManagedVectorStore(
|
||||
vector_store_id="configured",
|
||||
custom_llm_provider="milvus",
|
||||
litellm_params={"api_base": "https://configured-milvus:19530"},
|
||||
)
|
||||
registry = VectorStoreRegistry(vector_stores=[configured])
|
||||
registry.config_vector_store_ids = frozenset(("configured",))
|
||||
|
||||
registry.update_vector_store_in_registry(
|
||||
"configured",
|
||||
LiteLLM_ManagedVectorStore(
|
||||
vector_store_id="configured",
|
||||
custom_llm_provider="milvus",
|
||||
litellm_params={"api_base": "https://attacker.example"},
|
||||
),
|
||||
)
|
||||
registry.delete_vector_store_from_registry("configured")
|
||||
|
||||
assert registry.get_litellm_managed_vector_store_from_registry("configured") == configured
|
||||
|
||||
|
||||
def test_admin_persistence_strips_forged_marker_and_adds_server_marker():
|
||||
params = prepare_milvus_connection_for_persistence(
|
||||
custom_llm_provider="milvus/probe",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue