mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
fix: mutable-default in VectorStoreRegistry, ruff format, endpoint test
- Root-cause the shard failures on the previous commit: `__init__`'s `vector_stores: list = []` default was a shared mutable, so every `VectorStoreRegistry()` without an explicit list appended to the same underlying list — the four new tests leaked state into one another and passed locally only because they were the only registrants. Switch the default to `None` and materialize a per-instance list. - Apply `ruff format` to vector_store_registry.py (fixes `lint` job). - Add an endpoint-level regression test to raise codecov/patch: `test_list_vector_stores_returns_config_sourced_and_leaves_registry_intact` exercises the DB-empty + config-loaded branch of `list_vector_stores` and asserts the registry is NOT mutated after the call.
This commit is contained in:
parent
30cdbd4d9e
commit
97b9d9a0f9
2 changed files with 69 additions and 4 deletions
|
|
@ -94,8 +94,12 @@ 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):
|
||||
# Never bind a mutable default: the previous `= []` default caused every
|
||||
# VectorStoreRegistry() with no argument to share the same list, so a
|
||||
# mutation in one place (test, request handler, etc.) leaked into every
|
||||
# other caller.
|
||||
self.vector_stores: list[LiteLLM_ManagedVectorStore] = list(vector_stores) if vector_stores is not None else []
|
||||
self.vector_store_ids_to_vector_store_map: dict[str, LiteLLM_ManagedVectorStore] = {}
|
||||
# IDs of vector stores that were loaded from config.yaml / kustomization.
|
||||
# These have no DB row, so reconciliation code paths that treat "in memory
|
||||
|
|
@ -400,8 +404,7 @@ class VectorStoreRegistry:
|
|||
list-endpoint reconciliation on the very next reload.
|
||||
"""
|
||||
new_config_ids: set[str] = {
|
||||
(cfg.get("litellm_params") or {}).get("vector_store_id")
|
||||
for cfg in vector_stores_config
|
||||
(cfg.get("litellm_params") or {}).get("vector_store_id") for cfg in vector_stores_config
|
||||
}
|
||||
new_config_ids.discard(None) # ids missing here will re-raise below with a clearer error
|
||||
|
||||
|
|
|
|||
|
|
@ -3107,3 +3107,65 @@ class TestAzureAIAnalyzeNamedIndexClassification:
|
|||
user_api_key_dict=self._team_member("analyze", ["read"]),
|
||||
)
|
||||
assert result is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# /vector_store/list — config-sourced entries must survive DB reconciliation.
|
||||
# Regression test for: config.yaml-loaded stores were pruned from the
|
||||
# in-memory registry on every list call because they had no DB row.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_vector_stores_returns_config_sourced_and_leaves_registry_intact():
|
||||
from litellm.proxy.vector_store_endpoints.management_endpoints import list_vector_stores
|
||||
from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
|
||||
|
||||
registry = VectorStoreRegistry()
|
||||
registry.load_vector_stores_from_config(
|
||||
[
|
||||
{
|
||||
"vector_store_name": "docs-kb",
|
||||
"litellm_params": {
|
||||
"vector_store_id": "vs_from_config",
|
||||
"custom_llm_provider": "bedrock",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
original_registry = litellm.vector_store_registry
|
||||
try:
|
||||
litellm.vector_store_registry = registry
|
||||
|
||||
with (
|
||||
# Empty DB — the store exists only in config/memory.
|
||||
patch(
|
||||
"litellm.vector_stores.vector_store_registry.VectorStoreRegistry._get_vector_stores_from_db",
|
||||
new=AsyncMock(return_value=[]),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
|
||||
patch(
|
||||
"litellm.proxy.vector_store_endpoints.management_endpoints._check_vector_store_access",
|
||||
new=AsyncMock(return_value=True),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.vector_store_endpoints.management_endpoints.check_feature_access_for_user",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
):
|
||||
response = await list_vector_stores(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
token="sk-test",
|
||||
key_name="sk-...test",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
),
|
||||
)
|
||||
|
||||
returned_ids = [vs["vector_store_id"] for vs in response.data]
|
||||
assert "vs_from_config" in returned_ids
|
||||
# And, critically, the registry must NOT have been mutated —
|
||||
# config-loaded stores stay resident for downstream routing.
|
||||
assert any(vs.get("vector_store_id") == "vs_from_config" for vs in registry.vector_stores)
|
||||
finally:
|
||||
litellm.vector_store_registry = original_registry
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue