From 943ecff11d519fc4bf56b018273c1c21fa310726 Mon Sep 17 00:00:00 2001 From: sayonfortify Date: Fri, 19 Jun 2026 16:36:08 +0530 Subject: [PATCH] feat(langchain): emit streaming latency span attrs --- .../langchain/callback_handler.py | 47 +++++++- .../instrumentation/langchain/span_utils.py | 1 + .../test_streaming_latency_span_attrs.py | 113 ++++++++++++++++++ 3 files changed, 160 insertions(+), 1 deletion(-) create mode 100644 packages/opentelemetry-instrumentation-langchain/tests/test_streaming_latency_span_attrs.py diff --git a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py index 381c278780..50f069a01c 100644 --- a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py +++ b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py @@ -74,6 +74,11 @@ from opentelemetry.trace.status import Status, StatusCode from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE +FR_STREAMING_TIME_TO_FIRST_TOKEN_MS = ( + "fortifyroot.llm.streaming.time_to_first_token_ms" +) +FR_STREAMING_TIME_TO_GENERATE_MS = "fortifyroot.llm.streaming.time_to_generate_ms" + def _extract_class_name_from_serialized(serialized: Optional[dict[str, Any]]) -> str: """ @@ -596,6 +601,38 @@ def on_llm_start( else: set_llm_request(span, serialized, prompts, kwargs, self.spans[run_id]) + @dont_throw + def on_llm_new_token( + self, + token: str, + *, + run_id: UUID, + parent_run_id: Union[UUID, None] = None, + chunk: Union[GenerationChunk, ChatGenerationChunk, None] = None, + **kwargs: Any, + ) -> Any: + """Run when a streaming LLM emits a new token.""" + if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY): + return + if not token and chunk is None: + return + if run_id not in self.spans: + return + + span_holder = self.spans[run_id] + if span_holder.streaming_first_token_time is not None: + return + + first_token_time = time.time() + span_holder.streaming_first_token_time = first_token_time + span = span_holder.span + if span.is_recording(): + _set_span_attribute( + span, + FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, + round((first_token_time - span_holder.start_time) * 1000), + ) + @dont_throw def on_llm_end( self, @@ -698,7 +735,15 @@ def on_llm_end( set_chat_response(span, response) # Record duration before ending span - duration = time.time() - self.spans[run_id].start_time + end_time = time.time() + streaming_first_token_time = self.spans[run_id].streaming_first_token_time + if streaming_first_token_time is not None and span.is_recording(): + _set_span_attribute( + span, + FR_STREAMING_TIME_TO_GENERATE_MS, + round((end_time - streaming_first_token_time) * 1000), + ) + duration = end_time - self.spans[run_id].start_time vendor = span.attributes.get(GenAIAttributes.GEN_AI_SYSTEM, "Langchain") self.duration_histogram.record( duration, diff --git a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/span_utils.py b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/span_utils.py index 4198c974dd..04c0809dee 100644 --- a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/span_utils.py +++ b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/span_utils.py @@ -40,6 +40,7 @@ class SpanHolder: entity_path: str start_time: float = field(default_factory=time.time) request_model: Optional[str] = None + streaming_first_token_time: Optional[float] = None # ST-10 review-round-2 fix (2026-05-11): every context_api.attach() # performed for this span — span-context, suppression, metadata # association_properties — appended here in attach order. ``_end_span`` diff --git a/packages/opentelemetry-instrumentation-langchain/tests/test_streaming_latency_span_attrs.py b/packages/opentelemetry-instrumentation-langchain/tests/test_streaming_latency_span_attrs.py new file mode 100644 index 0000000000..39e80b713c --- /dev/null +++ b/packages/opentelemetry-instrumentation-langchain/tests/test_streaming_latency_span_attrs.py @@ -0,0 +1,113 @@ +from types import SimpleNamespace +from uuid import uuid4 + +from langchain_core.outputs import Generation, LLMResult +from opentelemetry import trace +from opentelemetry.instrumentation.langchain import callback_handler +from opentelemetry.instrumentation.langchain.callback_handler import ( + FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, + FR_STREAMING_TIME_TO_GENERATE_MS, + TraceloopCallbackHandler, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + + +class _StubHistogram: + def __init__(self): + self.records = [] + + def record(self, *args, **kwargs): + self.records.append((args, kwargs)) + + +def _install_noop_fortifyroot(monkeypatch): + monkeypatch.setitem( + __import__("sys").modules, + "opentelemetry.instrumentation.fortifyroot", + SimpleNamespace(emit_deferred_findings=lambda span: None), + ) + + +def _handler_and_exporter(): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + handler = TraceloopCallbackHandler( + provider.get_tracer("test"), + _StubHistogram(), + _StubHistogram(), + ) + return handler, exporter + + +def _llm_result(): + return LLMResult( + generations=[[Generation(text="hello")]], + llm_output={ + "model_name": "gpt-4o-mini", + "token_usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + ) + + +def test_streaming_callbacks_set_ttft_and_sttg_span_attrs(monkeypatch): + _install_noop_fortifyroot(monkeypatch) + handler, exporter = _handler_and_exporter() + run_id = uuid4() + + handler.on_llm_start( + serialized={"id": ["langchain_openai", "llms", "base", "OpenAI"]}, + prompts=["hello"], + run_id=run_id, + invocation_params={"model": "gpt-4o-mini"}, + ) + handler.spans[run_id].start_time = 100.0 + + times = [101.25, 103.0] + + def fake_time(): + if times: + return times.pop(0) + return 103.0 + + monkeypatch.setattr(callback_handler.time, "time", fake_time) + + handler.on_llm_new_token("hello", run_id=run_id) + handler.on_llm_end(_llm_result(), run_id=run_id) + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs[FR_STREAMING_TIME_TO_FIRST_TOKEN_MS] == 1250 + assert attrs[FR_STREAMING_TIME_TO_GENERATE_MS] == 1750 + + +def test_non_streaming_llm_end_does_not_set_streaming_latency_attrs(monkeypatch): + _install_noop_fortifyroot(monkeypatch) + handler, exporter = _handler_and_exporter() + run_id = uuid4() + + handler.on_llm_start( + serialized={"id": ["langchain_openai", "llms", "base", "OpenAI"]}, + prompts=["hello"], + run_id=run_id, + invocation_params={"model": "gpt-4o-mini"}, + ) + handler.spans[run_id].start_time = 100.0 + monkeypatch.setattr(callback_handler.time, "time", lambda: 103.0) + + handler.on_llm_end(_llm_result(), run_id=run_id) + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + attrs = spans[0].attributes + assert FR_STREAMING_TIME_TO_FIRST_TOKEN_MS not in attrs + assert FR_STREAMING_TIME_TO_GENERATE_MS not in attrs