|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the terms described in the LICENSE file in |
| 5 | +# the root directory of this source tree. |
| 6 | + |
| 7 | +import asyncio |
| 8 | +from importlib.metadata import version |
| 9 | +from typing import Any |
| 10 | + |
| 11 | +from pydantic import BaseModel |
| 12 | + |
| 13 | +from llama_stack.core.datatypes import StackRunConfig |
| 14 | +from llama_stack.core.external import load_external_apis |
| 15 | +from llama_stack.core.server.fastapi_router_registry import ( |
| 16 | + _ROUTER_FACTORIES, |
| 17 | + build_fastapi_router, |
| 18 | + get_router_routes, |
| 19 | +) |
| 20 | +from llama_stack.core.server.routes import get_all_api_routes |
| 21 | +from llama_stack.core.utils.config import redact_sensitive_fields |
| 22 | +from llama_stack.log import get_logger |
| 23 | +from llama_stack_api import ( |
| 24 | + Admin, |
| 25 | + Api, |
| 26 | + ApiFilter, |
| 27 | + HealthInfo, |
| 28 | + HealthResponse, |
| 29 | + HealthStatus, |
| 30 | + InspectProviderRequest, |
| 31 | + ListProvidersResponse, |
| 32 | + ListRoutesRequest, |
| 33 | + ListRoutesResponse, |
| 34 | + ProviderInfo, |
| 35 | + RouteInfo, |
| 36 | + VersionInfo, |
| 37 | +) |
| 38 | + |
| 39 | +logger = get_logger(name=__name__, category="core") |
| 40 | + |
| 41 | + |
| 42 | +class AdminImplConfig(BaseModel): |
| 43 | + run_config: StackRunConfig |
| 44 | + |
| 45 | + |
| 46 | +async def get_provider_impl(config, deps): |
| 47 | + impl = AdminImpl(config, deps) |
| 48 | + await impl.initialize() |
| 49 | + return impl |
| 50 | + |
| 51 | + |
| 52 | +class AdminImpl(Admin): |
| 53 | + def __init__(self, config: AdminImplConfig, deps): |
| 54 | + self.config = config |
| 55 | + self.deps = deps |
| 56 | + |
| 57 | + async def initialize(self) -> None: |
| 58 | + pass |
| 59 | + |
| 60 | + async def shutdown(self) -> None: |
| 61 | + logger.debug("AdminImpl.shutdown") |
| 62 | + pass |
| 63 | + |
| 64 | + # Provider management methods |
| 65 | + async def list_providers(self) -> ListProvidersResponse: |
| 66 | + run_config = self.config.run_config |
| 67 | + safe_config = StackRunConfig(**redact_sensitive_fields(run_config.model_dump())) |
| 68 | + providers_health = await self.get_providers_health() |
| 69 | + ret = [] |
| 70 | + for api, providers in safe_config.providers.items(): |
| 71 | + for p in providers: |
| 72 | + # Skip providers that are not enabled |
| 73 | + if p.provider_id is None: |
| 74 | + continue |
| 75 | + ret.append( |
| 76 | + ProviderInfo( |
| 77 | + api=api, |
| 78 | + provider_id=p.provider_id, |
| 79 | + provider_type=p.provider_type, |
| 80 | + config=p.config, |
| 81 | + health=providers_health.get(api, {}).get( |
| 82 | + p.provider_id, |
| 83 | + HealthResponse( |
| 84 | + status=HealthStatus.NOT_IMPLEMENTED, message="Provider does not implement health check" |
| 85 | + ), |
| 86 | + ), |
| 87 | + ) |
| 88 | + ) |
| 89 | + |
| 90 | + return ListProvidersResponse(data=ret) |
| 91 | + |
| 92 | + async def inspect_provider(self, request: InspectProviderRequest) -> ProviderInfo: |
| 93 | + all_providers = await self.list_providers() |
| 94 | + for p in all_providers.data: |
| 95 | + if p.provider_id == request.provider_id: |
| 96 | + return p |
| 97 | + |
| 98 | + raise ValueError(f"Provider {request.provider_id} not found") |
| 99 | + |
| 100 | + async def get_providers_health(self) -> dict[str, dict[str, HealthResponse]]: |
| 101 | + """Get health status for all providers. |
| 102 | +
|
| 103 | + Returns: |
| 104 | + Dict[str, Dict[str, HealthResponse]]: A dictionary mapping API names to provider health statuses. |
| 105 | + Each API maps to a dictionary of provider IDs to their health responses. |
| 106 | + """ |
| 107 | + providers_health: dict[str, dict[str, HealthResponse]] = {} |
| 108 | + |
| 109 | + # The timeout has to be long enough to allow all the providers to be checked, especially in |
| 110 | + # the case of the inference router health check since it checks all registered inference |
| 111 | + # providers. |
| 112 | + # The timeout must not be equal to the one set by health method for a given implementation, |
| 113 | + # otherwise we will miss some providers. |
| 114 | + timeout = 3.0 |
| 115 | + |
| 116 | + async def check_provider_health(impl: Any) -> tuple[str, HealthResponse] | None: |
| 117 | + # Skip special implementations (inspect/providers/admin) that don't have provider specs |
| 118 | + if not hasattr(impl, "__provider_spec__"): |
| 119 | + return None |
| 120 | + api_name = impl.__provider_spec__.api.name |
| 121 | + if not hasattr(impl, "health"): |
| 122 | + return ( |
| 123 | + api_name, |
| 124 | + HealthResponse( |
| 125 | + status=HealthStatus.NOT_IMPLEMENTED, message="Provider does not implement health check" |
| 126 | + ), |
| 127 | + ) |
| 128 | + |
| 129 | + try: |
| 130 | + health = await asyncio.wait_for(impl.health(), timeout=timeout) |
| 131 | + return api_name, health |
| 132 | + except TimeoutError: |
| 133 | + return ( |
| 134 | + api_name, |
| 135 | + HealthResponse( |
| 136 | + status=HealthStatus.ERROR, message=f"Health check timed out after {timeout} seconds" |
| 137 | + ), |
| 138 | + ) |
| 139 | + except Exception as e: |
| 140 | + return ( |
| 141 | + api_name, |
| 142 | + HealthResponse(status=HealthStatus.ERROR, message=f"Health check failed: {str(e)}"), |
| 143 | + ) |
| 144 | + |
| 145 | + # Create tasks for all providers |
| 146 | + tasks = [check_provider_health(impl) for impl in self.deps.values()] |
| 147 | + |
| 148 | + # Wait for all health checks to complete |
| 149 | + results = await asyncio.gather(*tasks) |
| 150 | + |
| 151 | + # Organize results by API and provider ID |
| 152 | + for result in results: |
| 153 | + if result is None: # Skip special implementations |
| 154 | + continue |
| 155 | + api_name, health_response = result |
| 156 | + providers_health[api_name] = health_response |
| 157 | + |
| 158 | + return providers_health |
| 159 | + |
| 160 | + # Inspect methods |
| 161 | + async def list_routes(self, request: ListRoutesRequest) -> ListRoutesResponse: |
| 162 | + run_config: StackRunConfig = self.config.run_config |
| 163 | + api_filter = request.api_filter |
| 164 | + |
| 165 | + # Helper function to determine if a route should be included based on api_filter |
| 166 | + # TODO: remove this once we've migrated all APIs to FastAPI routers |
| 167 | + def should_include_route(webmethod) -> bool: |
| 168 | + if api_filter is None: |
| 169 | + # Default: only non-deprecated APIs |
| 170 | + return not webmethod.deprecated |
| 171 | + elif api_filter == "deprecated": |
| 172 | + # Special filter: show deprecated routes regardless of their actual level |
| 173 | + return bool(webmethod.deprecated) |
| 174 | + else: |
| 175 | + # Filter by API level (non-deprecated routes only) |
| 176 | + return not webmethod.deprecated and webmethod.level == api_filter |
| 177 | + |
| 178 | + # Helper function to get provider types for an API |
| 179 | + def _get_provider_types(api: Api) -> list[str]: |
| 180 | + if api.value in ["providers", "inspect", "admin"]: |
| 181 | + return [] # These APIs don't have "real" providers - they're internal to the stack |
| 182 | + providers = run_config.providers.get(api.value, []) |
| 183 | + return [p.provider_type for p in providers] if providers else [] |
| 184 | + |
| 185 | + # Helper function to determine if a router route should be included based on api_filter |
| 186 | + def _should_include_router_route(route, router_prefix: str | None) -> bool: |
| 187 | + """Check if a router-based route should be included based on api_filter.""" |
| 188 | + # Check deprecated status |
| 189 | + route_deprecated = getattr(route, "deprecated", False) or False |
| 190 | + |
| 191 | + if api_filter is None: |
| 192 | + # Default: only non-deprecated routes |
| 193 | + return not route_deprecated |
| 194 | + elif api_filter == "deprecated": |
| 195 | + # Special filter: show deprecated routes regardless of their actual level |
| 196 | + return route_deprecated |
| 197 | + else: |
| 198 | + # Filter by API level (non-deprecated routes only) |
| 199 | + # Extract level from router prefix (e.g., "/v1" -> "v1") |
| 200 | + if router_prefix: |
| 201 | + prefix_level = router_prefix.lstrip("/") |
| 202 | + return not route_deprecated and prefix_level == api_filter |
| 203 | + return not route_deprecated |
| 204 | + |
| 205 | + ret = [] |
| 206 | + external_apis = load_external_apis(run_config) |
| 207 | + all_endpoints = get_all_api_routes(external_apis) |
| 208 | + |
| 209 | + # Process routes from APIs with FastAPI routers |
| 210 | + for api_name in _ROUTER_FACTORIES.keys(): |
| 211 | + api = Api(api_name) |
| 212 | + router = build_fastapi_router(api, None) # we don't need the impl here, just the routes |
| 213 | + if router: |
| 214 | + router_routes = get_router_routes(router) |
| 215 | + for route in router_routes: |
| 216 | + if _should_include_router_route(route, router.prefix): |
| 217 | + if route.methods is not None: |
| 218 | + available_methods = [m for m in route.methods if m != "HEAD"] |
| 219 | + if available_methods: |
| 220 | + ret.append( |
| 221 | + RouteInfo( |
| 222 | + route=route.path, |
| 223 | + method=available_methods[0], |
| 224 | + provider_types=_get_provider_types(api), |
| 225 | + ) |
| 226 | + ) |
| 227 | + |
| 228 | + # Process routes from legacy webmethod-based APIs |
| 229 | + for api, endpoints in all_endpoints.items(): |
| 230 | + # Skip APIs that have routers (already processed above) |
| 231 | + if api.value in _ROUTER_FACTORIES: |
| 232 | + continue |
| 233 | + |
| 234 | + # Always include provider, inspect, and admin APIs, filter others based on run config |
| 235 | + if api.value in ["providers", "inspect", "admin"]: |
| 236 | + ret.extend( |
| 237 | + [ |
| 238 | + RouteInfo( |
| 239 | + route=e.path, |
| 240 | + method=next(iter([m for m in e.methods if m != "HEAD"])), |
| 241 | + provider_types=[], # These APIs don't have "real" providers - they're internal to the stack |
| 242 | + ) |
| 243 | + for e, webmethod in endpoints |
| 244 | + if e.methods is not None and should_include_route(webmethod) |
| 245 | + ] |
| 246 | + ) |
| 247 | + else: |
| 248 | + providers = run_config.providers.get(api.value, []) |
| 249 | + if providers: # Only process if there are providers for this API |
| 250 | + ret.extend( |
| 251 | + [ |
| 252 | + RouteInfo( |
| 253 | + route=e.path, |
| 254 | + method=next(iter([m for m in e.methods if m != "HEAD"])), |
| 255 | + provider_types=[p.provider_type for p in providers], |
| 256 | + ) |
| 257 | + for e, webmethod in endpoints |
| 258 | + if e.methods is not None and should_include_route(webmethod) |
| 259 | + ] |
| 260 | + ) |
| 261 | + |
| 262 | + return ListRoutesResponse(data=ret) |
| 263 | + |
| 264 | + async def health(self) -> HealthInfo: |
| 265 | + return HealthInfo(status=HealthStatus.OK) |
| 266 | + |
| 267 | + async def version(self) -> VersionInfo: |
| 268 | + return VersionInfo(version=version("llama-stack")) |
0 commit comments