Upload 11 files

aworld/trace/__init__.py

@@ -0,0 +1,80 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+import traceback
+from typing import Sequence, Union
+from aworld.trace.context_manager import TraceManager
+from aworld.trace.constants import RunType
+from aworld.logs.util import logger
+from aworld.trace.config import configure, ObservabilityConfig
+
+
+def get_tool_name(tool_name: str,
+                  action: Union['ActionModel', Sequence['ActionModel']]) -> tuple[str, RunType]:
+    if tool_name == "mcp" and action:
+        try:
+            if isinstance(action, (list, tuple)):
+                action = action[0]
+            mcp_name = action.action_name.split("__")[0]
+            return (mcp_name, RunType.MCP)
+        except ValueError:
+            logger.warning(traceback.format_exc())
+            return (tool_name, RunType.MCP)
+    return (tool_name, RunType.TOOL)
+
+
+def get_span_name_from_message(message: 'aworld.core.event.base.Message') -> tuple[str, RunType]:
+    from aworld.core.event.base import Constants
+    span_name = (message.receiver or message.id)
+    if message.category == Constants.AGENT:
+        return (span_name, RunType.AGNET)
+    if message.category == Constants.TOOL:
+        action = message.payload
+        if isinstance(action, (list, tuple)):
+            action = action[0]
+        if action:
+            tool_name, run_type = get_tool_name(action.tool_name, action)
+            return (tool_name, run_type)
+        return (span_name, RunType.TOOL)
+    return (span_name, RunType.OTHER)
+
+
+def message_span(message: 'aworld.core.event.base.Message' = None, attributes: dict = None):
+    if message:
+        span_name, run_type = get_span_name_from_message(message)
+        message_span_attribute = {
+            "event.payload": str(message.payload),
+            "event.topic": message.topic or "",
+            "event.receiver": message.receiver or "",
+            "event.sender": message.sender or "",
+            "event.category": message.category,
+            "event.id": message.id,
+            "event.session_id": message.session_id
+        }
+        message_span_attribute.update(attributes or {})
+        return GLOBAL_TRACE_MANAGER.span(
+            span_name=f"{run_type.value.lower()}_event_{span_name}",
+            attributes=message_span_attribute,
+            run_type=run_type
+        )
+    else:
+        raise ValueError("message_span message is None")
+
+
+GLOBAL_TRACE_MANAGER: TraceManager = TraceManager()
+span = GLOBAL_TRACE_MANAGER.span
+func_span = GLOBAL_TRACE_MANAGER.func_span
+auto_tracing = GLOBAL_TRACE_MANAGER.auto_tracing
+get_current_span = GLOBAL_TRACE_MANAGER.get_current_span
+new_manager = GLOBAL_TRACE_MANAGER.get_current_span
+
+__all__ = [
+    "span",
+    "func_span",
+    "message_span",
+    "auto_tracing",
+    "get_current_span",
+    "new_manager",
+    "RunType",
+    "configure",
+    "ObservabilityConfig"
+]

aworld/trace/auto_trace.py

@@ -0,0 +1,192 @@
+import ast
+import re
+import sys
+import warnings
+from importlib.abc import Loader, MetaPathFinder
+from importlib.machinery import ModuleSpec
+from importlib.util import spec_from_loader
+from types import ModuleType
+from typing import TYPE_CHECKING, Sequence, Union, Callable, Iterator, TypeVar, Any, cast
+
+from aworld.trace.base import log_trace_error
+from .rewrite_ast import compile_source
+
+if TYPE_CHECKING:
+    from .context_manager import TraceManager
+
+
+class AutoTraceModule:
+    """A class that represents a module being imported that should maybe be traced automatically."""
+
+    def __init__(self, module_name: str) -> None:
+        self._module_name = module_name
+        """Fully qualified absolute name of the module being imported."""
+
+    def need_auto_trace(self, prefix: Union[str, Sequence[str]]) -> bool:
+        """
+        Check if the module name starts with the given prefix.
+        """
+        if isinstance(prefix, str):
+            prefix = (prefix,)
+        pattern = '|'.join([get_module_pattern(p) for p in prefix])
+        return bool(re.match(pattern, self._module_name))
+
+
+class TraceImportFinder(MetaPathFinder):
+    """A class that implements the `find_spec` method of the `MetaPathFinder` protocol."""
+
+    def __init__(self, trace_manager: "TraceManager", module_funcs: Callable[[AutoTraceModule], bool],
+                 min_duration_ns: int) -> None:
+        self._trace_manager = trace_manager
+        self._modules_filter = module_funcs
+        self._min_duration_ns = min_duration_ns
+
+    def _find_plain_specs(
+            self, fullname: str, path: Sequence[str] = None, target: ModuleType = None
+    ) -> Iterator[ModuleSpec]:
+        """Yield module specs returned by other finders on `sys.meta_path`."""
+        for finder in sys.meta_path:
+            # Skip this finder or any like it to avoid infinite recursion.
+            if isinstance(finder, TraceImportFinder):
+                continue
+
+            try:
+                plain_spec = finder.find_spec(fullname, path, target)
+            except Exception:  # pragma: no cover
+                continue
+
+            if plain_spec:
+                yield plain_spec
+
+    def find_spec(self, fullname: str, path: Sequence[str], target=None) -> None:
+        """Find the spec for the given module name."""
+
+        for plain_spec in self._find_plain_specs(fullname, path, target):
+            # Get module specs returned by other finders on `sys.meta_path`
+            get_source = getattr(plain_spec.loader, 'get_source', None)
+            if not callable(get_source):
+                continue
+            try:
+                source = cast(str, get_source(fullname))
+            except Exception:
+                continue
+
+            if not source:
+                continue
+
+            filename = plain_spec.origin
+            if not filename:
+                try:
+                    filename = cast('str | None', plain_spec.loader.get_filename(fullname))
+                except Exception:
+                    pass
+            filename = filename or f'<{fullname}>'
+
+            if not self._modules_filter(AutoTraceModule(fullname)):
+                return None
+
+            try:
+                tree = ast.parse(source)
+            except Exception:
+                # Invalid source code. Try another one.
+                continue
+
+            try:
+                execute = compile_source(tree, filename, fullname, self._trace_manager, self._min_duration_ns)
+            except Exception:  # pragma: no cover
+                log_trace_error()
+                return None
+
+            loader = AutoTraceLoader(plain_spec, execute)
+            return spec_from_loader(fullname, loader)
+
+
+class AutoTraceLoader(Loader):
+    """
+    A class that implements the `exec_module` method of the `Loader` protocol.
+    """
+
+    def __init__(self, plain_spec: ModuleSpec, execute: Callable[[dict[str, Any]], None]) -> None:
+        self._plain_spec = plain_spec
+        self._execute = execute
+
+    def exec_module(self, module: ModuleType):
+        """Execute a modified AST of the module's source code in the module's namespace.
+        """
+        self._execute(module.__dict__)
+
+    def create_module(self, spec: ModuleSpec):
+        return None
+
+    def get_code(self, _name: str):
+        """`python -m` uses the `runpy` module which calls this method instead of going through the normal protocol.
+        So return some code which can be executed with the module namespace.
+        Here `__loader__` will be this object, i.e. `self`.
+        source = '__loader__.execute(globals())'
+        return compile(source, '<string>', 'exec', dont_inherit=True)
+        """
+
+    def __getattr__(self, item: str):
+        """Forward some methods to the plain spec's loader (likely a `SourceFileLoader`) if they exist."""
+        if item in {'get_filename', 'is_package'}:
+            return getattr(self.plain_spec.loader, item)
+        raise AttributeError(item)
+
+
+def convert_to_modules_func(modules: Sequence[str]) -> Callable[[AutoTraceModule], bool]:
+    """Convert a sequence of module names to a function that checks if a module name starts with any of the given module names.
+    """
+    return lambda module: module.need_auto_trace(modules)
+
+
+def get_module_pattern(module: str):
+    """
+    Get the regex pattern for the given module name.
+    """
+
+    if not re.match(r'[\w.]+$', module, re.UNICODE):
+        return module
+    module = re.escape(module)
+    return rf'{module}($|\.)'
+
+
+def install_auto_tracing(trace_manager: "TraceManager",
+                         modules: Union[Sequence[str],
+                         Callable[[AutoTraceModule], bool]],
+                         min_duration_seconds: float
+                         ) -> None:
+    """
+    Automatically trace the execution of a function.
+    """
+    if isinstance(modules, Sequence):
+        module_funcs = convert_to_modules_func(modules)
+    else:
+        module_funcs = modules
+
+    if not callable(module_funcs):
+        raise TypeError('modules must be a list of strings or a callable')
+
+    for module in list(sys.modules.values()):
+        try:
+            auto_trace_module = AutoTraceModule(module.__name__)
+        except Exception:
+            continue
+
+        if module_funcs(auto_trace_module):
+            warnings.warn(f'The module {module.__name__!r} matches modules to trace, but it has already been imported. '
+                          f'Call `auto_tracing` earlier',
+                          stacklevel=2,
+                          )
+
+    min_duration_ns = int(min_duration_seconds * 1_000_000_000)
+    trace_manager = trace_manager.new_manager('auto_tracing')
+    finder = TraceImportFinder(trace_manager, module_funcs, min_duration_ns)
+    sys.meta_path.insert(0, finder)
+
+
+T = TypeVar('T')
+
+
+def not_auto_trace(x: T) -> T:
+    """Decorator to prevent a function/class from being traced by `auto_tracing`"""
+    return x

aworld/trace/base.py

@@ -0,0 +1,422 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Any, Iterator, Union, Sequence, Protocol, Iterable
+from enum import Enum
+from weakref import WeakSet
+from dataclasses import dataclass, field
+from aworld.logs.util import trace_logger
+
+
+class TraceProvider(ABC):
+
+    @abstractmethod
+    def get_tracer(
+            self,
+            name: str,
+            version: Optional[str] = None
+    ) -> "Tracer":
+        """Returns a `Tracer` for use by the given name.
+
+        This function may return different `Tracer` types (e.g. a no-op tracer
+        vs.  a functional tracer).
+
+        Args:
+            name: The uniquely identifiable name for instrumentation
+                scope, such as instrumentation library, package, module or class name.
+                ``__name__`` may not be used as this can result in
+                different tracer names if the tracers are in different files.
+                It is better to use a fixed string that can be imported where
+                needed and used consistently as the name of the tracer.
+
+                This should *not* be the name of the module that is
+                instrumented but the name of the module doing the instrumentation.
+                E.g., instead of ``"requests"``, use
+                ``"opentelemetry.instrumentation.requests"``.
+
+            version: Optional. The version string of the
+                instrumenting library.  Usually this should be the same as
+                ``importlib.metadata.version(instrumenting_library_name)``
+        """
+
+    @abstractmethod
+    def shutdown(self) -> None:
+        """Shuts down the provider and all its resources.
+        This method should be called when the application is shutting down.
+        """
+
+    @abstractmethod
+    def force_flush(self, timeout: Optional[float] = None) -> bool:
+        """Forces all the data to be sent to the backend.
+        This method should be called when the application is shutting down.
+        Args:
+            timeout: The maximum time to wait for the data to be sent.
+        Returns:
+            True if the data was sent successfully, False otherwise.
+        """
+
+    @abstractmethod
+    def get_current_span(self) -> Optional["Span"]:
+        """Returns the current span from the current context.
+        Returns:
+            The current span from the current context.
+        """
+
+
+class SpanType(Enum):
+    """Specifies additional details on how this span relates to its parent span.
+    """
+
+    #: Default value. Indicates that the span is used internally in the
+    # application.
+    INTERNAL = 0
+
+    #: Indicates that the span describes an operation that handles a remote
+    # request.
+    SERVER = 1
+
+    #: Indicates that the span describes a request to some remote service.
+    CLIENT = 2
+
+    #: Indicates that the span describes a producer sending a message to a
+    #: broker. Unlike client and server, there is usually no direct critical
+    #: path latency relationship between producer and consumer spans.
+    PRODUCER = 3
+
+    #: Indicates that the span describes a consumer receiving a message from a
+    #: broker. Unlike client and server, there is usually no direct critical
+    #: path latency relationship between producer and consumer spans.
+    CONSUMER = 4
+
+
+AttributeValueType = Union[
+    str,
+    bool,
+    int,
+    float,
+    Sequence[str],
+    Sequence[bool],
+    Sequence[int],
+    Sequence[float],
+]
+
+
+class Tracer(ABC):
+    """Handles span creation and in-process context propagation.
+    """
+
+    @abstractmethod
+    def start_span(
+            self,
+            name: str,
+            span_type: SpanType = SpanType.INTERNAL,
+            attributes: dict[str, AttributeValueType] = None,
+            start_time: Optional[int] = None,
+            record_exception: bool = True,
+            set_status_on_exception: bool = True,
+            trace_context: Optional["TraceContext"] = None,
+    ) -> "Span":
+        """Starts and returns a new Span.
+        Args:
+            name: The name of the span.
+            kind: The span's kind (relationship to parent). Note that is
+                meaningful even if there is no parent.
+            attributes: The span's attributes.
+            start_time: Sets the start time of a span
+            record_exception: Whether to record any exceptions raised within the
+                context as error event on the span.
+            set_status_on_exception: Only relevant if the returned span is used
+                in a with/context manager. Defines whether the span status will
+                be automatically set to ERROR when an uncaught exception is
+                raised in the span with block. The span status won't be set by
+                this mechanism if it was previously set manually.
+            trace_context: The trace context to use for the span. If not
+                provided, the current trace context will be used.
+        """
+
+    @abstractmethod
+    def start_as_current_span(
+            self,
+            name: str,
+            span_type: SpanType = SpanType.INTERNAL,
+            attributes: dict[str, AttributeValueType] = None,
+            start_time: Optional[int] = None,
+            record_exception: bool = True,
+            set_status_on_exception: bool = True,
+            end_on_exit: bool = True,
+            trace_context: Optional['TraceContext'] = None
+    ) -> Iterator["Span"]:
+        """Context manager for creating a new span and set it
+        as the current span in this tracer's context.
+
+        Example::
+
+            with tracer.start_as_current_span("one") as parent:
+                parent.add_event("parent's event")
+                with tracer.start_as_current_span("two") as child:
+                    child.add_event("child's event")
+                    trace.get_current_span()  # returns child
+                trace.get_current_span()      # returns parent
+            trace.get_current_span()          # returns previously active span
+
+            This can also be used as a decorator::
+            @tracer.start_as_current_span("name")
+            def function():
+
+        Args:
+            name: The name of the span to be created.
+            kind: The span's kind (relationship to parent). Note that is
+                meaningful even if there is no parent.
+            attributes: The span's attributes.
+            start_time: Sets the start time of a span
+            record_exception: Whether to record any exceptions raised within the
+                context as error event on the span.
+            set_status_on_exception: Only relevant if the returned span is used
+                in a with/context manager. Defines whether the span status will
+                be automatically set to ERROR when an uncaught exception is
+                raised in the span with block. The span status won't be set by
+                this mechanism if it was previously set manually.
+            end_on_exit: Whether to end the span automatically when leaving the
+                context manager.
+            trace_context: The trace context to use for the span. If not
+                provided, the current trace context will be used.
+    """
+
+
+class Span(ABC):
+    """A Span represents a single operation within a trace.
+    """
+
+    @abstractmethod
+    def end(self, end_time: Optional[int] = None) -> None:
+        """Sets the current time as the span's end time.
+
+        The span's end time is the wall time at which the operation finished.
+
+        Only the first call to `end` should modify the span, and
+        implementations are free to ignore or raise on further calls.
+        """
+
+    @abstractmethod
+    def set_attribute(self, key: str, value: Any) -> None:
+        """Sets an attribute on the Span.
+        Args:
+            key: The attribute key.
+            value: The attribute value.
+        """
+
+    @abstractmethod
+    def set_attributes(self, attributes: dict[str, Any]) -> None:
+        """Sets multiple attributes on the Span.
+        Args:
+            attributes: A dictionary of attributes to set.
+        """
+
+    @abstractmethod
+    def is_recording(self) -> bool:
+        """Returns whether this span will be recorded.
+        Returns true if this Span is active and recording information like attributes using set_attribute.
+        """
+
+    @abstractmethod
+    def record_exception(
+            self,
+            exception: BaseException,
+            attributes: dict[str, Any] = None,
+            timestamp: Optional[int] = None,
+            escaped: bool = False,
+    ) -> None:
+        """Records an exception in the span.
+        Args:
+            exception: The exception to record.
+            attributes: A dictionary of attributes to set on the exception event.
+            timestamp: The timestamp of the exception.
+            escaped: Whether the exception was escaped.
+        """
+
+    @abstractmethod
+    def get_trace_id(self) -> str:
+        """Returns the trace ID of the span.
+        Returns:
+            The trace ID of the span.
+        """
+
+    @abstractmethod
+    def get_span_id(self) -> str:
+        """Returns the ID of the span.
+        Returns:
+            The ID of the span.
+        """
+
+    def _add_to_open_spans(self) -> None:
+        """Add the current span to OPEN_SPANS."""
+        _OPEN_SPANS.add(self)
+
+    def _remove_from_open_spans(self) -> None:
+        """Remove the current span from OPEN_SPANS."""
+        _OPEN_SPANS.discard(self)
+
+
+class NoOpSpan(Span):
+    """No-op implementation of `Span`."""
+
+    def end(self, end_time: Optional[int] = None) -> None:
+        pass
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        pass
+
+    def set_attributes(self, attributes: dict[str, Any]) -> None:
+        pass
+
+    def is_recording(self) -> bool:
+        return False
+
+    def record_exception(
+            self,
+            exception: BaseException,
+            attributes: dict[str, Any] = None,
+            timestamp: Optional[int] = None,
+            escaped: bool = False,
+    ) -> None:
+        pass
+
+    def get_trace_id(self) -> str:
+        return ""
+
+    def get_span_id(self) -> str:
+        return ""
+
+
+class NoOpTracer(Tracer):
+    """No-op implementation of `Tracer`."""
+
+    def start_span(
+            self,
+            name: str,
+            span_type: SpanType = SpanType.INTERNAL,
+            attributes: dict[str, AttributeValueType] = None,
+            start_time: Optional[int] = None,
+            record_exception: bool = True,
+            set_status_on_exception: bool = True,
+            trace_context: Optional["TraceContext"] = None,
+    ) -> Span:
+        return NoOpSpan()
+
+    def start_as_current_span(
+            self,
+            name: str,
+            span_type: SpanType = SpanType.INTERNAL,
+            attributes: dict[str, AttributeValueType] = None,
+            start_time: Optional[int] = None,
+            record_exception: bool = True,
+            set_status_on_exception: bool = True,
+            end_on_exit: bool = True,
+            trace_context: Optional['TraceContext'] = None
+    ) -> Iterator[Span]:
+        yield NoOpSpan()
+
+
+class Carrier(Protocol):
+    """Carrier is a protocol that represents a carrier for trace context.
+    """
+
+    def get(self, key: str) -> Optional[str]:
+        """Returns the value of the given key from the carrier.
+        Args:
+            key: The key to get the value for.
+        Returns:
+            The value of the given key from the carrier.
+        """
+
+    def set(self, key: str, value: str) -> None:
+        """Sets the value of the given key in the carrier.
+        Args:
+            key: The key to set the value for.
+            value: The value to set.
+        """
+
+    def keys(self) -> Iterable[str]:
+        """Returns an iterable of keys in the carrier.
+        Returns:
+            An iterable of keys in the carrier.
+        """
+
+
+@dataclass(frozen=True)
+class TraceContext:
+    """TraceContext is a class that represents a trace context.
+    """
+    trace_id: str
+    span_id: str
+    version: str = "00"
+    trace_flags: str = "01"
+    attributes: dict[str, Any] = field(default_factory=dict)
+
+
+class Propagator(ABC):
+    """Propagator is a protocol that represents a propagator for trace context.
+    """
+
+    def _get_value(self, carrier: Carrier, name: str) -> str:
+        """
+        Get value from carrier.
+        Args:
+            carrier: The carrier to get value from.
+            name: The name of the value.
+        Returns:
+            The value of the name.
+        """
+        return carrier.get(name) or carrier.get('HTTP_' + name.upper().replace('-', '_'))
+
+    @abstractmethod
+    def extract(self, carrier: Carrier) -> Optional[TraceContext]:
+        """Extracts a trace context from the given carrier.
+        Args:
+            carrier: The carrier to extract the trace context from.
+        Returns:
+            The trace context extracted from the carrier.
+        """
+    @abstractmethod
+    def inject(self, trace_context: TraceContext, carrier: Carrier) -> None:
+        """Injects a trace context into the given carrier.
+        Args:
+            trace_context: The trace context to inject.
+            carrier: The carrier to inject the trace context into.
+        """
+
+
+_GLOBAL_TRACER_PROVIDER: Optional[TraceProvider] = None
+_OPEN_SPANS: WeakSet[Span] = WeakSet()
+
+
+def set_tracer_provider(provider: TraceProvider):
+    """
+    Set the global tracer provider.
+    """
+    global _GLOBAL_TRACER_PROVIDER
+    _GLOBAL_TRACER_PROVIDER = provider
+
+
+def get_tracer_provider() -> TraceProvider:
+    """
+    Get the global tracer provider.
+    """
+    global _GLOBAL_TRACER_PROVIDER
+    if _GLOBAL_TRACER_PROVIDER is None:
+        raise Exception("No tracer provider has been set.")
+    return _GLOBAL_TRACER_PROVIDER
+
+
+def get_tracer_provider_silent():
+    try:
+        return get_tracer_provider()
+    except Exception:
+        return None
+
+
+def log_trace_error():
+    """
+    Log an error with traceback information.
+    """
+    trace_logger.exception(
+        'This is logging the trace internal error.',
+    )

aworld/trace/config.py

@@ -0,0 +1,93 @@
+import os
+from pydantic import BaseModel
+from typing import Sequence, Optional
+from aworld.trace.span_cosumer import SpanConsumer
+from logging import Logger
+from aworld.logs.util import trace_logger
+from aworld.trace.context_manager import trace_configure
+from aworld.metrics.context_manager import MetricContext
+from aworld.logs.log import set_log_provider, instrument_logging
+from aworld.trace.instrumentation.uni_llmmodel import LLMModelInstrumentor
+from aworld.trace.instrumentation.eventbus import EventBusInstrumentor
+
+
+class ObservabilityConfig(BaseModel):
+    '''
+    Observability configuration
+    '''
+    class Config:
+        arbitrary_types_allowed = True
+    trace_provider: Optional[str] = "otlp"
+    trace_backends: Optional[Sequence[str]] = ["memory"]
+    trace_base_url: Optional[str] = None
+    trace_write_token: Optional[str] = None
+    trace_span_consumers: Optional[Sequence[SpanConsumer]] = None
+    # whether to start the trace service
+    trace_server_enabled: Optional[bool] = False
+    trace_server_port: Optional[int] = 7079
+    metrics_provider: Optional[str] = None
+    metrics_backend: Optional[str] = None
+    metrics_base_url: Optional[str] = None
+    metrics_write_token: Optional[str] = None
+    # whether to instrument system metrics
+    metrics_system_enabled: Optional[bool] = False
+    logs_provider: Optional[str] = None
+    logs_backend: Optional[str] = None
+    logs_base_url: Optional[str] = None
+    logs_write_token: Optional[str] = None
+    # The loggers that need to record the log as a span
+    logs_trace_instrumented_loggers: Sequence[Logger] = [trace_logger]
+
+
+def configure(config: ObservabilityConfig = None):
+    if config is None:
+        config = ObservabilityConfig()
+    _trace_configure(config)
+    _metrics_configure(config)
+    _log_configure(config)
+    LLMModelInstrumentor().instrument()
+    EventBusInstrumentor().instrument()
+
+
+def _trace_configure(config: ObservabilityConfig):
+    if not config.trace_base_url and config.trace_provider == "otlp":
+        if "logfire" in config.trace_backends:
+            config.trace_base_url = os.getenv("LOGFIRE_WRITE_TOKEN")
+        elif os.getenv("OTLP_TRACES_ENDPOINT"):
+            config.trace_base_url = os.getenv("OTLP_TRACES_ENDPOINT")
+            config.trace_backends.append("other_otlp")
+
+    trace_configure(
+        provider=config.trace_provider,
+        backends=config.trace_backends,
+        base_url=config.trace_base_url,
+        write_token=config.trace_write_token,
+        span_consumers=config.trace_span_consumers,
+        server_enabled=config.trace_server_enabled,
+        server_port=config.trace_server_port
+    )
+
+
+def _metrics_configure(config: ObservabilityConfig):
+    if config.metrics_provider and config.metrics_backend:
+        MetricContext.configure(
+            provider=config.metrics_provider,
+            backend=config.metrics_backend,
+            base_url=config.metrics_base_url,
+            write_token=config.metrics_write_token,
+            metrics_system_enabled=config.metrics_system_enabled
+        )
+
+
+def _log_configure(config: ObservabilityConfig):
+    if config.logs_provider and config.logs_backend:
+        if config.logs_backend == "logfire" and not config.logs_write_token:
+            config.logs_write_token = os.getenv("LOGFIRE_WRITE_TOKEN")
+        set_log_provider(provider=config.logs_provider,
+                         backend=config.logs_backend,
+                         base_url=config.logs_base_url,
+                         write_token=config.logs_write_token)
+
+    if config.logs_trace_instrumented_loggers:
+        for logger in config.logs_trace_instrumented_loggers:
+            instrument_logging(logger)

aworld/trace/constants.py

@@ -0,0 +1,26 @@
+from enum import Enum
+
+ATTRIBUTES_NAMESPACE = 'aworld'
+"""Namespace within OTEL attributes used by aworld."""
+
+ATTRIBUTES_MESSAGE_KEY = f'{ATTRIBUTES_NAMESPACE}.msg'
+"""The formatted message for a log."""
+
+ATTRIBUTES_MESSAGE_TEMPLATE_KEY = f'{ATTRIBUTES_NAMESPACE}.msg_template'
+"""The template for a log message."""
+
+ATTRIBUTES_MESSAGE_RUN_TYPE_KEY = f'{ATTRIBUTES_NAMESPACE}.run_type'
+"""The template for a log message."""
+
+MESSAGE_FORMATTED_VALUE_LENGTH_LIMIT = 128
+"""Maximum number of characters for formatted values in a trace message."""
+
+
+class RunType(Enum):
+    '''Span run type supported in the framework
+    '''
+    AGNET = "AGENT"
+    TOOL = "TOOL"
+    MCP = "MCP"
+    LLM = "LLM"
+    OTHER = "OTHER"

aworld/trace/context_manager.py

