Compare commits

...

12 Commits

Author SHA1 Message Date
Yuge Zhang aa305ccca9 . 2025-12-04 14:23:12 +08:00
Yuge Zhang 466d6e1501 merge operation into annotation 2025-12-04 14:20:02 +08:00
Yuge Zhang 57ab004f8c . 2025-12-04 09:19:45 +08:00
Yuge Zhang 628cbf8d29 update docs 2025-12-04 09:04:02 +08:00
Yuge Zhang 4ec6aa27bf add more tests 2025-12-04 08:50:31 +08:00
Yuge Zhang 4794b20871 add tests 2025-12-04 00:31:54 +08:00
Yuge Zhang 276cf930d8 add operation name 2025-12-04 00:05:35 +08:00
Yuge Zhang 9a2fd5966f . 2025-12-03 23:59:02 +08:00
Yuge Zhang 96f6fe3c27 minor fix 2025-12-03 23:54:56 +08:00
Yuge Zhang 2aeed81caf update operations 2025-12-03 18:35:08 +08:00
Yuge Zhang 226c1b4b7d update contributing 2025-12-03 18:32:24 +08:00
Yuge Zhang ba0b03a7f7 init AGL operation 2025-12-03 18:30:47 +08:00
8 changed files with 790 additions and 6 deletions
+2 -1
View File
@@ -1,6 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
from .annotation import emit_annotation
from .annotation import emit_annotation, operation
from .exception import emit_exception
from .message import emit_message, get_message_value
from .object import emit_object, get_object_value
@@ -16,6 +16,7 @@ from .reward import (
__all__ = [
"reward",
"operation",
"emit_reward",
"get_reward_value",
"get_rewards_from_span",
+319 -3
View File
@@ -1,15 +1,36 @@
# Copyright (c) Microsoft. All rights reserved.
"""Helpers for emitting annotation spans."""
"""Helpers for emitting annotation/operation spans."""
import asyncio
import functools
import inspect
import json
import logging
from typing import Any, Dict
from types import TracebackType
from typing import (
Any,
Callable,
ContextManager,
Dict,
Optional,
Tuple,
Type,
TypeVar,
Union,
cast,
overload,
)
from opentelemetry import trace
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.trace import Status, StatusCode
from agentlightning.semconv import AGL_ANNOTATION
from agentlightning.semconv import AGL_ANNOTATION, AGL_OPERATION, LightningSpanAttributes
from agentlightning.utils.otel import flatten_attributes, get_tracer
_FnType = TypeVar("_FnType", bound=Callable[..., Any])
logger = logging.getLogger(__name__)
@@ -46,3 +67,298 @@ def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> Reada
raise ValueError(f"Span is not a ReadableSpan: {span}")
return span
def _safe_json_dump(obj: Any) -> str:
"""Serialize an object to JSON, falling back to ``str(obj)`` if needed.
Args:
obj: Object to be serialized.
Returns:
The JSON-encoded string representation of the object, or its string
representation if JSON encoding fails.
"""
try:
return json.dumps(obj, default=str, ensure_ascii=False)
except Exception:
return str(obj)
class OperationContext:
"""Context manager and decorator for tracing operations.
This class manages an OpenTelemetry 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 :meth:`set_input` and :meth:`set_output`.
Attributes:
name: Human-readable span name.
initial_attributes: Attributes applied when the span is created.
tracer: OpenTelemetry tracer used to create spans.
span: The currently active span, if any.
"""
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: str = name
self.initial_attributes: Dict[str, Any] = attributes
self.propagate: bool = propagate
self.tracer: trace.Tracer = get_tracer(use_active_span_processor=propagate)
self.span: Optional[trace.Span] = None
self._ctx_token: Optional[ContextManager[Any]] = None
def __enter__(self) -> "OperationContext":
"""Enter the context manager and start a new span.
Returns:
The current :class:`OperationContext` instance with an active span.
"""
# 1. Start the span with initial attributes (JSON serialized)
sanitized_attrs = {
k: _safe_json_dump(v) if not isinstance(v, (str, int, float, bool)) else v
for k, v in self.initial_attributes.items()
}
self.span = self.tracer.start_span(self.name, attributes=sanitized_attrs)
self._ctx_token = trace.use_span(self.span, end_on_exit=True)
self._ctx_token.__enter__()
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.
Any exception raised inside the context is recorded on the span and the
span status is set to error.
Args:
exc_type: Exception type, if an exception occurred.
exc_val: Exception instance, if an exception occurred.
exc_tb: Traceback object, if an exception occurred.
"""
# 1. Record Exception if present
if exc_val and self.span:
self.span.record_exception(exc_val)
self.span.set_status(Status(StatusCode.ERROR, str(exc_val)))
# 2. Close span
if self._ctx_token:
self._ctx_token.__exit__(exc_type, exc_val, exc_tb)
def set_input(self, *args: Any, **kwargs: Any) -> None:
"""Record input arguments on the current span.
Positional arguments are stored under the ``input.args`` attribute,
and keyword arguments are stored under ``input.<name>`` 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.span:
return
if args:
self.span.set_attribute("input.args", _safe_json_dump(args))
if kwargs:
for k, v in kwargs.items():
self.span.set_attribute(f"input.{k}", _safe_json_dump(v))
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.span:
return
self.span.set_attribute("output", _safe_json_dump(output))
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)
def _record_auto_inputs(span: trace.Span, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> None:
"""Bind arguments to signature and log them on the span.
Args:
span: Span on which to record attributes.
args: Positional arguments passed to the wrapped callable.
kwargs: Keyword arguments passed to the wrapped callable.
"""
try:
bound = sig.bind(*args, **kwargs)
bound.apply_defaults()
for k, v in bound.arguments.items():
span.set_attribute(
f"{LightningSpanAttributes.OPERATION_INPUT.value}.{k}",
_safe_json_dump(v),
)
except Exception:
span.set_attribute(
f"{LightningSpanAttributes.OPERATION_INPUT.value}.args",
_safe_json_dump(args),
)
span.set_attribute(
f"{LightningSpanAttributes.OPERATION_INPUT.value}.kwargs",
_safe_json_dump(kwargs),
)
if asyncio.iscoroutinefunction(fn) or inspect.iscoroutinefunction(fn):
@functools.wraps(fn)
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
"""Async wrapper that traces the wrapped coroutine."""
# Reuse __enter__ logic via 'with self' would share state incorrectly
# across concurrent calls. We must create a new span per call.
# So we manually reimplement the span logic for the wrapper here.
sanitized_attrs = {
k: _safe_json_dump(v) if not isinstance(v, (str, int, float, bool)) else v
for k, v in self.initial_attributes.items()
}
with self.tracer.start_as_current_span(self.name, attributes=sanitized_attrs) as span:
span.set_attribute(LightningSpanAttributes.OPERATION_NAME.value, function_name)
_record_auto_inputs(span, args, kwargs)
try:
result = await fn(*args, **kwargs)
span.set_attribute(
LightningSpanAttributes.OPERATION_OUTPUT.value,
_safe_json_dump(result),
)
return result
except Exception as e:
span.record_exception(e)
span.set_status(Status(StatusCode.ERROR, str(e)))
raise
return cast(_FnType, async_wrapper)
else:
@functools.wraps(fn)
def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
"""Sync wrapper that traces the wrapped callable."""
sanitized_attrs = {
k: _safe_json_dump(v) if not isinstance(v, (str, int, float, bool)) else v
for k, v in self.initial_attributes.items()
}
with self.tracer.start_as_current_span(self.name, attributes=sanitized_attrs) as span:
span.set_attribute(LightningSpanAttributes.OPERATION_NAME.value, function_name)
_record_auto_inputs(span, args, kwargs)
try:
result = fn(*args, **kwargs)
span.set_attribute(
LightningSpanAttributes.OPERATION_OUTPUT.value,
_safe_json_dump(result),
)
return result
except Exception as e:
span.record_exception(e)
span.set_status(Status(StatusCode.ERROR, str(e)))
raise
return cast(_FnType, sync_wrapper)
@overload
def operation(fn: _FnType, *, propagate: bool = True, **additional_attributes: Any) -> _FnType: ...
@overload
def operation(*, propagate: bool = True, **additional_attributes: Any) -> OperationContext: ...
def operation(
fn: Optional[_FnType] = None,
*,
propagate: bool = True,
**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.
**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).
"""
# 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)
+14
View File
@@ -29,6 +29,11 @@ AGL_EXCEPTION = "agentlightning.exception"
Used by the exception emitter to record exception details.
"""
AGL_OPERATION = "agentlightning.operation"
"""Agent-lightning's standard span name for functions.
Wrap function or code-blocks as operations.
"""
AGL_VIRTUAL = "agentlightning.virtual"
"""Agent-lightning's standard span name for virtual operations.
@@ -84,6 +89,15 @@ class LightningSpanAttributes(Enum):
OBJECT_JSON = "agentlightning.object.json"
"""Attribute name for object serialized value (JSON) in object spans."""
OPERATION_NAME = "agentlightning.operation.name"
"""Attribute name for operation name in operation spans, normally the function name."""
OPERATION_INPUT = "agentlightning.operation.input"
"""Attribute name for operation input in operation spans."""
OPERATION_OUTPUT = "agentlightning.operation.output"
"""Attribute name for operation output in operation spans."""
class RewardAttributes(Enum):
"""Multi-dimensional reward attributes will look like:
+1 -1
View File
@@ -28,7 +28,7 @@ Documentation improvements are the easiest way to get started. You can find more
Bug fixes are the fastest way to get familiar with the codebase. To get started, you can:
- Browse the ["good first issue"](https://github.com/microsoft/agent-lightning/labels/good%20first%20issue) and ["bug"](https://github.com/microsoft/agent-lightning/labels/bug) labels; drop a comment before you start so we can mark it as taken.
- Browse the ["help wanted"](https://github.com/microsoft/agent-lightning/labels/help%20wanted) and ["bug"](https://github.com/microsoft/agent-lightning/labels/bug) labels; drop a comment before you start so we can mark it as taken.
- For fresh bugs, open an issue with reproduction steps, logs, and expected behavior before submitting a fix.
- Keep each pull request focused, ideally avoiding breaking API changes. Larger refactors should be discussed via RFC or maintainer sync.
+2
View File
@@ -22,6 +22,8 @@
## Emitter
::: agentlightning.operation
::: agentlightning.emit_annotation
::: agentlightning.emit_reward
+2
View File
@@ -24,6 +24,8 @@
::: agentlightning.litagent.decorator.prompt_rollout
::: agentlightning.emitter.annotation.OperationContext
## LLM Proxy
::: agentlightning.llm_proxy.ModelConfig
+54 -1
View File
@@ -211,6 +211,8 @@ While returning a single float for the final reward is sufficient for many algor
Agent-lightning provides an **emitter** module that allows you to record custom spans from within your agent's logic. Like many common operations (like LLM calls) that are automatically instrumented by [Tracer][agentlightning.Tracer], the emitter will also send a [Span][agentlightning.Span] that records an Agent-lightning-specific operation. Then algorithms can query and read those spans later. See [Working with Traces](./traces.md) for more details.
For multi-step routines (function calls, tools, or adapters) you can wrap code with [`operation`][agentlightning.operation], either as a decorator or a context manager,to capture inputs, outputs, and metadata on a dedicated `"agentlightning.operation"` span. This makes it easier to correlate downstream annotations (like rewards or messages) with the higher-level work that produced them.
You can find the emitter functions from [agentlightning.emitter](../reference/agent.md).
### Emitting Rewards, Messages, and More
@@ -221,7 +223,6 @@ Here are the primary emitter functions:
* [`emit_message(message: str)`][agentlightning.emit_message]: Records a simple log message as a span.
* [`emit_exception(exception: BaseException)`][agentlightning.emit_exception]: Records a Python exception, including its type, message, and stack trace.
* [`emit_object(obj: Any)`][agentlightning.emit_object]: Records any JSON-serializable object, perfect for structured data.
Let's see an example of an agent using these emitters to provide detailed feedback.
```python
@@ -256,3 +257,55 @@ def multi_step_agent(task: dict, prompt_template: PromptTemplate) -> float:
```
By using the emitter, you create a rich, detailed trace of your agent's execution. This data can be invaluable for debugging and is essential for advanced algorithms that can learn from more than just a single final score.
### Linking to Other Spans
Sometimes a span should explicitly point back to another span that produced the input it is working on (for example, linking a reward annotation to the `"agentlightning.operation"` span that generated a response). Agent-lightning encodes these relationships through flattened link attributes. The helper [`make_link_attributes`][agentlightning.utils.otel.make_link_attributes] converts a dictionary of keys—such as `trace_id`, `span_id`, or any custom attribute—into the `"agentlightning.link.*"` fields expected by the backend. Later on, [`query_linked_spans`][agentlightning.utils.otel.query_linked_spans] can be used to recover the original span(s) from those link descriptors.
```python
import opentelemetry.trace as trace_api
from agentlightning import emit_annotation, operation
from agentlightning.utils.otel import make_link_attributes, make_tag_attributes
with operation(conversation_id="chat-42") as op:
# ... perform the work ...
span_ctx = op.span.get_span_context()
link_attrs = make_link_attributes({
"conversation_id": "chat-42",
})
emit_annotation(
{
**link_attrs,
**make_tag_attributes(["reward", "good"]),
}
)
```
When analyzing in adapters, pass the extracted link models to [`query_linked_spans`][agentlightning.utils.otel.query_linked_spans] to retrieve the matching span(s):
```python
from agentlightning.utils.otel import extract_links_from_attributes, query_linked_spans
annotation_span = ... # Span from your trace store
operation_spans = [...] # list of spans you want to search
link_models = extract_links_from_attributes(annotation_span.attributes)
matches = query_linked_spans(operation_spans, link_models)
assert matches # Contains the original operation span
```
!!! tip "Correlating Rewards with LLM Requests"
[Tracer](./traces.md) instruments each request/response as its own span. You can link to the [`gen_ai.response.id`](https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-events/) attribute, which comes from the LLM response ID.
```python
from agentlightning import emit_reward
from agentlightning.utils.otel import make_link_attributes
result = call_llm(prompt)
reward_links = make_link_attributes({"gen_ai.response.id": result.id})
emit_reward(0.9, attributes=reward_links)
```
Later, use the same `gen_ai.response.id` key inside `query_linked_spans` to find the reward(s) that reference that specific LLM request span.
+396
View File
@@ -0,0 +1,396 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import json
from dataclasses import dataclass
from types import TracebackType
from typing import Any, Dict, List, Optional, Tuple, Type
import opentelemetry.trace as trace_api
import pytest
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.trace import Status, StatusCode
import agentlightning.emitter.annotation as annotation_module
from agentlightning.emitter.annotation import _safe_json_dump # pyright: ignore[reportPrivateUsage]
from agentlightning.emitter.annotation import (
OperationContext,
emit_annotation,
operation,
)
from agentlightning.semconv import AGL_ANNOTATION, AGL_OPERATION, LightningSpanAttributes
from agentlightning.utils.otel import extract_links_from_attributes, make_link_attributes, query_linked_spans
class RecordingSpan:
def __init__(self) -> None:
self.attributes: Dict[str, Any] = {}
self.recorded_exceptions: List[BaseException] = []
self.statuses: List[Status] = []
def set_attribute(self, key: str, value: Any) -> None:
self.attributes[key] = value
def record_exception(self, exc: BaseException) -> None:
self.recorded_exceptions.append(exc)
def set_status(self, status: Status) -> None:
self.statuses.append(status)
class DummySpanContextManager:
def __init__(self, span: RecordingSpan) -> None:
self.span = span
self.exit_calls: List[
Tuple[Optional[Type[BaseException]], Optional[BaseException], Optional[TracebackType]]
] = []
def __enter__(self) -> RecordingSpan:
return self.span
def __exit__(
self,
exc_type: Optional[Type[BaseException]],
exc_val: Optional[BaseException],
exc_tb: Optional[TracebackType],
) -> bool:
self.exit_calls.append((exc_type, exc_val, exc_tb))
return False
class DummyTracer:
def __init__(self, start_span_instance: Optional[RecordingSpan] = None) -> None:
self._start_span_instance = start_span_instance
self.start_span_calls: List[Tuple[str, Dict[str, Any]]] = []
self.start_as_current_span_calls: List[Tuple[str, Dict[str, Any], RecordingSpan]] = []
def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> RecordingSpan:
span = self._start_span_instance or RecordingSpan()
self.start_span_calls.append((name, dict(attributes or {})))
return span
def start_as_current_span(
self,
name: str,
attributes: Optional[Dict[str, Any]] = None,
) -> DummySpanContextManager:
span = RecordingSpan()
self.start_as_current_span_calls.append((name, dict(attributes or {}), span))
return DummySpanContextManager(span)
class DummyUseSpan:
def __init__(self) -> None:
self.calls: List[Tuple[RecordingSpan, bool]] = []
self.exit_calls: List[
Tuple[Optional[Type[BaseException]], Optional[BaseException], Optional[TracebackType]]
] = []
def __call__(self, span: RecordingSpan, end_on_exit: bool) -> DummyUseSpan:
self.calls.append((span, end_on_exit))
self._span = span
return self
def __enter__(self) -> None:
return None
def __exit__(
self,
exc_type: Optional[Type[BaseException]],
exc_val: Optional[BaseException],
exc_tb: Optional[TracebackType],
) -> bool:
self.exit_calls.append((exc_type, exc_val, exc_tb))
return False
@dataclass
class ComplexResult:
values: Tuple[int, ...]
marker: str
def test_safe_json_dump_handles_recursive_structures() -> None:
payload: List[Any] = []
payload.append(payload)
assert _safe_json_dump(payload) == "[[...]]"
def test_operation_context_records_inputs_and_outputs(monkeypatch: pytest.MonkeyPatch) -> None:
span = RecordingSpan()
tracer = DummyTracer(start_span_instance=span)
use_span = DummyUseSpan()
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
ctx = OperationContext("custom-span", {"meta": {"foo": 1}, "count": 2})
with ctx as op:
op.set_input({"payload": 1}, flag=True)
op.set_output({"success": True})
assert tracer.start_span_calls
start_name, start_attributes = tracer.start_span_calls[0]
assert start_name == "custom-span"
assert json.loads(start_attributes["meta"]) == {"foo": 1}
assert start_attributes["count"] == 2
assert json.loads(span.attributes["input.args"]) == [{"payload": 1}]
assert span.attributes["input.flag"] == "true"
assert json.loads(span.attributes["output"]) == {"success": True}
assert use_span.calls == [(span, True)]
def test_operation_context_set_input_supports_multiple_values(monkeypatch: pytest.MonkeyPatch) -> None:
span = RecordingSpan()
tracer = DummyTracer(start_span_instance=span)
use_span = DummyUseSpan()
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
ctx = OperationContext("ctx", {})
with ctx as op:
op.set_input(1, 2, data={"foo": ["bar"]}, flags=[True, False])
assert json.loads(span.attributes["input.args"]) == [1, 2]
assert json.loads(span.attributes["input.data"]) == {"foo": ["bar"]}
assert json.loads(span.attributes["input.flags"]) == [True, False]
def test_operation_context_records_non_serializable_output(monkeypatch: pytest.MonkeyPatch) -> None:
class Unserializable:
def __str__(self) -> str:
return "<Unserializable>"
span = RecordingSpan()
tracer = DummyTracer(start_span_instance=span)
use_span = DummyUseSpan()
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
ctx = OperationContext("ctx", {})
with ctx as op:
op.set_output(Unserializable())
assert json.loads(span.attributes["output"]) == "<Unserializable>"
def test_operation_context_records_exceptions(monkeypatch: pytest.MonkeyPatch) -> None:
span = RecordingSpan()
tracer = DummyTracer(start_span_instance=span)
use_span = DummyUseSpan()
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
ctx = OperationContext("custom-span", {})
with pytest.raises(RuntimeError):
with ctx:
raise RuntimeError("boom")
assert isinstance(span.recorded_exceptions[0], RuntimeError)
status = span.statuses[-1]
assert status.status_code == StatusCode.ERROR
assert status.description == "boom"
assert use_span.exit_calls[-1][1].args == ("boom",) # type: ignore
def test_operation_factory_context_records_inputs_and_outputs(monkeypatch: pytest.MonkeyPatch) -> None:
span = RecordingSpan()
tracer = DummyTracer(start_span_instance=span)
use_span = DummyUseSpan()
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
with operation(tags=["one", "two"]) as ctx:
ctx.set_input("alpha", meta={"score": 0.5})
ctx.set_output(["beta", "gamma"])
start_name, attrs = tracer.start_span_calls[0]
assert start_name == AGL_OPERATION
assert json.loads(attrs["tags"]) == ["one", "two"]
assert json.loads(span.attributes["input.args"]) == ["alpha"]
assert json.loads(span.attributes["input.meta"]) == {"score": 0.5}
assert json.loads(span.attributes["output"]) == ["beta", "gamma"]
def test_operation_factory_uses_standard_span_name(monkeypatch: pytest.MonkeyPatch) -> None:
span = RecordingSpan()
tracer = DummyTracer(start_span_instance=span)
use_span = DummyUseSpan()
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
with operation(user={"id": 5}) as ctx:
ctx.set_output("done")
assert tracer.start_span_calls
start_name, attrs = tracer.start_span_calls[0]
assert start_name == AGL_OPERATION
assert json.loads(attrs["user"]) == {"id": 5}
def test_operation_rejects_custom_span_names() -> None:
with pytest.raises(ValueError):
operation("custom-name") # type: ignore
def test_operation_decorator_sync_records_span_attributes(monkeypatch: pytest.MonkeyPatch) -> None:
tracer = DummyTracer()
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
@operation(category={"kind": "combine"})
def combine(data: Dict[str, int], *, meta: Dict[str, str]) -> Dict[str, Any]:
return {"joined": {**data, **meta}}
result = combine({"value": 1}, meta={"source": "unit"})
assert result == {"joined": {"value": 1, "source": "unit"}}
assert tracer.start_as_current_span_calls
span_name, span_attributes, span = tracer.start_as_current_span_calls[0]
assert span_name == AGL_OPERATION
assert json.loads(span_attributes["category"]) == {"kind": "combine"}
input_prefix = LightningSpanAttributes.OPERATION_INPUT.value
assert json.loads(span.attributes[f"{input_prefix}.data"]) == {"value": 1}
assert json.loads(span.attributes[f"{input_prefix}.meta"]) == {"source": "unit"}
assert span.attributes[LightningSpanAttributes.OPERATION_NAME.value] == "combine"
assert json.loads(span.attributes[LightningSpanAttributes.OPERATION_OUTPUT.value]) == result
def test_operation_decorator_handles_complex_signature(monkeypatch: pytest.MonkeyPatch) -> None:
tracer = DummyTracer()
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
@operation()
def complicated(
first: int,
/,
required: str,
default: int = 5,
*extra: int,
kwonly: str,
kwdefault: str = "fallback",
**rest: Any,
) -> ComplexResult:
return ComplexResult(values=(first, len(extra), len(rest)), marker=kwonly + kwdefault + required)
result = complicated(1, "req", 7, 8, 9, kwonly="x", kwdefault="y", tag="value")
span = tracer.start_as_current_span_calls[0][2]
input_prefix = LightningSpanAttributes.OPERATION_INPUT.value
assert json.loads(span.attributes[f"{input_prefix}.first"]) == 1
assert json.loads(span.attributes[f"{input_prefix}.required"]) == "req"
assert json.loads(span.attributes[f"{input_prefix}.default"]) == 7
assert json.loads(span.attributes[f"{input_prefix}.extra"]) == [8, 9]
assert json.loads(span.attributes[f"{input_prefix}.kwonly"]) == "x"
assert json.loads(span.attributes[f"{input_prefix}.kwdefault"]) == "y"
assert json.loads(span.attributes[f"{input_prefix}.rest"]) == {"tag": "value"}
assert span.attributes[LightningSpanAttributes.OPERATION_NAME.value] == "complicated"
assert json.loads(span.attributes[LightningSpanAttributes.OPERATION_OUTPUT.value]) == str(result)
def test_operation_decorator_records_exceptions(monkeypatch: pytest.MonkeyPatch) -> None:
tracer = DummyTracer()
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
@operation()
def fail(value: int) -> int:
raise ValueError("bad input")
with pytest.raises(ValueError):
fail(1)
span = tracer.start_as_current_span_calls[0][2]
assert isinstance(span.recorded_exceptions[0], ValueError)
status = span.statuses[-1]
assert status.status_code == StatusCode.ERROR
assert status.description == "bad input"
@pytest.mark.asyncio()
async def test_operation_async_wrapper_records_attributes(monkeypatch: pytest.MonkeyPatch) -> None:
tracer = DummyTracer()
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
@operation()
async def echo(payload: Dict[str, Any]) -> Dict[str, Any]:
return {"payload": payload}
result = await echo({"value": 3})
assert result == {"payload": {"value": 3}}
span = tracer.start_as_current_span_calls[0][2]
prefix = LightningSpanAttributes.OPERATION_INPUT.value
assert json.loads(span.attributes[f"{prefix}.payload"]) == {"value": 3}
assert json.loads(span.attributes[LightningSpanAttributes.OPERATION_OUTPUT.value]) == result
def test_operation_span_can_be_resolved_via_annotation_links(monkeypatch: pytest.MonkeyPatch) -> None:
provider = TracerProvider()
exporter = InMemorySpanExporter()
provider.add_span_processor(SimpleSpanProcessor(exporter))
tracer = provider.get_tracer(__name__)
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
@operation(conversation_id="conv-1")
def decorated(value: int) -> int:
return value + 1
assert decorated(41) == 42
spans = exporter.get_finished_spans()
operation_span = next(span for span in spans if span.name == AGL_OPERATION)
assert operation_span.attributes["conversation_id"] == "conv-1" # type: ignore
trace_id_hex = trace_api.format_trace_id(operation_span.context.trace_id) # type: ignore
span_id_hex = trace_api.format_span_id(operation_span.context.span_id) # type: ignore
link_attrs = make_link_attributes({"trace_id": trace_id_hex, "span_id": span_id_hex})
emit_annotation({**link_attrs, "note": "operation-follow-up"})
spans = exporter.get_finished_spans()
annotation_span = next(span for span in spans if span.name == AGL_ANNOTATION)
annotation_links = extract_links_from_attributes(dict(annotation_span.attributes or {}))
matches = query_linked_spans([operation_span], annotation_links)
assert matches == [operation_span]
def test_operation_honors_propagate_flag(monkeypatch: pytest.MonkeyPatch) -> None:
tracer = DummyTracer()
flags: List[bool] = []
use_span = DummyUseSpan()
def fake_get_tracer(use_active_span_processor: bool = True) -> DummyTracer:
flags.append(use_active_span_processor)
return tracer
monkeypatch.setattr(annotation_module, "get_tracer", fake_get_tracer)
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
@operation(propagate=False)
def decorated(value: int) -> int:
return value
assert decorated(7) == 7
with operation(propagate=False):
pass
assert flags == [False, False]