scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

lazy-load optional feature routers on first request (#26534)

Browse source 
Co-authored-by: Michael Riad Zaky <michaelr@Mac.localdomain>
This commit is contained in:
Michael-RZ-Berri 2026-04-28 17:04:40 -07:00 • committed by GitHub
parent f2747e8c75
commit 21ed38971d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 609 additions and 105 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

307
litellm/proxy/_lazy_features.py Normal file
View file

@ -0,0 +1,307 @@
"""
Lazy registration for optional feature routers. Each LAZY_FEATURES entry
imports its module only on the first request matching its path prefix,
saving ~700 MB at idle for deployments that don't use these features.
First hit pays the import cost (1-3 s for heavy modules); /openapi.json
omits each feature's routes until the feature is warmed.
"""
import asyncio
import importlib
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Callable, Tuple
from starlette.types import Receive, Scope, Send
from litellm._logging import verbose_proxy_logger
if TYPE_CHECKING:
from fastapi import FastAPI
def _include_router(attr_name: str = "router") -> Callable[["FastAPI", object], None]:
def _register(app: "FastAPI", module: object) -> None:
app.include_router(getattr(module, attr_name))
return _register
def _mount_app(
prefix: str, attr_name: str = "app"
) -> Callable[["FastAPI", object], None]:
def _register(app: "FastAPI", module: object) -> None:
app.mount(path=prefix, app=getattr(module, attr_name))
return _register
@dataclass(frozen=True)
class LazyFeature:
name: str
module_path: str
path_prefixes: Tuple[str, ...]
register_fn: Callable[["FastAPI", object], None] = field(
default_factory=lambda: _include_router("router")
)
# For routes whose path has a leading parameter (e.g. /{server}/authorize)
# — startswith can't match those, so the matcher also checks endswith.
path_suffixes: Tuple[str, ...] = ()
LAZY_FEATURES: Tuple[LazyFeature, ...] = (
LazyFeature(
name="guardrails",
module_path="litellm.proxy.guardrails.guardrail_endpoints",
path_prefixes=(
"/guardrails",
"/v2/guardrails",
"/apply_guardrail",
"/policies/usage",
),
),
LazyFeature(
name="policies",
module_path="litellm.proxy.management_endpoints.policy_endpoints",
# Trailing slash to avoid matching /policies/... (policy_engine).
path_prefixes=("/policy/", "/utils/test_policies_and_guardrails"),
),
LazyFeature(
name="policy_engine",
module_path="litellm.proxy.policy_engine.policy_endpoints",
path_prefixes=("/policies",),
),
LazyFeature(
name="policy_resolve",
module_path="litellm.proxy.policy_engine.policy_resolve_endpoints",
path_prefixes=("/policies/resolve", "/policies/attachments/estimate-impact"),
),
LazyFeature(
name="agents",
module_path="litellm.proxy.agent_endpoints.endpoints",
path_prefixes=("/v1/agents", "/agents", "/agent/"),
),
LazyFeature(
name="a2a",
module_path="litellm.proxy.agent_endpoints.a2a_endpoints",
path_prefixes=("/a2a", "/v1/a2a"),
),
LazyFeature(
name="vector_stores",
module_path="litellm.proxy.vector_store_endpoints.endpoints",
path_prefixes=("/v1/vector_stores", "/vector_stores", "/v1/indexes"),
),
LazyFeature(
name="vector_store_management",
module_path="litellm.proxy.vector_store_endpoints.management_endpoints",
# Trailing slash to avoid matching /vector_stores/... (vector_stores).
path_prefixes=("/vector_store/", "/v1/vector_store/"),
),
LazyFeature(
name="vector_store_files",
# Routes appear under both /v1/vector_stores/{id}/files and the
# un-versioned form, so both prefixes must trigger the load.
module_path="litellm.proxy.vector_store_files_endpoints.endpoints",
path_prefixes=("/v1/vector_stores", "/vector_stores"),
),
LazyFeature(
name="tools",
module_path="litellm.proxy.management_endpoints.tool_management_endpoints",
path_prefixes=("/v1/tool", "/tool"),
),
LazyFeature(
name="search_tools",
module_path="litellm.proxy.search_endpoints.search_tool_management",
path_prefixes=("/search_tools",),
),
# mcp_management owns most /v1/mcp/* admin routes; mcp_app is the mounted
# streaming sub-app at /mcp.
LazyFeature(
name="mcp_management",
module_path="litellm.proxy.management_endpoints.mcp_management_endpoints",
path_prefixes=("/v1/mcp/",),
),
LazyFeature(
# Also serves /.well-known/oauth-* (OAuth metadata discovery).
# No /mcp/oauth prefix here: the mounted /mcp sub-app would
# shadow it, and there are no actual routes there anyway.
name="mcp_byok_oauth",
module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints",
path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"),
),
LazyFeature(
# Serves OAuth dance endpoints (/authorize, /token, /callback,
# /register) plus several /.well-known/ discovery URLs at the proxy
# root — needed for MCP-over-OAuth flows even before /mcp is hit.
name="mcp_discoverable",
module_path="litellm.proxy._experimental.mcp_server.discoverable_endpoints",
path_prefixes=(
"/.well-known/oauth-",
"/.well-known/openid-configuration",
"/.well-known/jwks.json",
"/authorize",
"/token",
"/callback",
"/register",
),
# Catches the /{mcp_server_name}/authorize|token|register variants.
path_suffixes=("/authorize", "/token", "/register"),
),
LazyFeature(
name="mcp_rest",
module_path="litellm.proxy._experimental.mcp_server.rest_endpoints",
path_prefixes=("/mcp-rest",),
),
LazyFeature(
# Hardcoded /mcp matches BASE_MCP_ROUTE; importing the constant
# here would defeat lazy loading.
name="mcp_app",
module_path="litellm.proxy._experimental.mcp_server.server",
path_prefixes=("/mcp",),
register_fn=_mount_app("/mcp", attr_name="app"),
),
LazyFeature(
name="config_overrides",
module_path="litellm.proxy.management_endpoints.config_override_endpoints",
path_prefixes=("/config_overrides",),
),
LazyFeature(
name="realtime",
module_path="litellm.proxy.realtime_endpoints.endpoints",
path_prefixes=("/openai/v1/realtime", "/v1/realtime", "/realtime"),
),
LazyFeature(
name="anthropic_passthrough",
module_path="litellm.proxy.anthropic_endpoints.endpoints",
path_prefixes=("/v1/messages", "/anthropic", "/api/event_logging"),
),
LazyFeature(
name="anthropic_skills",
module_path="litellm.proxy.anthropic_endpoints.skills_endpoints",
path_prefixes=("/v1/skills", "/skills"),
),
LazyFeature(
name="langfuse_passthrough",
module_path="litellm.proxy.vertex_ai_endpoints.langfuse_endpoints",
path_prefixes=("/langfuse",),
),
LazyFeature(
name="evals",
module_path="litellm.proxy.openai_evals_endpoints.endpoints",
path_prefixes=("/v1/evals", "/evals"),
),
LazyFeature(
name="claude_code_marketplace",
module_path="litellm.proxy.anthropic_endpoints.claude_code_endpoints",
path_prefixes=("/claude-code",),
register_fn=_include_router("claude_code_marketplace_router"),
),
LazyFeature(
name="scim",
module_path="litellm.proxy.management_endpoints.scim.scim_v2",
path_prefixes=("/scim",),
register_fn=_include_router("scim_router"),
),
LazyFeature(
name="cloudzero",
module_path="litellm.proxy.spend_tracking.cloudzero_endpoints",
path_prefixes=("/cloudzero",),
),
LazyFeature(
name="vantage",
module_path="litellm.proxy.spend_tracking.vantage_endpoints",
path_prefixes=("/vantage",),
),
LazyFeature(
name="usage_ai",
module_path="litellm.proxy.management_endpoints.usage_endpoints",
path_prefixes=("/usage/ai",),
),
LazyFeature(
name="prompts",
module_path="litellm.proxy.prompts.prompt_endpoints",
path_prefixes=("/prompts", "/utils/dotprompt_json_converter"),
),
LazyFeature(
name="jwt_mappings",
module_path="litellm.proxy.management_endpoints.jwt_key_mapping_endpoints",
path_prefixes=("/jwt/key/mapping",),
),
LazyFeature(
name="compliance",
module_path="litellm.proxy.management_endpoints.compliance_endpoints",
path_prefixes=("/compliance",),
),
LazyFeature(
name="access_groups",
module_path="litellm.proxy.management_endpoints.access_group_endpoints",
path_prefixes=("/access_group", "/v1/access_group", "/v1/unified_access_group"),
),
)
class LazyFeatureMiddleware:
"""ASGI middleware that imports + registers a feature router on first
matching request. Idempotent; once loaded, subsequent requests skip."""
def __init__(
self,
app,
fastapi_app: "FastAPI",
features: Tuple[LazyFeature, ...] = LAZY_FEATURES,
):
self.app = app
self._fastapi_app = fastapi_app
self._features = features
self._loaded: set = set()
# Per-feature locks so independent features can load in parallel.
self._locks: dict = {}
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
# Short-circuit once every feature has loaded.
if scope["type"] in ("http", "websocket") and len(self._loaded) < len(
self._features
):
path = scope.get("path", "")
for feat in self._features:
if feat.module_path in self._loaded:
continue
if any(path.startswith(p) for p in feat.path_prefixes) or any(
path.endswith(s) for s in feat.path_suffixes
):
await self._load(feat)
await self.app(scope, receive, send)
async def _load(self, feat: LazyFeature) -> None:
lock = self._locks.setdefault(feat.module_path, asyncio.Lock())
async with lock:
if feat.module_path in self._loaded:
return
try:
# Import on a thread (heavy modules take 1-3 s). register_fn
# mutates app.router.routes, so it stays on the loop thread.
loop = asyncio.get_running_loop()
module = await loop.run_in_executor(
None, importlib.import_module, feat.module_path
)
feat.register_fn(self._fastapi_app, module)
self._loaded.add(feat.module_path)
self._fastapi_app.openapi_schema = None
verbose_proxy_logger.info(
"Lazy-loaded optional feature %r (module: %s)",
feat.name,
feat.module_path,
)
except Exception as exc:
# Mark loaded anyway so we don't retry on every request.
self._loaded.add(feat.module_path)
verbose_proxy_logger.warning(
"Failed to lazy-load optional feature %r (module: %s): %s. "
"This feature's endpoints will return 404 until restart.",
feat.name,
feat.module_path,
exc,
)
def attach_lazy_features(app: "FastAPI") -> None:
app.add_middleware(LazyFeatureMiddleware, fastapi_app=app)

126
litellm/proxy/proxy_server.py
View file

@ -235,37 +235,11 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import (
router as mcp_byok_oauth_router,
)
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
router as mcp_discoverable_endpoints_router,
)
from litellm.proxy._experimental.mcp_server.rest_endpoints import (
router as mcp_rest_endpoints_router,
)
from litellm.proxy._experimental.mcp_server.server import app as mcp_app
from litellm.proxy._experimental.mcp_server.tool_registry import (
global_mcp_tool_registry,
)
from litellm.proxy._types import *
from litellm.proxy.agent_endpoints.a2a_endpoints import router as a2a_router
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.agent_endpoints.endpoints import router as agent_endpoints_router
from litellm.proxy.agent_endpoints.model_list_helpers import (
append_agents_to_model_group,
append_agents_to_model_info,
)
from litellm.proxy._lazy_features import attach_lazy_features
from litellm.proxy.analytics_endpoints.analytics_endpoints import (
router as analytics_router,
)
from litellm.proxy.anthropic_endpoints.claude_code_endpoints import (
claude_code_marketplace_router,
)
from litellm.proxy.anthropic_endpoints.endpoints import router as anthropic_router
from litellm.proxy.anthropic_endpoints.skills_endpoints import (
router as anthropic_skills_router,
)
from litellm.proxy.auth.auth_checks import (
ExperimentalUIJWTToken,
get_team_object,
@ -328,7 +302,6 @@ from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router
from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router
from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config
from litellm.proxy.google_endpoints.endpoints import router as google_router
from litellm.proxy.guardrails.guardrail_endpoints import router as guardrails_router
from litellm.proxy.guardrails.init_guardrails import (
init_guardrails_v2,
initialize_guardrails,
@ -344,9 +317,6 @@ from litellm.proxy.hooks.prompt_injection_detection import (
from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger
from litellm.proxy.image_endpoints.endpoints import router as image_router
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.proxy.management_endpoints.access_group_endpoints import (
router as access_group_router,
)
from litellm.proxy.management_endpoints.budget_management_endpoints import (
router as budget_management_router,
)
@ -360,12 +330,6 @@ from litellm.proxy.management_endpoints.common_utils import (
_user_has_admin_privileges,
admin_can_invite_user,
)
from litellm.proxy.management_endpoints.compliance_endpoints import (
router as compliance_router,
)
from litellm.proxy.management_endpoints.config_override_endpoints import (
router as config_override_router,
)
from litellm.proxy.management_endpoints.cost_tracking_settings import (
router as cost_tracking_settings_router,
)
@ -379,9 +343,6 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import (
router as internal_user_router,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import user_update
from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import (
router as jwt_key_mapping_router,
)
from litellm.proxy.management_endpoints.key_management_endpoints import (
delete_verification_tokens,
duration_in_seconds,
@ -390,9 +351,6 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
from litellm.proxy.management_endpoints.key_management_endpoints import (
router as key_management_router,
)
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
router as mcp_management_router,
)
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
router as model_access_group_management_router,
)
@ -407,11 +365,9 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
from litellm.proxy.management_endpoints.organization_endpoints import (
router as organization_router,
)
from litellm.proxy.management_endpoints.policy_endpoints import router as policy_router
from litellm.proxy.management_endpoints.router_settings_endpoints import (
router as router_settings_router,
)
from litellm.proxy.management_endpoints.scim.scim_v2 import scim_router
from litellm.proxy.management_endpoints.tag_management_endpoints import (
router as tag_management_router,
)
@ -423,15 +379,11 @@ from litellm.proxy.management_endpoints.team_endpoints import (
update_team,
validate_membership,
)
from litellm.proxy.management_endpoints.tool_management_endpoints import (
router as tool_management_router,
)
from litellm.proxy.memory.memory_endpoints import router as memory_router
from litellm.proxy.management_endpoints.ui_sso import (
get_disabled_non_admin_personal_key_creation,
)
from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router
from litellm.proxy.management_endpoints.usage_endpoints import router as usage_ai_router
from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import (
router as user_agent_analytics_router,
)
@ -441,7 +393,6 @@ from litellm.proxy.middleware.in_flight_requests_middleware import (
)
from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware
from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router
from litellm.proxy.openai_evals_endpoints.endpoints import router as evals_router
from litellm.proxy.openai_files_endpoints.files_endpoints import (
router as openai_files_router,
)
@ -461,27 +412,16 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
router as pass_through_router,
)
from litellm.proxy.policy_engine.policy_endpoints import router as policy_crud_router
from litellm.proxy.policy_engine.policy_resolve_endpoints import (
router as policy_resolve_router,
)
from litellm.proxy.prompts.prompt_endpoints import router as prompts_router
from litellm.proxy.public_endpoints import router as public_endpoints_router
from litellm.proxy.rag_endpoints.endpoints import router as rag_router
from litellm.proxy.realtime_endpoints.endpoints import router as webrtc_router
from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router
from litellm.proxy.response_api_endpoints.endpoints import router as response_router
from litellm.proxy.route_llm_request import route_request
from litellm.proxy.search_endpoints.endpoints import router as search_router
from litellm.proxy.search_endpoints.search_tool_management import (
router as search_tool_management_router,
)
from litellm.proxy.spend_tracking.cloudzero_endpoints import router as cloudzero_router
from litellm.proxy.spend_tracking.spend_management_endpoints import (
router as spend_management_router,
)
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
from litellm.proxy.spend_tracking.vantage_endpoints import router as vantage_router
from litellm.proxy.types_utils.utils import get_instance_fn
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
router as ui_crud_endpoints_router,
@ -511,16 +451,6 @@ from litellm.proxy.utils import (
prefetch_config_params,
update_spend,
)
from litellm.proxy.vector_store_endpoints.endpoints import router as vector_store_router
from litellm.proxy.vector_store_endpoints.management_endpoints import (
router as vector_store_management_router,
)
from litellm.proxy.vector_store_files_endpoints.endpoints import (
router as vector_store_files_router,
)
from litellm.proxy.vertex_ai_endpoints.langfuse_endpoints import (
router as langfuse_router,
)
from litellm.proxy.video_endpoints.endpoints import router as video_router
from litellm.router import (
AssistantsTypedDict,
@ -3854,11 +3784,19 @@ class ProxyConfig:
## MCP TOOLS
mcp_tools_config = config.get("mcp_tools", None)
if mcp_tools_config:
from litellm.proxy._experimental.mcp_server.tool_registry import (
global_mcp_tool_registry,
)
global_mcp_tool_registry.load_tools_from_config(mcp_tools_config)
## AGENTS
agent_config = config.get("agent_list", None)
if agent_config:
from litellm.proxy.agent_endpoints.agent_registry import (
global_agent_registry,
)
global_agent_registry.load_agents_from_config(agent_config) # type: ignore
mcp_servers_config = config.get("mcp_servers", None)
@ -10576,6 +10514,10 @@ async def model_info_v2(
verbose_proxy_logger.debug("all_models: %s", all_models)
# Append A2A agents to models list
from litellm.proxy.agent_endpoints.model_list_helpers import (
append_agents_to_model_info,
)
all_models = await append_agents_to_model_info(
models=all_models,
user_api_key_dict=user_api_key_dict,
@ -11425,6 +11367,10 @@ async def model_group_info(
)
# Append A2A agents to model groups
from litellm.proxy.agent_endpoints.model_list_helpers import (
append_agents_to_model_group,
)
model_groups = await append_agents_to_model_group(
model_groups=model_groups,
user_api_key_dict=user_api_key_dict,
@ -14230,65 +14176,40 @@ app.include_router(container_router)
app.include_router(search_router)
app.include_router(image_router)
app.include_router(fine_tuning_router)
app.include_router(vector_store_router)
app.include_router(vector_store_management_router)
app.include_router(vector_store_files_router)
app.include_router(credential_router)
app.include_router(llm_passthrough_router)
app.include_router(webrtc_router)
app.include_router(mcp_management_router)
app.include_router(mcp_byok_oauth_router)
app.include_router(anthropic_router)
app.include_router(anthropic_skills_router)
app.include_router(evals_router)
app.include_router(claude_code_marketplace_router)
app.include_router(google_router)
app.include_router(langfuse_router)
app.include_router(pass_through_router)
app.include_router(health_router)
app.include_router(key_management_router)
app.include_router(internal_user_router)
app.include_router(team_router)
app.include_router(ui_sso_router)
app.include_router(scim_router)
app.include_router(organization_router)
app.include_router(customer_router)
app.include_router(spend_management_router)
app.include_router(cloudzero_router)
app.include_router(vantage_router)
app.include_router(caching_router)
app.include_router(analytics_router)
app.include_router(guardrails_router)
app.include_router(policy_router)
app.include_router(usage_ai_router)
app.include_router(policy_crud_router)
app.include_router(policy_resolve_router)
app.include_router(search_tool_management_router)
app.include_router(prompts_router)
app.include_router(callback_management_endpoints_router)
app.include_router(debugging_endpoints_router)
app.include_router(ui_crud_endpoints_router)
app.include_router(openai_files_router)
app.include_router(team_callback_router)
app.include_router(jwt_key_mapping_router)
app.include_router(budget_management_router)
app.include_router(model_management_router)
app.include_router(model_access_group_management_router)
app.include_router(tag_management_router)
app.include_router(tool_management_router)
app.include_router(memory_router)
app.include_router(cost_tracking_settings_router)
app.include_router(router_settings_router)
app.include_router(fallback_management_router)
app.include_router(cache_settings_router)
app.include_router(config_override_router)
app.include_router(user_agent_analytics_router)
app.include_router(enterprise_router)
app.include_router(ui_discovery_endpoints_router)
app.include_router(agent_endpoints_router)
app.include_router(compliance_router)
app.include_router(a2a_router)
app.include_router(access_group_router)
# Eager: /models/{name}:method overlaps with the OpenAI /models endpoint.
app.include_router(google_router)
attach_lazy_features(app)
async def _stream_mcp_asgi_response(
@ -14521,8 +14442,3 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request):
f"Error handling dynamic MCP route for {mcp_server_name}: {str(e)}"
)
raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
app.mount(path=BASE_MCP_ROUTE, app=mcp_app)
app.include_router(mcp_rest_endpoints_router)
app.include_router(mcp_discoverable_endpoints_router)

14
tests/proxy_unit_tests/test_proxy_routes.py
View file

@ -39,6 +39,20 @@ def test_routes_on_litellm_proxy():
this prevents accidentelly deleting /threads, or /batches etc
"""
# Force-load lazy features so the test sees the full route set. Continue
# on per-feature import failure — the assertion below still catches
# missing-route regressions.
import importlib
from litellm.proxy._lazy_features import LAZY_FEATURES
for feat in LAZY_FEATURES:
try:
module = importlib.import_module(feat.module_path)
feat.register_fn(app, module)
except Exception as exc:
print(f"warning: failed to force-load {feat.name}: {exc}")
_all_routes = []
for route in app.routes:

252
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -5471,3 +5471,255 @@ async def test_reseed_warms_cache_even_on_zero_db_spend():
finally:
ps.spend_counter_cache = orig_counter
ps.prisma_client = orig_prisma
# ---------------------------------------------------------------------------
# Lazy feature loading (LazyFeatureMiddleware) — verifies that optional
# routers are NOT imported at module load and ARE imported on first request
# to a matching path prefix. The same module isn't re-imported on subsequent
# requests.
# ---------------------------------------------------------------------------
import sys
class TestLazyFeatureRegistry:
"""Sanity checks on the registry shape — guards against accidental edits."""
def test_registry_entries_have_required_fields(self):
from litellm.proxy._lazy_features import LAZY_FEATURES, LazyFeature
assert len(LAZY_FEATURES) > 0
for feat in LAZY_FEATURES:
assert isinstance(feat, LazyFeature)
assert feat.name
assert feat.module_path
assert feat.path_prefixes
assert all(p.startswith("/") for p in feat.path_prefixes)
assert callable(feat.register_fn)
def test_registry_names_unique(self):
from litellm.proxy._lazy_features import LAZY_FEATURES
names = [f.name for f in LAZY_FEATURES]
assert len(names) == len(set(names)), "duplicate feature names"
class TestLazyFeaturesNotImportedAtStartup:
"""
The whole point of the refactor: gated feature modules must NOT be
present in `sys.modules` immediately after `proxy_server` imports.
"""
def test_heavy_modules_absent_at_startup(self):
# Force a fresh `proxy_server` import in a subprocess so other tests
# in this run (which may have triggered lazy loads via the TestClient)
# don't pollute the result.
import subprocess
check = (
"import sys; "
"from litellm.proxy.proxy_server import app; " # noqa: F401
"heavy = ["
"'litellm.proxy._experimental.mcp_server.rest_endpoints',"
"'litellm.proxy._experimental.mcp_server.server',"
"'litellm.proxy.management_endpoints.config_override_endpoints',"
"'litellm.proxy.guardrails.guardrail_endpoints',"
"'litellm.proxy.openai_evals_endpoints.endpoints',"
"]; "
"still_present = [m for m in heavy if m in sys.modules]; "
"print('PRESENT_AT_STARTUP:', still_present)"
)
result = subprocess.run(
[sys.executable, "-c", check],
capture_output=True,
text=True,
timeout=120,
)
# Last non-empty line of stdout (skip warnings printed before)
out_lines = [
line for line in result.stdout.strip().splitlines() if line.strip()
]
report = next((line for line in out_lines if "PRESENT_AT_STARTUP" in line), "")
assert report, f"no report emitted (stderr: {result.stderr[-500:]})"
assert (
"PRESENT_AT_STARTUP: []" in report
), f"expected no heavy modules at startup, got: {report}"
class TestLazyFeatureMiddleware:
"""Behavior of the middleware itself, exercised in isolation."""
@pytest.mark.asyncio
async def test_first_request_triggers_load_subsequent_does_not(self):
from fastapi import FastAPI
from litellm.proxy._lazy_features import (
LazyFeature,
LazyFeatureMiddleware,
)
loads = []
def fake_register(app, module):
loads.append(getattr(module, "__name__", "?"))
feat = LazyFeature(
name="dummy",
module_path="json", # any always-importable stdlib module
path_prefixes=("/dummy",),
register_fn=fake_register,
)
# Build a minimal ASGI receiver to satisfy the middleware contract
async def downstream(scope, receive, send):
# echo back; no-op handler
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b""})
target_app = FastAPI()
mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,))
async def receive():
return {"type": "http.request", "body": b"", "more_body": False}
sent: list = []
async def send(message):
sent.append(message)
# First request matching the prefix triggers register
await mw(
{"type": "http", "path": "/dummy/x", "method": "GET", "headers": []},
receive,
send,
)
assert loads == ["json"]
# Second matching request must NOT re-register
sent.clear()
await mw(
{"type": "http", "path": "/dummy/y", "method": "GET", "headers": []},
receive,
send,
)
assert loads == ["json"], "register_fn called twice for the same feature"
# Non-matching path must not trigger anything
await mw(
{"type": "http", "path": "/unrelated", "method": "GET", "headers": []},
receive,
send,
)
assert loads == ["json"]
@pytest.mark.asyncio
async def test_concurrent_first_requests_only_register_once(self):
"""
Two requests to the same prefix arriving in parallel must result in
exactly one `register_fn` invocation — the lock prevents the import +
register from racing with itself.
"""
from fastapi import FastAPI
from litellm.proxy._lazy_features import (
LazyFeature,
LazyFeatureMiddleware,
)
loads = []
def slow_register(app, module):
loads.append(getattr(module, "__name__", "?"))
feat = LazyFeature(
name="dummy_concurrent",
module_path="json",
path_prefixes=("/dummy_c",),
register_fn=slow_register,
)
async def downstream(scope, receive, send):
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b""})
target_app = FastAPI()
mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,))
async def receive():
return {"type": "http.request", "body": b"", "more_body": False}
sent: list = []
async def send(message):
sent.append(message)
async def hit():
await mw(
{
"type": "http",
"path": "/dummy_c/x",
"method": "GET",
"headers": [],
},
receive,
send,
)
await asyncio.gather(hit(), hit(), hit(), hit(), hit())
assert loads == [
"json"
], f"expected one registration despite concurrent first hits, got {loads}"
@pytest.mark.asyncio
async def test_failing_import_does_not_loop(self):
"""
If a feature's module can't be imported, the middleware should mark it
loaded anyway so subsequent requests don't repeatedly retry the failing
import (which would amplify the cost on every request).
"""
from fastapi import FastAPI
from litellm.proxy._lazy_features import (
LazyFeature,
LazyFeatureMiddleware,
)
attempts = []
def fail_register(app, module):
attempts.append("called")
raise RuntimeError("boom")
feat = LazyFeature(
name="failing",
module_path="json",
path_prefixes=("/fail",),
register_fn=fail_register,
)
async def downstream(scope, receive, send):
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b""})
target_app = FastAPI()
mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,))
async def receive():
return {"type": "http.request", "body": b"", "more_body": False}
sent: list = []
async def send(message):
sent.append(message)
for _ in range(3):
await mw(
{"type": "http", "path": "/fail/x", "method": "GET", "headers": []},
receive,
send,
)
assert attempts == [
"called"
], f"failing register_fn should be invoked once, not on every request; got {attempts}"

15
tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py
View file

@ -786,8 +786,23 @@ class TestVectorStoreManagementEndpointsExist:
- POST /vector_store/info
- POST /vector_store/update
"""
import importlib
from litellm.proxy._lazy_features import LAZY_FEATURES
from litellm.proxy.proxy_server import app
# Force-register the lazy vector_store_management routes so the
# assertions can find them.
already_registered = any(
getattr(r, "path", None) == "/vector_store/new" for r in app.routes
)
if not already_registered:
for feat in LAZY_FEATURES:
if feat.name == "vector_store_management":
module = importlib.import_module(feat.module_path)
feat.register_fn(app, module)
break
# Define expected endpoints
expected_endpoints = [
("POST", "/vector_store/new"),