mirror of
https://github.com/Stirling-Tools/Stirling-PDF.git
synced 2026-09-03 05:10:16 +03:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3fa7ba9ebe |
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Mapping
|
||||
from contextvars import ContextVar
|
||||
@@ -29,6 +30,8 @@ from posthog.client import Client as PostHogClient
|
||||
from stirling.config import AppSettings
|
||||
from stirling.models import UserId
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Per-request user ID, set by middleware from the X-User-Id header.
|
||||
# When not set, PostHog generates a random ID and marks the event as personless.
|
||||
current_user_id: ContextVar[UserId | None] = ContextVar("current_user_id", default=None)
|
||||
@@ -127,17 +130,27 @@ class PostHogSpanProcessor(SpanProcessor):
|
||||
pass
|
||||
|
||||
def on_end(self, span: ReadableSpan) -> None:
|
||||
attrs = dict(span.attributes or {})
|
||||
if attrs.get(GEN_AI_OPERATION_NAME) != GenAiOperationNameValues.CHAT.value:
|
||||
return
|
||||
# OpenTelemetry calls this synchronously from ``Span.end()`` with no
|
||||
# exception handling of its own (see SynchronousMultiSpanProcessor),
|
||||
# and pydantic-ai ends the model span inside a ``with`` block. So any
|
||||
# exception raised here propagates straight into the model request and
|
||||
# breaks the agent run. A telemetry defect must never take down the
|
||||
# assistant or silently halt emission for every subsequent span, so we
|
||||
# translate and capture defensively and swallow (but log) any failure.
|
||||
try:
|
||||
attrs = dict(span.attributes or {})
|
||||
if attrs.get(GEN_AI_OPERATION_NAME) != GenAiOperationNameValues.CHAT.value:
|
||||
return
|
||||
|
||||
properties = self._build_generation_properties(span, attrs)
|
||||
self._maybe_emit_trace_event(span, attrs, properties)
|
||||
self._client.capture(
|
||||
distinct_id=current_user_id.get(),
|
||||
event="$ai_generation",
|
||||
properties=properties,
|
||||
)
|
||||
properties = self._build_generation_properties(span, attrs)
|
||||
self._maybe_emit_trace_event(span, attrs, properties)
|
||||
self._client.capture(
|
||||
distinct_id=current_user_id.get(),
|
||||
event="$ai_generation",
|
||||
properties=properties,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to emit $ai_generation for span; dropping this event")
|
||||
|
||||
def _build_generation_properties(self, span: ReadableSpan, attrs: Mapping[str, Any]) -> dict[str, object]:
|
||||
"""Build the $ai_generation event properties from span data."""
|
||||
@@ -231,6 +244,19 @@ class PostHogSpanProcessor(SpanProcessor):
|
||||
return True
|
||||
|
||||
|
||||
def _log_delivery_error(error: Exception, _batch: Any) -> None:
|
||||
"""PostHog client ``on_error`` callback.
|
||||
|
||||
The PostHog SDK delivers events from a background consumer thread and its
|
||||
``capture`` method is decorated ``@no_throw``, so without this hook a
|
||||
failing (or dead) delivery pipeline is completely invisible: events simply
|
||||
stop arriving with nothing in our logs. Surfacing the error lets a halt be
|
||||
noticed in minutes instead of sitting silently at zero. Kept lightweight so
|
||||
a transient upload error can't spam at capture volume.
|
||||
"""
|
||||
logger.warning("PostHog event delivery failed: %s", error)
|
||||
|
||||
|
||||
def setup_posthog_tracking(settings: AppSettings) -> TracerProvider | None:
|
||||
"""Configure OpenTelemetry with a PostHog span processor for LLM analytics.
|
||||
|
||||
@@ -238,11 +264,17 @@ def setup_posthog_tracking(settings: AppSettings) -> TracerProvider | None:
|
||||
or None when tracking is disabled.
|
||||
"""
|
||||
if not settings.posthog_enabled or not settings.posthog_api_key:
|
||||
logger.info("PostHog LLM tracking is disabled (posthog_enabled=%s)", settings.posthog_enabled)
|
||||
return None
|
||||
|
||||
client = PostHogClient(project_api_key=settings.posthog_api_key, host=settings.posthog_host)
|
||||
client = PostHogClient(
|
||||
project_api_key=settings.posthog_api_key,
|
||||
host=settings.posthog_host,
|
||||
on_error=_log_delivery_error,
|
||||
)
|
||||
processor = PostHogSpanProcessor(client)
|
||||
|
||||
provider = TracerProvider()
|
||||
provider.add_span_processor(processor)
|
||||
logger.info("PostHog LLM tracking enabled (host=%s)", settings.posthog_host)
|
||||
return provider
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
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
|
||||
Reference in New Issue
Block a user