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:
KK-MCP 2026-08-25 22:08:03 -07:00
parent 97b9d9a0f9
commit 2088ea17a2
2 changed files with 49 additions and 35 deletions

View file

@ -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",

View file

@ -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.