mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-17 10:59:48 +00:00
134 lines
4.6 KiB
Python
134 lines
4.6 KiB
Python
# 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.
|
|
|
|
import asyncio
|
|
import inspect
|
|
from functools import wraps
|
|
from typing import Any, AsyncGenerator, Callable, Type, TypeVar
|
|
|
|
from llama_stack.providers.utils.telemetry import tracing
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def trace_protocol(cls: Type[T]) -> Type[T]:
|
|
"""
|
|
A class decorator that automatically traces all methods in a protocol/base class
|
|
and its inheriting classes. Supports sync methods, async methods, and async generators.
|
|
"""
|
|
|
|
def trace_method(method: Callable) -> Callable:
|
|
is_async = asyncio.iscoroutinefunction(method)
|
|
is_async_gen = inspect.isasyncgenfunction(method)
|
|
|
|
@wraps(method)
|
|
async def async_gen_wrapper(
|
|
self: Any, *args: Any, **kwargs: Any
|
|
) -> AsyncGenerator:
|
|
class_name = self.__class__.__name__
|
|
method_name = f"{class_name}.{method.__name__}"
|
|
|
|
args_repr = [repr(arg) for arg in args]
|
|
kwargs_repr = [f"{k}={repr(v)}" for k, v in kwargs.items()]
|
|
signature = ", ".join(args_repr + kwargs_repr)
|
|
|
|
with tracing.span(
|
|
f"{class_name}.{method_name}",
|
|
{
|
|
"class": class_name,
|
|
"method": method_name,
|
|
"signature": signature,
|
|
"type": "async_generator",
|
|
},
|
|
) as span:
|
|
output = []
|
|
try:
|
|
async for item in method(self, *args, **kwargs):
|
|
output.append(item)
|
|
yield item
|
|
except Exception as e:
|
|
raise
|
|
finally:
|
|
span.set_attribute("output", output)
|
|
|
|
@wraps(method)
|
|
async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
|
|
class_name = self.__class__.__name__
|
|
method_name = f"{class_name}.{method.__name__}"
|
|
|
|
args_repr = [repr(arg) for arg in args]
|
|
kwargs_repr = [f"{k}={repr(v)}" for k, v in kwargs.items()]
|
|
signature = ", ".join(args_repr + kwargs_repr)
|
|
|
|
with tracing.span(
|
|
f"{class_name}.{method_name}",
|
|
{
|
|
"class": class_name,
|
|
"method": method_name,
|
|
"signature": signature,
|
|
"type": "async",
|
|
},
|
|
):
|
|
try:
|
|
result = await method(self, *args, **kwargs)
|
|
return result
|
|
except Exception as e:
|
|
raise
|
|
|
|
@wraps(method)
|
|
def sync_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
|
|
class_name = self.__class__.__name__
|
|
method_name = f"{class_name}.{method.__name__}"
|
|
|
|
args_repr = [repr(arg) for arg in args]
|
|
kwargs_repr = [f"{k}={repr(v)}" for k, v in kwargs.items()]
|
|
signature = ", ".join(args_repr + kwargs_repr)
|
|
|
|
with tracing.span(
|
|
f"{class_name}.{method_name}",
|
|
{
|
|
"class": class_name,
|
|
"method": method_name,
|
|
"signature": signature,
|
|
"type": "sync",
|
|
},
|
|
):
|
|
try:
|
|
result = method(self, *args, **kwargs)
|
|
return result
|
|
except Exception as e:
|
|
raise
|
|
|
|
if is_async_gen:
|
|
return async_gen_wrapper
|
|
elif is_async:
|
|
return async_wrapper
|
|
else:
|
|
return sync_wrapper
|
|
|
|
# Trace all existing methods in the base class
|
|
for name, method in vars(cls).items():
|
|
if inspect.isfunction(method) and not name.startswith("__"):
|
|
setattr(cls, name, trace_method(method))
|
|
|
|
# Store the original __init_subclass__ if it exists
|
|
original_init_subclass = getattr(cls, "__init_subclass__", None)
|
|
|
|
# Define a new __init_subclass__ to handle child classes
|
|
def __init_subclass__(cls_child, **kwargs): # noqa: N807
|
|
# Call original __init_subclass__ if it exists
|
|
if original_init_subclass:
|
|
original_init_subclass(**kwargs)
|
|
|
|
# Trace all methods defined in the child class
|
|
for name, method in vars(cls_child).items():
|
|
if inspect.isfunction(method) and not name.startswith("__"):
|
|
setattr(cls_child, name, trace_method(method))
|
|
|
|
# Set the new __init_subclass__
|
|
cls.__init_subclass__ = classmethod(__init_subclass__)
|
|
|
|
return cls
|