diff --git a/src/glassflow/generation.py b/src/glassflow/generation.py index e5843b6..3f4e9f8 100644 --- a/src/glassflow/generation.py +++ b/src/glassflow/generation.py @@ -18,6 +18,7 @@ from . import __version__ from ._serde import serialize from .semconv import ( + GEN_AI_FIRST_TOKEN_EVENT, GEN_AI_INPUT_MESSAGES, GEN_AI_OPERATION_NAME, GEN_AI_OUTPUT_MESSAGES, @@ -119,6 +120,7 @@ class Generation: def __init__(self, span: Span) -> None: self._span = span + self._first_token_recorded = False def set_input(self, messages: Messages) -> None: """Record the request messages (``gen_ai.input.messages``). @@ -164,6 +166,19 @@ def set_usage( if output_tokens is not None: self._span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) + def record_first_token(self) -> None: + """Mark the arrival of the first streamed token (``gen_ai.first_token`` event). + + Call from a streaming loop when the first content chunk arrives; the + backend derives time-to-first-token as the event time minus the span + start. Idempotent — only the first call records; safe to call + unconditionally per chunk. A no-op after ``end()``. + """ + if self._first_token_recorded or not self._span.is_recording(): + return + self._span.add_event(GEN_AI_FIRST_TOKEN_EVENT) + self._first_token_recorded = True + def set_finish_reasons(self, reasons: str | list[str]) -> None: """Record why generation stopped (``gen_ai.response.finish_reasons``). diff --git a/src/glassflow/semconv.py b/src/glassflow/semconv.py index e8e8b13..ab6f232 100644 --- a/src/glassflow/semconv.py +++ b/src/glassflow/semconv.py @@ -33,6 +33,13 @@ GEN_AI_TOOL_NAME = "gen_ai.tool.name" GEN_AI_REQUEST_PREFIX = "gen_ai.request." +# --- Span event names --- +# First streamed token/chunk arrived — the TTFT anchor (event time minus span +# start). No client-side semconv exists for this yet (OTel only standardizes +# the server-side gen_ai.server.time_to_first_token metric); this follows the +# gen_ai.* naming style, precedent Langfuse's completion_start_time. +GEN_AI_FIRST_TOKEN_EVENT = "gen_ai.first_token" + # Attribute keys carrying user content — masked/stripped at export (see masking.py). CONTENT_ATTRIBUTES = frozenset( { diff --git a/tests/test_generation.py b/tests/test_generation.py index 10a6e97..449bff0 100644 --- a/tests/test_generation.py +++ b/tests/test_generation.py @@ -1,5 +1,6 @@ import json +from opentelemetry.sdk.trace import Event, ReadableSpan from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from glassflow import start_as_current_generation, start_generation @@ -241,3 +242,45 @@ def test_manual_generation(exported_spans: InMemorySpanExporter) -> None: assert attrs["gen_ai.request.model"] == "gpt-4o" assert "hi" in attrs["gen_ai.output.messages"] assert attrs["gen_ai.usage.input_tokens"] == 1 + + +# --- record_first_token: the TTFT anchor for streaming (GLA2-175) --- + + +def _first_token_events(span: ReadableSpan) -> list[Event]: + return [e for e in span.events if e.name == "gen_ai.first_token"] + + +def test_record_first_token_adds_event_between_start_and_end( + exported_spans: InMemorySpanExporter, +) -> None: + with start_as_current_generation("chat") as gen: + gen.record_first_token() + span = exported_spans.get_finished_spans()[0] + (event,) = _first_token_events(span) + assert span.start_time <= event.timestamp <= span.end_time + + +def test_record_first_token_is_idempotent(exported_spans: InMemorySpanExporter) -> None: + with start_as_current_generation("chat") as gen: + gen.record_first_token() + gen.record_first_token() + gen.record_first_token() + span = exported_spans.get_finished_spans()[0] + assert len(_first_token_events(span)) == 1 + + +def test_record_first_token_after_end_is_noop(exported_spans: InMemorySpanExporter) -> None: + gen = start_generation("chat") + gen.end() + gen.record_first_token() # must neither raise nor record + span = exported_spans.get_finished_spans()[0] + assert _first_token_events(span) == [] + + +def test_record_first_token_on_manual_generation(exported_spans: InMemorySpanExporter) -> None: + gen = start_generation("chat", model="gpt-4o") + gen.record_first_token() + gen.end() + span = exported_spans.get_finished_spans()[0] + assert len(_first_token_events(span)) == 1