microsoft--agent-lightning/agentlightning/tracer/weave.py

# Copyright (c) Microsoft. All rights reserved.

from __future__ import annotations

import asyncio
import concurrent.futures as futures
import logging
import os
import re
import weakref
from contextlib import asynccontextmanager, contextmanager
from datetime import datetime
from typing import (
    Any,
    AsyncIterator,
    Callable,
    Dict,
    Iterator,
    List,
    Optional,
    cast,
)

import weave
from opentelemetry.semconv.attributes import exception_attributes
from weave.trace.call import Call
from weave.trace.settings import UserSettings
from weave.trace.weave_client import WeaveClient
from weave.trace_server import trace_server_interface as tsi
from weave.wandb_interface.context import set_wandb_api_context

from agentlightning.instrumentation.weave import InMemoryWeaveTraceServer, instrument_weave, uninstrument_weave
from agentlightning.semconv import LightningResourceAttributes, LightningSpanAttributes
from agentlightning.store.base import LightningStore
from agentlightning.types import (
    Attributes,
    OtelResource,
    Span,
    SpanContext,
    SpanCoreFields,
    SpanRecordingContext,
    StatusCode,
    TraceStatus,
)
from agentlightning.utils.id import generate_id
from agentlightning.utils.otel import (
    filter_and_unflatten_attributes,
    flatten_attributes,
    format_exception_attributes,
    sanitize_attributes,
)

from .base import Tracer, with_active_tracer_context

logger = logging.getLogger(__name__)


def op_name_to_func_name(op_name: str) -> str:
    """Convert a Weave operation name to a function name.

    Weave operation names look like this: `weave:///xxx/agentlightning.tracer.weave/op/openai.chat.completions.create:019b10be-...-44d74272569c`
    """
    match = re.search(r"/([^/:]+):", op_name)
    if match:
        return match.group(1)
    else:
        return op_name


def random_project_name() -> str:
    return "agl/weave-" + generate_id(12)


def get_timestamp_or_throw(date: Optional[datetime], field_name: str) -> float:
    if date is None:
        raise ValueError(f"{field_name} is required but not set")
    return date.timestamp()


class WeaveSpanRecordingContext(SpanRecordingContext):
    """Universal interface for recording operations on a Weave call."""

    def __init__(self, call: Call) -> None:
        self._call = call

    def record_exception(self, exception: BaseException) -> None:
        self._call.exception = str(exception)
        self.record_status("ERROR", str(exception))
        self.record_attributes(format_exception_attributes(exception))

    def _get_input_from_attributes(self, attributes: Attributes) -> Dict[str, Any]:
        if LightningSpanAttributes.OPERATION_INPUT.value in attributes:
            # This can be a very rare case. If it happens, we can just let it throw.
            return cast(Dict[str, Any], attributes[LightningSpanAttributes.OPERATION_INPUT.value])
        else:
            filtered_attributes = filter_and_unflatten_attributes(
                attributes, LightningSpanAttributes.OPERATION_INPUT.value
            )
            if isinstance(filtered_attributes, list):
                return {str(i): v for i, v in enumerate(filtered_attributes)}
            else:
                return filtered_attributes

    def _get_output_from_attributes(self, attributes: Attributes) -> Any:
        if LightningSpanAttributes.OPERATION_OUTPUT.value in attributes:
            return attributes[LightningSpanAttributes.OPERATION_OUTPUT.value]
        else:
            return filter_and_unflatten_attributes(attributes, LightningSpanAttributes.OPERATION_OUTPUT.value)

    def record_attributes(self, attributes: Attributes) -> None:
        input_attributes = self._get_input_from_attributes(attributes)
        if input_attributes:
            self._call.inputs.update(input_attributes)

        output_attributes = self._get_output_from_attributes(attributes)
        if output_attributes:
            if self._call.output is not None:
                logger.warning(f"Output is already set. It will be overridden: {self._call.output}")
            self._call.output = output_attributes

        if LightningSpanAttributes.OPERATION_NAME.value in attributes:
            logger.error(
                f"Cannot record operation name as an attribute. It will be skipped: {attributes[LightningSpanAttributes.OPERATION_NAME.value]}"
            )

        # The rest of the attributes are recorded as summary.
        for key, value in attributes.items():
            if (
                not key == LightningSpanAttributes.OPERATION_INPUT.value
                and not key.startswith(LightningSpanAttributes.OPERATION_INPUT.value + ".")
                and not key == LightningSpanAttributes.OPERATION_OUTPUT.value
                and not key.startswith(LightningSpanAttributes.OPERATION_OUTPUT.value + ".")
                and not key == LightningSpanAttributes.OPERATION_NAME.value
            ):
                if self._call.summary is None:
                    self._call.summary = {}
                self._call.summary[key] = value

    def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
        if status_code == "ERROR":
            if not description:
                raise ValueError("Description is required when status code is ERROR")
            self._call.exception = description
        elif status_code == "OK":
            self._call.exception = None
        # Do nothing for other status codes.

    def finalize(self) -> None:
        # Do nothing
        pass

    def get_recorded_span(self) -> SpanCoreFields:
        return SpanCoreFields(
            name=self._call.op_name,
            attributes=flatten_attributes(self._call.attributes or {}),
            start_time=self._call.started_at.timestamp() if self._call.started_at else None,
            end_time=self._call.ended_at.timestamp() if self._call.ended_at else None,
            status=TraceStatus(
                status_code="OK" if self._call.exception is None else "ERROR", description=self._call.exception
            ),
        )


