Streaming

Author: loukratz-bv

qtype/interpreter/base/stream_emitter.py

@@ -35,6 +35,8 @@ async def process_message(self, message: FlowMessage):
 
 from typing import Any
 
+from opentelemetry import trace
+
 from qtype.interpreter.types import (
     ErrorEvent,
     ReasoningStreamDeltaEvent,
@@ -54,6 +56,24 @@ async def process_message(self, message: FlowMessage):
 from qtype.semantic.model import Step
 
 
+def get_current_telemetry_metadata() -> dict[str, Any]:
+    """
+    Get current OpenTelemetry span context as metadata dict.
+
+    Returns:
+        Dictionary with span_id and trace_id if span is recording,
+        empty dict otherwise
+    """
+    span = trace.get_current_span()
+    if span and span.is_recording():
+        ctx = span.get_span_context()
+        return {
+            "span_id": format(ctx.span_id, "016x"),
+            "trace_id": format(ctx.trace_id, "032x"),
+        }
+    return {}
+
+
 class TextStreamContext:
     """
     Async context manager for text streaming.
@@ -83,7 +103,11 @@ async def __aenter__(self) -> TextStreamContext:
         """Emit TextStreamStartEvent when entering context."""
         if self.on_stream_event:
             await self.on_stream_event(
-                TextStreamStartEvent(step=self.step, stream_id=self.stream_id)
+                TextStreamStartEvent(
+                    step=self.step,
+                    stream_id=self.stream_id,
+                    metadata=get_current_telemetry_metadata(),
+                )
             )
         return self
 
@@ -96,7 +120,11 @@ async def __aexit__(
         """Emit TextStreamEndEvent when exiting context."""
         if self.on_stream_event:
             await self.on_stream_event(
-                TextStreamEndEvent(step=self.step, stream_id=self.stream_id)
+                TextStreamEndEvent(
+                    step=self.step,
+                    stream_id=self.stream_id,
+                    metadata=get_current_telemetry_metadata(),
+                )
             )
         return False
 
@@ -113,6 +141,7 @@ async def delta(self, text: str) -> None:
                     step=self.step,
                     stream_id=self.stream_id,
                     delta=text,
+                    metadata=get_current_telemetry_metadata(),
                 )
             )
 
@@ -158,7 +187,9 @@ async def __aexit__(
         if self._started and self.on_stream_event:
             await self.on_stream_event(
                 ReasoningStreamEndEvent(
-                    step=self.step, stream_id=self.stream_id
+                    step=self.step,
+                    stream_id=self.stream_id,
+                    metadata=get_current_telemetry_metadata(),
                 )
             )
         return False
@@ -173,11 +204,15 @@ async def delta(self, text: str) -> None:
             text: The incremental reasoning content to append to the stream
         """
         if self.on_stream_event:
+            metadata = get_current_telemetry_metadata()
+
             # Emit start event on first delta
             if not self._started:
                 await self.on_stream_event(
                     ReasoningStreamStartEvent(
-                        step=self.step, stream_id=self.stream_id
+                        step=self.step,
+                        stream_id=self.stream_id,
+                        metadata=metadata,
                     )
                 )
                 self._started = True
@@ -187,6 +222,7 @@ async def delta(self, text: str) -> None:
                     step=self.step,
                     stream_id=self.stream_id,
                     delta=text,
+                    metadata=metadata,
                 )
             )
 
@@ -218,7 +254,11 @@ def __init__(
     async def __aenter__(self) -> StepBoundaryContext:
         """Emit StepStartEvent when entering context."""
         if self.on_stream_event:
-            await self.on_stream_event(StepStartEvent(step=self.step))
+            await self.on_stream_event(
+                StepStartEvent(
+                    step=self.step, metadata=get_current_telemetry_metadata()
+                )
+            )
         return self
 
     async def __aexit__(
@@ -229,7 +269,11 @@ async def __aexit__(
     ) -> bool:
         """Emit StepEndEvent when exiting context."""
         if self.on_stream_event:
-            await self.on_stream_event(StepEndEvent(step=self.step))
+            await self.on_stream_event(
+                StepEndEvent(
+                    step=self.step, metadata=get_current_telemetry_metadata()
+                )
+            )
         return False
 
 
@@ -277,6 +321,7 @@ async def __aenter__(self) -> ToolExecutionContext:
                     tool_call_id=self.tool_call_id,
                     tool_name=self.tool_name,
                     tool_input=self.tool_input,
+                    metadata=get_current_telemetry_metadata(),
                 )
             )
         return self
