chore: import upstream snapshot with attribution

agentlightning/types/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .core import *
+from .resources import *
+from .tracer import *

agentlightning/types/core.py

@@ -0,0 +1,553 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Core data models shared across Agent Lightning components."""
+
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    Iterator,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Protocol,
+    Sequence,
+    SupportsIndex,
+    TypedDict,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+
+from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import BaseModel, Field, model_validator
+
+from .tracer import Span, SpanCoreFields
+
+if TYPE_CHECKING:
+    from agentlightning.litagent import LitAgent
+    from agentlightning.runner.base import Runner
+    from agentlightning.tracer.base import Tracer
+
+__all__ = [
+    "Triplet",
+    "RolloutLegacy",
+    "Task",
+    "TaskInput",
+    "TaskIfAny",
+    "RolloutRawResultLegacy",
+    "RolloutRawResult",
+    "RolloutMode",
+    "GenericResponse",
+    "ParallelWorkerBase",
+    "Dataset",
+    "AttemptStatus",
+    "RolloutStatus",
+    "RolloutConfig",
+    "Rollout",
+    "Attempt",
+    "AttemptedRollout",
+    "EnqueueRolloutRequest",
+    "Hook",
+    "Worker",
+    "WorkerStatus",
+    "PaginatedResult",
+    "FilterOptions",
+    "SortOptions",
+    "FilterField",
+]
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+class Triplet(BaseModel):
+    """Single interaction turn captured during reinforcement learning."""
+
+    prompt: Any
+    response: Any
+    reward: Optional[float] = None
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class RolloutLegacy(BaseModel):
+    """Legacy reporting payload exchanged with the deprecated HTTP server.
+
+    !!! warning "Deprecated"
+        Use [`Rollout`][agentlightning.Rollout] instead.
+    """
+
+    rollout_id: str
+
+    # Echoing the input task
+    task: Optional[Task] = None
+
+    # Primary, high-level feedback
+    final_reward: Optional[float] = None
+
+    # Structured, sequential feedback for RL-style optimization
+    triplets: Optional[List[Triplet]] = None
+
+    # Optional, rich-context data for deep analysis
+    trace: Optional[List[Dict[str, Any]]] = Field(
+        default=None,
+        description="A list of spans that conform to the OpenTelemetry JSON format. "
+        "Users of the opentelemetry-sdk can generate this by calling "
+        "json.loads(readable_span.to_json()).",
+    )
+    logs: Optional[List[str]] = None
+
+    # A bucket for any other relevant information
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+RolloutStatus = Literal[
+    "queuing",  # initial status
+    "preparing",  # after the trace is claimed
+    "running",  # after receiving the first trace
+    "failed",  # crashed
+    "succeeded",  # status OK
+    "cancelled",  # cancelled by user (or watchdog)
+    "requeuing",  # retrying
+]
+"""The status of a rollout."""
+
+AttemptStatus = Literal[
+    # A status is essentially a process.
+    # It should not have scheduling/management statuses like "queuing" or "cancelled".
+    "preparing",
+    "running",
+    "failed",
+    "succeeded",
+    "unresponsive",  # the worker has not reported results for a while
+    "timeout",  # the worker has been emitting new logs, but have been working on the task for too long
+]
+"""The status of an attempt."""
+
+RolloutMode = Literal["train", "val", "test"]
+"""Possible rollout modes."""
+
+
+class Attempt(BaseModel):
+    """Execution attempt for a rollout, including metadata for retries."""
+
+    rollout_id: str
+    """The rollout which this attempt belongs to."""
+    attempt_id: str
+    """The universal id for current attempt."""
+    sequence_id: int
+    """The sequence number of the attempt, starting from 1."""
+    start_time: float
+    """The time when the attempt has started."""
+    end_time: Optional[float] = None
+    """The time when the attempt has ended."""
+    status: AttemptStatus = "preparing"
+    """The status of the attempt."""
+    worker_id: Optional[str] = None
+    """The rollout worker which is executing this attempt."""
+
+    last_heartbeat_time: Optional[float] = None
+    """The last time when the worker has reported progress (i.e., a span)."""
+
+    metadata: Optional[Dict[str, Any]] = None
+    """A bucket for any other relevant information."""
+
+
+class RolloutConfig(BaseModel):
+    """Configuration controlling rollout retries and timeouts."""
+
+    timeout_seconds: Optional[float] = None
+    """The timeout for the rollout, in seconds. None indicates no timeout."""
+    unresponsive_seconds: Optional[float] = None
+    """The unresponsive timeout for the rollout, in seconds. None indicates no unresponsive timeout."""
+    max_attempts: int = Field(default=1, ge=1)
+    """The maximum number of attempts for the rollout, including the first attempt."""
+    retry_condition: List[AttemptStatus] = Field(default_factory=cast(Callable[[], List[AttemptStatus]], list))
+    """The list of statuses that should trigger a retry."""
+
+
+class Rollout(BaseModel):
+    rollout_id: str
+    """Unique identifier for the rollout."""
+
+    input: TaskInput
+    """Task input used to generate the rollout."""
+
+    # Time to track the lifecycle of the rollout
+    start_time: float
+    """Timestamp when the rollout started."""
+    end_time: Optional[float] = None
+    """Timestamp when the rollout ended."""
+
+    mode: Optional[RolloutMode] = None
+    """Execution mode such as `"train"`, `"val"` or `"test"`. See [`RolloutMode`][agentlightning.RolloutMode]."""
+    resources_id: Optional[str] = None
+    """Identifier of the resources required to execute the rollout."""
+
+    status: RolloutStatus = "queuing"
+    """Latest status emitted by the controller."""
+
+    config: RolloutConfig = Field(default_factory=RolloutConfig)
+    """Retry and timeout configuration associated with the rollout."""
+
+    metadata: Optional[Dict[str, Any]] = None
+    """Additional metadata attached to the rollout."""
+
+
+class AttemptedRollout(Rollout):
+    """Rollout paired with the currently active attempt."""
+
+    attempt: Attempt
+    """The attempt that is currently processing the rollout."""
+
+    @model_validator(mode="after")
+    def check_consistency(self) -> AttemptedRollout:
+        if self.attempt.rollout_id != self.rollout_id:
+            raise ValueError("Inconsistent rollout_id between Rollout and Attempt")
+        return self
+
+
+class EnqueueRolloutRequest(BaseModel):
+    """Payload describing a rollout to be queued via [`enqueue_rollout`][agentlightning.LightningStore.enqueue_rollout].
+
+    A subset of fields from [`Rollout`][agentlightning.Rollout] used for queuing new rollouts.
+    """
+
+    input: TaskInput
+    """Task input used to generate the rollout."""
+    mode: Optional[RolloutMode] = None
+    """Execution mode such as `"train"`, `"val"` or `"test"`. See [`RolloutMode`][agentlightning.RolloutMode]."""
+    resources_id: Optional[str] = None
+    """Identifier of the resources required to execute the rollout."""
+    config: Optional[RolloutConfig] = None
+    """Retry and timeout configuration associated with the rollout."""
+    metadata: Optional[Dict[str, Any]] = None
+    """Additional metadata attached to the rollout."""
+
+
+WorkerStatus = Literal["idle", "busy", "unknown"]
+
+
+class Worker(BaseModel):
+    """Worker information. This is actually the same as Runner info."""
+
+    worker_id: str
+    """The ID of the worker."""
+    status: WorkerStatus = "unknown"
+    """The status of the worker."""
+    heartbeat_stats: Optional[Dict[str, Any]] = None
+    """Statistics about the worker's heartbeat."""
+    last_heartbeat_time: Optional[float] = None
+    """The last time when the worker has reported the stats."""
+    last_dequeue_time: Optional[float] = None
+    """The last time when the worker has tried to dequeue a rollout."""
+    last_busy_time: Optional[float] = None
+    """The last time when the worker has started an attempt and became busy."""
+    last_idle_time: Optional[float] = None
+    """The last time when the worker has triggered the end of an attempt and became idle."""
+    current_rollout_id: Optional[str] = None
+    """The ID of the current rollout that the worker is processing."""
+    current_attempt_id: Optional[str] = None
+    """The ID of the current attempt that the worker is processing."""
+
+
+TaskInput = Any
+"""Task input type. Accepts arbitrary payloads."""
+
+
+class Task(BaseModel):
+    """Rollout request served to client agents.
+
+    !!! warning "Deprecated"
+        The legacy HTTP client/server stack still uses this model. Prefer
+        [`LightningStore`][agentlightning.LightningStore] APIs for new workflows.
+    """
+
+    rollout_id: str
+    input: TaskInput
+
+    mode: Optional[RolloutMode] = None
+    resources_id: Optional[str] = None
+
+    # Optional fields for tracking task lifecycle
+    create_time: Optional[float] = None
+    last_claim_time: Optional[float] = None
+    num_claims: Optional[int] = None
+
+    # Allow additional metadata fields
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class TaskIfAny(BaseModel):
+    """A task or indication that no task is available.
+
+    !!! warning "Deprecated"
+        Use [`LightningStore`][agentlightning.LightningStore] APIs for new workflows.
+    """
+
+    is_available: bool
+    """Indication that a task is available."""
+    task: Optional[Task] = None
+
+
+RolloutRawResultLegacy = Union[None, float, List[Triplet], List[Dict[str, Any]], List[ReadableSpan], RolloutLegacy]
+"""Legacy rollout result type.
+
+!!! warning "Deprecated"
+    Use [`RolloutRawResult`][agentlightning.RolloutRawResult] instead.
+"""
+
+RolloutRawResult = Union[
+    None,  # nothing (relies on tracer)
+    float,  # only final reward
+    List[ReadableSpan],  # constructed OTEL spans by user
+    List[Span],  # constructed Span objects by user
+    List[SpanCoreFields],  # constructed SpanCoreFields objects by user
+]
+"""Rollout result type.
+
+Possible return values of [`rollout`][agentlightning.LitAgent.rollout].
+"""
+
+
+class GenericResponse(BaseModel):
+    """Generic server response used by compatibility endpoints.
+
+    !!! warning "Deprecated"
+        This response is no longer used by the new
+        [`LightningStore`][agentlightning.LightningStore] APIs.
+
+    Attributes:
+        status: Status string describing the result of the request.
+        message: Optional human readable explanation.
+        data: Arbitrary payload serialized as JSON.
+    """
+
+    status: str = "success"
+    message: Optional[str] = None
+    data: Optional[Dict[str, Any]] = None
+
+
+class ParallelWorkerBase:
+    """Base class for workloads executed across multiple worker processes.
+
+    The lifecycle is orchestrated by the main process:
+
+    * [`init()`][agentlightning.ParallelWorkerBase.init] prepares shared state.
+    * Each worker calls [`init_worker()`][agentlightning.ParallelWorkerBase.init_worker] during start-up.
+    * [`run()`][agentlightning.ParallelWorkerBase.run] performs the parallel workload.
+    * Workers call [`teardown_worker()`][agentlightning.ParallelWorkerBase.teardown_worker] before exiting.
+    * The main process finalizes through [`teardown()`][agentlightning.ParallelWorkerBase.teardown].
+
+    Subclasses must implement [`run()`][agentlightning.ParallelWorkerBase.run]
+    and can override other lifecycle hooks.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the base class. This method can be overridden by subclasses."""
+        self.worker_id: Optional[int] = None
+
+    def init(self, *args: Any, **kwargs: Any) -> None:
+        """Initialize before spawning the workers. This method can be overridden by subclasses."""
+        pass
+
+    def init_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        """Initialize the worker. This method can be overridden by subclasses."""
+        self.worker_id = worker_id
+
+    def run(self, *args: Any, **kwargs: Any) -> Any:
+        """Run the workload. This method can be overridden by subclasses."""
+        pass
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        """Teardown the worker. This method can be overridden by subclasses."""
+        pass
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        """Teardown after the workers have exited. This method can be overridden by subclasses."""
+        pass
+
+
+class Dataset(Protocol, Generic[T_co]):
+    """The general interface for a dataset.
+
+    It's currently implemented as a protocol, having a similar interface to `torch.utils.data.Dataset`.
+    You don't have to inherit from this class; you can use a simple list if you want to.
+    """
+
+    def __getitem__(self, index: SupportsIndex, /) -> T_co: ...
+
+    def __len__(self) -> int: ...
+
+
+class Hook(ParallelWorkerBase):
+    """Base class for defining hooks in the agent runner's lifecycle."""
+
+    async def on_trace_start(
+        self, *, agent: LitAgent[Any], runner: Runner[Any], tracer: Tracer, rollout: Rollout
+    ) -> None:
+        """Hook called immediately after the tracer enters the trace context but before the rollout begins.
+
+        Args:
+            agent: The [`LitAgent`][agentlightning.LitAgent] instance associated with the runner.
+            runner: The [`Runner`][agentlightning.Runner] managing the rollout.
+            tracer: The [`Tracer`][agentlightning.Tracer] instance associated with the runner.
+            rollout: The [`Rollout`][agentlightning.Rollout] object that will be processed.
+
+        Subclasses can override this method to implement custom logic such as logging,
+        metric collection, or resource setup. By default, this is a no-op.
+        """
+
+    async def on_trace_end(
+        self, *, agent: LitAgent[Any], runner: Runner[Any], tracer: Tracer, rollout: Rollout
+    ) -> None:
+        """Hook called immediately after the rollout completes but before the tracer exits the trace context.
+
+        Args:
+            agent: The [`LitAgent`][agentlightning.LitAgent] instance associated with the runner.
+            runner: The [`Runner`][agentlightning.Runner] managing the rollout.
+            tracer: The [`Tracer`][agentlightning.Tracer] instance associated with the runner.
+            rollout: The [`Rollout`][agentlightning.Rollout] object that has been processed.
+
+        Subclasses can override this method to implement custom logic such as logging,
+        metric collection, or resource cleanup. By default, this is a no-op.
+        """
+
+    async def on_rollout_start(self, *, agent: LitAgent[Any], runner: Runner[Any], rollout: Rollout) -> None:
+        """Hook called immediately before a rollout *attempt* begins.
+
+        Args:
+            agent: The [`LitAgent`][agentlightning.LitAgent] instance associated with the runner.
+            runner: The [`Runner`][agentlightning.Runner] managing the rollout.
+            rollout: The [`Rollout`][agentlightning.Rollout] object that will be processed.
+
+        Subclasses can override this method to implement custom logic such as
+        logging, metric collection, or resource setup. By default, this is a
+        no-op.
+        """
+
+    async def on_rollout_end(
+        self,
+        *,
+        agent: LitAgent[Any],
+        runner: Runner[Any],
+        rollout: Rollout,
+        spans: Union[List[ReadableSpan], List[Span]],
+    ) -> None:
+        """Hook called after a rollout *attempt* completes.
+
+        Args:
+            agent: The [`LitAgent`][agentlightning.LitAgent] instance associated with the runner.
+            runner: The [`Runner`][agentlightning.Runner] managing the rollout.
+            rollout: The [`Rollout`][agentlightning.Rollout] object that has been processed.
+            spans: The spans that have been added to the store.
+
+        Subclasses can override this method for cleanup or additional
+        logging. By default, this is a no-op.
+        """
+
+
+class FilterField(TypedDict, total=False):
+    """An operator dict for a single field."""
+
+    exact: Any
+    within: Sequence[Any]
+    contains: str
+
+
+FilterOptions = Mapping[
+    Union[str, Literal["_aggregate", "_must"]],
+    Union[FilterField, Literal["and", "or"], Mapping[str, FilterField]],
+]
+"""A mapping of field name -> operator dict.
+
+Each operator dict can contain:
+
+- "exact": value for exact equality.
+- "within": iterable of allowed values.
+- "contains": substring to search for in string fields.
+
+The filter can also have a special field called "_aggregate" that can be used to specify the logic
+to combine the results of the filters:
+
+- "and": all conditions must match. This is the default value if not specified.
+- "or": at least one condition must match.
+
+All conditions within a field and between different fields are
+stored in a unified pool and combined using `_aggregate`.
+
+The filter can also have a special group called "_must", which is a mapping of filters that must all match,
+no matter whether the aggregate logic is "and" or "or".
+
+Example:
+
+```json
+{
+    "_aggregate": "or",
+    "_must": {
+        "city": {"exact": "New York"},
+        "timezone": {"within": ["America/New_York", "America/Los_Angeles"]},
+    },
+    "status": {"exact": "active"},
+    "id": {"within": [1, 2, 3]},
+    "name": {"contains": "foo"},
+}
+```
+"""
+
+
+class SortOptions(TypedDict):
+    """Options for sorting the collection."""
+
+    name: str
+    """The name of the field to sort by."""
+    order: Literal["asc", "desc"]
+    """The order to sort by."""
+
+
+T_item = TypeVar("T_item")
+
+
+class PaginatedResult(BaseModel, Sequence[T_item]):
+    """Result of a paginated query.
+
+    Behaves like a sequence, but also carries pagination metadata (limit, offset, total).
+    """
+
+    items: Sequence[T_item]
+    """Items in the result."""
+    limit: int
+    """Limit of the result."""
+    offset: int
+    """Offset of the result."""
+    total: int
+    """Total number of items in the collection."""
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    @overload
+    def __getitem__(self, index: int) -> T_item: ...
+
+    @overload
+    def __getitem__(self, index: slice) -> Sequence[T_item]: ...
+
+    def __getitem__(self, index: Union[int, slice]) -> Union[T_item, Sequence[T_item]]:
+        return self.items[index]
+
+    # Overriding __iter__ enables list(paginated_result) to work as expected,
+    # but changes Pydantic's default dict iteration behavior (which would otherwise
+    # iterate over field names).
+    def __iter__(self) -> Iterator[T_item]:  # type: ignore
+        return iter(self.items)
+
+    def __repr__(self) -> str:
+        first_item_repr = repr(self.items[0]) if self.items else "empty"
+        items_repr = f"[{first_item_repr}, ...]" if len(self.items) > 1 else first_item_repr
+        slice_repr = f"{self.offset}:" if self.limit == -1 else f"{self.offset}:{self.offset + self.limit}"
+        return f"<PaginatedResult ({slice_repr} of {self.total}) {items_repr}>"

