feat(telemetry:major): End to End Testing, Metric Capture, SQL Alchemy Injection

llama_stack/core/library_client.py

@@ -32,7 +32,7 @@ from termcolor import cprint
 
 from llama_stack.core.build import print_pip_install_help
 from llama_stack.core.configure import parse_and_maybe_upgrade_config
-from llama_stack.core.datatypes import Api, BuildConfig, BuildProvider, DistributionSpec
+from llama_stack.core.datatypes import BuildConfig, BuildProvider, DistributionSpec
 from llama_stack.core.request_headers import (
     PROVIDER_DATA_VAR,
     request_provider_data_context,
@@ -49,7 +49,6 @@ from llama_stack.core.utils.context import preserve_contexts_async_generator
 from llama_stack.core.utils.exec import in_notebook
 from llama_stack.log import get_logger
 
-
 logger = get_logger(name=__name__, category="core")
 
 T = TypeVar("T")

llama_stack/core/server/server.py

@@ -59,7 +59,6 @@ from llama_stack.core.utils.context import preserve_contexts_async_generator
 from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import Api
 
-
 from .auth import AuthenticationMiddleware
 from .quota import QuotaMiddleware
 
@@ -232,9 +231,7 @@ def create_dynamic_typed_route(func: Any, method: str, route: str) -> Callable:
 
             try:
                 if is_streaming:
-                    gen = preserve_contexts_async_generator(
-                        sse_generator(func(**kwargs)), [PROVIDER_DATA_VAR]
-                    )
+                    gen = preserve_contexts_async_generator(sse_generator(func(**kwargs)), [PROVIDER_DATA_VAR])
                     return StreamingResponse(gen, media_type="text/event-stream")
                 else:
                     value = func(**kwargs)
@@ -278,7 +275,7 @@ def create_dynamic_typed_route(func: Any, method: str, route: str) -> Callable:
             ]
         )
 
-    setattr(route_handler, "__signature__", sig.replace(parameters=new_params))
+    route_handler.__signature__ = sig.replace(parameters=new_params)
     return route_handler
 
 
@@ -405,6 +402,7 @@ def create_app() -> StackApp:
 
     if Api.telemetry in impls:
         impls[Api.telemetry].fastapi_middleware(app)
+        impls[Api.telemetry].sqlalchemy_instrumentation()
 
     # Load external APIs if configured
     external_apis = load_external_apis(config)

llama_stack/core/stack.py

@@ -359,7 +359,6 @@ class Stack:
         await refresh_registry_once(impls)
         self.impls = impls
 
-
     # safely access impls without raising an exception
     def get_impls(self) -> dict[Api, Any]:
         if self.impls is None:

llama_stack/core/telemetry/__init__.py

@@ -1,4 +1,10 @@
+# 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.
+# the root directory of this source tree.
+
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

llama_stack/core/telemetry/telemetry.py

@@ -4,46 +4,19 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 from abc import abstractmethod
+
 from fastapi import FastAPI
 from pydantic import BaseModel
-from typing import Any
 
 
 class TelemetryProvider(BaseModel):
     """
     TelemetryProvider standardizes how telemetry is provided to the application.
     """
+
     @abstractmethod
     def fastapi_middleware(self, app: FastAPI, *args, **kwargs):
         """
         Injects FastAPI middleware that instruments the application for telemetry.
         """
         ...
-
-    @abstractmethod
-    def custom_trace(self, name: str, *args, **kwargs) -> Any:
-        """
-        Creates a custom trace.
-        """
-        ...
-    
-    @abstractmethod
-    def record_count(self, name: str, *args, **kwargs):
-        """
-        Increments a counter metric.
-        """
-        ...
-    
-    @abstractmethod
-    def record_histogram(self, name: str, *args, **kwargs):
-        """
-        Records a histogram metric.
-        """
-        ...
-    
-    @abstractmethod
-    def record_up_down_counter(self, name: str, *args, **kwargs):
-        """
-        Records an up/down counter metric.
-        """
-        ...

llama_stack/core/telemetry/tracing.py

@@ -1,20 +0,0 @@
-# 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 abc import abstractmethod
-from fastapi import FastAPI
-from pydantic import BaseModel
-
-
-class TelemetryProvider(BaseModel):
-    """
-    TelemetryProvider standardizes how telemetry is provided to the application.
-    """
-    @abstractmethod
-    def fastapi_middleware(self, app: FastAPI, *args, **kwargs):
-        """
-        Injects FastAPI middleware that instruments the application for telemetry.
-        """
-        ...