feat(telemetry): GenAI OpenTelemetry spans on _predict and _stream

.github/workflows/ci.yml

@@ -52,6 +52,7 @@ jobs:
       - uses: ./.github/actions/setup-python-uv
         with:
           python-version: ${{ inputs.python-version || '3.12' }}
+      - run: uv sync --extra otel
       - run: uv run mypy -p celeste && uv run mypy tests/
 
   security:

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
 
 [project.optional-dependencies]
 gcp = ["google-auth[requests]>=2.0.0"]
+otel = ["opentelemetry-api>=1.30"]
 
 [project.urls]
 Homepage = "https://withceleste.ai"
@@ -121,6 +122,7 @@ known-first-party = ["celeste"]
 [tool.ruff.lint.per-file-ignores]
 "tests/*" = ["D"]  # No docstrings required in tests
 "**/client.py" = ["ANN401"]  # Allow Any return types in mixin client classes
+"**/telemetry.py" = ["ANN401"]  # Allow Any: mirrors OTel API surface for forwarding
 
 [tool.mypy]
 python_version = "3.12"

src/celeste/client.py

@@ -9,6 +9,7 @@
 import httpx
 from pydantic import BaseModel, ConfigDict, Field
 
+from celeste import telemetry
 from celeste.auth import Authentication
 from celeste.core import Modality, Protocol, Provider
 from celeste.exceptions import (
@@ -212,29 +213,44 @@ async def _predict(
         Returns:
             Output of the parameterized type.
         """
-        inputs, parameters = self._validate_artifacts(inputs, **parameters)
-        request_body = self._build_request(inputs, extra_body=extra_body, **parameters)
-        response_data = await self._make_request(
-            request_body, endpoint=endpoint, extra_headers=extra_headers, **parameters
-        )
-        content = self._parse_content(response_data)
-        content = self._transform_output(content, **parameters)
-        tool_calls = self._parse_tool_calls(response_data)
-        reasoning, signature = self._parse_reasoning(response_data)
-        kwargs: dict[str, Any] = {}
-        if reasoning is not None:
-            kwargs["reasoning"] = reasoning
-        if signature:
-            kwargs["signature"] = signature
-        output = self._output_class()(
-            content=content,
-            usage=self._get_usage(response_data),
-            finish_reason=self._get_finish_reason(response_data),
-            metadata=self._build_metadata(response_data),
-            tool_calls=tool_calls,
-            **kwargs,
-        )
-        return output
+        with telemetry.tracer.start_as_current_span(
+            telemetry.span_name(self.modality, self.model),
+            attributes=telemetry.request_attributes(
+                model=self.model,
+                provider=self.provider,
+                protocol=self.protocol,
+                modality=self.modality,
+            ),
+        ) as span:
+            inputs, parameters = self._validate_artifacts(inputs, **parameters)
+            request_body = self._build_request(
+                inputs, extra_body=extra_body, **parameters
+            )
+            response_data = await self._make_request(
+                request_body,
+                endpoint=endpoint,
+                extra_headers=extra_headers,
+                **parameters,
+            )
+            content = self._parse_content(response_data)
+            content = self._transform_output(content, **parameters)
+            tool_calls = self._parse_tool_calls(response_data)
+            reasoning, signature = self._parse_reasoning(response_data)
+            kwargs: dict[str, Any] = {}
+            if reasoning is not None:
+                kwargs["reasoning"] = reasoning
+            if signature:
+                kwargs["signature"] = signature
+            output = self._output_class()(
+                content=content,
+                usage=self._get_usage(response_data),
+                finish_reason=self._get_finish_reason(response_data),
+                metadata=self._build_metadata(response_data),
+                tool_calls=tool_calls,
+                **kwargs,
+            )
+            span.set_attributes(telemetry.output_attributes(output))
+            return output
 
     def _parse_tool_calls(self, response_data: dict[str, Any]) -> list[ToolCall]:
         """Parse tool calls from response. Override in providers that support tools."""
@@ -281,13 +297,26 @@ def _stream(
         request_body = self._build_request(
             inputs, extra_body=extra_body, streaming=True, **parameters
         )
+        span = telemetry.tracer.start_span(
+            telemetry.span_name(self.modality, self.model),
+            attributes={
+                **telemetry.request_attributes(
+                    model=self.model,
+                    provider=self.provider,
+                    protocol=self.protocol,
+                    modality=self.modality,
+                ),
+                "gen_ai.request.stream": True,
+            },
+        )
         sse_iterator = self._make_stream_request(
             request_body,
             endpoint=endpoint,
             extra_headers=extra_headers,
             **parameters,
         )
         sse_iterator = enrich_stream_errors(sse_iterator, self._handle_error_response)
+        sse_iterator = telemetry.trace_stream(sse_iterator, span)
         return stream_class(
             sse_iterator,
             transform_output=self._transform_output,

src/celeste/telemetry.py

@@ -0,0 +1,170 @@
+"""OpenTelemetry GenAI telemetry for Celeste."""
+
+import time
+from collections.abc import AsyncIterator, Iterator
+from contextlib import contextmanager
+from typing import Any
+
+from celeste.core import Modality, Protocol, Provider
+from celeste.io import Output
+from celeste.models import Model
+
+_PROVIDER_NAME_MAP: dict[Provider, str] = {
+    Provider.OPENAI: "openai",
+    Provider.ANTHROPIC: "anthropic",
+    Provider.GOOGLE: "gcp.gemini",
+    Provider.MISTRAL: "mistral_ai",
+    Provider.XAI: "x_ai",
+    Provider.GROQ: "groq",
+    Provider.DEEPSEEK: "deepseek",
+    Provider.PERPLEXITY: "perplexity",
+    Provider.COHERE: "cohere",
+}
+
+_MODALITY_OPERATION: dict[Modality, str] = {
+    Modality.TEXT: "chat",
+    Modality.EMBEDDINGS: "embeddings",
+}
+
+
+class _NoOpSpan:
+    """Span stand-in used when ``opentelemetry-api`` is not installed."""
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        """Discard the attribute."""
+
+    def set_attributes(self, attributes: dict[str, Any]) -> None:
+        """Discard the attribute mapping."""
+
+    def set_status(self, *args: Any, **kwargs: Any) -> None:
+        """Discard the status update."""
+
+    def record_exception(self, *args: Any, **kwargs: Any) -> None:
+        """Discard the exception."""
+
+    def end(self) -> None:
+        """Discard the end signal."""
+
+
+class _NoOpTracer:
+    """Tracer stand-in used when ``opentelemetry-api`` is not installed."""
+
+    @contextmanager
+    def start_as_current_span(self, name: str, **kwargs: Any) -> Iterator[_NoOpSpan]:
+        """Yield a no-op span as a context manager."""
+        yield _NoOpSpan()
+
+    def start_span(self, name: str, **kwargs: Any) -> _NoOpSpan:
+        """Return a detached no-op span."""
+        return _NoOpSpan()
+
+
+@contextmanager
+def _noop_use_span(span: Any, end_on_exit: bool = False) -> Iterator[Any]:
+    """No-op stand-in for ``opentelemetry.trace.use_span``."""
+    yield span
+
+
+try:
+    from opentelemetry import trace as _otel_trace
+
+    tracer: Any = _otel_trace.get_tracer("celeste")
+    use_span: Any = _otel_trace.use_span
+except ImportError:
+    tracer = _NoOpTracer()
+    use_span = _noop_use_span
+
+
+def request_attributes(
+    *,
+    model: Model,
+    provider: Provider | None,
+    protocol: Protocol | None,
+    modality: Modality,
+) -> dict[str, Any]:
+    """Build GenAI request attributes for a span.
+
+    Off-spec providers/modalities get ``celeste.*`` keys so the data is
+    queryable without polluting the ``gen_ai.*`` namespace.
+
+    Args:
+        model: Resolved model.
+        provider: Resolved provider, if any.
+        protocol: Resolved protocol, if any.
+        modality: Operation modality.
+
+    Returns:
+        Attribute dict suitable for passing to ``start_as_current_span``.
+    """
+    attrs: dict[str, Any] = {"gen_ai.request.model": model.id}
+    provider_name = _PROVIDER_NAME_MAP.get(provider) if provider is not None else None
+    if provider_name is not None:
+        attrs["gen_ai.provider.name"] = provider_name
+    elif protocol is not None:
+        attrs["gen_ai.provider.name"] = protocol.value
+    operation = _MODALITY_OPERATION.get(modality)
+    if operation is not None:
+        attrs["gen_ai.operation.name"] = operation
+    else:
+        attrs["celeste.modality"] = modality.value
+    return attrs
+
+
+def span_name(modality: Modality, model: Model) -> str:
+    """Build the span name per GenAI semconv ``"{operation} {model}"``.
+
+    Falls back to ``"celeste.{modality} {model}"`` for modalities the spec
+    does not yet cover.
+    """
+    operation = _MODALITY_OPERATION.get(modality)
+    if operation is not None:
+        return f"{operation} {model.id}"
+    return f"celeste.{modality.value} {model.id}"
+
+
+def output_attributes(output: Output[Any]) -> dict[str, Any]:
+    """Extract GenAI response attributes from a typed Output."""
+    attrs: dict[str, Any] = {}
+    input_tokens = getattr(output.usage, "input_tokens", None)
+    output_tokens = getattr(output.usage, "output_tokens", None)
+    if isinstance(input_tokens, int):
+        attrs["gen_ai.usage.input_tokens"] = input_tokens
+    if isinstance(output_tokens, int):
+        attrs["gen_ai.usage.output_tokens"] = output_tokens
+    if output.finish_reason is not None and output.finish_reason.reason is not None:
+        attrs["gen_ai.response.finish_reasons"] = (output.finish_reason.reason,)
+    return attrs
+
+
+async def trace_stream(
+    sse_iterator: AsyncIterator[dict[str, Any]],
+    span: Any,
+) -> AsyncIterator[dict[str, Any]]:
+    """Wrap an SSE iterator to record TTFC and end the span at iteration's end.
+
+    OTel's ``use_span(span, end_on_exit=True)`` context manager makes the
+    span current, auto-records exceptions, sets ERROR status on exception,
+    and ends the span on exit (including ``GeneratorExit`` from consumer
+    abandonment). No try/except needed.
+    """
+    with use_span(span, end_on_exit=True):
+        started = time.monotonic()
+        seen_first = False
+        async for event in sse_iterator:
+            if not seen_first:
+                seen_first = True
+                span.set_attribute(
+                    "gen_ai.response.time_to_first_chunk",
+                    time.monotonic() - started,
+                )
+            yield event
+
+
+__all__ = [
+    "output_attributes",
+    "request_attributes",
+    "span_name",
+    "trace_stream",
+    "tracer",
+    "use_span",
+]