# Copyright (c) Microsoft. All rights reserved. """Helpers for emitting annotation/operation spans.""" import asyncio import functools import inspect import logging from types import TracebackType from typing import ( Any, Callable, ContextManager, Dict, Optional, Tuple, Type, TypeVar, Union, cast, overload, ) from agentlightning.semconv import AGL_ANNOTATION, AGL_OPERATION, LightningSpanAttributes from agentlightning.tracer.base import get_active_tracer from agentlightning.tracer.dummy import DummyTracer from agentlightning.types import SpanCoreFields, SpanRecordingContext, TraceStatus from agentlightning.utils.otel import check_attributes_sanity, flatten_attributes, sanitize_attributes _FnType = TypeVar("_FnType", bound=Callable[..., Any]) logger = logging.getLogger(__name__) def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> SpanCoreFields: """Emit a new annotation span. This is the underlying implementation of [`emit_reward`][agentlightning.emit_reward]. Annotation spans are used to annotate a specific event or a part of rollout. See [semconv][agentlightning.semconv] for conventional annotation keys in Agent-lightning. If annotations contain nested dicts, they will be flattened before emitting. Complex objects will lead to emitting failures. Args: annotation: Dictionary containing annotation key-value pairs. Representatives are rewards, tags, and metadata. propagate: Whether to propagate the span to tracers automatically. """ annotation_attributes = flatten_attributes(annotation, expand_leaf_lists=False) check_attributes_sanity(annotation_attributes) sanitized_attributes = sanitize_attributes(annotation_attributes) logger.debug("Emitting annotation span with keys %s", sanitized_attributes.keys()) if propagate: tracer = get_active_tracer() if tracer is None: raise RuntimeError("No active tracer found. Cannot emit annotation span.") else: tracer = DummyTracer() return tracer.create_span( name=AGL_ANNOTATION, attributes=sanitized_attributes, status=TraceStatus(status_code="OK"), ) class OperationContext: """Context manager and decorator for tracing operations. This class manages a tracer-backed span for a logical unit of work. It can be used either: * As a decorator, in which case inputs and outputs are inferred automatically from the wrapped function's signature. * As a context manager, in which case inputs and outputs can be recorded explicitly via [`set_input`][agentlightning.emitter.annotation.OperationContext.set_input] and [`set_output`][agentlightning.emitter.annotation.OperationContext.set_output]. Attributes: name: Human-readable span name. initial_attributes: Attributes applied when the span is created. tracer: Tracer implementation used to create spans. """ def __init__(self, name: str, attributes: Dict[str, Any], propagate: bool = True) -> None: """Initialize a new operation context. Args: name: Human-readable name of the span. attributes: Initial attributes attached to the span. Values are JSON-serialized where necessary. propagate: Whether the span should be sent to active exporters. """ self.name = name self.initial_attributes = flatten_attributes(attributes, expand_leaf_lists=False) self.propagate = propagate if propagate: tracer = get_active_tracer() if tracer is None: raise RuntimeError("No active tracer found. Cannot trace operation spans.") self.tracer = tracer else: self.tracer = DummyTracer() self._ctx_manager: Optional[ContextManager[SpanRecordingContext]] = None self._recording_context: Optional[SpanRecordingContext] = None self._span: Optional[SpanCoreFields] = None def __enter__(self) -> "OperationContext": """Enter the context manager and start a new span. Returns: The current :class:`OperationContext` instance with an active span. """ sanitized_attrs = sanitize_attributes(self.initial_attributes) self._ctx_manager = self.tracer.operation_context(self.name, attributes=sanitized_attrs) recording_context = self._ctx_manager.__enter__() self._recording_context = recording_context return self def __exit__( self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType], ) -> None: """Exit the context manager and finish the span.""" if self._ctx_manager: self._ctx_manager.__exit__(exc_type, exc_val, exc_tb) if self._recording_context: self._span = self._recording_context.get_recorded_span() self._ctx_manager = None self._recording_context = None def span(self) -> SpanCoreFields: """Get the span that was created by this context manager.""" if self._span is None: raise RuntimeError("Span is not ready yet.") return self._span def set_input(self, *args: Any, **kwargs: Any) -> None: """Record input arguments on the current span. Positional arguments are stored under the `input.args.` attributes, and keyword arguments are stored under `input.` attributes. This is intended for use inside a `with operation(...) as op` block. Args: *args: Positional arguments to record. **kwargs: Keyword arguments to record. """ if not self._recording_context: raise RuntimeError("No recording context found. Cannot set input.") prefix = LightningSpanAttributes.OPERATION_INPUT.value attributes: Dict[str, Any] = {} if args: for idx, value in enumerate(args): flattened = flatten_attributes({str(idx): value}) for nested_key, nested_value in flattened.items(): attributes[f"{prefix}.args.{nested_key}"] = nested_value if kwargs: for key, value in kwargs.items(): flattened = flatten_attributes({key: value}) for nested_key, nested_value in flattened.items(): attributes[f"{prefix}.{nested_key}"] = nested_value if attributes: self._recording_context.record_attributes(sanitize_attributes(attributes)) def set_output(self, output: Any) -> None: """Record the output value on the current span. This is intended for use inside a `with operation(...) as op` block. Args: output: The output value to record. """ if not self._recording_context: raise RuntimeError("No recording context found. Cannot set output.") flattened = flatten_attributes({LightningSpanAttributes.OPERATION_OUTPUT.value: output}) self._recording_context.record_attributes(sanitize_attributes(flattened)) def __call__(self, fn: _FnType) -> _FnType: """Wrap a callable so its execution is traced in a span. When used as a decorator, a new span is created for each call to the wrapped function. The bound arguments are recorded as input attributes, the return value is recorded as an output attribute, and any exception is recorded and marks the span as an error. Args: fn: The function or coroutine function to wrap. Returns: The wrapped callable. """ function_name = fn.__name__ sig = inspect.signature(fn) sanitized_init_attrs = sanitize_attributes( {LightningSpanAttributes.OPERATION_NAME.value: function_name, **self.initial_attributes} ) def _record_auto_inputs( recording_ctx: SpanRecordingContext, args: Tuple[Any, ...], kwargs: Dict[str, Any] ) -> None: """Bind arguments to signature and log them on the span.""" attributes: Dict[str, Any] = {} try: bound = sig.bind(*args, **kwargs) bound.apply_defaults() for name, value in bound.arguments.items(): parameter = sig.parameters.get(name) if parameter and parameter.kind is inspect.Parameter.VAR_POSITIONAL: attr_prefix = f"{LightningSpanAttributes.OPERATION_INPUT.value}.{name}" for idx, item in enumerate(value): flattened = flatten_attributes({str(idx): item}) for nested_key, nested_value in flattened.items(): attributes[f"{attr_prefix}.{nested_key}"] = nested_value else: flattened = flatten_attributes({name: value}) for nested_key, nested_value in flattened.items(): attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.{nested_key}"] = nested_value except Exception: if args: for idx, value in enumerate(args): flattened = flatten_attributes({str(idx): value}) for nested_key, nested_value in flattened.items(): attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.args.{nested_key}"] = ( nested_value ) if kwargs: flattened = flatten_attributes({"kwargs": kwargs}) for nested_key, nested_value in flattened.items(): attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.{nested_key}"] = nested_value if attributes: recording_ctx.record_attributes(sanitize_attributes(attributes)) def _record_auto_outputs(recording_ctx: SpanRecordingContext, result: Any) -> None: """Record the output value on the span.""" flattened = flatten_attributes({LightningSpanAttributes.OPERATION_OUTPUT.value: result}) recording_ctx.record_attributes(sanitize_attributes(flattened)) if inspect.iscoroutinefunction(fn) or ( # For backwards compatibility. hasattr(asyncio, "iscoroutinefunction") and asyncio.iscoroutinefunction(fn) # type: ignore ): @functools.wraps(fn) async def async_wrapper(*args: Any, **kwargs: Any) -> Any: """Async wrapper that traces the wrapped coroutine.""" with self.tracer.operation_context(self.name, attributes=sanitized_init_attrs) as recording_ctx: _record_auto_inputs(recording_ctx, args, kwargs) result = await fn(*args, **kwargs) _record_auto_outputs(recording_ctx, result) return result return cast(_FnType, async_wrapper) else: @functools.wraps(fn) def sync_wrapper(*args: Any, **kwargs: Any) -> Any: """Sync wrapper that traces the wrapped callable.""" with self.tracer.operation_context(self.name, attributes=sanitized_init_attrs) as recording_ctx: _record_auto_inputs(recording_ctx, args, kwargs) result = fn(*args, **kwargs) _record_auto_outputs(recording_ctx, result) return result return cast(_FnType, sync_wrapper) @overload def operation( fn: _FnType, *, propagate: bool = True, name: Optional[str] = None, **additional_attributes: Any ) -> _FnType: ... @overload def operation( *, propagate: bool = True, name: Optional[str] = None, **additional_attributes: Any ) -> OperationContext: ... @overload def operation(fn: _FnType, *, name: Optional[str] = None, **additional_attributes: Any) -> _FnType: ... @overload def operation(*, name: Optional[str] = None, **additional_attributes: Any) -> OperationContext: ... @overload def operation(fn: _FnType, **additional_attributes: Any) -> _FnType: ... @overload def operation(**additional_attributes: Any) -> OperationContext: ... def operation( fn: Optional[_FnType] = None, *, propagate: bool = True, name: Optional[str] = None, **additional_attributes: Any, ) -> Union[_FnType, OperationContext]: """Entry point for tracking operations. This helper can be used either as a decorator or as a context manager. The span name is fixed to [`AGL_OPERATION`][agentlightning.semconv.AGL_OPERATION]; custom span names are not supported. Any keyword arguments are recorded as span attributes. Usage as a decorator: ```python @operation def func(...): ... @operation(category="compute") def func(...): ... ``` Usage as a context manager: ```python with operation(user_id=123) as op: op.set_input(data=data) # ... do work ... op.set_output(result) ``` Args: fn: When used as `@operation`, this is the wrapped function. When used as `operation(**attrs)`, this should be omitted (or left as `None`) and only keyword attributes are provided. propagate: Whether spans should use the active span processor. When False, spans will stay local and not be exported. name: Optional alias that populates [`LightningSpanAttributes.OPERATION_NAME`][agentlightning.semconv.LightningSpanAttributes.OPERATION_NAME] when `additional_attributes` does not already define it. **additional_attributes: Additional span attributes to attach at creation time. Returns: Either a wrapped callable (when used as a decorator) or an [`OperationContext`][agentlightning.emitter.annotation.OperationContext] (when used as a context manager factory). """ if name is not None: if LightningSpanAttributes.OPERATION_NAME.value in additional_attributes: raise ValueError("Cannot specify both `name` and `additional_attributes.operation_name`.") additional_attributes[LightningSpanAttributes.OPERATION_NAME.value] = name # Case 1: Used as @operation (bare decorator or with attributes) if callable(fn): # Create context with fixed name, then immediately wrap the function return OperationContext(AGL_OPERATION, additional_attributes, propagate=propagate)(fn) # Case 2: Used as operation(...) / with operation(...) # Custom span names are intentionally not supported; use AGL_OPERATION. if fn is not None: raise ValueError("Custom span names are intentionally not supported when used as a context manager.") return OperationContext(AGL_OPERATION, additional_attributes, propagate=propagate)