agentlightning/types/resources.py

@@ -0,0 +1,204 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+"""Typed representations of tunable resources shared between Agent Lightning components."""
+
+import inspect
+import logging
+from typing import (
+    Annotated,
+    Any,
+    Dict,
+    Literal,
+    Optional,
+    Union,
+)
+
+from pydantic import BaseModel, Field
+
+from .core import AttemptedRollout
+
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "Resource",
+    "LLM",
+    "ProxyLLM",
+    "PromptTemplate",
+    "ResourceUnion",
+    "NamedResources",
+    "ResourcesUpdate",
+]
+
+
+class Resource(BaseModel):
+    """Base class for tunable resources distributed to executors."""
+
+    resource_type: Any
+    """Alias of the resource type."""
+
+
+class LLM(Resource):
+    """Resource that identifies an LLM endpoint and its configuration."""
+
+    resource_type: Literal["llm"] = "llm"
+    endpoint: str
+    """The URL of the LLM API endpoint."""
+    model: str
+    """The identifier for the model to be used (e.g., 'gpt-4o')."""
+    api_key: Optional[str] = None
+    """Optional secret used to authenticate requests."""
+    sampling_parameters: Dict[str, Any] = Field(default_factory=dict)
+    """A dictionary of hyperparameters for model inference, such as temperature, top_p, etc."""
+
+    def get_base_url(self, *args: Any, **kwargs: Any) -> str:
+        """Return the base URL consumed by OpenAI-compatible clients.
+
+        Users are encouraged to use `get_base_url(rollout_id, attempt_id)` to get
+        the LLM endpoint instead of accessing `.endpoint` directly.
+        """
+        return self.endpoint
+
+
+class ProxyLLM(LLM):
+    """LLM resource that rewrites endpoints through [`LLMProxy`][agentlightning.LLMProxy].
+
+    The proxy injects rollout- and attempt-specific routing information into the
+    endpoint so that downstream services can attribute requests correctly.
+    """
+
+    resource_type: Literal["proxy_llm"] = "proxy_llm"  # type: ignore
+    _initialized: bool = False
+
+    def model_post_init(self, __context: Any) -> None:
+        """Mark initialization as complete after Pydantic finishes setup."""
+        super().model_post_init(__context)
+        object.__setattr__(self, "_initialized", True)
+
+    def __getattribute__(self, name: str) -> Any:
+        """Emit a warning when `endpoint` is accessed directly after initialization."""
+        # Check if we're accessing endpoint after initialization and not from base_url
+        if name == "endpoint":
+            try:
+                initialized = object.__getattribute__(self, "_initialized")
+            except AttributeError:
+                initialized = False
+
+            if initialized:
+                # Check the call stack to see if we're being called from base_url
+                frame = inspect.currentframe()
+                if frame and frame.f_back:
+                    caller_name = frame.f_back.f_code.co_name
+                    if caller_name != "get_base_url":
+                        logger.warning(
+                            "Accessing 'endpoint' directly on ProxyLLM is discouraged. "
+                            "Use 'get_base_url(rollout_id, attempt_id)' instead to get the properly formatted endpoint."
+                        )
+        return super().__getattribute__(name)
+
+    def with_attempted_rollout(self, rollout: AttemptedRollout) -> LLM:
+        """Bake rollout metadata into a concrete [`LLM`][agentlightning.LLM] instance."""
+        return LLM(
+            endpoint=self.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
+            model=self.model,
+            sampling_parameters=self.sampling_parameters,
+            api_key=self.api_key,
+        )
+
+    def get_base_url(self, rollout_id: Optional[str], attempt_id: Optional[str]) -> str:
+        """Return the routed endpoint for a specific rollout/attempt pair.
+
+        Args:
+            rollout_id: Identifier of the rollout making the request.
+            attempt_id: Identifier of the attempt within that rollout.
+
+        Returns:
+            Fully qualified endpoint including rollout metadata.
+
+        Raises:
+            ValueError: If exactly one of ``rollout_id`` or ``attempt_id`` is provided.
+        """
+        if rollout_id is None and attempt_id is None:
+            return self.endpoint
+
+        if not (isinstance(rollout_id, str) and isinstance(attempt_id, str)):
+            raise ValueError("rollout_id and attempt_id must be strings or all be empty")
+
+        prefix = self.endpoint
+        if prefix.endswith("/"):
+            prefix = prefix[:-1]
+        if prefix.endswith("/v1"):
+            prefix = prefix[:-3]
+            has_v1 = True
+        else:
+            has_v1 = False
+        # Now the prefix should look like "http://localhost:11434"
+
+        # Append the rollout and attempt id to the prefix
+        prefix = prefix + f"/rollout/{rollout_id}/attempt/{attempt_id}"
+        if has_v1:
+            prefix = prefix + "/v1"
+        return prefix
+
+
+class PromptTemplate(Resource):
+    """Resource describing a reusable prompt template."""
+
+    resource_type: Literal["prompt_template"] = "prompt_template"
+    template: str
+    """The template string. The format depends on the engine."""
+    engine: Literal["jinja", "f-string", "poml"]
+    """The templating engine to use for rendering the prompt."""
+
+    def format(self, **kwargs: Any) -> str:
+        """Format the prompt using keyword arguments.
+
+        !!! warning
+            Only the `f-string` engine is supported for now.
+        """
+        if self.engine == "f-string":
+            return self.template.format(**kwargs)
+        else:
+            raise NotImplementedError(
+                "Formatting prompt templates for non-f-string engines with format() helper is not supported yet."
+            )
+
+
+# Use discriminated union for proper deserialization
+# TODO: migrate to use a registry
+ResourceUnion = Annotated[Union[LLM, ProxyLLM, PromptTemplate], Field(discriminator="resource_type")]
+NamedResources = Dict[str, ResourceUnion]
+"""Mapping from resource names to their configured instances.
+
+Examples:
+    ```python
+    resources: NamedResources = {
+        "main_llm": LLM(
+            endpoint="http://localhost:8080",
+            model="llama3",
+            sampling_parameters={"temperature": 0.7, "max_tokens": 100},
+        ),
+        "system_prompt": PromptTemplate(
+            template="You are a helpful assistant.",
+            engine="f-string",
+        ),
+    }
+    ```
+"""
+
+
+class ResourcesUpdate(BaseModel):
+    """Update payload broadcast to clients when resources change."""
+
+    resources_id: str
+    """Identifier used to version the resources."""
+    create_time: float
+    """Timestamp of the creation time of the resources."""
+    update_time: float
+    """Timestamp of the last update time of the resources."""
+    version: int
+    """Version of the resources."""
+    resources: NamedResources
+    """Mapping of resource names to their definitions."""

