mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
swagger: stub-inject unloaded lazy features and warm on dropdown expand
This commit is contained in:
parent
0f8dd28542
commit
adcab435a4
2 changed files with 73 additions and 2 deletions
|
|
@ -8,8 +8,9 @@ omits each feature's routes until the feature is warmed.
|
|||
|
||||
import asyncio
|
||||
import importlib
|
||||
import sys
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Callable, Tuple
|
||||
from typing import TYPE_CHECKING, Callable, Dict, Tuple
|
||||
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
|
|
@ -46,6 +47,9 @@ class LazyFeature:
|
|||
# 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, ...] = ()
|
||||
# Keep the stub injected even after load — for mounted ASGI sub-apps
|
||||
# whose routes don't appear in the parent app's openapi spec.
|
||||
persistent_swagger_stub: bool = False
|
||||
|
||||
|
||||
LAZY_FEATURES: Tuple[LazyFeature, ...] = (
|
||||
|
|
@ -158,6 +162,7 @@ LAZY_FEATURES: Tuple[LazyFeature, ...] = (
|
|||
module_path="litellm.proxy._experimental.mcp_server.server",
|
||||
path_prefixes=("/mcp",),
|
||||
register_fn=_mount_app("/mcp", attr_name="app"),
|
||||
persistent_swagger_stub=True,
|
||||
),
|
||||
LazyFeature(
|
||||
name="config_overrides",
|
||||
|
|
@ -305,3 +310,33 @@ class LazyFeatureMiddleware:
|
|||
|
||||
def attach_lazy_features(app: "FastAPI") -> None:
|
||||
app.add_middleware(LazyFeatureMiddleware, fastapi_app=app)
|
||||
|
||||
|
||||
def inject_lazy_stubs(schema: Dict) -> Dict:
|
||||
"""Stub openapi entries for unloaded features so Swagger renders sections."""
|
||||
paths = schema.setdefault("paths", {})
|
||||
for feat in LAZY_FEATURES:
|
||||
if feat.module_path in sys.modules and not feat.persistent_swagger_stub:
|
||||
continue
|
||||
prefix = feat.path_prefixes[0]
|
||||
if prefix in paths:
|
||||
continue
|
||||
paths[prefix] = {
|
||||
"get": {
|
||||
"tags": [feat.name],
|
||||
"summary": feat.name,
|
||||
"responses": {"200": {"description": "OK"}},
|
||||
}
|
||||
}
|
||||
return schema
|
||||
|
||||
|
||||
def lazy_tag_to_prefix() -> Dict[str, str]:
|
||||
"""feature.name -> first prefix, used by the Swagger warmup JS plugin.
|
||||
Excludes persistent-stub features (mounted sub-apps) — warming them
|
||||
triggers a streaming hit and no useful new routes appear."""
|
||||
return {
|
||||
feat.name: feat.path_prefixes[0]
|
||||
for feat in LAZY_FEATURES
|
||||
if not feat.persistent_swagger_stub
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1033,6 +1033,11 @@ def get_openapi_schema():
|
|||
|
||||
openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema)
|
||||
|
||||
# Stub unloaded lazy features so they appear as Swagger sections.
|
||||
from litellm.proxy._lazy_features import inject_lazy_stubs
|
||||
|
||||
openapi_schema = inject_lazy_stubs(openapi_schema)
|
||||
|
||||
# Fix Swagger UI execute path error when server_root_path is set
|
||||
if server_root_path:
|
||||
openapi_schema["servers"] = [{"url": "/" + server_root_path.strip("/")}]
|
||||
|
|
@ -1059,6 +1064,11 @@ def custom_openapi():
|
|||
|
||||
openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema)
|
||||
|
||||
# Stub unloaded lazy features so they appear as Swagger sections.
|
||||
from litellm.proxy._lazy_features import inject_lazy_stubs
|
||||
|
||||
openapi_schema = inject_lazy_stubs(openapi_schema)
|
||||
|
||||
# Fix Swagger UI execute path error when server_root_path is set
|
||||
if server_root_path:
|
||||
openapi_schema["servers"] = [{"url": "/" + server_root_path.strip("/")}]
|
||||
|
|
@ -1486,14 +1496,40 @@ def mount_swagger_ui():
|
|||
|
||||
app.mount("/swagger", StaticFiles(directory=swagger_directory), name="swagger")
|
||||
|
||||
# On dropdown expand: one-time fetch to the prefix (triggers lazy load),
|
||||
# then spec re-download so real routes replace the stub. Raw JS (no
|
||||
# <script> tag) since it's injected inside the existing inline script.
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
from litellm.proxy._lazy_features import lazy_tag_to_prefix
|
||||
|
||||
_lazy_plugin_js = (
|
||||
"const TAG_TO_PREFIX = "
|
||||
+ json.dumps(lazy_tag_to_prefix())
|
||||
+ ";const warmedTags = new Set();const LazyLoadPlugin = () => ({"
|
||||
"statePlugins:{layout:{wrapActions:{show:(ori,sys)=>(...args)=>{"
|
||||
"const thing=args[0];let tag=null;"
|
||||
"if(Array.isArray(thing)){for(const t of thing)if(TAG_TO_PREFIX[t])tag=t;}"
|
||||
"if(tag&&!warmedTags.has(tag)){warmedTags.add(tag);"
|
||||
"fetch(TAG_TO_PREFIX[tag]).finally(()=>setTimeout(()=>sys.specActions.download(),800));}"
|
||||
"return ori(...args);}}}}});"
|
||||
)
|
||||
|
||||
def swagger_monkey_patch(*args, **kwargs):
|
||||
return get_swagger_ui_html(
|
||||
response = get_swagger_ui_html(
|
||||
*args,
|
||||
**kwargs,
|
||||
swagger_js_url=f"{custom_root_path_swagger_path}/swagger-ui-bundle.js",
|
||||
swagger_css_url=f"{custom_root_path_swagger_path}/swagger-ui.css",
|
||||
swagger_favicon_url=f"{custom_root_path_swagger_path}/favicon.png",
|
||||
)
|
||||
body = response.body.decode("utf-8")
|
||||
body = body.replace(
|
||||
"const ui = SwaggerUIBundle({",
|
||||
_lazy_plugin_js + "const ui = SwaggerUIBundle({plugins:[LazyLoadPlugin],",
|
||||
1,
|
||||
)
|
||||
return HTMLResponse(content=body)
|
||||
|
||||
applications.get_swagger_ui_html = swagger_monkey_patch
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue