mantisdk 0.1.0__py3-none-any.whl

mantisdk/runner/base.py

@@ -0,0 +1,182 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Abstract runner interface for executing agent tasks."""
+
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Generic, Iterator, Optional, Sequence, TypeVar
+
+from mantisdk.execution.events import ExecutionEvent
+from mantisdk.litagent import LitAgent
+from mantisdk.store.base import LightningStore
+from mantisdk.types import Hook, NamedResources, ParallelWorkerBase, Rollout, RolloutMode
+
+if TYPE_CHECKING:
+    from mantisdk.execution.events import ExecutionEvent
+
+
+T_task = TypeVar("T_task")
+
+logger = logging.getLogger(__name__)
+
+
+class Runner(ParallelWorkerBase, Generic[T_task]):
+    """Abstract base class for long-running agent executors.
+
+    Runner implementations coordinate [`LitAgent`][mantisdk.LitAgent]
+    instances, acquire work from a [`LightningStore`][mantisdk.LightningStore],
+    and emit [`Rollout`][mantisdk.Rollout] objects. Subclasses decide how
+    to schedule work (polling, streaming, etc.) while this base class provides a
+    minimal lifecycle contract.
+    """
+
+    def init(self, agent: LitAgent[T_task], **kwargs: Any) -> None:
+        """Prepare the runner to execute tasks for `agent`.
+
+        This method is called only once during the setup for all workers, not for each worker.
+
+        Args:
+            agent: Agent instance providing task-specific logic.
+            **kwargs: Optional runner-specific configuration.
+
+        Raises:
+            NotImplementedError: Subclasses must supply the initialization
+                routine.
+        """
+        raise NotImplementedError()
+
+    def init_worker(self, worker_id: int, store: LightningStore, **kwargs: Any) -> None:
+        """Configure worker-local state before processing tasks.
+
+        This method is called for **each** worker during the setup.
+
+        Args:
+            worker_id: Unique identifier for this worker process or thread.
+            store: Shared [`LightningStore`][mantisdk.LightningStore]
+                backing task coordination.
+            **kwargs: Optional worker-specific configuration.
+
+        Raises:
+            NotImplementedError: Subclasses must prepare per-worker resources.
+        """
+        raise NotImplementedError()
+
+    def run(self, *args: Any, **kwargs: Any) -> None:
+        """Deprecated synchronous entry point.
+
+        Use [`iter()`][mantisdk.Runner.iter] or [`step()`][mantisdk.Runner.step] instead.
+
+        Raises:
+            RuntimeError: Always raised to direct callers to
+                [iter()][mantisdk.Runner.iter] or
+                [step()][mantisdk.Runner.step].
+        """
+        raise RuntimeError("The behavior of run() of Runner is undefined. Use iter() or step() instead.")
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        """Release resources acquired during [`init()`][mantisdk.Runner.init].
+
+        Raises:
+            NotImplementedError: Subclasses must implement the shutdown routine.
+        """
+        raise NotImplementedError()
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        """Release per-worker resources allocated by [`init_worker()`][mantisdk.Runner.init_worker].
+
+        Args:
+            worker_id: Identifier of the worker being torn down.
+
+        Raises:
+            NotImplementedError: Subclasses must implement the shutdown routine.
+        """
+        raise NotImplementedError()
+
+    @contextmanager
+    def run_context(
+        self,
+        *,
+        agent: LitAgent[T_task],
+        store: LightningStore,
+        hooks: Optional[Sequence[Hook]] = None,
+        worker_id: Optional[int] = None,
+    ) -> Iterator[Runner[T_task]]:
+        """Initialize and tear down a runner within a simple context manager.
+
+        The helper is primarily intended for debugging runner implementations
+        outside of a full [`Trainer`][mantisdk.Trainer] stack.
+
+        Args:
+            agent: Agent executed by this runner.
+            store: Backing [`LightningStore`][mantisdk.LightningStore].
+                If you don't have one, you can easily create one with
+                [`InMemoryLightningStore`][mantisdk.InMemoryLightningStore].
+            hooks: Optional sequence of hooks recognised by the runner.
+                Not all runners support hooks.
+            worker_id: Override the worker identifier used during setup. Defaults
+                to `0`.
+        """
+        _initialized: bool = False
+        _worker_initialized: bool = False
+        try:
+            self.init(agent=agent, hooks=hooks)
+            _initialized = True
+            self.init_worker(worker_id=0, store=store)
+            _worker_initialized = True
+            yield self
+        finally:
+            try:
+                if _worker_initialized:
+                    self.teardown_worker(worker_id=worker_id if worker_id is not None else 0)
+            except Exception:
+                logger.error("Error during runner worker teardown", exc_info=True)
+
+            try:
+                if _initialized:
+                    self.teardown()
+            except Exception:
+                logger.error("Error during runner teardown", exc_info=True)
+
+    async def iter(self, *, event: Optional[ExecutionEvent] = None) -> None:
+        """Run the runner, continuously iterating over tasks in the store.
+
+        This method runs in a loop, polling the store for new tasks and executing
+        them until interrupted by the event or when no more tasks are available.
+
+        Args:
+            event: Cooperative stop signal. When set, the runner should complete
+                the current unit of work and exit the loop.
+
+        Raises:
+            NotImplementedError: Subclasses provide the iteration behavior.
+        """
+        raise NotImplementedError()
+
+    async def step(
+        self,
+        input: T_task,
+        *,
+        resources: Optional[NamedResources] = None,
+        mode: Optional[RolloutMode] = None,
+        event: Optional[ExecutionEvent] = None,
+    ) -> Rollout:
+        """Execute a single task with the given input.
+
+        This method provides fine-grained control for executing individual tasks
+        directly, bypassing the store's task queue.
+
+        Args:
+            input: Task payload consumed by the agent.
+            resources: Optional named resources scoped to this invocation.
+            mode: Optional rollout mode such as `"train"` or `"eval"`.
+            event: Cooperative stop signal for long-running tasks.
+
+        Returns:
+            Completed rollout produced by the agent.
+
+        Raises:
+            NotImplementedError: Subclasses provide the execution behavior.
+        """
+        raise NotImplementedError()

mantisdk/runner/legacy.py

@@ -0,0 +1,309 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import logging
+import time
+from typing import Any, Dict, List, Optional, cast
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from mantisdk.adapter import TracerTraceToTriplet
+from mantisdk.client import MantisdkClient
+from mantisdk.litagent import LitAgent
+from mantisdk.litagent.litagent import is_v0_1_rollout_api
+from mantisdk.tracer.base import Tracer
+from mantisdk.types import RolloutLegacy, RolloutRawResultLegacy, Span, SpanLike, Triplet
+
+from .base import Runner
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "LegacyAgentRunner",
+]
+
+
+class LegacyAgentRunner(Runner[Any]):
+    """Manages the agent's execution loop and integrates with AgentOps.
+
+    This class orchestrates the interaction between the agent (`LitAgent`) and
+    the server (`MantisdkClient`). It handles polling for tasks, executing
+    the agent's logic, and reporting results back to the server. If enabled,
+    it will also automatically trace each rollout using AgentOps.
+
+    Attributes:
+        agent: The `LitAgent` instance containing the agent's logic.
+        client: The `MantisdkClient` for server communication.
+        tracer: The tracer instance for this runner/worker.
+        worker_id: An optional identifier for the worker process.
+        max_tasks: The maximum number of tasks to process before stopping.
+    """
+
+    def __init__(
+        self,
+        agent: LitAgent[Any],
+        client: MantisdkClient,
+        tracer: Tracer,
+        triplet_exporter: TracerTraceToTriplet,
+        worker_id: Optional[int] = None,
+        max_tasks: Optional[int] = None,
+    ):
+        super().__init__()
+        self.agent = agent
+        self.client = client
+        self.tracer = tracer
+        self.triplet_exporter = triplet_exporter
+
+        # Worker-specific attributes
+        self.worker_id = worker_id
+        self.max_tasks = max_tasks
+
+    # These methods are overridden by Runner, getting them back to old behavior.
+    def init(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def init_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        self.worker_id = worker_id
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def _log_prefix(self, rollout_id: Optional[str] = None) -> str:
+        """Generates a standardized log prefix for the current worker."""
+        if self.worker_id is not None:
+            if rollout_id:
+                return f"[Worker {self.worker_id} | RolloutLegacy {rollout_id}]"
+            else:
+                return f"[Worker {self.worker_id}]"
+        if rollout_id:
+            return f"[RolloutLegacy {rollout_id}]"
+        return "[Default Worker]"
+
+    def _to_rollout_object(
+        self,
+        result: RolloutRawResultLegacy,
+        rollout_id: str,
+    ) -> RolloutLegacy:
+        """Standardizes the agent's return value into a RolloutLegacy object.
+
+        Args:
+            result: The output from the agent's rollout method.
+            rollout_id: The unique identifier for the current task.
+
+        Returns:
+            A standardized `RolloutLegacy` object for reporting to the server.
+        """
+        trace: Any = None
+        final_reward: Optional[float] = None
+        triplets: Optional[List[Triplet]] = None
+        trace_spans: Optional[List[SpanLike]] = None
+
+        # Handle different types of results from the agent
+        # Case 1: result is a float (final reward)
+        if isinstance(result, float):
+            final_reward = result
+        # Case 2: result is a list of Triplets
+        if isinstance(result, list) and all(isinstance(t, Triplet) for t in result):
+            triplets = result  # type: ignore
+        # Case 3.1: result is a list of ReadableSpan (OpenTelemetry spans)
+        if isinstance(result, list) and all(isinstance(t, (ReadableSpan)) for t in result):
+            trace_spans = result  # type: ignore
+            trace = [json.loads(readable_span.to_json()) for readable_span in trace_spans]  # type: ignore
+        # Case 3.2: result is a list of Span (Mantisdk spans)
+        if isinstance(result, list) and all(isinstance(t, Span) for t in result):
+            trace_spans = result  # type: ignore
+            trace = [span.model_dump() for span in trace_spans]  # type: ignore
+        # Case 4: result is a list of dict (trace JSON)
+        if isinstance(result, list) and all(isinstance(t, dict) for t in result):
+            trace = result
+        # Case 5: result is a RolloutLegacy object
+        if isinstance(result, RolloutLegacy):
+            final_reward = result.final_reward
+            triplets = result.triplets
+            trace = result.trace
+
+        # If the agent has tracing enabled, use the tracer's last trace if not already set
+        if self.tracer and (trace is None or trace_spans is None):
+            trace_spans = self.tracer.get_last_trace()  # type: ignore
+            if trace_spans:
+                trace = [cast(Span, span).model_dump() for span in trace_spans]
+
+        # Always extract triplets from the trace using TracerTraceToTriplet
+        if trace_spans:
+            triplets = self.triplet_exporter(trace_spans)  # type: ignore
+
+        # If the agent has triplets, use the last one for final reward if not set
+        if triplets and triplets[-1].reward is not None and final_reward is None:
+            final_reward = triplets[-1].reward
+
+        # Create the RolloutLegacy object with standardized fields
+        result_dict: Dict[str, Any] = {
+            "rollout_id": rollout_id,
+        }
+        if final_reward is not None:
+            result_dict["final_reward"] = final_reward
+        if triplets is not None:
+            result_dict["triplets"] = triplets
+        if trace is not None:
+            result_dict["trace"] = trace
+
+        if isinstance(result, RolloutLegacy):
+            return result.model_copy(update=result_dict)
+        return RolloutLegacy(**result_dict)
+
+    def run(self) -> bool:  # type: ignore
+        """Poll the task and rollout once synchronously."""
+        self.agent.set_runner(self)  # Ensure the agent has a reference to this runner
+
+        task = self.client.poll_next_task()
+        if task is None:
+            logger.info(f"{self._log_prefix()} Poll returned no task. Exiting.")
+            return False
+        rollout_id = task.rollout_id
+
+        resources_id = task.resources_id
+        resources_update = None
+        if resources_id:
+            resources_update = self.client.get_resources_by_id(resources_id)
+        else:
+            logger.debug(f"{self._log_prefix(rollout_id)} No 'resources_id'. Fetching latest resources.")
+            resources_update = self.client.get_latest_resources()
+        if not resources_update:
+            logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
+            return False
+
+        rollout_obj = RolloutLegacy(rollout_id=task.rollout_id, task=task)  # Default empty rollout
+
+        try:
+            try:
+                self.agent.on_rollout_start(task, self, self.tracer)
+            except Exception:
+                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")
+
+            with self.tracer._trace_context_sync(name=f"rollout_{rollout_id}"):  # pyright: ignore[reportPrivateUsage]
+                start_time = time.time()
+                rollout_method = self.agent.training_rollout if task.mode == "train" else self.agent.validation_rollout
+                # Pass the task input, not the whole task object
+                if is_v0_1_rollout_api(rollout_method):
+                    result = cast(
+                        RolloutRawResultLegacy,
+                        rollout_method(
+                            task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources  # type: ignore
+                        ),
+                    )  # type: ignore
+                else:
+                    result = rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)  # type: ignore
+                rollout_obj = self._to_rollout_object(result, task.rollout_id)  # type: ignore
+                end_time = time.time()
+                logger.info(
+                    f"{self._log_prefix(rollout_id)} Completed in "
+                    f"{end_time - start_time:.2f}s. Triplet length: "
+                    f"{len(rollout_obj.triplets) if rollout_obj.triplets is not None else 'N/A'}. "
+                    f"Reward: {rollout_obj.final_reward}"
+                )
+
+        except Exception:
+            logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
+        finally:
+            try:
+                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)  # type: ignore
+            except Exception:
+                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
+            self.client.post_rollout(rollout_obj)
+
+        return True
+
+    def iter(self) -> int:  # type: ignore
+        """Executes the synchronous polling and rollout loop."""
+        num_tasks_processed = 0
+        logger.info(f"{self._log_prefix()} Started sync rollouts (max: {self.max_tasks or 'unlimited'}).")
+
+        while self.max_tasks is None or num_tasks_processed < self.max_tasks:
+            if self.run():
+                num_tasks_processed += 1
+
+            if num_tasks_processed % 10 == 0 or num_tasks_processed == 1:
+                logger.info(f"{self._log_prefix()} Progress: {num_tasks_processed}/{self.max_tasks or 'unlimited'}")
+
+        logger.info(f"{self._log_prefix()} Finished sync rollouts. Processed {num_tasks_processed} tasks.")
+        return num_tasks_processed
+
+    async def run_async(self) -> bool:
+        """Poll the task and rollout once."""
+        self.agent.set_runner(self)  # Ensure the agent has a reference to this runner
+
+        task = await self.client.poll_next_task_async()
+        if task is None:
+            logger.info(f"{self._log_prefix()} Poll returned no task. Exiting.")
+            return False
+        rollout_id = task.rollout_id
+
+        resources_id = task.resources_id
+        resources_update = None
+        if resources_id:
+            resources_update = await self.client.get_resources_by_id_async(resources_id)
+        else:
+            logger.debug(f"{self._log_prefix(rollout_id)} No 'resources_id'. Fetching latest resources.")
+            resources_update = await self.client.get_latest_resources_async()
+        if not resources_update:
+            logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
+            return False
+
+        rollout_obj = RolloutLegacy(rollout_id=task.rollout_id, task=task)  # Default empty rollout
+
+        try:
+            try:
+                self.agent.on_rollout_start(task, self, self.tracer)
+            except Exception:
+                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")
+
+            async with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
+                start_time = time.time()
+                rollout_method = (
+                    self.agent.training_rollout_async if task.mode == "train" else self.agent.validation_rollout_async
+                )
+                # Pass the task input, not the whole task object
+                if is_v0_1_rollout_api(rollout_method):
+                    result = cast(
+                        RolloutRawResultLegacy,
+                        await rollout_method(
+                            task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources  # type: ignore
+                        ),
+                    )  # type: ignore
+                else:
+                    result = await rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)  # type: ignore
+                rollout_obj = self._to_rollout_object(result, task.rollout_id)  # type: ignore
+                end_time = time.time()
+                logger.info(
+                    f"{self._log_prefix(rollout_id)} Completed in "
+                    f"{end_time - start_time:.2f}s. Triplet length: "
+                    f"{len(rollout_obj.triplets) if rollout_obj.triplets is not None else 'N/A'}. "
+                    f"Reward: {rollout_obj.final_reward}"
+                )
+        except Exception:
+            logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
+        finally:
+            try:
+                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)  # type: ignore
+            except Exception:
+                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
+            await self.client.post_rollout_async(rollout_obj)
+
+        return True
+
+    async def iter_async(self) -> int:
+        """Executes the asynchronous polling and rollout loop."""
+        num_tasks_processed = 0
+        logger.info(f"{self._log_prefix()} Started async rollouts (max: {self.max_tasks or 'unlimited'}).")
+
+        while self.max_tasks is None or num_tasks_processed < self.max_tasks:
+            if await self.run_async():
+                num_tasks_processed += 1
+
+            if num_tasks_processed % 10 == 0 or num_tasks_processed == 1:
+                logger.info(f"{self._log_prefix()} Progress: {num_tasks_processed}/{self.max_tasks or 'unlimited'}")
+        logger.info(f"{self._log_prefix()} Finished async rollouts. Processed {num_tasks_processed} tasks.")
+        return num_tasks_processed

mantisdk/semconv.py

@@ -0,0 +1,170 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Semantic conventions for Mantisdk spans.
+
+Conventions in this file are added on demand. We generally DO NOT add
+new semantic conventions unless it's absolutely needed for certain algorithms or scenarios.
+"""
+
+from enum import Enum
+
+from pydantic import BaseModel
+
+AGL_ANNOTATION = "mantisdk.annotation"
+"""Mantisdk's standard span name for annotations.
+
+Annotations are minimal span units for rewards, tags, and metadatas.
+They are used to "annotate" a specific event or a part of rollout.
+"""
+
+AGL_MESSAGE = "mantisdk.message"
+"""Mantisdk's standard span name for messages and logs."""
+
+AGL_OBJECT = "mantisdk.object"
+"""Mantisdk's standard span name for customized objects."""
+
+AGL_EXCEPTION = "mantisdk.exception"
+"""Mantisdk's standard span name for exceptions.
+
+Used by the exception emitter to record exception details.
+"""
+
+AGL_OPERATION = "mantisdk.operation"
+"""Mantisdk's standard span name for functions.
+Wrap function or code-blocks as operations.
+"""
+
+AGL_REWARD = "mantisdk.reward"
+"""Mantisdk's standard span name for reward operations."""
+
+AGL_VIRTUAL = "mantisdk.virtual"
+"""Mantisdk's standard span name for virtual operations.
+
+Mostly used in adapter when needing to represent the root or intermediate operations.
+"""
+
+
+class LightningResourceAttributes(Enum):
+    """Resource attribute names used in Mantisdk spans."""
+
+    ROLLOUT_ID = "mantisdk.rollout_id"
+    """Resource name for rollout ID in Mantisdk spans."""
+
+    ATTEMPT_ID = "mantisdk.attempt_id"
+    """Resource name for attempt ID in Mantisdk spans."""
+
+    SPAN_SEQUENCE_ID = "mantisdk.span_sequence_id"
+    """Resource name for span sequence ID in Mantisdk spans."""
+
+    TRACER_NAME = "mantisdk.tracer.name"
+    """Which tracer is used to create this span."""
+
+    JOB_ID = "mantisdk.job_id"
+    """Resource name for job ID in Mantisdk spans (for non-rollout spans)."""
+
+    SPAN_TYPE = "mantisdk.span_type"
+    """Type of span: 'rollout' for rollout-context spans, 'job' for job-level spans."""
+
+
+class LightningSpanAttributes(Enum):
+    """Attribute names that commonly appear in Mantisdk spans.
+
+    Exception types can't be found here because they are defined in OpenTelemetry's official semantic conventions.
+    """
+
+    REWARD = "mantisdk.reward"
+    """Attribute prefix for rewards-related data in reward spans.
+
+    It should be used as a prefix. For example, "mantisdk.reward.0.value" can
+    be used to track a specific metric. See [RewardAttributes][mantisdk.semconv.RewardAttributes].
+    """
+
+    LINK = "mantisdk.link"
+    """Attribute name for linking the current span to another span or other objects like requests/responses."""
+
+    TAG = "mantisdk.tag"
+    """Attribute name for tagging spans with customized strings."""
+
+    MESSAGE_BODY = "mantisdk.message.body"
+    """Attribute name for message text in message spans."""
+
+    OBJECT_TYPE = "mantisdk.object.type"
+    """Attribute name for object type (full qualified name) in object spans.
+
+    I think builtin types like str, int, bool, list, dict are self-explanatory and
+    should also be qualified to use here.
+    """
+
+    OBJECT_LITERAL = "mantisdk.object.literal"
+    """Attribute name for object literal value in object spans (for str, int, bool, ...)."""
+
+    OBJECT_JSON = "mantisdk.object.json"
+    """Attribute name for object serialized value (JSON) in object spans."""
+
+    OPERATION_NAME = "mantisdk.operation.name"
+    """Attribute name for operation name in operation spans, normally the function name."""
+
+    OPERATION_INPUT = "mantisdk.operation.input"
+    """Attribute name for operation input in operation spans."""
+
+    OPERATION_OUTPUT = "mantisdk.operation.output"
+    """Attribute name for operation output in operation spans."""
+
+
+class RewardAttributes(Enum):
+    """Multi-dimensional reward attributes will look like:
+
+    ```json
+    {"mantisdk.reward.0.name": "efficiency", "mantisdk.reward.0.value": 0.75}
+    ```
+
+    The first reward in the reward list will automatically be the primary reward.
+    If the reward list has greater than 1, it shall be a multi-dimensional case.
+    """
+
+    REWARD_NAME = "name"
+    """Key for each dimension in multi-dimensional reward spans."""
+
+    REWARD_VALUE = "value"
+    """Value for each dimension in multi-dimensional reward spans."""
+
+
+class RewardPydanticModel(BaseModel):
+    """A stricter implementation of RewardAttributes used in otel helpers."""
+
+    name: str
+    """Name of the reward dimension."""
+
+    value: float
+    """Value of the reward dimension."""
+
+
+class LinkAttributes(Enum):
+    """Standard link types used in Mantisdk spans.
+
+    The link is more powerful than [OpenTelemetry link](https://opentelemetry.io/docs/specs/otel/trace/api/#link)
+    in that it supports linking to a queryset of spans.
+    It can even link to span object that hasn't been emitted yet.
+    """
+
+    KEY_MATCH = "key_match"
+    """Linking to spans with matching attribute keys.
+
+    `trace_id` and `span_id` are reserved and will be used to link to specific spans directly.
+
+    For example, it can be `gen_ai.response.id` if intended to be link to a chat completion response span.
+    Or it can be `span_id` to link to a specific span by its ID.
+    """
+
+    VALUE_MATCH = "value_match"
+    """Linking to spans with corresponding attribute values on those keys."""
+
+
+class LinkPydanticModel(BaseModel):
+    """A stricter implementation of LinkAttributes used in otel helpers."""
+
+    key_match: str
+    """The attribute key to match on the target spans."""
+
+    value_match: str
+    """The attribute value to match on the target spans."""