mantisdk 0.1.0__py3-none-any.whl

mantisdk/tracer/agentops.py

@@ -0,0 +1,242 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+import os
+import warnings
+from contextlib import asynccontextmanager, contextmanager
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Iterator, List, Optional
+
+import agentops
+import agentops.sdk.core
+import opentelemetry.trace as trace_api
+from agentops.sdk.core import TracingCore
+from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
+from opentelemetry.trace.status import StatusCode
+
+from mantisdk.instrumentation import instrument_all, uninstrument_all
+from mantisdk.store.base import LightningStore
+from mantisdk.utils.otel import get_span_processors, get_tracer_provider
+
+from .base import with_active_tracer_context
+from .otel import LightningSpanProcessor, OtelTracer
+
+if TYPE_CHECKING:
+    from agentops.integration.callbacks.langchain import LangchainCallbackHandler
+
+
+logger = logging.getLogger(__name__)
+
+
+class AgentOpsTracer(OtelTracer):
+    """Traces agent execution using AgentOps.
+
+    This tracer provides functionality to capture execution details using the
+    AgentOps library. It manages the AgentOps client initialization, server setup,
+    and integration with the OpenTelemetry tracing ecosystem.
+
+    Attributes:
+        agentops_managed: Whether to automatically manage `agentops`.
+                          When set to true, tracer calls `agentops.init()`
+                          automatically and launches an agentops endpoint locally.
+                          If not, you are responsible for calling and using it
+                          before using the tracer.
+        instrument_managed: Whether to automatically manage instrumentation.
+                            When set to false, you will manage the instrumentation
+                            yourself and the tracer might not work as expected.
+        daemon: Whether the AgentOps server runs as a daemon process.
+                Only applicable if `agentops_managed` is True.
+    """
+
+    def __init__(self, *, agentops_managed: bool = True, instrument_managed: bool = True, daemon: bool = True):
+        super().__init__()
+        self._lightning_span_processor: Optional[LightningSpanProcessor] = None
+        self.agentops_managed = agentops_managed
+        self.instrument_managed = instrument_managed
+        self.daemon = daemon
+
+        if not self.agentops_managed:
+            logger.warning("agentops_managed=False. You are responsible for AgentOps setup.")
+        if not self.instrument_managed:
+            logger.warning("instrument_managed=False. You are responsible for all instrumentation.")
+
+    def instrument(self, worker_id: int):
+        instrument_all()
+
+    def uninstrument(self, worker_id: int):
+        uninstrument_all()
+
+    def _initialize_tracer_provider(self, worker_id: int):
+        logger.info(f"[Worker {worker_id}] Setting up AgentOps tracer...")  # worker_id included in process name
+
+        if self.instrument_managed:
+            self.instrument(worker_id)
+            logger.info(f"[Worker {worker_id}] Instrumentation applied.")
+
+        if self.agentops_managed:
+            os.environ.setdefault("AGENTOPS_API_KEY", "dummy")
+            if not agentops.get_client().initialized:
+                agentops.init(auto_start_session=False)  # type: ignore
+                logger.info(f"[Worker {worker_id}] AgentOps client initialized.")
+            else:
+                logger.warning(f"[Worker {worker_id}] AgentOps client was already initialized. Skip initialization.")
+
+        span_processors = get_span_processors(self._get_tracer_provider(), LightningSpanProcessor)
+        if len(span_processors) > 0:
+            logger.warning(
+                "LightningSpanProcessor already present in TracerProvider. You might have called init_worker() multiple times."
+                "Mantisdk will try to reuse the existing LightningSpanProcessor."
+            )
+            if len(span_processors) > 1:
+                logger.error("More than one LightningSpanProcessors present in TracerProvider. This should not happen.")
+            self._lightning_span_processor = span_processors[0]
+        else:
+            self._lightning_span_processor = LightningSpanProcessor()
+            self._get_tracer_provider().add_span_processor(self._lightning_span_processor)  # type: ignore
+
+    def teardown_worker(self, worker_id: int) -> None:
+        super().teardown_worker(worker_id)
+
+        if self.instrument_managed:
+            self.uninstrument(worker_id)
+            logger.info(f"[Worker {worker_id}] Instrumentation removed.")
+
+        # NOTE: The teardown doesn't try to remove the LightningSpanProcessor from the TracerProvider.
+        # Currently there is no stable way to fully restore the AgentOps state to the initial state.
+
+    @with_active_tracer_context
+    @asynccontextmanager
+    async def trace_context(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> AsyncGenerator[trace_api.Tracer, None]:
+        """
+        Starts a new tracing context. This should be used as a context manager.
+
+        Args:
+            name: Optional name for the tracing context.
+            store: Optional store to add the spans to.
+            rollout_id: Optional rollout ID to add the spans to.
+            attempt_id: Optional attempt ID to add the spans to.
+
+        Yields:
+            The OpenTelemetry tracer instance to collect spans.
+        """
+        if store is not None:
+            warnings.warn(
+                "store is deprecated in favor of init_worker(). It will be removed in the future.",
+                DeprecationWarning,
+                stacklevel=3,
+            )
+        else:
+            store = self._store
+        with self._trace_context_sync(name=name, store=store, rollout_id=rollout_id, attempt_id=attempt_id) as tracer:
+            yield tracer
+
+    @contextmanager
+    def _trace_context_sync(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> Iterator[trace_api.Tracer]:
+        """Implementation of `trace_context` for synchronous execution."""
+        if not self._lightning_span_processor:
+            raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
+        tracer_provider = self._get_tracer_provider()
+
+        kwargs: dict[str, Any] = {}
+        if name is not None:
+            kwargs["trace_name"] = name
+        elif rollout_id is not None:
+            kwargs["trace_name"] = rollout_id
+        if store is not None and rollout_id is not None and attempt_id is not None:
+            if store.capabilities.get("otlp_traces", False) is True:
+                logger.debug(f"Tracing to LightningStore rollout_id={rollout_id}, attempt_id={attempt_id}")
+                self._enable_native_otlp_exporter(store, rollout_id, attempt_id)
+            else:
+                self._disable_native_otlp_exporter()
+            ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
+            with ctx:
+                # AgentOps end_trace and start_trace must live inside the lightning span processor context.
+                # Otherwise some traces might not be recorded.
+                with self._agentops_trace_context(rollout_id, attempt_id, kwargs):
+                    yield trace_api.get_tracer(__name__, tracer_provider=tracer_provider)
+        elif store is None and rollout_id is None and attempt_id is None:
+            self._disable_native_otlp_exporter()
+            with self._lightning_span_processor:
+                with self._agentops_trace_context(None, None, kwargs):
+                    yield trace_api.get_tracer(__name__, tracer_provider=tracer_provider)
+        else:
+            raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
+
+    @contextmanager
+    def _agentops_trace_context(self, rollout_id: Optional[str], attempt_id: Optional[str], kwargs: dict[str, Any]):
+        trace = agentops.start_trace(**kwargs)
+        status = StatusCode.OK  # type: ignore
+        try:
+            yield
+        except Exception as e:
+            # This will catch errors in user code.
+            status = StatusCode.ERROR  # type: ignore
+            logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}: {e}")
+            raise  # should reraise the error here so that runner can handle it
+        finally:
+            agentops.end_trace(trace, end_state=status)  # type: ignore
+
+    def get_langchain_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
+        """
+        Get the Langchain callback handler for integrating with Langchain.
+
+        Args:
+            tags: Optional list of tags to apply to the Langchain callback handler.
+
+        Returns:
+            An instance of the Langchain callback handler.
+        """
+        import agentops
+        from agentops.integration.callbacks.langchain import LangchainCallbackHandler
+
+        tags = tags or []
+        client_instance = agentops.get_client()
+        api_key = None
+        if client_instance.initialized:
+            api_key = client_instance.config.api_key
+        else:
+            logger.warning(
+                "AgentOps client not initialized when creating LangchainCallbackHandler. API key may be missing."
+            )
+        return LangchainCallbackHandler(api_key=api_key, tags=tags)
+
+    get_langchain_callback_handler = get_langchain_handler  # alias
+
+    def _get_tracer_provider(self) -> TracerProviderImpl:
+        try:
+            # new versions
+            instance = agentops.sdk.core.tracer
+            if instance.provider is None:
+                raise RuntimeError("AgentOps TracerProvider is not initialized.")
+
+            if get_tracer_provider() is not instance.provider:
+                logger.error(
+                    "Mismatch between global singleton TracerProvider and AgentOps TracerProvider. "
+                    "AgentOps might not work properly."
+                )
+
+            if not isinstance(instance.provider, TracerProviderImpl):  # type: ignore
+                raise RuntimeError("Unsupported TracerProvider type for AgentOps instrumentation.")
+
+            self._tracer_provider = instance.provider
+            return self._tracer_provider
+        except AttributeError:
+            # old versions
+            instance = TracingCore.get_instance()  # type: ignore
+            self._tracer_provider = instance._provider  # type: ignore
+            return self._tracer_provider  # type: ignore

mantisdk/tracer/base.py

@@ -0,0 +1,287 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import functools
+import logging
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional, TypeVar
+
+from mantisdk.store.base import LightningStore
+from mantisdk.types import Attributes, ParallelWorkerBase, Span, SpanCoreFields, SpanRecordingContext, TraceStatus
+
+if TYPE_CHECKING:
+    from langchain_core.callbacks.base import BaseCallbackHandler  # type: ignore
+
+logger = logging.getLogger(__name__)
+
+
+T = TypeVar("T")
+
+
+_active_tracer: Optional[Tracer] = None
+
+T_func = Callable[..., Awaitable[Any]]
+
+
+class Tracer(ParallelWorkerBase):
+    """
+    An abstract base class for tracers.
+
+    This class defines a standard interface for tracing code execution,
+    capturing the resulting spans, and providing them for analysis. It is
+    designed to be backend-agnostic, allowing for different implementations
+    (e.g., for AgentOps, OpenTelemetry, Docker, etc.).
+
+    The primary interaction pattern is through the [`trace_context`][mantisdk.Tracer.trace_context]
+    context manager, which ensures that traces are properly started and captured,
+    even in the case of exceptions.
+
+    A typical workflow:
+
+    ```python
+    tracer = YourTracerImplementation()
+
+    try:
+        async with tracer.trace_context(name="my_traced_task"):
+            # ... code to be traced ...
+            await run_my_agent_logic()
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+    # Retrieve the trace data after the context block
+    spans: list[ReadableSpan] = tracer.get_last_trace()
+
+    # Process the trace data
+    if trace_tree:
+        rl_triplets = TracerTraceToTriplet().adapt(spans)
+        # ... do something with the triplets
+    ```
+    """
+
+    _store: Optional[LightningStore] = None
+
+    def init_worker(self, worker_id: int, store: Optional[LightningStore] = None) -> None:
+        """Initialize the tracer for a worker.
+
+        Args:
+            worker_id: The ID of the worker.
+            store: The store to add the spans to. If it's provided, traces will be added to the store when tracing.
+        """
+        super().init_worker(worker_id)
+        self._store = store
+
+    def trace_context(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> AsyncContextManager[Any]:
+        """
+        Starts a new tracing context. This should be used as a context manager.
+
+        The implementation should handle the setup and teardown of the tracing
+        for the enclosed code block. It must ensure that any spans generated
+        within the `with` block are collected and made available via
+        [`get_last_trace`][mantisdk.Tracer.get_last_trace].
+
+        Args:
+            name: The name for the root span of this trace context.
+            store: The store to add the spans to. Deprecated in favor of passing store to init_worker().
+            rollout_id: The rollout ID to add the spans to.
+            attempt_id: The attempt ID to add the spans to.
+        """
+        raise NotImplementedError()
+
+    def _trace_context_sync(
+        self,
+        name: Optional[str] = None,
+        *,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> ContextManager[Any]:
+        """Internal API for CI backward compatibility."""
+        raise NotImplementedError()
+
+    def get_last_trace(self) -> List[Span]:
+        """
+        Retrieves the raw list of captured spans from the most recent trace.
+
+        Returns:
+            A list of [`Span`][mantisdk.Span] objects collected during the last trace.
+        """
+        raise NotImplementedError()
+
+    def trace_run(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+        """
+        A convenience wrapper to trace the execution of a single synchronous function.
+
+        Deprecated in favor of customizing Runners.
+
+        Args:
+            func: The synchronous function to execute and trace.
+            *args: Positional arguments to pass to the function.
+            **kwargs: Keyword arguments to pass to the function.
+
+        Returns:
+            The return value of the function.
+        """
+        with self._trace_context_sync(name=func.__name__):
+            return func(*args, **kwargs)
+
+    def create_span(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        timestamp: Optional[float] = None,
+        status: Optional[TraceStatus] = None,
+    ) -> SpanCoreFields:
+        """Notify the tracer that a span should be created here.
+
+        It uses a fire-and-forget approach and doesn't wait for the span to be created.
+
+        Args:
+            name: The name of the span.
+            attributes: The attributes of the span.
+            timestamp: The timestamp of the span.
+            status: The status of the span.
+
+        Returns:
+            The core fields of the span.
+        """
+        raise NotImplementedError()
+
+    def operation_context(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+    ) -> ContextManager[SpanRecordingContext]:
+        """Start to record an operation to a span.
+
+        Args:
+            name: The name of the operation.
+            attributes: The attributes of the operation.
+            start_time: The start time of the operation.
+            end_time: The end time of the operation.
+
+        Returns:
+            A [`SpanRecordingContext`][mantisdk.SpanRecordingContext] for recording the operation on the span.
+        """
+        raise NotImplementedError()
+
+    async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> Any:
+        """
+        A convenience wrapper to trace the execution of a single asynchronous function.
+
+        Deprecated in favor of customizing Runners.
+
+        Args:
+            func: The asynchronous function to execute and trace.
+            *args: Positional arguments to pass to the function.
+            **kwargs: Keyword arguments to pass to the function.
+
+        Returns:
+            The return value of the function.
+        """
+        async with self.trace_context(name=func.__name__):
+            return await func(*args, **kwargs)
+
+    def get_langchain_handler(self) -> Optional[BaseCallbackHandler]:  # type: ignore
+        """Get a handler to install in langchain agent callback.
+
+        Agents are expected to use this handler in their agents to enable tracing.
+        """
+        logger.warning(f"{self.__class__.__name__} does not provide a LangChain callback handler.")
+        return None
+
+    @contextmanager
+    def lifespan(self, store: Optional[LightningStore] = None):
+        """A context manager to manage the lifespan of the tracer.
+
+        This can be used to set up and tear down any necessary resources
+        for the tracer, useful for debugging purposes.
+
+        Args:
+            store: The store to add the spans to. If it's provided, traces will be added to the store when tracing.
+        """
+        has_init = False
+        has_init_worker = False
+        try:
+            self.init()
+            has_init = True
+
+            self.init_worker(0, store)
+            has_init_worker = True
+
+            yield
+
+        finally:
+            if has_init_worker:
+                self.teardown_worker(0)
+            if has_init:
+                self.teardown()
+
+
+def set_active_tracer(tracer: Tracer):
+    """Set the active tracer for the current process.
+
+    Args:
+        tracer: The tracer to set as active.
+    """
+    global _active_tracer
+    if _active_tracer is not None:
+        raise ValueError("An active tracer is already set. Cannot set a new one.")
+    _active_tracer = tracer
+
+
+def clear_active_tracer():
+    """Clear the active tracer for the current process."""
+    global _active_tracer
+    _active_tracer = None
+
+
+def get_active_tracer() -> Optional[Tracer]:
+    """Get the active tracer for the current process.
+
+    Returns:
+        The active tracer, or None if no tracer is active.
+    """
+    global _active_tracer
+    return _active_tracer
+
+
+class _ActiveTracerAsyncCM(AsyncContextManager[T]):
+    def __init__(self, tracer: Tracer, inner: AsyncContextManager[T]):
+        self._tracer = tracer
+        self._inner = inner
+
+    async def __aenter__(self) -> T:
+        set_active_tracer(self._tracer)  # will raise if nested
+        try:
+            return await self._inner.__aenter__()
+        except Exception:
+            clear_active_tracer()
+            raise
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> Optional[bool]:
+        try:
+            return await self._inner.__aexit__(*args, **kwargs)
+        finally:
+            clear_active_tracer()
+
+
+def with_active_tracer_context(
+    func: Callable[..., AsyncContextManager[T]],
+) -> Callable[..., AsyncContextManager[T]]:
+    """Decorate a method returning an AsyncContextManager so tracer is active for the whole `async with`."""
+
+    @functools.wraps(func)
+    def wrapper(self: Tracer, *args: Any, **kwargs: Any) -> AsyncContextManager[T]:
+        cm = func(self, *args, **kwargs)
+        return _ActiveTracerAsyncCM(self, cm)
+
+    return wrapper

mantisdk/tracer/dummy.py

@@ -0,0 +1,106 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+import time
+from contextlib import contextmanager
+from typing import (
+    Iterator,
+    Optional,
+)
+
+from mantisdk.types import (
+    Attributes,
+    SpanCoreFields,
+    SpanRecordingContext,
+    StatusCode,
+    TraceStatus,
+)
+from mantisdk.utils.otel import format_exception_attributes
+
+from .base import Tracer
+
+logger = logging.getLogger(__name__)
+
+
+class DummySpanRecordingContext(SpanRecordingContext):
+    """Context for recording operations on a dummy span, not dependent on any backend tracer."""
+
+    def __init__(self, name: str, attributes: Optional[Attributes] = None, start_time: Optional[float] = None) -> None:
+        self.name = name
+        self.attributes = attributes or {}
+        self.start_time = start_time or time.time()
+        self.end_time = None
+        self.status = TraceStatus(status_code="OK")
+
+    def record_exception(self, exception: BaseException) -> None:
+        self.record_status("ERROR", str(exception))
+        self.record_attributes(format_exception_attributes(exception))
+
+    def record_attributes(self, attributes: Attributes) -> None:
+        self.attributes.update(attributes)
+
+    def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
+        self.status = TraceStatus(status_code=status_code, description=description)
+
+    def finalize(self, end_time: Optional[float] = None) -> None:
+        self.end_time = end_time or time.time()
+
+    def get_recorded_span(self) -> SpanCoreFields:
+        if self.end_time is None:
+            raise ValueError("End time is not set. Call finalize() first.")
+        return SpanCoreFields(
+            name=self.name,
+            attributes=self.attributes,
+            start_time=self.start_time,
+            end_time=self.end_time,
+            status=self.status,
+        )
+
+
+class DummyTracer(Tracer):
+    """A dummy tracer that does not trace anything, but it is compatible with the emitter API.
+
+    It doesn't rely on any backend tracer, and also doesn't use any stores.
+    """
+
+    def create_span(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        timestamp: Optional[float] = None,
+        status: Optional[TraceStatus] = None,
+    ) -> SpanCoreFields:
+        if attributes is None:
+            attributes = {}
+        if timestamp is None:
+            timestamp = time.time()
+        if status is None:
+            status = TraceStatus(status_code="OK")
+        return SpanCoreFields(
+            name=name,
+            attributes=attributes,
+            start_time=timestamp,
+            end_time=timestamp,
+            status=status,
+        )
+
+    @contextmanager
+    def operation_context(
+        self,
+        name: str,
+        attributes: Optional[Attributes] = None,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+    ) -> Iterator[DummySpanRecordingContext]:
+        start_time = start_time or time.time()
+        recording_context = DummySpanRecordingContext(name, attributes, start_time)
+        try:
+            yield recording_context
+        except Exception as exc:
+            recording_context.record_exception(exc)
+            recording_context.record_status("ERROR", str(exc))
+            raise
+        finally:
+            recording_context.finalize(end_time)