Add an introspection "Api.inspect" API

llama_stack/distribution/datatypes.py

@@ -17,6 +17,53 @@ LLAMA_STACK_BUILD_CONFIG_VERSION = "v1"
 LLAMA_STACK_RUN_CONFIG_VERSION = "v1"
 
 
+RoutingKey = Union[str, List[str]]
+
+
+class GenericProviderConfig(BaseModel):
+    provider_type: str
+    config: Dict[str, Any]
+
+
+class RoutableProviderConfig(GenericProviderConfig):
+    routing_key: RoutingKey
+
+
+class PlaceholderProviderConfig(BaseModel):
+    """Placeholder provider config for API whose provider are defined in routing_table"""
+
+    providers: List[str]
+
+
+# Example: /inference, /safety
+class AutoRoutedProviderSpec(ProviderSpec):
+    provider_type: str = "router"
+    config_class: str = ""
+
+    docker_image: Optional[str] = None
+    routing_table_api: Api
+    module: str
+    provider_data_validator: Optional[str] = Field(
+        default=None,
+    )
+
+    @property
+    def pip_packages(self) -> List[str]:
+        raise AssertionError("Should not be called on AutoRoutedProviderSpec")
+
+
+# Example: /models, /shields
+@json_schema_type
+class RoutingTableProviderSpec(ProviderSpec):
+    provider_type: str = "routing_table"
+    config_class: str = ""
+    docker_image: Optional[str] = None
+
+    inner_specs: List[ProviderSpec]
+    module: str
+    pip_packages: List[str] = Field(default_factory=list)
+
+
 @json_schema_type
 class DistributionSpec(BaseModel):
     description: Optional[str] = Field(

llama_stack/distribution/distribution.py

@@ -46,6 +46,8 @@ def get_provider_registry() -> Dict[Api, Dict[str, ProviderSpec]]:
     for api in stack_apis():
         if api in routing_table_apis:
             continue
+        if api == Api.inspect:
+            continue
 
         name = api.name.lower()
         module = importlib.import_module(f"llama_stack.providers.registry.{name}")

llama_stack/distribution/inspect.py

@@ -0,0 +1,54 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Dict, List
+from llama_stack.apis.inspect import *  # noqa: F403
+
+
+from llama_stack.distribution.distribution import get_provider_registry
+from llama_stack.distribution.server.endpoints import get_all_api_endpoints
+from llama_stack.providers.datatypes import *  # noqa: F403
+
+
+def is_passthrough(spec: ProviderSpec) -> bool:
+    return isinstance(spec, RemoteProviderSpec) and spec.adapter is None
+
+
+class DistributionInspectImpl(Inspect):
+    def __init__(self):
+        pass
+
+    async def list_providers(self) -> Dict[str, List[ProviderInfo]]:
+        ret = {}
+        all_providers = get_provider_registry()
+        for api, providers in all_providers.items():
+            ret[api.value] = [
+                ProviderInfo(
+                    provider_type=p.provider_type,
+                    description="Passthrough" if is_passthrough(p) else "",
+                )
+                for p in providers.values()
+            ]
+
+        return ret
+
+    async def list_routes(self) -> Dict[str, List[RouteInfo]]:
+        ret = {}
+        all_endpoints = get_all_api_endpoints()
+
+        for api, endpoints in all_endpoints.items():
+            ret[api.value] = [
+                RouteInfo(
+                    route=e.route,
+                    method=e.method,
+                    providers=[],
+                )
+                for e in endpoints
+            ]
+        return ret
+
+    async def health(self) -> HealthInfo:
+        return HealthInfo(status="OK")

llama_stack/distribution/resolver.py

@@ -3,6 +3,7 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
+import importlib
 
 from typing import Any, Dict, List, Set
 
@@ -11,7 +12,8 @@ from llama_stack.distribution.distribution import (
     builtin_automatically_routed_apis,
     get_provider_registry,
 )
-from llama_stack.distribution.utils.dynamic import instantiate_provider
+from llama_stack.distribution.inspect import DistributionInspectImpl
+from llama_stack.distribution.utils.dynamic import instantiate_class_type
 
 
 async def resolve_impls_with_routing(run_config: StackRunConfig) -> Dict[Api, Any]:
@@ -57,7 +59,6 @@ async def resolve_impls_with_routing(run_config: StackRunConfig) -> Dict[Api, An
         if info.router_api.value not in apis_to_serve:
             continue
 
-        print("router_api", info.router_api)
         if info.router_api.value not in run_config.routing_table:
             raise ValueError(f"Routing table for `{source_api.value}` is not provided?")
 
@@ -104,6 +105,14 @@ async def resolve_impls_with_routing(run_config: StackRunConfig) -> Dict[Api, An
 
         impls[api] = impl
 
+    impls[Api.inspect] = DistributionInspectImpl()
+    specs[Api.inspect] = InlineProviderSpec(
+        api=Api.inspect,
+        provider_type="__distribution_builtin__",
+        config_class="",
+        module="",
+    )
+
     return impls, specs
 
 
@@ -127,3 +136,60 @@ def topological_sort(providers: List[ProviderSpec]) -> List[ProviderSpec]:
             dfs(a, visited, stack)
 
     return [by_id[x] for x in stack]
+
+
+# returns a class implementing the protocol corresponding to the Api
+async def instantiate_provider(
+    provider_spec: ProviderSpec,
+    deps: Dict[str, Any],
+    provider_config: Union[GenericProviderConfig, RoutingTable],
+):
+    module = importlib.import_module(provider_spec.module)
+
+    args = []
+    if isinstance(provider_spec, RemoteProviderSpec):
+        if provider_spec.adapter:
+            method = "get_adapter_impl"
+        else:
+            method = "get_client_impl"
+
+        assert isinstance(provider_config, GenericProviderConfig)
+        config_type = instantiate_class_type(provider_spec.config_class)
+        config = config_type(**provider_config.config)
+        args = [config, deps]
+    elif isinstance(provider_spec, AutoRoutedProviderSpec):
+        method = "get_auto_router_impl"
+
+        config = None
+        args = [provider_spec.api, deps[provider_spec.routing_table_api], deps]
+    elif isinstance(provider_spec, RoutingTableProviderSpec):
+        method = "get_routing_table_impl"
+
+        assert isinstance(provider_config, List)
+        routing_table = provider_config
+
+        inner_specs = {x.provider_type: x for x in provider_spec.inner_specs}
+        inner_impls = []
+        for routing_entry in routing_table:
+            impl = await instantiate_provider(
+                inner_specs[routing_entry.provider_type],
+                deps,
+                routing_entry,
+            )
+            inner_impls.append((routing_entry.routing_key, impl))
+
+        config = None
+        args = [provider_spec.api, inner_impls, routing_table, deps]
+    else:
+        method = "get_provider_impl"
+
+        assert isinstance(provider_config, GenericProviderConfig)
+        config_type = instantiate_class_type(provider_spec.config_class)
+        config = config_type(**provider_config.config)
+        args = [config, deps]
+
+    fn = getattr(module, method)
+    impl = await fn(*args)
+    impl.__provider_spec__ = provider_spec
+    impl.__provider_config__ = config
+    return impl

llama_stack/distribution/server/endpoints.py

@@ -11,12 +11,14 @@ from pydantic import BaseModel
 
 from llama_stack.apis.agents import Agents
 from llama_stack.apis.inference import Inference
+from llama_stack.apis.inspect import Inspect
 from llama_stack.apis.memory import Memory
 from llama_stack.apis.memory_banks import MemoryBanks
 from llama_stack.apis.models import Models
 from llama_stack.apis.safety import Safety
 from llama_stack.apis.shields import Shields
 from llama_stack.apis.telemetry import Telemetry
+
 from llama_stack.providers.datatypes import Api
 
 
@@ -38,6 +40,7 @@ def get_all_api_endpoints() -> Dict[Api, List[ApiEndpoint]]:
         Api.models: Models,
         Api.shields: Shields,
         Api.memory_banks: MemoryBanks,
+        Api.inspect: Inspect,
     }
 
     for api, protocol in protocols.items():

llama_stack/distribution/server/server.py

@@ -15,7 +15,6 @@ from collections.abc import (
     AsyncIterator as AsyncIteratorABC,
 )
 from contextlib import asynccontextmanager
-from http import HTTPStatus
 from ssl import SSLError
 from typing import Any, AsyncGenerator, AsyncIterator, Dict, get_type_hints, Optional
 
@@ -26,7 +25,6 @@ import yaml
 from fastapi import Body, FastAPI, HTTPException, Request, Response
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse, StreamingResponse
-from fastapi.routing import APIRoute
 from pydantic import BaseModel, ValidationError
 from termcolor import cprint
 from typing_extensions import Annotated
@@ -287,15 +285,6 @@ def main(
 
     app = FastAPI()
 
-    # Health check is added to enable deploying the docker container image on Kubernetes which require
-    # a health check that can return 200 for readiness and liveness check
-    class HealthCheck(BaseModel):
-        status: str = "OK"
-
-    @app.get("/healthcheck", status_code=HTTPStatus.OK, response_model=HealthCheck)
-    async def healthcheck():
-        return HealthCheck(status="OK")
-
     impls, specs = asyncio.run(resolve_impls_with_routing(config))
     if Api.telemetry in impls:
         setup_logger(impls[Api.telemetry])
@@ -307,6 +296,7 @@ def main(
     else:
         apis_to_serve = set(impls.keys())
 
+    apis_to_serve.add(Api.inspect)
     for api_str in apis_to_serve:
         api = Api(api_str)
 
@@ -340,14 +330,11 @@ def main(
                     )
                 )
 
-    for route in app.routes:
-        if isinstance(route, APIRoute):
-            cprint(
-                f"Serving {next(iter(route.methods))} {route.path}",
-                "white",
-                attrs=["bold"],
-            )
+        cprint(f"Serving API {api_str}", "white", attrs=["bold"])
+        for endpoint in endpoints:
+            cprint(f" {endpoint.method.upper()} {endpoint.route}", "white")
 
+    print("")
     app.exception_handler(RequestValidationError)(global_exception_handler)
     app.exception_handler(Exception)(global_exception_handler)
     signal.signal(signal.SIGINT, handle_sigint)

llama_stack/distribution/utils/dynamic.py

@@ -5,69 +5,9 @@
 # the root directory of this source tree.
 
 import importlib
-from typing import Any, Dict
-
-from llama_stack.distribution.datatypes import *  # noqa: F403
 
 
 def instantiate_class_type(fully_qualified_name):
     module_name, class_name = fully_qualified_name.rsplit(".", 1)
     module = importlib.import_module(module_name)
     return getattr(module, class_name)
-
-
-# returns a class implementing the protocol corresponding to the Api
-async def instantiate_provider(
-    provider_spec: ProviderSpec,
-    deps: Dict[str, Any],
-    provider_config: Union[GenericProviderConfig, RoutingTable],
-):
-    module = importlib.import_module(provider_spec.module)
-
-    args = []
-    if isinstance(provider_spec, RemoteProviderSpec):
-        if provider_spec.adapter:
-            method = "get_adapter_impl"
-        else:
-            method = "get_client_impl"
-
-        assert isinstance(provider_config, GenericProviderConfig)
-        config_type = instantiate_class_type(provider_spec.config_class)
-        config = config_type(**provider_config.config)
-        args = [config, deps]
-    elif isinstance(provider_spec, AutoRoutedProviderSpec):
-        method = "get_auto_router_impl"
-
-        config = None
-        args = [provider_spec.api, deps[provider_spec.routing_table_api], deps]
-    elif isinstance(provider_spec, RoutingTableProviderSpec):
-        method = "get_routing_table_impl"
-
-        assert isinstance(provider_config, List)
-        routing_table = provider_config
-
-        inner_specs = {x.provider_type: x for x in provider_spec.inner_specs}
-        inner_impls = []
-        for routing_entry in routing_table:
-            impl = await instantiate_provider(
-                inner_specs[routing_entry.provider_type],
-                deps,
-                routing_entry,
-            )
-            inner_impls.append((routing_entry.routing_key, impl))
-
-        config = None
-        args = [provider_spec.api, inner_impls, routing_table, deps]
-    else:
-        method = "get_provider_impl"
-
-        assert isinstance(provider_config, GenericProviderConfig)
-        config_type = instantiate_class_type(provider_spec.config_class)
-        config = config_type(**provider_config.config)
-        args = [config, deps]
-
-    fn = getattr(module, method)
-    impl = await fn(*args)
-    impl.__provider_spec__ = provider_spec
-    impl.__provider_config__ = config
-    return impl