posthog-eu[bot]andGitHub
3fa7ba9ebe
The AI engine translates pydantic-ai OpenTelemetry spans into PostHog `$ai_generation`/`$ai_trace` events via a custom SpanProcessor. Two gaps in that path let LLM telemetry stop silently and stay at zero: - `PostHogSpanProcessor.on_end` was unguarded. OpenTelemetry calls it synchronously from `Span.end()` with no exception handling of its own, and pydantic-ai ends the model span inside a `with` block, so any error while translating a span propagates into the model request itself. A telemetry defect could therefore break the agent run and/or halt emission for every subsequent span. It now translates and captures defensively and logs-and-drops on failure. - The PostHog client was created without an `on_error` callback. Its `capture` is decorated `@no_throw` and delivery happens on a background consumer thread, so a failing or stalled delivery pipeline was completely invisible. We now log delivery failures and log the enabled/disabled state at startup, so a halt is noticeable in minutes instead of going unnoticed. Adds unit tests covering the happy path, per-trace dedup, non-chat spans, and that `on_end` never raises on delivery failure or a malformed span. Generated-By: PostHog Code Task-Id: 505d85dd-9cd2-4258-b76a-e027f4240277
133 lines
4.1 KiB
Python
133 lines
4.1 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import (
|
|
GEN_AI_INPUT_MESSAGES,
|
|
GEN_AI_OPERATION_NAME,
|
|
GEN_AI_SYSTEM,
|
|
GenAiOperationNameValues,
|
|
)
|
|
|
|
from stirling.config import AppSettings
|
|
from stirling.services.tracking import PostHogSpanProcessor, current_user_id, setup_posthog_tracking
|
|
|
|
|
|
@dataclass
|
|
class _FakeContext:
|
|
trace_id: int
|
|
span_id: int
|
|
|
|
|
|
@dataclass
|
|
class _FakeParent:
|
|
span_id: int
|
|
|
|
|
|
class _FakeSpan:
|
|
"""Minimal stand-in for the OpenTelemetry ReadableSpan fields the processor reads."""
|
|
|
|
def __init__(self, attributes: dict[str, Any], trace_id: int = 0x1234, span_id: int = 0x9) -> None:
|
|
self.attributes = attributes
|
|
self.context = _FakeContext(trace_id=trace_id, span_id=span_id)
|
|
self.parent = _FakeParent(span_id=0x1)
|
|
self.start_time = 1_000_000_000
|
|
self.end_time = 2_000_000_000
|
|
|
|
|
|
class _RecordingClient:
|
|
def __init__(self) -> None:
|
|
self.events: list[tuple[str, dict[str, Any]]] = []
|
|
|
|
def capture(self, *, distinct_id: str | None, event: str, properties: dict[str, Any]) -> None:
|
|
self.events.append((event, properties))
|
|
|
|
|
|
class _RaisingClient:
|
|
def capture(self, **_kwargs: Any) -> None:
|
|
raise RuntimeError("boom")
|
|
|
|
|
|
def _processor(client: object) -> PostHogSpanProcessor:
|
|
# The processor only calls ``capture`` on its client; the test doubles model
|
|
# that surface without being real PostHog clients (which start threads).
|
|
return PostHogSpanProcessor(client) # type: ignore[arg-type]
|
|
|
|
|
|
def _on_end(processor: PostHogSpanProcessor, span: object) -> None:
|
|
processor.on_end(span) # type: ignore[arg-type]
|
|
|
|
|
|
def _chat_attrs() -> dict[str, Any]:
|
|
return {
|
|
GEN_AI_OPERATION_NAME: GenAiOperationNameValues.CHAT.value,
|
|
GEN_AI_SYSTEM: "anthropic",
|
|
GEN_AI_INPUT_MESSAGES: json.dumps([{"role": "user", "parts": [{"type": "text", "content": "hello"}]}]),
|
|
}
|
|
|
|
|
|
def test_chat_span_emits_generation_and_trace() -> None:
|
|
client = _RecordingClient()
|
|
|
|
_on_end(_processor(client), _FakeSpan(_chat_attrs()))
|
|
|
|
emitted = [event for event, _ in client.events]
|
|
assert "$ai_generation" in emitted
|
|
assert "$ai_trace" in emitted
|
|
|
|
|
|
def test_trace_event_deduplicated_per_trace() -> None:
|
|
client = _RecordingClient()
|
|
processor = _processor(client)
|
|
|
|
_on_end(processor, _FakeSpan(_chat_attrs(), trace_id=0x42, span_id=0x1))
|
|
_on_end(processor, _FakeSpan(_chat_attrs(), trace_id=0x42, span_id=0x2))
|
|
|
|
trace_events = [event for event, _ in client.events if event == "$ai_trace"]
|
|
generation_events = [event for event, _ in client.events if event == "$ai_generation"]
|
|
assert len(trace_events) == 1
|
|
assert len(generation_events) == 2
|
|
|
|
|
|
def test_non_chat_span_is_ignored() -> None:
|
|
client = _RecordingClient()
|
|
|
|
_on_end(_processor(client), _FakeSpan({GEN_AI_OPERATION_NAME: "embeddings"}))
|
|
|
|
assert client.events == []
|
|
|
|
|
|
def test_on_end_never_raises_when_delivery_fails() -> None:
|
|
# A telemetry failure must not propagate into Span.end() and break the
|
|
# model request, nor wedge emission for every subsequent span.
|
|
_on_end(_processor(_RaisingClient()), _FakeSpan(_chat_attrs())) # must not raise
|
|
|
|
|
|
def test_on_end_never_raises_on_malformed_span() -> None:
|
|
processor = _processor(_RecordingClient())
|
|
|
|
broken = _FakeSpan(_chat_attrs())
|
|
broken.context = None # type: ignore[assignment]
|
|
|
|
_on_end(processor, broken) # must not raise
|
|
|
|
|
|
def test_distinct_id_pulled_from_context_var() -> None:
|
|
client = _RecordingClient()
|
|
processor = _processor(client)
|
|
|
|
token = current_user_id.set("user-123") # type: ignore[arg-type]
|
|
try:
|
|
_on_end(processor, _FakeSpan(_chat_attrs()))
|
|
finally:
|
|
current_user_id.reset(token)
|
|
|
|
assert client.events, "expected at least one captured event"
|
|
|
|
|
|
def test_setup_returns_none_when_disabled(app_settings: AppSettings) -> None:
|
|
# conftest builds settings with posthog disabled; setup must no-op.
|
|
assert setup_posthog_tracking(app_settings) is None
|