mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
address CI: reduce C901 complexity, fix TypedDict test access
- Extract two helpers from load_vector_stores_from_config: _build_managed_vector_store_from_config_entry (per-entry construction + required-param validation) and _drop_config_entries_removed_from_new_config (stale marker eviction). Puts the main method back under the C901 strict-budget limit (was +1 over). - Endpoint test: LiteLLM_ManagedVectorStoreListResponse is a TypedDict; use dict-style access on the response body instead of attribute access.
This commit is contained in:
parent
97b9d9a0f9
commit
2088ea17a2
2 changed files with 49 additions and 35 deletions
|
|
@ -93,6 +93,40 @@ class VectorStoreIndexRegistry:
|
|||
return vector_stores_from_db
|
||||
|
||||
|
||||
def _build_managed_vector_store_from_config_entry(vector_store_config: dict) -> LiteLLM_ManagedVectorStore:
|
||||
"""Construct a LiteLLM_ManagedVectorStore from one config.yaml entry.
|
||||
|
||||
Extracted from load_vector_stores_from_config to keep that method under the
|
||||
C901 complexity budget; also enforces the two required params (vector_store_id
|
||||
and custom_llm_provider) in one place with clear error messages.
|
||||
"""
|
||||
litellm_vector_store_config = LiteLLM_VectorStoreConfig(**vector_store_config)
|
||||
vector_store_name = litellm_vector_store_config.get("vector_store_name")
|
||||
vector_store_litellm_params: dict[str, Any] = litellm_vector_store_config.get("litellm_params") or {}
|
||||
|
||||
vector_store_id = vector_store_litellm_params.get("vector_store_id")
|
||||
if vector_store_id is None:
|
||||
raise ValueError(
|
||||
f"vector_store_id is required for initializing vector store, got vector_store_id={vector_store_id}"
|
||||
)
|
||||
custom_llm_provider = vector_store_litellm_params.get("custom_llm_provider")
|
||||
if custom_llm_provider is None:
|
||||
raise ValueError(
|
||||
f"custom_llm_provider is required for initializing vector store, got custom_llm_provider={custom_llm_provider}"
|
||||
)
|
||||
|
||||
return LiteLLM_ManagedVectorStore(
|
||||
vector_store_id=vector_store_id,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_params=vector_store_litellm_params,
|
||||
vector_store_name=vector_store_name,
|
||||
vector_store_description=vector_store_litellm_params.get("vector_store_description"),
|
||||
vector_store_metadata=vector_store_litellm_params.get("vector_store_metadata"),
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
|
||||
class VectorStoreRegistry:
|
||||
def __init__(self, vector_stores: list[LiteLLM_ManagedVectorStore] | None = None):
|
||||
# Never bind a mutable default: the previous `= []` default caused every
|
||||
|
|
@ -394,6 +428,14 @@ class VectorStoreRegistry:
|
|||
vector_store_ids.extend(tool["vector_store_ids"])
|
||||
return vector_store_ids
|
||||
|
||||
def _drop_config_entries_removed_from_new_config(self, new_config_ids: set[str]) -> None:
|
||||
"""Evict any previously-config-loaded store whose id isn't in the new config,
|
||||
so config removals stop being protected from list-endpoint reconciliation."""
|
||||
for stale_id in self.config_loaded_vector_store_ids - new_config_ids:
|
||||
# delete_vector_store_from_registry also discards from
|
||||
# config_loaded_vector_store_ids, keeping the two in sync.
|
||||
self.delete_vector_store_from_registry(stale_id)
|
||||
|
||||
def load_vector_stores_from_config(self, vector_stores_config: list[dict]):
|
||||
"""
|
||||
Loads vector stores from the litellm proxy config.yaml.
|
||||
|
|
@ -407,44 +449,14 @@ class VectorStoreRegistry:
|
|||
(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
|
||||
|
||||
stale_ids: Final = self.config_loaded_vector_store_ids - new_config_ids
|
||||
for stale_id in stale_ids:
|
||||
# delete_vector_store_from_registry also discards from
|
||||
# config_loaded_vector_store_ids, keeping the two in sync.
|
||||
self.delete_vector_store_from_registry(stale_id)
|
||||
self._drop_config_entries_removed_from_new_config(new_config_ids)
|
||||
|
||||
for vector_store_config in vector_stores_config:
|
||||
# cast to VectorStoreConfig
|
||||
litellm_vector_store_config = LiteLLM_VectorStoreConfig(**vector_store_config)
|
||||
vector_store_name = litellm_vector_store_config.get("vector_store_name")
|
||||
vector_store_litellm_params: dict[str, Any] = litellm_vector_store_config.get("litellm_params") or {}
|
||||
|
||||
vector_store_id = vector_store_litellm_params.get("vector_store_id")
|
||||
if vector_store_id is None:
|
||||
raise ValueError(
|
||||
f"vector_store_id is required for initializing vector store, got vector_store_id={vector_store_id}"
|
||||
)
|
||||
custom_llm_provider = vector_store_litellm_params.get("custom_llm_provider")
|
||||
if custom_llm_provider is None:
|
||||
raise ValueError(
|
||||
f"custom_llm_provider is required for initializing vector store, got custom_llm_provider={custom_llm_provider}"
|
||||
)
|
||||
|
||||
litellm_managed_vector_store = LiteLLM_ManagedVectorStore(
|
||||
vector_store_id=vector_store_id,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_params=vector_store_litellm_params,
|
||||
vector_store_name=vector_store_name,
|
||||
vector_store_description=vector_store_litellm_params.get("vector_store_description"),
|
||||
vector_store_metadata=vector_store_litellm_params.get("vector_store_metadata"),
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
managed = _build_managed_vector_store_from_config_entry(vector_store_config)
|
||||
# add_vector_store_to_registry is idempotent by id, so a reload
|
||||
# that re-lists an existing store won't create duplicates.
|
||||
self.add_vector_store_to_registry(litellm_managed_vector_store)
|
||||
self.config_loaded_vector_store_ids.add(vector_store_id)
|
||||
self.add_vector_store_to_registry(managed)
|
||||
self.config_loaded_vector_store_ids.add(managed["vector_store_id"])
|
||||
|
||||
verbose_logger.debug(
|
||||
"all loaded vector stores = %s",
|
||||
|
|
|
|||
|
|
@ -3162,7 +3162,9 @@ async def test_list_vector_stores_returns_config_sourced_and_leaves_registry_int
|
|||
),
|
||||
)
|
||||
|
||||
returned_ids = [vs["vector_store_id"] for vs in response.data]
|
||||
# LiteLLM_ManagedVectorStoreListResponse is a TypedDict — use dict access.
|
||||
response_data = response["data"] if isinstance(response, dict) else response.data
|
||||
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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue