microsoft--agent-lightning/agentlightning/tracer/base.py

# Copyright (c) Microsoft. All rights reserved.

from __future__ import annotations

import functools
import logging
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional, TypeVar

from agentlightning.store.base import LightningStore
from agentlightning.types import Attributes, ParallelWorkerBase, Span, SpanCoreFields, SpanRecordingContext, TraceStatus

if TYPE_CHECKING:
    from langchain_core.callbacks.base import BaseCallbackHandler  # type: ignore

logger = logging.getLogger(__name__)


T = TypeVar("T")


_active_tracer: Optional[Tracer] = None

T_func = Callable[..., Awaitable[Any]]


class Tracer(ParallelWorkerBase):
    """
    An abstract base class for tracers.

    This class defines a standard interface for tracing code execution,
    capturing the resulting spans, and providing them for analysis. It is
    designed to be backend-agnostic, allowing for different implementations
    (e.g., for AgentOps, OpenTelemetry, Docker, etc.).

    The primary interaction pattern is through the [`trace_context`][agentlightning.Tracer.trace_context]
    context manager, which ensures that traces are properly started and captured,
    even in the case of exceptions.

    A typical workflow:

    ```python
    tracer = YourTracerImplementation()

    try:
        async with tracer.trace_context(name="my_traced_task"):
            # ... code to be traced ...
            await run_my_agent_logic()
    except Exception as e:
        print(f"An error occurred: {e}")

    # Retrieve the trace data after the context block
    spans: list[ReadableSpan] = tracer.get_last_trace()

    # Process the trace data
    if trace_tree:
        rl_triplets = TracerTraceToTriplet().adapt(spans)
        # ... do something with the triplets
    ```
    """

    _store: Optional[LightningStore] = None

    def init_worker(self, worker_id: int, store: Optional[LightningStore] = None) -> None:
        """Initialize the tracer for a worker.

        Args:
            worker_id: The ID of the worker.
            store: The store to add the spans to. If it's provided, traces will be added to the store when tracing.
        """
        super().init_worker(worker_id)
        self._store = store

    def trace_context(
        self,
        name: Optional[str] = None,
        *,
        store: Optional[LightningStore] = None,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
    ) -> AsyncContextManager[Any]:
        """
        Starts a new tracing context. This should be used as a context manager.

        The implementation should handle the setup and teardown of the tracing
        for the enclosed code block. It must ensure that any spans generated
        within the `with` block are collected and made available via
        [`get_last_trace`][agentlightning.Tracer.get_last_trace].

        Args:
            name: The name for the root span of this trace context.
            store: The store to add the spans to. Deprecated in favor of passing store to init_worker().
            rollout_id: The rollout ID to add the spans to.
            attempt_id: The attempt ID to add the spans to.
        """
        raise NotImplementedError()

    def _trace_context_sync(
        self,
        name: Optional[str] = None,
        *,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
    ) -> ContextManager[Any]:
        """Internal API for CI backward compatibility."""
        raise NotImplementedError()

    def get_last_trace(self) -> List[Span]:
        """
        Retrieves the raw list of captured spans from the most recent trace.

        Returns:
            A list of [`Span`][agentlightning.Span] objects collected during the last trace.
        """
        raise NotImplementedError()

    def trace_run(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
        """
        A convenience wrapper to trace the execution of a single synchronous function.

        Deprecated in favor of customizing Runners.

        Args:
            func: The synchronous function to execute and trace.
            *args: Positional arguments to pass to the function.
            **kwargs: Keyword arguments to pass to the function.

        Returns:
            The return value of the function.
        """
        with self._trace_context_sync(name=func.__name__):
            return func(*args, **kwargs)

    def create_span(
        self,
        name: str,
        attributes: Optional[Attributes] = None,
        timestamp: Optional[float] = None,
        status: Optional[TraceStatus] = None,
    ) -> SpanCoreFields:
        """Notify the tracer that a span should be created here.

        It uses a fire-and-forget approach and doesn't wait for the span to be created.

        Args:
            name: The name of the span.
            attributes: The attributes of the span.
            timestamp: The timestamp of the span.
            status: The status of the span.

        Returns:
            The core fields of the span.
        """
        raise NotImplementedError()

    def operation_context(
        self,
        name: str,
        attributes: Optional[Attributes] = None,
        start_time: Optional[float] = None,
        end_time: Optional[float] = None,
    ) -> ContextManager[SpanRecordingContext]:
        """Start to record an operation to a span.

        Args:
            name: The name of the operation.
            attributes: The attributes of the operation.
            start_time: The start time of the operation.
            end_time: The end time of the operation.

        Returns:
            A [`SpanRecordingContext`][agentlightning.SpanRecordingContext] for recording the operation on the span.
        """
        raise NotImplementedError()

    async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> Any:
        """
        A convenience wrapper to trace the execution of a single asynchronous function.

        Deprecated in favor of customizing Runners.

        Args:
            func: The asynchronous function to execute and trace.
            *args: Positional arguments to pass to the function.
            **kwargs: Keyword arguments to pass to the function.

        Returns:
            The return value of the function.
        """
        async with self.trace_context(name=func.__name__):
            return await func(*args, **kwargs)

    def get_langchain_handler(self) -> Optional[BaseCallbackHandler]:  # type: ignore
        """Get a handler to install in langchain agent callback.

        Agents are expected to use this handler in their agents to enable tracing.
        """
        logger.warning(f"{self.__class__.__name__} does not provide a LangChain callback handler.")
        return None

    @contextmanager
    def lifespan(self, store: Optional[LightningStore] = None):
        """A context manager to manage the lifespan of the tracer.

        This can be used to set up and tear down any necessary resources
        for the tracer, useful for debugging purposes.

        Args:
            store: The store to add the spans to. If it's provided, traces will be added to the store when tracing.
        """
        has_init = False
        has_init_worker = False
        try:
            self.init()
            has_init = True

            self.init_worker(0, store)
            has_init_worker = True

            yield

        finally:
            if has_init_worker:
                self.teardown_worker(0)
            if has_init:
                self.teardown()


def set_active_tracer(tracer: Tracer):
    """Set the active tracer for the current process.

    Args:
        tracer: The tracer to set as active.
    """
    global _active_tracer
    if _active_tracer is not None:
        raise ValueError("An active tracer is already set. Cannot set a new one.")
    _active_tracer = tracer


def clear_active_tracer():
    """Clear the active tracer for the current process."""
    global _active_tracer
    _active_tracer = None


def get_active_tracer() -> Optional[Tracer]:
    """Get the active tracer for the current process.

    Returns:
        The active tracer, or None if no tracer is active.
    """
    global _active_tracer
    return _active_tracer


class _ActiveTracerAsyncCM(AsyncContextManager[T]):
    def __init__(self, tracer: Tracer, inner: AsyncContextManager[T]):
        self._tracer = tracer
        self._inner = inner

    async def __aenter__(self) -> T:
        set_active_tracer(self._tracer)  # will raise if nested
        try:
            return await self._inner.__aenter__()
        except Exception:
            clear_active_tracer()
            raise

    async def __aexit__(self, *args: Any, **kwargs: Any) -> Optional[bool]:
        try:
            return await self._inner.__aexit__(*args, **kwargs)
        finally:
            clear_active_tracer()


def with_active_tracer_context(
    func: Callable[..., AsyncContextManager[T]],
) -> Callable[..., AsyncContextManager[T]]:
    """Decorate a method returning an AsyncContextManager so tracer is active for the whole `async with`."""

    @functools.wraps(func)
    def wrapper(self: Tracer, *args: Any, **kwargs: Any) -> AsyncContextManager[T]:
        cm = func(self, *args, **kwargs)
        return _ActiveTracerAsyncCM(self, cm)

    return wrapper