mlflow--mlflow/mlflow/semantic_kernel/tracing_utils.py

import logging
from typing import Any

from opentelemetry import context as otel_context_api
from opentelemetry import trace as otel_trace
from opentelemetry.trace import get_current_span
from semantic_kernel.contents.chat_history import ChatHistory
from semantic_kernel.contents.kernel_content import KernelContent
from semantic_kernel.contents.streaming_content_mixin import StreamingContentMixin
from semantic_kernel.functions import FunctionResult
from semantic_kernel.utils.telemetry.model_diagnostics import (
    gen_ai_attributes as model_gen_ai_attributes,
)
from semantic_kernel.utils.telemetry.model_diagnostics.decorators import (
    CHAT_COMPLETION_OPERATION,
    TEXT_COMPLETION_OPERATION,
)
from semantic_kernel.utils.telemetry.model_diagnostics.function_tracer import (
    OPERATION_NAME as FUNCTION_OPERATION_NAME,
)

import mlflow
from mlflow.entities import SpanType
from mlflow.entities.span import LiveSpan
from mlflow.tracing.constant import SpanAttributeKey, TokenUsageKey
from mlflow.tracing.utils import (
    construct_full_inputs,
    get_mlflow_span_for_otel_span,
)

_OPERATION_TO_SPAN_TYPE = {
    CHAT_COMPLETION_OPERATION: SpanType.CHAT_MODEL,
    TEXT_COMPLETION_OPERATION: SpanType.LLM,
    FUNCTION_OPERATION_NAME: SpanType.TOOL,
    # https://github.com/microsoft/semantic-kernel/blob/d5ee6aa1c176a4b860aba72edaa961570874661b/python/semantic_kernel/utils/telemetry/agent_diagnostics/decorators.py#L22
    "invoke_agent": SpanType.AGENT,
}

# NB: Streaming operation names were removed in Semantic Kernel 1.38.0
try:
    from semantic_kernel.utils.telemetry.agent_diagnostics.decorators import (
        CHAT_STREAMING_COMPLETION_OPERATION,
        TEXT_STREAMING_COMPLETION_OPERATION,
    )

    _OPERATION_TO_SPAN_TYPE[CHAT_STREAMING_COMPLETION_OPERATION] = SpanType.CHAT_MODEL
    _OPERATION_TO_SPAN_TYPE[TEXT_STREAMING_COMPLETION_OPERATION] = SpanType.LLM
except ImportError:
    pass

_logger = logging.getLogger(__name__)


def semantic_kernel_diagnostics_wrapper(original, *args, **kwargs) -> None:
    """
    Wrapper for Semantic Kernel's model diagnostics decorators.

    This wrapper is used to record the inputs and outputs to the span, because
    Semantic Kernel's Otel span do not record the inputs and outputs.
    """
    full_kwargs = construct_full_inputs(original, *args, **kwargs)
    current_span = full_kwargs.get("current_span") or get_current_span()
    mlflow_span = get_mlflow_span_for_otel_span(current_span)

    if not mlflow_span:
        _logger.debug("Span is not found or recording. Skipping error handling.")
        return original(*args, **kwargs)

    if prompt := full_kwargs.get("prompt"):
        # Wrapping _set_completion_input
        # https://github.com/microsoft/semantic-kernel/blob/d5ee6aa1c176a4b860aba72edaa961570874661b/python/semantic_kernel/utils/telemetry/model_diagnostics/decorators.py#L369
        mlflow_span.set_inputs(_parse_content(prompt))

    if completions := full_kwargs.get("completions"):
        # Wrapping _set_completion_response
        # https://github.com/microsoft/semantic-kernel/blob/d5ee6aa1c176a4b860aba72edaa961570874661b/
        mlflow_span.set_outputs({"messages": [_parse_content(c) for c in completions]})

    if error := full_kwargs.get("error"):
        # Wrapping _set_completion_error
        # https://github.com/microsoft/semantic-kernel/blob/d5ee6aa1c176a4b860aba72edaa961570874661b/python/semantic_kernel/utils/telemetry/model_diagnostics/decorators.py#L452
        mlflow_span.record_exception(error)

    return original(*args, **kwargs)


