feat: Bring Your Own API (BYOA) (#2228)

# What does this PR do?

Prototype on a new feature to allow new APIs to be plugged in Llama
Stack. Opened for early feedback on the approach and test appetite on
the functionality.

@ashwinb @raghotham open for early feedback, thanks!

---------

Signed-off-by: Sébastien Han <seb@redhat.com>
Co-authored-by: Ashwin Bharambe <ashwin.bharambe@gmail.com>

llama_stack/distribution/server/routes.py

@@ -12,10 +12,10 @@ from typing import Any
 from aiohttp import hdrs
 from starlette.routing import Route
 
+from llama_stack.apis.datatypes import Api, ExternalApiSpec
 from llama_stack.apis.tools import RAGToolRuntime, SpecialToolGroup
 from llama_stack.apis.version import LLAMA_STACK_API_VERSION
 from llama_stack.distribution.resolver import api_protocol_map
-from llama_stack.providers.datatypes import Api
 from llama_stack.schema_utils import WebMethod
 
 EndpointFunc = Callable[..., Any]
@@ -32,10 +32,12 @@ def toolgroup_protocol_map():
     }
 
 
-def get_all_api_routes() -> dict[Api, list[tuple[Route, WebMethod]]]:
+def get_all_api_routes(
+    external_apis: dict[Api, ExternalApiSpec] | None = None,
+) -> dict[Api, list[tuple[Route, WebMethod]]]:
     apis = {}
 
-    protocols = api_protocol_map()
+    protocols = api_protocol_map(external_apis)
     toolgroup_protocols = toolgroup_protocol_map()
     for api, protocol in protocols.items():
         routes = []
@@ -74,8 +76,8 @@ def get_all_api_routes() -> dict[Api, list[tuple[Route, WebMethod]]]:
     return apis
 
 
-def initialize_route_impls(impls: dict[Api, Any]) -> RouteImpls:
-    api_to_routes = get_all_api_routes()
+def initialize_route_impls(impls, external_apis: dict[Api, ExternalApiSpec] | None = None) -> RouteImpls:
+    api_to_routes = get_all_api_routes(external_apis)
     route_impls: RouteImpls = {}
 
     def _convert_path_to_regex(path: str) -> str:

llama_stack/distribution/server/server.py

@@ -40,6 +40,7 @@ from llama_stack.distribution.datatypes import (
     StackRunConfig,
 )
 from llama_stack.distribution.distribution import builtin_automatically_routed_apis
+from llama_stack.distribution.external import ExternalApiSpec, load_external_apis
 from llama_stack.distribution.request_headers import PROVIDER_DATA_VAR, User, request_provider_data_context
 from llama_stack.distribution.resolver import InvalidProviderError
 from llama_stack.distribution.server.routes import (
@@ -282,9 +283,10 @@ def create_dynamic_typed_route(func: Any, method: str, route: str) -> Callable:
 
 
 class TracingMiddleware:
-    def __init__(self, app, impls):
+    def __init__(self, app, impls, external_apis: dict[str, ExternalApiSpec]):
         self.app = app
         self.impls = impls
+        self.external_apis = external_apis
         # FastAPI built-in paths that should bypass custom routing
         self.fastapi_paths = ("/docs", "/redoc", "/openapi.json", "/favicon.ico", "/static")
 
@@ -301,7 +303,7 @@ class TracingMiddleware:
             return await self.app(scope, receive, send)
 
         if not hasattr(self, "route_impls"):
-            self.route_impls = initialize_route_impls(self.impls)
+            self.route_impls = initialize_route_impls(self.impls, self.external_apis)
 
         try:
             _, _, route_path, webmethod = find_matching_route(
@@ -486,7 +488,9 @@ def main(args: argparse.Namespace | None = None):
     else:
         setup_logger(TelemetryAdapter(TelemetryConfig(), {}))
 
-    all_routes = get_all_api_routes()
+    # Load external APIs if configured
+    external_apis = load_external_apis(config)
+    all_routes = get_all_api_routes(external_apis)
 
     if config.apis:
         apis_to_serve = set(config.apis)
@@ -505,7 +509,10 @@ def main(args: argparse.Namespace | None = None):
         api = Api(api_str)
 
         routes = all_routes[api]
-        impl = impls[api]
+        try:
+            impl = impls[api]
+        except KeyError as e:
+            raise ValueError(f"Could not find provider implementation for {api} API") from e
 
         for route, _ in routes:
             if not hasattr(impl, route.name):
@@ -536,7 +543,7 @@ def main(args: argparse.Namespace | None = None):
     app.exception_handler(Exception)(global_exception_handler)
 
     app.__llama_stack_impls__ = impls
-    app.add_middleware(TracingMiddleware, impls=impls)
+    app.add_middleware(TracingMiddleware, impls=impls, external_apis=external_apis)
 
     import uvicorn