@@ -299,6 +344,7 @@ async def __aexit__(
                     step=self.step,
                     tool_call_id=self.tool_call_id,
                     error_message=str(exc_val),
+                    metadata=get_current_telemetry_metadata(),
                 )
             )
             self._completed = True
@@ -319,6 +365,7 @@ async def complete(self, output: Any) -> None:
                     step=self.step,
                     tool_call_id=self.tool_call_id,
                     tool_output=output,
+                    metadata=get_current_telemetry_metadata(),
                 )
             )
         self._completed = True
@@ -338,6 +385,7 @@ async def error(self, error_message: str) -> None:
                     step=self.step,
                     tool_call_id=self.tool_call_id,
                     error_message=error_message,
+                    metadata=get_current_telemetry_metadata(),
                 )
             )
         self._completed = True
@@ -387,6 +435,18 @@ def __init__(
         self.step = step
         self.on_stream_event = on_stream_event
 
+    def _get_span_context(self) -> tuple[str | None, str | None]:
+        """
+        Get current OpenTelemetry span context.
+
+        Returns:
+            Tuple of (span_id, trace_id) as hex strings, or (None, None)
+
+        Deprecated: Use get_current_telemetry_metadata() instead.
+        """
+        metadata = get_current_telemetry_metadata()
+        return metadata.get("span_id"), metadata.get("trace_id")
+
     def text_stream(self, stream_id: str) -> TextStreamContext:
         """
         Create a context manager for text streaming.
@@ -459,7 +519,11 @@ async def status(self, message: str) -> None:
         """
         if self.on_stream_event:
             await self.on_stream_event(
-                StatusEvent(step=self.step, message=message)
+                StatusEvent(
+                    step=self.step,
+                    message=message,
+                    metadata=get_current_telemetry_metadata(),
+                )
             )
 
     async def error(self, error_message: str) -> None:
@@ -471,5 +535,9 @@ async def error(self, error_message: str) -> None:
         """
         if self.on_stream_event:
             await self.on_stream_event(
-                ErrorEvent(step=self.step, error_message=error_message)
+                ErrorEvent(
+                    step=self.step,
+                    error_message=error_message,
+                    metadata=get_current_telemetry_metadata(),
+                )
             )

qtype/interpreter/stream/chat/converter.py

@@ -377,11 +377,22 @@ async def stream_events():
     # Create converter for stateful event-to-chunk conversion
     converter = StreamEventConverter()
 
+    # Track telemetry metadata from events
+    telemetry_metadata: dict[str, str] = {}
+
     # Process events and convert to chunks
     async for event in event_stream:
         if event is None:
             break  # End of stream
 
+        # Extract telemetry metadata from first event that has it
+        if not telemetry_metadata and event.metadata:
+            if "span_id" in event.metadata and "trace_id" in event.metadata:
+                telemetry_metadata = {
+                    "span_id": event.metadata["span_id"],
+                    "trace_id": event.metadata["trace_id"],
+                }
+
         # Convert event to chunks and yield as SSE
         for chunk in converter.convert(event):
             yield (
@@ -390,8 +401,15 @@ async def stream_events():
                 f"\n\n"
             )
 
-    # End message stream with optional metadata
-    finish_chunk = FinishChunk(messageMetadata=output_metadata)  # type: ignore[arg-type]
+    # Merge telemetry metadata with output_metadata for FinishChunk
+    final_metadata = {**telemetry_metadata}
+    if output_metadata:
+        final_metadata.update(output_metadata)
+
+    # End message stream with metadata (includes telemetry)
+    finish_chunk = FinishChunk(
+        messageMetadata=final_metadata if final_metadata else None
+    )
     yield (
         f"data: "
         f"{finish_chunk.model_dump_json(by_alias=True, exclude_none=True)}"

qtype/interpreter/types.py

@@ -38,7 +38,23 @@ def __bool__(self) -> bool:
 # and can be converted to Vercel UI chunks for frontend display
 
 
-class TextStreamStartEvent(BaseModel):
+class BaseStreamEvent(BaseModel):
+    """
+    Base class for all stream events.
+
+    Provides common metadata field for telemetry and other contextual data.
+    The metadata dict typically contains:
+    - span_id: OpenTelemetry span ID (16 hex chars)
+    - trace_id: OpenTelemetry trace ID (32 hex chars)
+    """
+
+    metadata: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Metadata for telemetry and context tracking",
+    )
+
+
+class TextStreamStartEvent(BaseStreamEvent):
     """Signals the start of incremental text streaming.
 
     Use this when beginning to stream LLM-generated content or other
@@ -55,7 +71,7 @@ class TextStreamStartEvent(BaseModel):
     )
 
 