async def patched_kernel_entry_point(original, self, *args, **kwargs):
    with mlflow.start_span(
        name=f"{self.__class__.__name__}.{original.__name__}",
        span_type=SpanType.AGENT,
    ) as mlflow_span:
        inputs = construct_full_inputs(original, self, *args, **kwargs)
        mlflow_span.set_inputs(_parse_content(inputs))

        # Attach the MLflow span to the global OTel context so that Semantic Kernel's
        # internal OTel spans (e.g., execute_tool, chat.completions) will inherit the
        # same trace_id and be properly linked as child spans.
        global_ctx = otel_trace.set_span_in_context(mlflow_span._span)
        token = otel_context_api.attach(global_ctx)
        try:
            result = await original(self, *args, **kwargs)
        finally:
            otel_context_api.detach(token)

        mlflow_span.set_outputs(_parse_content(result))

    return result


def _parse_content(value: Any) -> Any:
    """
    Parse the message content objects in Semantic Kernel into a more readable format.

    Those objects are Pydantic models, but includes many noisy fields that are not
    useful for debugging and hard to read. The base KernelContent class has a to_dict()
    method that converts them into more readable format (role, content), so we use that.
    """
    if isinstance(value, dict) and (chat_history := value.get("chat_history")):
        value = _parse_content(chat_history)
    elif isinstance(value, ChatHistory):
        # Record chat history as a list of messages for better readability
        value = {"messages": [_parse_content(m) for m in value.messages]}
    elif isinstance(value, (KernelContent, StreamingContentMixin)):
        value = value.to_dict()
    elif isinstance(value, FunctionResult):
        # Extract "value" field from the FunctionResult object
        value = _parse_content(value.value)
    elif isinstance(value, list):
        value = [_parse_content(item) for item in value]
    return value


def set_span_type(mlflow_span: LiveSpan) -> str:
    """Determine the span type based on the operation."""
    span_type = SpanType.UNKNOWN
    if operation := mlflow_span.get_attribute(model_gen_ai_attributes.OPERATION):
        span_type = _OPERATION_TO_SPAN_TYPE.get(operation, SpanType.UNKNOWN)

    mlflow_span.set_span_type(span_type)


def set_token_usage(mlflow_span: LiveSpan) -> None:
    """Set token usage attributes on the MLflow span."""
    input_tokens = mlflow_span.get_attribute(model_gen_ai_attributes.INPUT_TOKENS)
    output_tokens = mlflow_span.get_attribute(model_gen_ai_attributes.OUTPUT_TOKENS)

    usage_dict = {}
    if input_tokens is not None:
        usage_dict[TokenUsageKey.INPUT_TOKENS] = input_tokens
    if output_tokens is not None:
        usage_dict[TokenUsageKey.OUTPUT_TOKENS] = output_tokens

    if input_tokens is not None or output_tokens is not None:
        total_tokens = (input_tokens or 0) + (output_tokens or 0)
        usage_dict[TokenUsageKey.TOTAL_TOKENS] = total_tokens

    if usage_dict:
        mlflow_span.set_attribute(SpanAttributeKey.CHAT_USAGE, usage_dict)


def set_model(mlflow_span: LiveSpan) -> None:
    """Set model name and provider attributes on the MLflow span."""
    if model := mlflow_span.get_attribute(model_gen_ai_attributes.MODEL):
        mlflow_span.set_attribute(SpanAttributeKey.MODEL, model)
    if provider := mlflow_span.get_attribute(model_gen_ai_attributes.SYSTEM):
        mlflow_span.set_attribute(SpanAttributeKey.MODEL_PROVIDER, provider)