diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 36b7abae..651f3d13 100644 --- a/.github/workflows/ci.yml +++ b/.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: diff --git a/pyproject.toml b/pyproject.toml index f82578e4..f12b87dd 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/celeste/client.py b/src/celeste/client.py index a8e9f49b..0f8f5af9 100644 --- a/src/celeste/client.py +++ b/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,6 +297,18 @@ 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, @@ -288,6 +316,7 @@ def _stream( **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, diff --git a/src/celeste/telemetry.py b/src/celeste/telemetry.py new file mode 100644 index 00000000..8b60abc9 --- /dev/null +++ b/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", +]