@@ -0,0 +1,316 @@
+import types
+import inspect
+from typing import Union, Optional, Any, Type, Sequence, Callable, Iterable
+from aworld.trace.base import (
+    AttributeValueType,
+    NoOpSpan,
+    Span, Tracer,
+    NoOpTracer,
+    get_tracer_provider,
+    get_tracer_provider_silent,
+    log_trace_error
+)
+from aworld.trace.span_cosumer import SpanConsumer
+from aworld.version_gen import __version__
+from aworld.trace.auto_trace import AutoTraceModule, install_auto_tracing
+from aworld.trace.stack_info import get_user_stack_info
+from aworld.trace.constants import (
+    ATTRIBUTES_MESSAGE_KEY,
+    ATTRIBUTES_MESSAGE_RUN_TYPE_KEY,
+    ATTRIBUTES_MESSAGE_TEMPLATE_KEY,
+    RunType
+)
+from aworld.trace.msg_format import (
+    chunks_formatter,
+    warn_formatting,
+    FStringAwaitError,
+    KnownFormattingError,
+    warn_fstring_await
+)
+from aworld.trace.function_trace import trace_func
+from .opentelemetry.opentelemetry_adapter import configure_otlp_provider
+from aworld.logs.util import logger
+
+
+def trace_configure(provider: str = "otlp",
+                    backends: Sequence[str] = None,
+                    base_url: str = None,
+                    write_token: str = None,
+                    span_consumers: Optional[Sequence[SpanConsumer]] = None,
+                    **kwargs
+                    ) -> None:
+    """
+    Configure the trace provider.
+    Args:
+        provider: The trace provider to use.
+        backends: The trace backends to use.
+        base_url: The base URL of the trace backend.
+        write_token: The write token of the trace backend.
+        span_consumers: The span consumers to use.
+        **kwargs: Additional arguments to pass to the trace provider.
+    Returns:
+        None
+    """
+    exist_provider = get_tracer_provider_silent()
+    if exist_provider:
+        logger.info("Trace provider already configured, shutting down...")
+        exist_provider.shutdown()
+    if provider == "otlp":
+        configure_otlp_provider(
+            backends=backends, base_url=base_url, write_token=write_token, span_consumers=span_consumers, **kwargs)
+    else:
+        raise ValueError(f"Unknown trace provider: {provider}")
+
+
+class TraceManager:
+    """
+    TraceManager is a class that provides a way to trace the execution of a function.
+    """
+
+    def __init__(self, tracer_name: str = None) -> None:
+        self._tracer_name = tracer_name or "aworld"
+        self._version = __version__
+
+    def _create_auto_span(self,
+                          name: str,
+                          attributes: dict[str, AttributeValueType] = None
+                          ) -> Span:
+        """
+        Create a auto trace span with the given name and attributes.
+        """
+        return self._create_context_span(name, attributes)
+
+    def _create_context_span(self,
+                             name: str,
+                             attributes: dict[str, AttributeValueType] = None) -> Span:
+        try:
+            tracer = get_tracer_provider().get_tracer(
+                name=self._tracer_name, version=self._version)
+            return ContextSpan(span_name=name, tracer=tracer, attributes=attributes)
+        except Exception:
+            return ContextSpan(span_name=name, tracer=NoOpTracer(), attributes=attributes)
+
+    def get_current_span(self) -> Span:
+        """
+        Get the current span.
+        """
+        try:
+            return get_tracer_provider().get_current_span()
+        except Exception:
+            return NoOpSpan()
+
+    def new_manager(self, tracer_name_suffix: str = None) -> "TraceManager":
+        """
+        Create a new TraceManager with the given tracer name suffix.
+        """
+        tracer_name = self._tracer_name if not tracer_name_suffix else f"{self._tracer_name}.{tracer_name_suffix}"
+        return TraceManager(tracer_name=tracer_name)
+
+    def auto_tracing(self,
+                     modules: Union[Sequence[str], Callable[[AutoTraceModule], bool]],
+                     min_duration: float) -> None:
+        """
+        Automatically trace the execution of a function.
+        Args:
+            modules: A list of module names or a callable that takes a `AutoTraceModule` and returns a boolean.
+            min_duration: The minimum duration of a function to be traced.
+        Returns:
+            None
+        """
+        install_auto_tracing(self, modules, min_duration)
+
+    def span(self,
+             msg_template: str = "",
+             attributes: dict[str, AttributeValueType] = None,
+             *,
+             span_name: str = None,
+             run_type: RunType = RunType.OTHER) -> "ContextSpan":
+
+        try:
+            attributes = attributes or {}
+            stack_info = get_user_stack_info()
+            merged_attributes = {**stack_info, **attributes}
+            # Retrieve stack information of user code and add it to the attributes
+
+            if any(c in msg_template for c in ('{', '}')):
+                fstring_frame = inspect.currentframe().f_back
+            else:
+                fstring_frame = None
+            log_message, extra_attrs, msg_template = format_span_msg(
+                msg_template,
+                merged_attributes,
+                fstring_frame=fstring_frame,
+            )
+            merged_attributes[ATTRIBUTES_MESSAGE_KEY] = log_message
+            merged_attributes.update(extra_attrs)
+            merged_attributes[ATTRIBUTES_MESSAGE_TEMPLATE_KEY] = msg_template
+            merged_attributes[ATTRIBUTES_MESSAGE_RUN_TYPE_KEY] = run_type.value
+            span_name = span_name or msg_template
+
+            return self._create_context_span(span_name, merged_attributes)
+
+        except Exception:
+            log_trace_error()
+            return ContextSpan(span_name=span_name, tracer=NoOpTracer(), attributes=attributes)
+
+    def func_span(self,
+                  msg_template: Union[str, Callable] = None,
+                  *,
+                  attributes: dict[str, AttributeValueType] = None,
+                  span_name: str = None,
+                  extract_args: Union[bool, Iterable[str]] = False,
+                  **kwargs) -> Callable:
+        """
+        A decorator that traces the execution of a function.
+        Args:
+            msg_template: The message template to use.
+            attributes: The attributes to add to the span.
+            span_name: The name of the span.
+            extract_args: Whether to extract arguments from the function call.
+            **kwargs: Additional attributes to add to the span.
+        Returns:
+            A decorator that traces the execution of a function.
+        """
+        if callable(msg_template):
+            #  @trace_func
+            #  def foo():
+            return self.func_span()(msg_template)
+
+        attributes = attributes or {}
+        attributes.update(kwargs)
+        return trace_func(self, msg_template, attributes, span_name, extract_args)
+
+
+class ContextSpan(Span):
+    """A context manager that wraps an existing `Span` object.
+    This class provides a way to use a `Span` object as a context manager.
+    When the context manager is entered, it returns the `Span` itself.
+    When the context manager is exited, it calls `end` on the `Span`.
+    Args:
+        span: The `Span` object to wrap.
+    """
+
+    def __init__(self,
+                 span_name: str,
+                 tracer: Tracer,
+                 attributes: dict[str, AttributeValueType] = None) -> None:
+        self._span_name = span_name
+        self._tracer = tracer
+        self._attributes = attributes
+        self._span: Span = None
+        self._coro_context = None
+
+    def _start(self):
+        if self._span is not None:
+            return
+
+        self._span = self._tracer.start_span(
+            name=self._span_name,
+            attributes=self._attributes,
+        )
+
+    def __enter__(self) -> "Span":
+        self._start()
+        return self
+
+    def __exit__(
+            self,
+            exc_type: Optional[Type[BaseException]],
+            exc_val: Optional[BaseException],
+            traceback: Optional[Any],
+    ) -> None:
+        """Ends context manager and calls `end` on the `Span`."""
+        self._handle_exit(exc_type, exc_val, traceback)
+
+    async def __aenter__(self) -> "Span":
+        self._start()
+
+        return self
+
+    async def __aexit__(
+            self,
+            exc_type: Optional[Type[BaseException]],
+            exc_val: Optional[BaseException],
+            traceback: Optional[Any],
+    ) -> None:
+        self._handle_exit(exc_type, exc_val, traceback)
+
+    def _handle_exit(
+            self,
+            exc_type: Optional[Type[BaseException]],
+            exc_val: Optional[BaseException],
+            traceback: Optional[Any],
+    ) -> None:
+        try:
+            if self._span and self._span.is_recording() and isinstance(exc_val, BaseException):
+                self._span.record_exception(exc_val, escaped=True)
+        except ValueError as e:
+            logger.warning(f"Failed to record_exception: {e}")
+        finally:
+            if self._span:
+                self._span.end()
+
+    def end(self, end_time: Optional[int] = None) -> None:
+        if self._span:
+            self._span.end(end_time)
+
+    def set_attribute(self, key: str, value: AttributeValueType) -> None:
+        if self._span:
+            self._span.set_attribute(key, value)
+
+    def set_attributes(self, attributes: dict[str, AttributeValueType]) -> None:
+        if self._span:
+            self._span.set_attributes(attributes)
+
+    def is_recording(self) -> bool:
+        if self._span:
+            return self._span.is_recording()
+        return False
+
+    def record_exception(
+            self,
+            exception: BaseException,
+            attributes: dict[str, Any] = None,
+            timestamp: Optional[int] = None,
+            escaped: bool = False,
+    ) -> None:
+        if self._span:
+            self._span.record_exception(
+                exception, attributes, timestamp, escaped)
+
+    def get_trace_id(self) -> str:
+        if self._span:
+            return self._span.get_trace_id()
+
+    def get_span_id(self) -> str:
+        if self._span:
+            return self._span.get_span_id()
+
+
+def format_span_msg(
+        format_string: str,
+        kwargs: dict[str, Any],
+        fstring_frame: types.FrameType = None,
+) -> tuple[str, dict[str, Any], str]:
+    """ Returns
+    1. The formatted message.
+    2. A dictionary of extra attributes to add to the span/log.
+         These can come from evaluating values in f-strings.
+    3. The final message template, which may differ from `format_string` if it was an f-string.
+    """
+    try:
+        chunks, extra_attrs, new_template = chunks_formatter.chunks(
+            format_string,
+            kwargs,
+            fstring_frame=fstring_frame
+        )
+        return ''.join(chunk['v'] for chunk in chunks), extra_attrs, new_template
+    except KnownFormattingError as e:
+        warn_formatting(str(e) or str(e.__cause__))
+    except FStringAwaitError as e:
+        warn_fstring_await(str(e))
+    except Exception:
+        log_trace_error()
+
+    # Formatting failed, so just use the original format string as the message.
+    return format_string, {}, format_string

aworld/trace/function_trace.py

@@ -0,0 +1,166 @@
+import inspect
+import contextlib
+import functools
+from typing import TYPE_CHECKING, Callable, Any, Union, Iterable
+from aworld.trace.base import (
+    AttributeValueType
+)
+
+from aworld.trace.stack_info import get_filepath_attribute
+from aworld.trace.constants import (
+    ATTRIBUTES_MESSAGE_TEMPLATE_KEY
+)
+
+if TYPE_CHECKING:
+    from aworld.trace.context_manager import TraceManager, ContextSpan
+
+
+def trace_func(trace_manager: "TraceManager",
+               msg_template: str = None,
+               attributes: dict[str, AttributeValueType] = None,
+               span_name: str = None,
+               extract_args: Union[bool, Iterable[str]] = False):
+    """A decorator that traces the execution of a function.
+
+    Args:
+        trace_manager: The trace manager to use.
+        msg_template: The message template to use.
+        attributes: The attributes to use.
+        span_name: The span name to use.
+        extract_args: Whether to extract arguments from the function call.
+
+    Returns:
+        The decorated function.
+    """
+
+    def decorator(func: Callable) -> Callable:
+        func_meta = get_function_meta(func, msg_template)
+        func_meta.update(attributes or {})
+        final_span_name = span_name or func_meta.get(ATTRIBUTES_MESSAGE_TEMPLATE_KEY) or func.__name__
+
+        if inspect.isgeneratorfunction(func):
+            def wrapper(*args, **kwargs):
+                with open_func_span(trace_manager, func_meta, final_span_name,
+                                    get_func_args(func, extract_args, *args, **kwargs)):
+                    for item in func(*args, **kwargs):
+                        yield item
+        elif inspect.isasyncgenfunction(func):
+            async def wrapper(*args, **kwargs):
+                with open_func_span(trace_manager, func_meta, final_span_name,
+                                    get_func_args(func, extract_args, *args, **kwargs)):
+                    async for item in func(*args, **kwargs):
+                        yield item
+        elif inspect.iscoroutinefunction(func):
+            async def wrapper(*args, **kwargs):
+                with open_func_span(trace_manager, func_meta, final_span_name,
+                                    get_func_args(func, extract_args, *args, **kwargs)):
+                    return await func(*args, **kwargs)
+        else:
+            def wrapper(*args, **kwargs):
+                with open_func_span(trace_manager, func_meta, final_span_name,
+                                    get_func_args(func, extract_args, *args, **kwargs)):
+                    return func(*args, **kwargs)
+
+        wrapper = functools.wraps(func)(wrapper)  # type: ignore
+        return wrapper
+
+    return decorator
+
+
+def open_func_span(trace_manager: "TraceManager",
+                   func_meta: dict[str, AttributeValueType],
+                   span_name: str,
+                   func_args: dict[str, AttributeValueType]):
+    """Open a function span.
+
+    Args:
+        func_meta: The function meta information.
+        span_name: The span name.
+
+    Returns:
+        The function span.
+    """
+    func_meta.update(func_args)
+    return trace_manager._create_auto_span(name=span_name, attributes=func_meta)
+
+
+def get_func_args(func: Callable,
+                  extract_args: Union[bool, Iterable[str]] = False,
+                  *args,
+                  **kwargs):
+    """Get the arguments of a function.
+
+    Args:
+        func: The function to get the arguments of.
+        extract_args: Whether to extract arguments from the function call.
+        *args: The positional arguments.
+        **kwargs: The keyword arguments.
+
+    Returns:
+        The arguments of the function.
+    """
+    func_sig = inspect.signature(func)
+    if func_sig.parameters:
+        func_args = func_sig.bind(*args, **kwargs).arguments
+        if extract_args is not False:
+            if isinstance(extract_args, bool):
+                extract_args = func_sig.parameters.keys()
+            func_args = {k: v for k, v in func_args.items() if k in extract_args}
+        return func_args
+    return {}
+
+
+def get_function_meta(func: Any,
+                      msg_template: str = None) -> dict[str, AttributeValueType]:
+    """Get the meta information of a function.\
+
+    Args:
+        func: The function to get the meta information of.
+        msg_template: The message template to use.
+
+    Returns:
+        The meta information of the function.
+    """
+    func = inspect.unwrap(func)
+    if not inspect.isfunction(func) and hasattr(func, '__call__'):
+        func = func.__call__
+        func = inspect.unwrap(func)
+
+    func_name = getattr(func, '__qualname__', getattr(func, '__name__', build_func_name(func)))
+    if not msg_template:
+        try:
+            msg_template = f'Calling {inspect.getmodule(func).__name__}.{func_name}'  # type: ignore
+        except Exception:  # pragma: no cover
+            msg_template = f'Calling {func_name}'
+    meta: dict[str, AttributeValueType] = {
+        'code.function': func_name,
+        ATTRIBUTES_MESSAGE_TEMPLATE_KEY: msg_template,
+    }
+    with contextlib.suppress(Exception):
+        meta['code.lineno'] = func.__code__.co_firstlineno
+    with contextlib.suppress(Exception):
+        # get code.filepath
+        meta.update(get_filepath_attribute(inspect.getsourcefile(func)))
+
+    func_sig = inspect.signature(func)
+    if func_sig.parameters:
+        meta['func.args'] = [str(param) for param in func_sig.parameters.values()
+                             if param.name != 'self']
+    return meta
+
+
+def build_func_name(func: Any) -> str:
+    """Build the function name.
+
+    Args:
+        func: The function to build the name of.
+
+    Returns:
+        The function name.
+    """
+    try:
+        result = repr(func)
+    except Exception:
+        result = f'<{type(func).__name__} object>'
+
+    return result

aworld/trace/msg_format.py

@@ -0,0 +1,403 @@
+import ast
+import inspect
+import sys
+import types
+import warnings
+import executing
+from functools import lru_cache
+from string import Formatter
+from types import CodeType
+from typing import Any, Literal, TypeVar
+from typing_extensions import NotRequired, TypedDict
+from .constants import MESSAGE_FORMATTED_VALUE_LENGTH_LIMIT
+from .stack_info import get_user_frame_and_stacklevel
+
+Truncatable = TypeVar('Truncatable', str, bytes, 'list[Any]', 'tuple[Any, ...]')
+
+class LiteralChunk(TypedDict):
+    t: Literal['lit']
+    v: str
+
+
+class ArgChunk(TypedDict):
+    t: Literal['arg']
+    v: str
+    spec: NotRequired[str]
+
+
+class KnownFormattingError(Exception):
+    """An error raised when there's something wrong with a format string or the field values.
+
+    In other words this should correspond to errors that would be raised when using `str.format`,
+    and generally indicate a user error, most likely that they weren't trying to pass a template string at all.
+    """
+
+
+class FStringAwaitError(Exception):
+    """An error raised when an await expression is found in an f-string.
+
+    This is a specific case that can't be handled by f-string introspection and requires
+    pre-evaluating the await expression before logging.
+    """
+
+
+class FormattingFailedWarning(UserWarning):
+    pass
+
+class InspectArgumentsFailedWarning(Warning):
+    pass
+
+class ChunksFormatter(Formatter):
+    def chunks(
+        self,
+        format_string: str,
+        kwargs: dict[str, Any],
+        *,
+        fstring_frame: types.FrameType = None,
+    ) -> tuple[list[LiteralChunk | ArgChunk], dict[str, Any], str]:
+        # Returns
+        # 1. A list of chunks
+        # 2. A dictionary of extra attributes to add to the span/log.
+        #      These can come from evaluating values in f-strings,
+        #      or from noting scrubbed values.
+        # 3. The final message template, which may differ from `format_string` if it was an f-string.
+        if fstring_frame:
+            result = self._fstring_chunks(kwargs, fstring_frame)
+            if result:  # returns None if faile
+                return result
+
+        chunks = self._vformat_chunks(
+            format_string,
+            kwargs=kwargs
+        )
+        # When there's no f-string magic, there's no changes in the template string.
+        return chunks, {}, format_string
+
+    def _fstring_chunks(
+        self,
+        kwargs: dict[str, Any],
+        frame: types.FrameType,
+    ) -> tuple[list[LiteralChunk | ArgChunk], dict[str, Any], str]:
+        # `frame` is the frame of the method that's being called by the user
+        # called_code = frame.f_code
+        frame = frame.f_back or frame  # type: ignore
+        assert frame is not None
+        # This is where the magic happens. It has caching.
+        ex = executing.Source.executing(frame)
+
+        call_node = ex.node
+        if call_node is None:  # type: ignore[reportUnnecessaryComparison]
+            # `executing` failed to find a node.
+            # This shouldn't happen in most cases, but it's best not to rely on it always working.
+            if not ex.source.text:
+                # This is a very likely cause.
+                # There's nothing we could possibly do to make magic work here,
+                # and it's a clear case where the user should turn the magic off.
+                warn_inspect_arguments(
+                    'No source code available. '
+                    'This happens when running in an interactive shell, '
+                    'using exec(), or running .pyc files without the source .py files.',
+                    get_stacklevel(frame),
+                )
+                return None
+
+            msg = '`executing` failed to find a node.'
+            if sys.version_info[:2] < (3, 11):  # pragma: no cover
+                # inspect_arguments is only on by default for 3.11+ for this reason.
+                # The AST modifications made by auto-tracing
+                # mean that the bytecode doesn't match the source code seen by `executing`.
+                # In 3.11+, a different algorithm is used by `executing` which can deal with this.
+                msg += ' This may be caused by a combination of using Python < 3.11 and auto-tracing.'
+
+            # Try a simple fallback heuristic to find the node which should work in most cases.
+            main_nodes: list[ast.AST] = []
+            for statement in ex.statements:
+                if isinstance(statement, ast.With):
+                    # Only look at the 'header' of a with statement, not its body.
+                    main_nodes += statement.items
+                else:
+                    main_nodes.append(statement)
+            call_nodes = [
+                node
+                for main_node in main_nodes
+                for node in ast.walk(main_node)
+                if isinstance(node, ast.Call)
+                if node.args or node.keywords
+            ]
+            if len(call_nodes) != 1:
+                warn_inspect_arguments(msg, get_stacklevel(frame))
+                return None
+
+            [call_node] = call_nodes
+
+        if not isinstance(call_node, ast.Call):  # pragma: no cover
+            # Very unlikely.
+            warn_inspect_arguments(
+                '`executing` unexpectedly identified a non-Call node.',
+                get_stacklevel(frame),
+            )
+            return None
+        
+        if call_node.args:
+            arg_node = call_node.args[0]
+        else:
+            # Very unlikely.
+            warn_inspect_arguments(
+                "Couldn't identify the `msg_template` argument in the call.",
+                get_stacklevel(frame),
+            )
+            return None
+
+        if not isinstance(arg_node, ast.JoinedStr):
+            # Not an f-string, not a problem.
+            # Just use normal formatting.
+            return None
+
+        # We have an f-string AST node.
+        # Now prepare the namespaces that we will use to evaluate the components.
+        global_vars = frame.f_globals
+        local_vars = {**frame.f_locals, **kwargs}
+
+        # Now for the actual formatting!
+        result: list[LiteralChunk | ArgChunk] = []
+
+        # We construct the message template (i.e. the span name) from the AST.
+        # We don't use the source code of the f-string because that gets messy
+        # if there's escaped quotes or implicit joining of adjacent strings.
+        new_template = ''
+
+        extra_attrs: dict[str, Any] = {}
+        for node_value in arg_node.values:
+            if isinstance(node_value, ast.Constant):
+                # These are the parts of the f-string not enclosed by `{}`, e.g. 'foo ' in f'foo {bar}'
+                value: str = node_value.value
+                result.append({'v': value, 't': 'lit'})
+                new_template += value
+            else:
+                # These are the parts of the f-string enclosed by `{}`, e.g. 'bar' in f'foo {bar}'
+                assert isinstance(node_value, ast.FormattedValue)
+
+                # This is cached.
+                source, value_code, formatted_code = compile_formatted_value(node_value, ex.source)
+
+                # Note that this doesn't include:
+                # - The format spec, e.g. `:0.2f`
+                # - The conversion, e.g. `!r`
+                # - The '=' sign within the braces, e.g. `{bar=}`.
+                #     The AST represents f'{bar = }' as f'bar = {bar}' which is how the template will look.
+                new_template += '{' + source + '}'
+
+                # The actual value of the expression.
+                value = eval(value_code, global_vars, local_vars)
+                extra_attrs[source] = value
+
+                # Format the value according to the format spec, converting to a string.
+                formatted = eval(formatted_code, global_vars, {**local_vars, '@fvalue': value})
+                formatted = self._clean_value(formatted)
+                result.append({'v': formatted, 't': 'arg'})
+
+        return result, extra_attrs, new_template
+
+    def _vformat_chunks(
+        self,
+        format_string: str,
+        kwargs: dict[str, Any],
+        *,
+        recursion_depth: int = 2,
+    ) -> list[LiteralChunk | ArgChunk]:
+        """Copied from `string.Formatter._vformat` https://github.com/python/cpython/blob/v3.11.4/Lib/string.py#L198-L247 then altered."""
+        if recursion_depth < 0:
+            raise KnownFormattingError('Max format spec recursion exceeded')
+        result: list[LiteralChunk | ArgChunk] = []
+        # We currently don't use positional arguments
+        args = ()
+
+        for literal_text, field_name, format_spec, conversion in self.parse(format_string):
+            # output the literal text
+            if literal_text:
+                result.append({'v': literal_text, 't': 'lit'})
+
+            # if there's a field, output it
+            if field_name is not None:
+                # this is some markup, find the object and do
+                #  the formatting
+                if field_name == '':
+                    raise KnownFormattingError('Empty curly brackets `{}` are not allowed. A field name is required.')
+
+                # ADDED BY US:
+                if field_name.endswith('='):
+                    if result and result[-1]['t'] == 'lit':
+                        result[-1]['v'] += field_name
+                    else:
+                        result.append({'v': field_name, 't': 'lit'})
+                    field_name = field_name[:-1]
+
+                # given the field_name, find the object it references
+                #  and the argument it came from
+                try:
+                    obj, _arg_used = self.get_field(field_name, args, kwargs)
+                except IndexError:
+                    raise KnownFormattingError('Numeric field names are not allowed.')
+                except KeyError as exc1:
+                    if str(exc1) == repr(field_name):
+                        raise KnownFormattingError(f'The field {{{field_name}}} is not defined.') from exc1
+
+                    try:
+                        # field_name is something like 'a.b' or 'a[b]'
+                        # Evaluating that expression failed, so now just try getting the whole thing from kwargs.
+                        # In particular, OTEL attributes with dots in their names are normal and handled here.
+                        obj = kwargs[field_name]
+                    except KeyError as exc2:
+                        # e.g. neither 'a' nor 'a.b' is defined
+                        raise KnownFormattingError(f'The fields {exc1} and {exc2} are not defined.') from exc2
+                except Exception as exc:
+                    raise KnownFormattingError(f'Error getting field {{{field_name}}}: {exc}') from exc
+
+                # do any conversion on the resulting object
+                if conversion is not None:
+                    try:
+                        obj = self.convert_field(obj, conversion)
+                    except Exception as exc:
+                        raise KnownFormattingError(f'Error converting field {{{field_name}}}: {exc}') from exc
+
+                # expand the format spec, if needed
+                format_spec_chunks = self._vformat_chunks(
+                    format_spec or '', kwargs, recursion_depth=recursion_depth - 1
+                )
+                format_spec = ''.join(chunk['v'] for chunk in format_spec_chunks)
+
+                try:
+                    value = self.format_field(obj, format_spec)
+                except Exception as exc:
+                    raise KnownFormattingError(f'Error formatting field {{{field_name}}}: {exc}') from exc
+                value = self._clean_value(value)
+                d: ArgChunk = {'v': value, 't': 'arg'}
+                if format_spec:
+                    d['spec'] = format_spec
+                result.append(d)
+
+        return result
+
+    def _clean_value(self, value: str) -> str:
+        return truncate_sequence(seq=value, max_length=MESSAGE_FORMATTED_VALUE_LENGTH_LIMIT, middle='...')
+
+def warn_inspect_arguments(msg: str, stacklevel: int):
+    """Warn about an error in inspecting arguments.
+    This is a separate function so that it can be called from multiple places.
+    """
+    msg = (
+        'Failed to introspect calling code. '
+        'Falling back to normal message formatting '
+        'which may result in loss of information if using an f-string. '
+        'The problem was:\n'
+    ) + msg
+    warnings.warn(msg, InspectArgumentsFailedWarning, stacklevel=stacklevel)
+
+
+def get_stacklevel(frame: types.FrameType):
+    """Get a stacklevel which can be passed to warn_inspect_arguments
+    which points at the given frame, where the f-string was found.
+    """
+    current_frame = inspect.currentframe()
+    stacklevel = 0
+    while current_frame:  # pragma: no branch
+        if current_frame == frame:
+            break
+        stacklevel += 1
+        current_frame = current_frame.f_back
+    return stacklevel
+
+@lru_cache
+def compile_formatted_value(node: ast.FormattedValue, ex_source: executing.Source) -> tuple[str, CodeType, CodeType]:
+    """Returns three things that can be expensive to compute.
+
+    1. Source code corresponding to the node value (excluding the format spec).
+    2. A compiled code object which can be evaluated to calculate the value.
+    3. Another code object which formats the value.
+    """
+    source = get_node_source_text(node.value, ex_source)
+
+    # Check if the expression contains await before attempting to compile
+    for sub_node in ast.walk(node.value):
+        if isinstance(sub_node, ast.Await):
+            raise FStringAwaitError(source)
+
+    value_code = compile(source, '<fvalue1>', 'eval')
+    expr = ast.Expression(
+        ast.JoinedStr(
+            values=[
+                # Similar to the original FormattedValue node,
+                # but replace the actual expression with a simple variable lookup
+                # so that it the expression doesn't need to be evaluated again.
+                # Use @ in the variable name so that it can't possibly conflict
+                # with a normal variable.
+                # The value of this variable will be provided in the eval() call
+                # and will come from evaluating value_code above.
+                ast.FormattedValue(
+                    value=ast.Name(id='@fvalue', ctx=ast.Load()),
+                    conversion=node.conversion,
+                    format_spec=node.format_spec,
+                )
+            ]
+        )
+    )
+    ast.fix_missing_locations(expr)
+    formatted_code = compile(expr, '<fvalue2>', 'eval')
+    return source, value_code, formatted_code
+
+def get_node_source_text(node: ast.AST, ex_source: executing.Source):
+    """Returns some Python source code representing `node`.
+
+    Preferably the actual original code given by `ast.get_source_segment`,
+    but falling back to `ast.unparse(node)` if the former is incorrect.
+    This happens sometimes due to Python bugs (especially for older Python versions)
+    in the source positions of AST nodes inside f-strings.
+    """
+    # ast.unparse is not available in Python 3.8, which is why inspect_arguments is forbidden in 3.8.
+    source_unparsed = ast.unparse(node)
+    source_segment = ast.get_source_segment(ex_source.text, node) or ''
+    try:
+        # Verify that the source segment is correct by checking that the AST is equivalent to what we have.
+        source_segment_unparsed = ast.unparse(ast.parse(source_segment, mode='eval'))
+    except Exception:  # probably SyntaxError, but ast.parse can raise other exceptions too
+        source_segment_unparsed = ''
+    return source_segment if source_unparsed == source_segment_unparsed else source_unparsed
+
+
+def truncate_sequence(seq: Truncatable, *, max_length: int, middle: Truncatable) -> Truncatable:
+    """Return a sequence at with `len()` at most `max_length`, with `middle` in the middle if truncated."""
+    if len(seq) <= max_length:
+        return seq
+    remaining_length = max_length - len(middle)
+    half = remaining_length // 2
+    return seq[:half] + middle + seq[-half:]
+
+def warn_at_user_stacklevel(msg: str, category: type[Warning]):
+    """Warn at the user's stack level.
+    """
+    _frame, stacklevel = get_user_frame_and_stacklevel()
+    warnings.warn(msg, stacklevel=stacklevel, category=category)
+
+def warn_formatting(msg: str):
+    """Warn about a formatting error.   
+    """
+    warn_at_user_stacklevel(
+        f'\n'
+        f'    Ensure you are either:\n'
+        '      (1) passing an f-string directly, or\n'
+        '      (2) passing a literal `str.format`-style template, not a preformatted string.\n'
+        f'    The problem was: {msg}',
+        category=FormattingFailedWarning,
+    )
+
+def warn_fstring_await(msg: str):
+    """Warn about an await expression in an f-string.
+    """
+    warn_at_user_stacklevel(
+        f'\n'
+        f'    Cannot evaluate await expression in f-string. Pre-evaluate the expression before logging.\n'
+        f'    The problematic f-string value was: {msg}',
+        category=FormattingFailedWarning,
+    )
+
+chunks_formatter = ChunksFormatter()

aworld/trace/rewrite_ast.py

@@ -0,0 +1,259 @@
+from __future__ import annotations
+
+import ast
+import uuid
+import time
+from pathlib import Path
+from collections import deque
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable, ContextManager, cast
+
+from aworld.trace.base import AttributeValueType
+from aworld.trace.constants import ATTRIBUTES_MESSAGE_TEMPLATE_KEY
+
+if TYPE_CHECKING:
+    from .context_manager import TraceManager
+    from .auto_trace import not_auto_trace
+
+
+def compile_source(
+        tree: ast.AST, filename: str, module_name: str, trace_manager: TraceManager, min_duration_ns: int
+) -> Callable[[dict[str, Any]], None]:
+    """Compile a modified AST of the module's source code in the module's namespace.
+
+    Returns a function which accepts module globals and executes the compiled code.
+
+    The modified AST wraps the body of every function definition in `with context_factories[index]():`.
+    `context_factories` is added to the module's namespace as `aworld_<uuid>`.
+    `index` is a different constant number for each function definition.
+    """
+
+    context_factories_var_name = f'aworld_{uuid.uuid4().hex}'
+    # The variable name for storing context_factors in the module's namespace.
+
+    context_factories: list[Callable[[], ContextManager[Any]]] = []
+    tree = rewrite_ast(tree, filename, context_factories_var_name, module_name, trace_manager, context_factories,
+                       min_duration_ns)
+    assert isinstance(tree, ast.Module)  # for type checking
+    # dont_inherit=True is necessary to prevent the module from inheriting the __future__ import from this module.
+    code = compile(tree, filename, 'exec', dont_inherit=True)
+
+    def execute(globs: dict[str, Any]):
+        globs[context_factories_var_name] = context_factories
+        exec(code, globs, globs)
+
+    return execute
+
+
+def rewrite_ast(
+        tree: ast.AST,
+        filename: str,
+        context_factories_var_name: str,
+        module_name: str,
+        trace_manager: TraceManager,
+        context_factories: list[Callable[[], ContextManager[Any]]],
+        min_duration_ns: int,
+) -> ast.AST:
+    transformer = AutoTraceTransformer(
+        context_factories_var_name, filename, module_name, trace_manager, context_factories, min_duration_ns
+    )
+    return transformer.visit(tree)
+
+
+class AutoTraceTransformer(ast.NodeTransformer):
+    """Trace all encountered functions except those explicitly marked with `@no_auto_trace`."""
+
+    def __init__(
+            self,
+            context_factories_var_name: str,
+            filename: str,
+            module_name: str,
+            trace_manager: TraceManager,
+            context_factories: list[Callable[[], ContextManager[Any]]],
+            min_duration_ns: int,
+    ):
+        self._context_factories_var_name = context_factories_var_name
+        self._filename = filename
+        self._module_name = module_name
+        self._trace_manager = trace_manager
+        self._context_factories = context_factories
+        self._min_duration_ns = min_duration_ns
+        self._qualname_stack: list[str] = []
+
+    def visit_ClassDef(self, node: ast.ClassDef):
+        """Visit a class definition and rewrite its methods."""
+
+        if self.check_not_auto_trace(node):
+            return node
+
+        self._qualname_stack.append(node.name)
+        node = cast(ast.ClassDef, self.generic_visit(node))
+        self._qualname_stack.pop()
+        return node
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST:
+        """Visit a function definition and rewrite it."""
+
+        if self.check_not_auto_trace(node):
+            return node
+
+        self._qualname_stack.append(node.name)
+        qualname = '.'.join(self._qualname_stack)
+        self._qualname_stack.append('<locals>')
+        self.generic_visit(node)
+        self._qualname_stack.pop()  # <locals>
+        self._qualname_stack.pop()  # node.name
+        return self.rewrite_function(node, qualname)
+
+    def check_not_auto_trace(self, node: ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef) -> bool:
+        """Return true if the node has a `@not_auto_trace` decorator."""
+        return any(
+            (
+                    isinstance(node, ast.Name)
+                    and node.id == not_auto_trace.__name__
+                # or (
+                #     isinstance(node, ast.Attribute)
+                #     and node.attr == not_auto_trace.__name__
+                #     and isinstance(node.value, ast.Name)
+                #     and node.value.id == xxx.__name__
+                # )
+            )
+            for node in node.decorator_list
+        )
+
+    def rewrite_function(self, node: ast.FunctionDef | ast.AsyncFunctionDef, qualname: str) -> ast.AST:
+        """Rewrite a function definition to trace its execution."""
+
+        if has_yield(node):
+            return node
+
+        body = node.body.copy()
+        new_body: list[ast.stmt] = []
+        if (
+                body
+                and isinstance(body[0], ast.Expr)
+                and isinstance(body[0].value, ast.Constant)
+                and isinstance(body[0].value.value, str)
+        ):
+            new_body.append(body.pop(0))
+
+        if not body or (
+                len(body) == 1
+                and (
+                        isinstance(body[0], ast.Pass)
+                        or (isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant))
+                )
+        ):
+            return node
+
+        span = ast.With(
+            items=[
+                ast.withitem(
+                    context_expr=self.trace_context_method_call_node(node, qualname),
+                )
+            ],
+            body=body,
+            type_comment=node.type_comment,
+        )
+        new_body.append(span)
+
+        return ast.fix_missing_locations(
+            ast.copy_location(
+                type(node)(  # type: ignore
+                    name=node.name,
+                    args=node.args,
+                    body=new_body,
+                    decorator_list=node.decorator_list,
+                    returns=node.returns,
+                    type_comment=node.type_comment,
+                ),
+                node,
+            )
+        )
+
+    def trace_context_method_call_node(self, node: ast.FunctionDef | ast.AsyncFunctionDef, qualname: str) -> ast.Call:
+        """Return a method call to `context_factories[index]()`."""
+
+        index = len(self._context_factories)
+        span_factory = partial(
+            self._trace_manager._create_auto_span,  # type: ignore
+            *self.build_create_auto_span_args(qualname, node.lineno),
+        )
+        if self._min_duration_ns > 0:
+
+            timer = time.time_ns
+            min_duration = self._min_duration_ns
+
+            # This needs to be as fast as possible since it's the cost of auto-tracing a function
+            # that never actually gets instrumented because its calls are all faster than `min_duration`.
+            class MeasureTime:
+                __slots__ = 'start'
+
+                def __enter__(_self):
+                    _self.start = timer()
+
+                def __exit__(_self, *_):
+                    # the first call exceeding min_ruration will not be tracked, and subsequent calls will only be tracked
+                    if timer() - _self.start >= min_duration:
+                        self._context_factories[index] = span_factory
+
+            self._context_factories.append(MeasureTime)
+        else:
+            self._context_factories.append(span_factory)
+
+        # This node means:
+        #   context_factories[index]()
+        # where `context_factories` is a global variable with the name `self._context_factories_var_name`
+        # pointing to the `self.context_factories` list.
+        return ast.Call(
+            func=ast.Subscript(
+                value=ast.Name(id=self._context_factories_var_name, ctx=ast.Load()),
+                slice=ast.Index(value=ast.Constant(value=index)),  # type: ignore
+                ctx=ast.Load(),
+            ),
+            args=[],
+            keywords=[],
+        )
+
+    def build_create_auto_span_args(self, qualname: str, lineno: int) -> tuple[str, dict[str, AttributeValueType]]:
+        """Build the arguments for `create_auto_span`."""
+
+        stack_info = {
+            'code.filepath': get_filepath(self._filename),
+            'code.lineno': lineno,
+            'code.function': qualname,
+        }
+        attributes: dict[str, AttributeValueType] = {**stack_info}  # type: ignore
+
+        msg_template = f'Calling {self._module_name}.{qualname}'
+        attributes[ATTRIBUTES_MESSAGE_TEMPLATE_KEY] = msg_template
+
+        span_name = msg_template
+
+        return span_name, attributes
+
+
+def has_yield(node: ast.AST):
+    """Return true if the node has a yield statement."""
+
+    queue = deque([node])
+    while queue:
+        node = queue.popleft()
+        for child in ast.iter_child_nodes(node):
+            if isinstance(child, (ast.Yield, ast.YieldFrom)):
+                return True
+            if not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
+                queue.append(child)
+
+
+def get_filepath(file: str):
+    """Return a dict with the filepath attribute."""
+
+    path = Path(file)
+    if path.is_absolute():
+        try:
+            path = path.relative_to(Path('.').resolve())
+        except ValueError:  # pragma: no cover
+            # happens if filename path is not within CWD
+            pass
+    return str(path)