-class TextStreamDeltaEvent(BaseModel):
+class TextStreamDeltaEvent(BaseStreamEvent):
     """Carries an incremental chunk of text content.
 
     Use this for streaming LLM responses or other incremental text.
@@ -72,7 +88,7 @@ class TextStreamDeltaEvent(BaseModel):
     delta: str = Field(description="Incremental text content to append")
 
 
-class TextStreamEndEvent(BaseModel):
+class TextStreamEndEvent(BaseStreamEvent):
     """Signals the completion of incremental text streaming.
 
     Use this to mark the end of a text stream. After this event,
@@ -88,7 +104,7 @@ class TextStreamEndEvent(BaseModel):
     )
 
 
-class ReasoningStreamStartEvent(BaseModel):
+class ReasoningStreamStartEvent(BaseStreamEvent):
     """Signals the start of incremental reasoning streaming.
 
     Use this when an agent begins outputting reasoning/thinking steps.
@@ -105,7 +121,7 @@ class ReasoningStreamStartEvent(BaseModel):
     )
 
 
-class ReasoningStreamDeltaEvent(BaseModel):
+class ReasoningStreamDeltaEvent(BaseStreamEvent):
     """Carries an incremental chunk of reasoning content.
 
     Use this for streaming agent reasoning/thinking steps.
@@ -122,7 +138,7 @@ class ReasoningStreamDeltaEvent(BaseModel):
     delta: str = Field(description="Incremental reasoning content to append")
 
 
-class ReasoningStreamEndEvent(BaseModel):
+class ReasoningStreamEndEvent(BaseStreamEvent):
     """Signals the completion of incremental reasoning streaming.
 
     Use this to mark the end of a reasoning stream. After this event,
@@ -138,7 +154,7 @@ class ReasoningStreamEndEvent(BaseModel):
     )
 
 
-class StatusEvent(BaseModel):
+class StatusEvent(BaseStreamEvent):
     """Reports a complete status message from a step.
 
     Use this for non-streaming status updates like:
@@ -155,7 +171,7 @@ class StatusEvent(BaseModel):
     message: str = Field(description="Complete status message to display")
 
 
-class StepStartEvent(BaseModel):
+class StepStartEvent(BaseStreamEvent):
     """Marks the beginning of a logical step boundary.
 
     Use this to group related events together visually in the UI.
@@ -168,7 +184,7 @@ class StepStartEvent(BaseModel):
     step: Step
 
 
-class StepEndEvent(BaseModel):
+class StepEndEvent(BaseStreamEvent):
     """Marks the end of a logical step boundary.
 
     Use this to close a step boundary opened by StepStartEvent.
@@ -180,7 +196,7 @@ class StepEndEvent(BaseModel):
     step: Step
 
 
-class ToolExecutionStartEvent(BaseModel):
+class ToolExecutionStartEvent(BaseStreamEvent):
     """Signals the start of tool execution.
 
     Use this when a tool is about to be invoked, either by an LLM
@@ -198,7 +214,7 @@ class ToolExecutionStartEvent(BaseModel):
     )
 
 
-class ToolExecutionEndEvent(BaseModel):
+class ToolExecutionEndEvent(BaseStreamEvent):
     """Signals the completion of tool execution.
 
     Use this when a tool has finished executing successfully.
@@ -214,7 +230,7 @@ class ToolExecutionEndEvent(BaseModel):
     tool_output: Any = Field(description="Output returned by the tool")
 
 
-class ToolExecutionErrorEvent(BaseModel):
+class ToolExecutionErrorEvent(BaseStreamEvent):
     """Signals that tool execution failed.
 
     Use this when a tool encounters an error during execution.
@@ -230,7 +246,7 @@ class ToolExecutionErrorEvent(BaseModel):
     error_message: str = Field(description="Description of the error")
 
 
-class ErrorEvent(BaseModel):
+class ErrorEvent(BaseStreamEvent):
     """Signals a general error occurred during step execution.
 
     Use this for errors that aren't specific to tool execution.