class WeaveTracerManagedTraceServer(InMemoryWeaveTraceServer):
    """A managed trace server for WeaveTracer."""

    def __init__(
        self,
        partial_call_callback: Callable[[Dict[str, Any]], None],
        complete_call_callback: Callable[[tsi.CallSchema], None],
    ):
        super().__init__()
        self.partial_call_callback = partial_call_callback
        self.complete_call_callback = complete_call_callback
        self._calls_already_invoked: set[str] = set()

    def trigger_callbacks(self, call_id: str) -> None:
        with self._call_threading_lock:
            if call_id in self.calls:
                if call_id not in self._calls_already_invoked:
                    self._calls_already_invoked.add(call_id)
                    self.complete_call_callback(self.calls[call_id])
                else:
                    logger.info(f"Call {call_id} has callback already invoked. Skipping.")
            elif call_id in self.partial_calls:
                self.partial_call_callback(self.partial_calls[call_id])
            else:
                logger.error(f"Call {call_id} not found in partial_calls or calls")

    def call_start(self, req: tsi.CallStartReq) -> tsi.CallStartRes:
        try:
            ret = super().call_start(req)
            self.trigger_callbacks(ret.id)
            return ret
        except Exception:
            logger.exception(f"Error calling call_start: {req}", exc_info=True)
            raise

    def call_end(self, req: tsi.CallEndReq) -> tsi.CallEndRes:
        try:
            ret = super().call_end(req)
            self.trigger_callbacks(req.end.id)
            return ret
        except Exception:
            logger.exception(f"Error calling call_end: {req}", exc_info=True)
            raise

    def clear(self) -> None:
        self._calls_already_invoked.clear()


