spans.py

"""Span helper utilities and OTel semantic-convention attribute keys.

Use semconv-defined attribute names where one exists. The full GenAI registry
is at https://opentelemetry.io/docs/specs/semconv/registry/attributes/gen-ai/
and the DB registry at https://opentelemetry.io/docs/specs/semconv/database/.

Semconv-stable attribute keys live as module constants so a typo at the
call site is a NameError, not a silently-different attribute.
"""

from __future__ import annotations

from contextlib import contextmanager
from typing import TYPE_CHECKING

from opentelemetry import trace

if TYPE_CHECKING:
    from collections.abc import Iterator, Mapping

# ---------------------------------------------------------------------------
# GenAI semantic convention attribute keys
# https://opentelemetry.io/docs/specs/semconv/registry/attributes/gen-ai/
# ---------------------------------------------------------------------------

# GenAI — conversation & request
GENAI_CONVERSATION_ID: str = "gen_ai.conversation.id"
GENAI_REQUEST_MODEL: str = "gen_ai.request.model"
GENAI_OPERATION_NAME: str = "gen_ai.operation.name"
GENAI_PROVIDER_NAME: str = "gen_ai.provider.name"

# GenAI — usage (token counts)
GENAI_USAGE_INPUT_TOKENS: str = "gen_ai.usage.input_tokens"
GENAI_USAGE_OUTPUT_TOKENS: str = "gen_ai.usage.output_tokens"

# GenAI — tool calling
GENAI_TOOL_NAME: str = "gen_ai.tool.name"
GENAI_TOOL_CALL_ARGUMENTS: str = "gen_ai.tool.call.arguments"
GENAI_TOOL_CALL_RESULT: str = "gen_ai.tool.call.result"

# DB semantic convention attribute keys
# https://opentelemetry.io/docs/specs/semconv/database/
DB_QUERY_TEXT: str = "db.query.text"
DB_RESPONSE_RETURNED_ROWS: str = "db.response.returned_rows"


# ---------------------------------------------------------------------------
# Span helpers
# ---------------------------------------------------------------------------


@contextmanager
def agent_span(
    name: str,
    attributes: Mapping[str, str | int | float | bool] | None = None,
) -> Iterator[trace.Span]:
    """Create a span and yield it for further mutation.

    Args:
        name: Human-readable span name.
        attributes: Initial key-value attributes to set on the span.
    """
    tracer = trace.get_tracer(__name__)
    with tracer.start_as_current_span(name) as span:
        if attributes:
            for key, value in attributes.items():
                span.set_attribute(key, value)
        yield span


def set_span_attributes(
    span: trace.Span,
    **kwargs: str | int | float | bool | None,
) -> None:
    """Set multiple attributes on a span, filtering ``None`` values.

    Args:
        span: The span to annotate.
        **kwargs: Attribute key-value pairs. ``None`` values are silently
            skipped.
    """
    for key, value in kwargs.items():
        if value is not None:
            span.set_attribute(key, value)