agentlightning/types/tracer.py

@@ -0,0 +1,512 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import time
+
+"""Data models that mirror OpenTelemetry spans for Agent Lightning."""
+
+import json
+from enum import Enum
+from typing import Any, Dict, List, Literal, Optional, Protocol, Sequence, Union
+
+from opentelemetry import trace as trace_api
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import Event as OtelEvent
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
+from opentelemetry.trace.status import Status as OtelStatus
+from pydantic import BaseModel, ConfigDict
+
+from agentlightning.semconv import AGL_VIRTUAL
+
+__all__ = [
+    "AttributeValue",
+    "Attributes",
+    "TraceState",
+    "SpanContext",
+    "TraceStatus",
+    "Event",
+    "Link",
+    "OtelResource",
+    "Span",
+    "SpanNames",
+    "SpanAttributeNames",
+    "SpanLike",
+    "StatusCode",
+    "SpanCoreFields",
+    "SpanRecordingContext",
+]
+
+
+def convert_timestamp(timestamp: Optional[int]) -> Optional[float]:
+    """Normalize OpenTelemetry timestamps to seconds.
+
+    Args:
+        timestamp: Timestamp expressed either in seconds or nanoseconds.
+
+    Returns:
+        Timestamp in seconds when `timestamp` is provided; otherwise `None`.
+    """
+    if not timestamp:
+        return None
+    return timestamp / 1_000_000_000 if timestamp > 1e12 else timestamp
+
+
+def extract_extra_fields(src: Any, excluded_fields: List[str]) -> Dict[str, Any]:
+    """Capture custom attributes from an OpenTelemetry object.
+
+    Args:
+        src: Object that exposes a `__dict__` of potential attributes.
+        excluded_fields: Attribute names that should be removed from the output.
+
+    Returns:
+        Dictionary containing JSON-serializable representations of the remaining fields.
+    """
+    excluded_fields_set = set(excluded_fields) | set(["_" + k for k in excluded_fields])
+    # Exclude the function fields
+    excluded_fields_set |= set(src.__class__.__dict__.keys())
+    stripped_dict = {k.lstrip("_"): v for k, v in src.__dict__.items()}
+    candidates = {k: v for k, v in stripped_dict.items() if k not in excluded_fields_set and not k.startswith("_")}
+    # This should strip or flatten the unserializable fields
+    candidates_serialized = json.dumps(candidates, default=str)
+    return json.loads(candidates_serialized)
+
+
+AttributeValue = Union[
+    str,
+    bool,
+    int,
+    float,
+    Sequence[str],
+    Sequence[bool],
+    Sequence[int],
+    Sequence[float],
+]
+"""Possible values for OpenTelemetry attributes."""
+Attributes = Dict[str, AttributeValue]
+"""Mapping from attribute names to their values. Same as OpenTelemetry `Attributes` type."""
+TraceState = Dict[str, str]
+"""Mapping from trace state key to its value. Same as OpenTelemetry `TraceState` type."""
+StatusCode = Literal["UNSET", "OK", "ERROR"]
+"""The status code of the span."""
+
+
+class SpanContext(BaseModel):
+    """Pydantic representation of `opentelemetry.trace.SpanContext` values."""
+
+    trace_id: str
+    """The trace ID of the span."""
+    span_id: str
+    """The span ID of the span."""
+    is_remote: bool
+    """Whether the span is remote."""
+    trace_state: TraceState
+    """Mapping from trace state key to its value."""
+
+    model_config = ConfigDict(extra="allow")
+
+    @classmethod
+    def from_opentelemetry(cls, src: trace_api.SpanContext) -> "SpanContext":
+        """Construct a [`SpanContext`][agentlightning.SpanContext] from OpenTelemetry data."""
+
+        return cls(
+            trace_id=trace_api.format_trace_id(src.trace_id),
+            span_id=trace_api.format_span_id(src.span_id),
+            is_remote=src.is_remote,
+            trace_state={k: v for k, v in src.trace_state.items()} if src.trace_state else {},
+            **extract_extra_fields(src, ["trace_id", "span_id", "is_remote", "trace_state"]),
+        )
+
+
+class TraceStatus(BaseModel):
+    """Serializable variant of `opentelemetry.trace.Status`."""
+
+    status_code: StatusCode
+    """The status code of the span. Same as OpenTelemetry `Status.status_code` type."""
+    description: Optional[str] = None
+    """The description of the span. Same as OpenTelemetry `Status.description` type."""
+
+    model_config = ConfigDict(extra="allow")
+
+    @classmethod
+    def from_opentelemetry(cls, src: OtelStatus) -> "TraceStatus":
+        """Create a [`TraceStatus`][agentlightning.TraceStatus] from OpenTelemetry metadata."""
+
+        return cls(
+            status_code=src.status_code.name,
+            description=src.description,
+            **extract_extra_fields(src, ["status_code", "description"]),
+        )
+
+
+class Event(BaseModel):
+    """Serializable representation of OpenTelemetry `Event` values."""
+
+    name: str
+    """The name of the event."""
+    attributes: Attributes
+    """Mapping from attribute names to their values. Same as OpenTelemetry `Attributes` type."""
+    timestamp: Optional[float] = None
+    """The timestamp of the event. Same as OpenTelemetry `Event.timestamp` type."""
+
+    model_config = ConfigDict(extra="allow")
+
+    @classmethod
+    def from_opentelemetry(cls, src: OtelEvent) -> "Event":
+        """Create an [`Event`][agentlightning.Event] from an OpenTelemetry event."""
+
+        return cls(
+            name=src.name,
+            attributes=dict(src.attributes) if src.attributes else {},
+            timestamp=convert_timestamp(src.timestamp),
+            **extract_extra_fields(src, ["name", "attributes", "timestamp"]),
+        )
+
+
+class Link(BaseModel):
+    """Serializable representation of OpenTelemetry `Link` values."""
+
+    context: SpanContext
+    """The context of the link."""
+    attributes: Optional[Attributes] = None
+    """Optional attributes."""
+
+    model_config = ConfigDict(extra="allow")
+
+    @classmethod
+    def from_opentelemetry(cls, src: trace_api.Link) -> "Link":
+        """Create a [`Link`][agentlightning.Link] from an OpenTelemetry link."""
+
+        return cls(
+            context=SpanContext.from_opentelemetry(src.context),
+            attributes=dict(src.attributes) if src.attributes else None,
+            **extract_extra_fields(src, ["context", "attributes"]),
+        )
+
+
+class OtelResource(BaseModel):
+    """Serializable representation of OpenTelemetry `Resource` values.
+
+    Named as `OtelResource` to avoid confusion with the [`Resource`][agentlightning.Resource] class.
+    Users will very rarely need to construct this class directly. Most of the times,
+    they deal with the [`Resource`][agentlightning.Resource] class instead, which describes
+    a very different concept.
+    """
+
+    attributes: Attributes
+    """Mapping from attribute names to their values. Same as OpenTelemetry `Attributes` type."""
+    schema_url: str
+    """The schema URL of the resource."""
+
+    @classmethod
+    def from_opentelemetry(cls, src: Resource) -> "OtelResource":
+        """Create a [`Resource`][agentlightning.Resource] from an OpenTelemetry resource."""
+
+        return cls(
+            attributes=dict(src.attributes) if src.attributes else {},
+            schema_url=src.schema_url if src.schema_url else "",
+            **extract_extra_fields(src, ["attributes", "schema_url"]),
+        )
+
+
+class SpanCoreFields(BaseModel):
+    """Core fields of a span. Used by span creators who don't care about the full span model.
+
+    If the spans are managed by some OTel tracer provider, it's not advised to create spans via this path.
+    """
+
+    name: str
+    """The name of the span."""
+    status: TraceStatus
+    """The status of the span."""
+    attributes: Attributes
+    """The attributes of the span."""
+    start_time: Optional[float]
+    """The start time of the span."""
+    end_time: Optional[float]
+    """The end time of the span."""
+
+
+class SpanRecordingContext(Protocol):
+    """Context for recording operations on a span. It doesn't have to finalize the span; the caller will do it."""
+
+    def record_exception(self, exception: BaseException) -> None:
+        """Record an exception on the span."""
+        raise NotImplementedError()
+
+    def record_attributes(self, attributes: Attributes) -> None:
+        """Record attributes on the span."""
+        raise NotImplementedError()
+
+    def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
+        """Record the status of the span."""
+        raise NotImplementedError()
+
+    def get_recorded_span(self) -> SpanCoreFields:
+        """Get the recording of the span."""
+        raise NotImplementedError()
+
+
+class Span(BaseModel):
+    """Agent Lightning's canonical span model used for persistence and analytics.
+
+    The model captures the most relevant fields from
+    `opentelemetry.sdk.trace.ReadableSpan` instances while preserving unmodeled
+    attributes in Pydantic `BaseModel`'s extra storage. This keeps the serialized format
+    stable even as upstream OpenTelemetry types evolve.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    rollout_id: str
+    """The rollout which this span belongs to."""
+    attempt_id: str
+    """The attempt which this span belongs to."""
+    sequence_id: int
+    """The ID to make spans ordered within a single attempt."""
+
+    # Current ID (in hex, formatted via trace_api.format_*)
+    trace_id: str  # one rollout can have traces coming from multiple places
+    """The trace ID of the span. One rollout/attempt can have multiple traces.
+    This ID comes from the OpenTelemetry trace ID generator.
+    """
+    span_id: str
+    """The span ID of the span. This ID comes from the OpenTelemetry span ID generator."""
+    parent_id: Optional[str]
+    """The parent span ID of the span."""
+
+    # Core ReadableSpan fields
+    name: str
+    """The name of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
+    status: TraceStatus
+    """The status of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
+    attributes: Attributes
+    """The attributes of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
+    events: List[Event]
+    """The events of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
+    links: List[Link]
+    """The links of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
+
+    # Timestamps
+    start_time: Optional[float]
+    """The start time of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
+    end_time: Optional[float]
+    """The end time of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
+
+    # Other parsable fields
+    context: Optional[SpanContext]
+    """The context of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
+    parent: Optional[SpanContext]
+    """The parent context of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
+    resource: OtelResource
+    """The resource of the span. See [OpenTelemetry docs](https://opentelemetry.io/docs/concepts/signals/traces/)."""
+
+    # Preserve other fields in the readable span as extra fields
+    # Make sure that are json serializable (so no bytes, complex objects, ...)
+
+    @classmethod
+    def from_opentelemetry(
+        cls,
+        src: ReadableSpan,
+        rollout_id: str,
+        attempt_id: str,
+        sequence_id: int,
+    ) -> "Span":
+        """Convert an OpenTelemetry span into the Agent Lightning data model.
+
+        Args:
+            src: Span captured by OpenTelemetry.
+            rollout_id: Identifier for the rollout that produced the span.
+            attempt_id: Identifier of the attempt within the rollout.
+            sequence_id: Monotonically increasing identifier assigned to the span.
+
+        Returns:
+            Parsed [`Span`][agentlightning.Span] instance suitable for persistence.
+        """
+        context = src.get_span_context()
+        if context is None:
+            trace_id = span_id = 0
+        else:
+            trace_id = context.trace_id
+            span_id = context.span_id
+        return cls(
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
+            sequence_id=sequence_id,
+            trace_id=trace_api.format_trace_id(trace_id),
+            span_id=trace_api.format_span_id(span_id),
+            parent_id=(trace_api.format_span_id(src.parent.span_id) if src.parent else None),
+            name=src.name,
+            status=TraceStatus.from_opentelemetry(src.status),
+            attributes=dict(src.attributes) if src.attributes else {},
+            events=[Event.from_opentelemetry(event) for event in src.events] if src.events else [],
+            links=[Link.from_opentelemetry(link) for link in src.links] if src.links else [],
+            start_time=convert_timestamp(src.start_time),
+            end_time=convert_timestamp(src.end_time),
+            context=SpanContext.from_opentelemetry(context) if context else None,
+            parent=(SpanContext.from_opentelemetry(src.parent) if src.parent else None),
+            resource=OtelResource.from_opentelemetry(src.resource),
+            **extract_extra_fields(
+                src,
+                [
+                    "name",
+                    "context",
+                    "parent",
+                    "resource",
+                    "attributes",
+                    "events",
+                    "links",
+                    "start_time",
+                    "end_time",
+                    "status",
+                    "span_processor",
+                    "rollout_id",
+                    "attempt_id",
+                    "trace_id",
+                    "span_id",
+                    "parent_id",
+                ],
+            ),
+        )
+
+    @classmethod
+    def from_attributes(
+        cls,
+        *,
+        attributes: Attributes,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+        sequence_id: Optional[int] = None,
+        name: Optional[str] = None,
+        trace_id: Optional[str] = None,
+        span_id: Optional[str] = None,
+        parent_id: Optional[str] = None,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+        resource: Optional[OtelResource] = None,
+        status: Optional[TraceStatus] = None,
+    ) -> "Span":
+        """Build a synthetic span from raw attributes.
+        Different from the [`from_opentelemetry`][agentlightning.Span.from_opentelemetry] method,
+        all parameters other than `attributes` are optional and will be generated if not provided.
+
+        Args:
+            attributes: Span attributes to persist.
+            rollout_id: Optional rollout identifier associated with the span.
+            attempt_id: Optional attempt identifier associated with the span.
+            sequence_id: Optional sequence number to preserve ordering.
+            name: Optional human-readable span name.
+            trace_id: Custom trace identifier. When omitted, a random identifier is generated.
+            span_id: Custom span identifier. When omitted, a random identifier is generated.
+            parent_id: Optional parent span identifier.
+            start_time: Span start timestamp in seconds.
+            end_time: Span end timestamp in seconds.
+            resource: Explicit resource information to attach to the span.
+            status: Optional status of the span.
+
+        Returns:
+            [`Span`][agentlightning.Span] populated with the provided attributes.
+        """
+
+        id_generator = RandomIdGenerator()
+        trace_id = trace_id or trace_api.format_trace_id(id_generator.generate_trace_id())
+        span_id = span_id or trace_api.format_span_id(id_generator.generate_span_id())
+
+        return cls(
+            rollout_id=rollout_id or "",
+            attempt_id=attempt_id or "",
+            sequence_id=sequence_id or 0,
+            trace_id=trace_id,
+            span_id=span_id,
+            parent_id=parent_id,
+            start_time=start_time,
+            end_time=end_time,
+            context=SpanContext(
+                trace_id=trace_id,
+                span_id=span_id,
+                is_remote=False,
+                trace_state={},
+            ),
+            name=name or AGL_VIRTUAL,
+            resource=resource or OtelResource(attributes={}, schema_url=""),
+            attributes=attributes,
+            status=status or TraceStatus(status_code="OK"),
+            events=[],
+            links=[],
+            parent=(
+                SpanContext(
+                    trace_id=trace_id,
+                    span_id=parent_id,
+                    is_remote=False,
+                    trace_state={},
+                )
+                if parent_id
+                else None
+            ),
+        )
+
+    @classmethod
+    def from_core_fields(
+        cls,
+        core: SpanCoreFields,
+        *,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+        sequence_id: Optional[int] = None,
+    ) -> Span:
+        """Build a span from a core span.
+
+        Args:
+            core: Core span to build from.
+            rollout_id: Optional rollout identifier associated with the span.
+            attempt_id: Optional attempt identifier associated with the span.
+            sequence_id: Optional sequence number to preserve ordering.
+
+        Returns:
+            [`Span`][agentlightning.Span] populated with the provided attributes.
+        """
+        return cls.from_attributes(
+            attributes=core.attributes,
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
+            sequence_id=sequence_id,
+            name=core.name,
+            start_time=core.start_time or time.time(),
+            end_time=core.end_time,
+            status=core.status,
+        )
+
+
+class SpanNames(str, Enum):
+    """Enumerated span names recognised by Agent-lightning. Deprecated in favor of [semconv][agentlightning.semconv]."""
+
+    REWARD = "agentlightning.reward"
+    """The name of the reward span."""
+    MESSAGE = "agentlightning.message"
+    """The name of the message span."""
+    OBJECT = "agentlightning.object"
+    """The name of the object span."""
+    EXCEPTION = "agentlightning.exception"
+    """The name of the exception span."""
+    VIRTUAL = "agentlightning.virtual"
+    """The name of the virtual span. It represents derived spans without concrete operations."""
+    ROLLOUT_ID = "agentlightning.rollout_id"
+    """The name of the rollout ID."""
+    ATTEMPT_ID = "agentlightning.attempt_id"
+    """The name of the attempt ID."""
+    SPAN_SEQUENCE_ID = "agentlightning.span_sequence_id"
+    """The name of the span sequence ID."""
+
+
+class SpanAttributeNames(str, Enum):
+    """Canonical attribute names written by Agent Lightning emitters. Deprecated in favor of [semconv][agentlightning.semconv]."""
+
+    MESSAGE = "message"
+    """The name of the message attribute."""
+    OBJECT = "object"
+    """The name of the object attribute."""
+
+
+SpanLike = Union[ReadableSpan, Span]
+"""Union type of OpenTelemetry `ReadableSpan` and Agent-lightning [`Span`][agentlightning.Span]."""