aworld/trace/span_cosumer.py

@@ -0,0 +1,43 @@
+from abc import ABC, abstractmethod
+from typing import Sequence
+from aworld.trace.base import Span
+
+
+class SpanConsumer(ABC):
+    """SpanConsumer is a protocol that represents a consumer for spans.
+    """
+    @abstractmethod
+    def consume(self, spans: Sequence[Span]) -> None:
+        """Consumes a span.
+        Args:
+            spans: The span to consume.
+        """
+
+
+_SPAN_CONSUMER_REGISTRY = {}
+
+
+def register_span_consumer(default_kwargs=None) -> None:
+    """Registers a span consumer.
+    Args:
+        default_kwargs: A dictionary of default keyword arguments to pass to the span consumer.
+    """
+
+    default_kwargs = default_kwargs or {}
+
+    def decorator(cls):
+        _SPAN_CONSUMER_REGISTRY[cls.__name__] = (cls, default_kwargs)
+        return cls
+
+    return decorator
+
+
+def get_span_consumers() -> Sequence[SpanConsumer]:
+    """Returns a list of span consumers.
+    Returns:
+        A list of span consumers.
+    """
+    return [
+        cls(**kwargs)
+        for cls, kwargs in _SPAN_CONSUMER_REGISTRY.values()
+    ]

aworld/trace/stack_info.py

@@ -0,0 +1,91 @@
+import inspect
+import sys
+import aworld.trace as atrace
+from types import CodeType, FrameType
+from typing import Optional, TypedDict, Union
+from functools import lru_cache
+from pathlib import Path
+
+StackInfo = TypedDict('StackInfo', {'code.filepath': str, 'code.lineno': int, 'code.function': str}, total=False)
+
+NON_USER_CODE_PREFIXES: tuple[str, ...] = ()
+
+def add_non_user_code_prefix(path: Union[str, Path]) -> None:
+    global NON_USER_CODE_PREFIXES
+    path = str(Path(path).absolute())
+    NON_USER_CODE_PREFIXES += (path,)  
+
+add_non_user_code_prefix(Path(inspect.__file__).parent)
+add_non_user_code_prefix(Path(atrace.__file__).parent)
+
+def get_user_stack_info() -> StackInfo:
+    """Get the stack info for the first calling frame in user code.
+
+    See is_user_code for details.
+    Returns an empty dict if no such frame is found.
+    """
+    frame, _stacklevel = get_user_frame_and_stacklevel()
+    if frame:
+        return get_stack_info_from_frame(frame)
+    return {}
+
+
+def get_user_frame_and_stacklevel() -> tuple[Optional[FrameType], int]:
+    """Get the first calling frame in user code and a corresponding stacklevel that can be passed to `warnings.warn`.
+
+    See is_user_code for details.
+    Returns `(None, 0)` if no such frame is found.
+    """
+    frame = inspect.currentframe()
+    stacklevel = 0
+    while frame:
+        if is_user_code(frame.f_code):
+            return frame, stacklevel
+        frame = frame.f_back
+        stacklevel += 1
+    return None, 0
+
+def get_stack_info_from_frame(frame: FrameType) -> StackInfo:
+    return {
+        **get_code_object_info(frame.f_code),
+        'code.lineno': frame.f_lineno,
+    }
+
+@lru_cache(maxsize=2048)
+def get_code_object_info(code: CodeType) -> StackInfo:
+    result = get_filepath_attribute(code.co_filename)
+    if code.co_name != '<module>':  # pragma: no branch
+        result['code.function'] = code.co_qualname if sys.version_info >= (3, 11) else code.co_name
+    result['code.lineno'] = code.co_firstlineno
+    return result
+
+def get_filepath_attribute(file: str) -> StackInfo:
+    path = Path(file)
+    if path.is_absolute():
+        try:
+            path = path.relative_to(Path('.').resolve())
+        except ValueError:  # pragma: no cover
+            # happens if filename path is not within CWD
+            pass
+    return {'code.filepath': str(path)}
+
+@lru_cache(maxsize=8192)
+def is_user_code(code: CodeType) -> bool:
+    """Check if the code object is from user code.
+
+    A code object is not user code if:
+    - It is from a file in
+        - the standard library
+        - site-packages (specifically wherever opentelemetry is installed)
+        - an unknown location (e.g. a dynamically generated code object) indicated by a filename starting with '<'
+
+    - It is a list/dict/set comprehension.
+        These are artificial frames only created before Python 3.12,
+        and they are always called directly from the enclosing function so it makes sense to skip them.
+        On the other hand, generator expressions and lambdas might be called far away from where they are defined.
+    """
+    return not (
+        str(Path(code.co_filename).absolute()).startswith(NON_USER_CODE_PREFIXES)
+        or code.co_filename.startswith('<')
+        or code.co_name in ('<listcomp>', '<dictcomp>', '<setcomp>')
+    )