wehub-resource-sync
53 lines
1.5 KiB
Python
53 lines
1.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
from pathlib import Path
|
|
from typing import List, Literal, Optional
|
|
|
|
# Third Party
|
|
from fastapi import APIRouter
|
|
|
|
# First Party
|
|
from lmcache.v1.utils.router_discovery import discover_api_routers
|
|
|
|
APICategory = Literal["common", "vllm", "controller"]
|
|
|
|
|
|
class APIRegistry:
|
|
"""
|
|
Automatically discovers and registers API routes by category
|
|
|
|
Categories:
|
|
- common: APIs that work for all components (metrics, logs, etc.)
|
|
- vllm: APIs specific to vLLM scheduler/worker
|
|
- controller: APIs specific to LMCache controller
|
|
"""
|
|
|
|
def __init__(self, app):
|
|
self.app = app
|
|
self.router = APIRouter()
|
|
|
|
def register_all_apis(self, categories: Optional[List[APICategory]] = None):
|
|
"""
|
|
Discover and register API modules from specified categories
|
|
|
|
Args:
|
|
categories: List of categories to register.
|
|
If None, registers all categories.
|
|
"""
|
|
if categories is None:
|
|
categories = ["common", "vllm", "controller"]
|
|
|
|
package_path = Path(__file__).parent
|
|
package_name = __package__
|
|
|
|
for category in categories:
|
|
category_path = package_path / category
|
|
if not category_path.exists():
|
|
continue
|
|
|
|
category_package = f"{package_name}.{category}"
|
|
for r in discover_api_routers(category_path, category_package):
|
|
self.router.include_router(r)
|
|
|
|
self.app.include_router(self.router)
|