class WeaveTracer(Tracer):
    """Tracer implementation using Weave for telemetry and trace logging.

    This replaces AgentOpsTracer with a Weave-based manual trace context. It tracks:

    - Function/method calls
    - Input/Output data
    - Exceptions

    and logs them to Weave Cloud (W&B backend) or optionally bypasses the network for testing.
    """

    def __init__(
        self,
        *,
        project_name: str | None = None,
        weave_user_settings: UserSettings | None = None,
        instrument_managed: bool = True,
    ):
        """Initialize a WeaveTracer instance.

        Args:
            project_name: Optional project name for Weave; defaults to the current module name.
            weave_user_settings: Optional UserSettings for Weave.
            instrument_managed: Whether to patch the Weave/W&B integration to bypass actual network calls for testing.
        """
        super().__init__()
        self.project_name = project_name
        self.instrument_managed = instrument_managed
        self.weave_user_settings = weave_user_settings or UserSettings(use_server_cache=False)

        self._store: Optional[LightningStore] = None
        self._server = WeaveTracerManagedTraceServer(
            partial_call_callback=self.partial_call_callback, complete_call_callback=self.complete_call_callback
        )

        self._default_sequence_counter: int = 0
        self._calls: Dict[str, tsi.CallSchema] = {}  # call_id -> call
        self._spans: List[Span] = []  # spans in the current trace
        self._rollout_id: Optional[str] = None
        self._attempt_id: Optional[str] = None
        self._partial_call_futures: Dict[str, asyncio.Future[int] | futures.Future[int]] = {}
        self._complete_call_futures: List[asyncio.Future[None] | futures.Future[None]] = []
        self._loop: weakref.ReferenceType[asyncio.AbstractEventLoop] | None = None

    def instrument(self, worker_id: int):
        instrument_weave(self._server)

    def uninstrument(self, worker_id: int):
        uninstrument_weave()

    def init_worker(self, worker_id: int, store: Optional[LightningStore] = None):
        """
        Initialize the tracer for a worker thread/process.

        Args:
            worker_id: Identifier of the worker.
            store: Optional LightningStore for storing spans.
        """
        super().init_worker(worker_id, store)
        logger.info(f"[Worker {worker_id}] Setting up Weave tracer...")
        self._store = store

        # Optionally patch network calls to bypass real Weave/W&B endpoints
        if self.instrument_managed:
            self.instrument(worker_id)

        # If WANDB_API_KEY is not set, we need to initialize Weave with a hack
        if not os.getenv("WANDB_API_KEY"):
            logger.info("WANDB_API_KEY is not set. Initializing Weave a mock context.")
            set_wandb_api_context("agl", api_key=None, headers=None, cookies=None)
        else:
            logger.debug("WANDB_API_KEY is set. Weave will be initialized automatically.")

        weave_client = weave.get_client()
        if self.project_name is None:
            self.project_name = random_project_name()

        if weave_client is not None:
            logger.warning("Weave client was already initialized. Reentrant calls are at your own risk.")
            if weave_client.project == self.project_name:
                logger.error(
                    f"Weave client was already initialized for the same project '{self.project_name}'. It's very likely that weave won't work correctly."
                )

        # Init no matter what
        try:
            weave.init(project_name=self.project_name, settings=self.weave_user_settings)
            logger.info(f"[Worker {worker_id}] Weave client initialized.")
        except Exception as exc:
            raise RuntimeError(f"Failed to initialize Weave for project '{self.project_name}'") from exc

    def teardown_worker(self, worker_id: int):
        """
        Clean up tracer resources for the worker.

        Args:
            worker_id: Identifier of the worker.
        """
        super().teardown_worker(worker_id)

        if self.instrument_managed:
            self.uninstrument(worker_id)
            logger.info(f"[Worker {worker_id}] Instrumentation removed.")

    @with_active_tracer_context
    @asynccontextmanager
    async def trace_context(
        self,
        name: Optional[str] = None,
        *,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
        **kwargs: Any,
    ) -> AsyncIterator[Any]:
        """Asynchronous implementation of the tracing context.

        Args:
            name: Optional operation name.
            rollout_id: Optional rollout ID.
            attempt_id: Optional attempt ID.

        Raises:
            ValueError: If store, rollout_id, and attempt_id are inconsistently provided.
            RuntimeError: If Weave is not installed or client is uninitialized.
        """

        if rollout_id is not None and attempt_id is not None:
            self._rollout_id = rollout_id
            self._attempt_id = attempt_id
        elif rollout_id is None and attempt_id is None:
            logger.info("No rollout_id or attempt_id provided. Skipping writing to store.")
            self._rollout_id = self._attempt_id = None
        else:
            raise ValueError("rollout_id and attempt_id must be either both provided or both None")

        await self._init_trace_context()

        weave_client = self._get_weave_client()

        if weave_client.server is not self._server:
            logger.error(
                "Weave client is not using the correct trace server. You might have multiple WeaveTracer instances running in the same process. "
                f"Expected {self._server}, got {weave_client.server}"
            )

        arg_op = name or weave_client.project
        arg_inputs: dict[str, str] = {}
        if rollout_id is not None:
            arg_inputs[LightningResourceAttributes.ROLLOUT_ID.value] = rollout_id
        if attempt_id is not None:
            arg_inputs[LightningResourceAttributes.ATTEMPT_ID.value] = attempt_id

        try:
            # Create a new trace call object in Weave
            trace_call = weave_client.create_call(  # pyright: ignore[reportUnknownMemberType]
                op=arg_op, inputs=arg_inputs
            )

            try:
                yield trace_call
                # Finish trace even if no exception
                weave_client.finish_call(trace_call)  # pyright: ignore[reportUnknownMemberType]
            except Exception as exc:
                # Finish trace and log any exception
                weave_client.finish_call(trace_call, exception=exc)  # pyright: ignore[reportUnknownMemberType]
                logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}, error={exc}")
                raise

        finally:
            try:
                weave_client.flush()
                # It's possible that the call end futures are from a dedicated Weave thread pool,
                await asyncio.gather(*[asyncio.wrap_future(future) for future in self._complete_call_futures])

            finally:
                # Mandatory cleanup
                self._rollout_id = None
                self._attempt_id = None
                self._server.clear()

    def create_span(
        self,
        name: str,
        attributes: Optional[Attributes] = None,
        timestamp: Optional[float] = None,
        status: Optional[TraceStatus] = None,
    ) -> SpanCoreFields:
        if timestamp is not None:
            logger.warning("Weave doesn't support customizing the start time of a call. Timestamp is ignored.")
        weave_client = self._get_weave_client()
        trace_call = weave_client.create_call(  # pyright: ignore[reportUnknownMemberType]
            op=name,
            attributes=attributes,
            inputs={},
        )
        # Immediately finish the call
        weave_client.finish_call(trace_call)  # pyright: ignore[reportUnknownMemberType]
        # We don't wait for the call to be propagated to the server.
        start_time = trace_call.started_at.timestamp() if trace_call.started_at else None
        end_time = trace_call.ended_at.timestamp() if trace_call.ended_at else None
        trace_status = (
            TraceStatus(status_code="OK")
            if trace_call.exception is None
            else TraceStatus(status_code="ERROR", description=trace_call.exception)
        )
        return SpanCoreFields(
            name=name,
            attributes=flatten_attributes(trace_call.attributes or {}),
            start_time=start_time,
            end_time=end_time,
            status=trace_status,
        )

    @contextmanager
    def operation_context(
        self,
        name: str,
        attributes: Optional[Attributes] = None,
        start_time: Optional[float] = None,
        end_time: Optional[float] = None,
    ) -> Iterator[SpanRecordingContext]:
        if start_time is not None:
            logger.warning("Weave doesn't support customizing the start time of a call. Timestamp is ignored.")
        if end_time is not None:
            logger.warning("Weave doesn't support customizing the end time of a call. Timestamp is ignored.")
        weave_client = self._get_weave_client()
        trace_call = weave_client.create_call(  # pyright: ignore[reportUnknownMemberType]
            op=name,
            attributes=attributes,
            inputs={},
        )
        recording_context = WeaveSpanRecordingContext(trace_call)
        try:
            yield recording_context
        except Exception as exc:
            recording_context.record_exception(exc)
            raise
        finally:
            weave_client.finish_call(trace_call)  # pyright: ignore[reportUnknownMemberType]

    async def _init_trace_context(self) -> None:
        """Initialize the trace context."""
        self._spans.clear()
        self._calls.clear()
        self._partial_call_futures.clear()
        self._complete_call_futures.clear()
        self._loop = weakref.ref(asyncio.get_running_loop())

    def _get_weave_client(self) -> WeaveClient:
        """Get the Weave client."""
        weave_client = weave.get_client()
        if not weave_client:
            raise RuntimeError("Weave client is not initialized. Call init_worker() first.")
        return weave_client

    def _ensure_loop(self) -> tuple[asyncio.AbstractEventLoop, bool]:
        """Returns a usable event loop and a boolean indicating whether it's the current running loop.

        Prefer using the main loop if it's possible. Otherwise, use the current running loop.
        """
        # Get the current running loop
        try:
            running_loop = asyncio.get_running_loop()
        except RuntimeError:
            running_loop = None

        # Get the main loop, which can be a different loop
        if self._loop is not None:
            main_loop = self._loop()
        else:
            main_loop = None

        if main_loop is not None:
            return main_loop, id(main_loop) == id(running_loop)
        elif running_loop is not None:
            return running_loop, True
        else:
            raise RuntimeError("No running event loop found. This should not happen.")

    def get_last_trace(self) -> List[Span]:
        return self._spans

    def partial_call_callback(self, request_content: Dict[str, Any]) -> None:
        call_id = request_content.get("id")
        if call_id is None:
            raise ValueError("Call ID is required even for partial calls")

        if call_id in self._partial_call_futures:
            raise ValueError(f"Call {call_id} already has a start future")

        # The callback must possibly be called from a dedicated Weave thread pool,
        # but it should be executed on the main event loop.
        try:
            loop, is_current_loop = self._ensure_loop()
            if is_current_loop:
                task = loop.create_task(self.partial_call_handler(request_content))
            else:
                # Schedule the task on the dedicated loop
                task = asyncio.run_coroutine_threadsafe(self.partial_call_handler(request_content), loop)
            self._partial_call_futures[call_id] = task
        except Exception as exc:
            logger.exception(f"Error creating call start task: {exc}", exc_info=True)

    def complete_call_callback(self, call: tsi.CallSchema) -> None:
        try:
            loop, is_current_loop = self._ensure_loop()
            if is_current_loop:
                task = loop.create_task(self.complete_call_handler(call))
            else:
                # Schedule the task on the dedicated loop
                task = asyncio.run_coroutine_threadsafe(self.complete_call_handler(call), loop)
            self._complete_call_futures.append(task)
        except Exception as exc:
            logger.exception(f"Error creating call finish task: {exc}", exc_info=True)

    async def _get_next_sequence_id(self) -> int:
        """Get the next sequence ID for a span.

        Use store to get the next sequence ID if available, otherwise use a default counter.
        """
        if self._rollout_id and self._attempt_id and self._store:
            return await self._store.get_next_span_sequence_id(self._rollout_id, self._attempt_id)
        else:
            self._default_sequence_counter += 1
            return self._default_sequence_counter

    async def partial_call_handler(self, request_content: Dict[str, Any]) -> int:
        """Handler called when a Weave Call starts.

        Args:
            request_content: The partial Weave Call object.

        Returns:
            The sequence ID for the call.
        """
        sequence_id = await self._get_next_sequence_id()
        return sequence_id

    async def complete_call_handler(self, call: tsi.CallSchema) -> None:
        """Handler called when a Weave Call finishes.

        Converts the call (including nested children) into spans and stores them in LightningStore.
        """
        # Make sure the corresponding call_start_future is complete
        if call.id in self._partial_call_futures:
            sequence_id = await asyncio.wrap_future(self._partial_call_futures[call.id])
            del self._partial_call_futures[call.id]
        else:
            # Fetch a new sequence ID as the call_start is somehow missing
            if call.id in self._calls:
                logger.warning(
                    f"Call {call.id} is already in calls. The call is already completed. Overwriting the call."
                )
            else:
                logger.warning(f"Call {call.id} has no start future. Fetching a new sequence ID.")
            sequence_id = await self._get_next_sequence_id()

        self._calls[call.id] = call

        span = await self.convert_call_to_span(call, self._rollout_id, self._attempt_id, sequence_id)
        self._spans.append(span)
        if self._store and self._rollout_id and self._attempt_id:
            try:
                await self._store.add_span(span)
            except Exception as exc:
                logger.exception(f"Error adding span to store: {exc}")

    async def convert_call_to_span(
        self,
        call: tsi.CallSchema,
        rollout_id: Optional[str] = None,
        attempt_id: Optional[str] = None,
        sequence_id: Optional[int] = None,
    ) -> Span:
        """Convert a Weave Call (with nested children) into a Agent-lightning Span.

        `rollout_id` and `attempt_id` are required to attach the spans to the store.

        Args:
            call: The Weave Call object.
            rollout_id: Optional rollout ID to attach to spans.
            attempt_id: Optional attempt ID to attach to spans.
            sequence_id: Optional sequence ID to attach to spans.

        Returns:
            List of converted spans.
        """
        rollout_id = rollout_id or "rollout-dummy"
        attempt_id = attempt_id or "attempt-dummy"
        sequence_id = sequence_id or 0

        start_ts: float = call.started_at.timestamp()
        end_ts: Optional[float] = call.ended_at.timestamp() if call.ended_at else None

        if call.exception:
            status = TraceStatus(status_code="ERROR", description=call.exception)
        else:
            status = TraceStatus(status_code="OK")

        attributes: Dict[str, Any] = {
            LightningSpanAttributes.OPERATION_NAME.value: call.op_name,
            # op_name can be possibly overridden by the attributes.
            **call.attributes,
        }
        if call.inputs:
            attributes[LightningSpanAttributes.OPERATION_INPUT.value] = call.inputs
        if call.output:
            attributes[LightningSpanAttributes.OPERATION_OUTPUT.value] = call.output
        if call.summary:
            # attributes can be possibly overridden by the summary.
            attributes.update(call.summary)
        if call.exception:
            attributes[exception_attributes.EXCEPTION_MESSAGE] = call.exception

        sanitized_attributes = sanitize_attributes(flatten_attributes(attributes, expand_leaf_lists=False))

        context = SpanContext(
            trace_id=call.trace_id,
            span_id=call.id,
            is_remote=False,
            trace_state={},
        )

        # Get context for parent
        if call.parent_id:
            parent_call = self._calls.get(call.parent_id)
            if parent_call:
                parent_context = SpanContext(
                    trace_id=parent_call.trace_id,
                    span_id=parent_call.id,
                    is_remote=False,
                    trace_state={},
                )
            else:
                parent_context = None
        else:
            parent_context = None

        # Build the Span object
        return Span(
            rollout_id=rollout_id,
            attempt_id=attempt_id,
            sequence_id=sequence_id,
            trace_id=call.trace_id,
            span_id=call.id,
            parent_id=call.parent_id,
            name=op_name_to_func_name(call.op_name),
            status=status,
            attributes=sanitized_attributes,
            events=[],  # Weave calls do not generate events
            links=[],  # Weave calls do not generate links
            start_time=start_ts,
            end_time=end_ts,
            context=context,
            parent=parent_context,
            resource=OtelResource(
                attributes={
                    LightningResourceAttributes.ROLLOUT_ID.value: rollout_id,
                    LightningResourceAttributes.ATTEMPT_ID.value: attempt_id,
                    LightningResourceAttributes.SPAN_SEQUENCE_ID.value: sequence_id,
                    LightningResourceAttributes.TRACER_NAME.value: "weave",
                },
                schema_url="",
            ),
        )