From 01ac402decacc4712919399556064740f7478a0b Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Fri, 23 Oct 2026 23:15:39 +0800 Subject: [PATCH 1/2] Separate incremental stream accumulation from provider parsing and spans --- src/stream.rs | 1591 +-------------------------------- src/stream/accumulator.rs | 157 ++++ src/stream/instrumentation.rs | 162 ++++ src/stream/provider.rs | 227 +++++ src/stream/tests.rs | 363 ++++++++ src/stream/types.rs | 598 +++++++++++++ src/stream/update.rs | 27 + tests/stream_lifecycle.rs | 135 +++ 8 files changed, 1698 insertions(+), 1562 deletions(-) create mode 100644 src/stream/accumulator.rs create mode 100644 src/stream/instrumentation.rs create mode 100644 src/stream/provider.rs create mode 100644 src/stream/tests.rs create mode 100644 src/stream/types.rs create mode 100644 src/stream/update.rs create mode 100644 tests/stream_lifecycle.rs diff --git a/src/stream.rs b/src/stream.rs index 26a9bf1..f8bd6cd 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -1,1587 +1,54 @@ -//! Stream aggregation for LLM streaming responses. -//! -//! This module provides `BraintrustStream`, a wrapper that aggregates streaming -//! chunks into a final response value, following the JS/Python SDK pattern. -//! -//! It also provides `wrap_stream_with_span` for wrapping streams with span logging. - -use std::collections::HashMap; -use std::fmt; -use std::pin::Pin; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::sync::Arc; -use std::task::{Context, Poll}; -use std::time::{Duration, Instant}; - -/// Accumulates partial tool call information from streaming chunks. -/// -/// In OpenAI streaming, each chunk may contain partial tool call data. -/// This struct accumulates the id, type, function name, and raw arguments -/// across multiple chunks, then parses the complete arguments string -/// as JSON when aggregation is complete. -#[derive(Clone, Debug, Default)] -struct StreamToolCallAccumulator { - id: String, - call_type: String, - function_name: String, - raw_arguments: String, -} - -impl StreamToolCallAccumulator { - /// Creates a `ToolCall` from the accumulated data. - /// - /// The raw arguments string is parsed as JSON to extract the function - /// name and arguments, falling back to the raw string if parsing fails. - pub fn build(self) -> crate::stream::ToolCall { - let args: serde_json::Value = serde_json::from_str(&self.raw_arguments).unwrap_or_default(); - - // Extract function name from JSON object, falling back to stored name - let function_name = if let Some(obj) = args.as_object() { - obj.get("name") - .and_then(|v| v.as_str()) - .unwrap_or(&self.function_name) - .to_string() - } else { - self.function_name.clone() - }; - - // Extract arguments string from JSON object, falling back to raw string - let arguments = if let Some(obj) = args.as_object() { - obj.get("arguments") - .and_then(|v| v.as_str()) - .unwrap_or(&self.raw_arguments) - .to_string() - } else { - self.raw_arguments.clone() - }; - - let function = crate::stream::FunctionCall::new(function_name, arguments); - crate::stream::ToolCall::new(self.id, self.call_type, function) - } -} +//! Incremental aggregation of OpenAI-compatible streaming responses. +mod accumulator; +mod instrumentation; +mod provider; +#[cfg(test)] +mod tests; +mod types; +mod update; +use accumulator::Accumulator; use anyhow::Result; -use futures::Stream; -use serde::{Deserialize, Serialize}; +pub use instrumentation::wrap_stream_with_span; use serde_json::Value; -use tokio::sync::Mutex; - -use crate::span::{SpanHandle, SpanLog}; -use crate::types::{usage_metrics_to_map, UsageMetrics}; - -// ============================================================================= -// Error types for stream output builders (empty for now, extensible later) -// ============================================================================= - -/// Error type for ToolCall builder validation. -#[derive(Debug, Clone, PartialEq)] -#[non_exhaustive] -pub enum ToolCallBuilderError {} - -impl fmt::Display for ToolCallBuilderError { - fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { - match *self {} - } -} - -impl std::error::Error for ToolCallBuilderError {} - -/// Error type for FunctionCall builder validation. -#[derive(Debug, Clone, PartialEq)] -#[non_exhaustive] -pub enum FunctionCallBuilderError {} - -impl fmt::Display for FunctionCallBuilderError { - fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { - match *self {} - } -} - -impl std::error::Error for FunctionCallBuilderError {} - -/// Error type for ChatMessage builder validation. -#[derive(Debug, Clone, PartialEq)] -#[non_exhaustive] -pub enum ChatMessageBuilderError {} - -impl fmt::Display for ChatMessageBuilderError { - fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { - match *self {} - } -} - -impl std::error::Error for ChatMessageBuilderError {} - -/// Error type for OutputChoice builder validation. -#[derive(Debug, Clone, PartialEq)] -#[non_exhaustive] -pub enum OutputChoiceBuilderError {} - -impl fmt::Display for OutputChoiceBuilderError { - fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { - match *self {} - } -} - -impl std::error::Error for OutputChoiceBuilderError {} - -/// Error type for StreamMetadata builder validation. -#[derive(Debug, Clone, PartialEq)] -#[non_exhaustive] -pub enum StreamMetadataBuilderError {} - -impl fmt::Display for StreamMetadataBuilderError { - fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { - match *self {} - } -} - -impl std::error::Error for StreamMetadataBuilderError {} - -/// Error type for FinalizedStream builder validation. -#[derive(Debug, Clone, PartialEq)] -#[non_exhaustive] -pub enum FinalizedStreamBuilderError {} - -impl fmt::Display for FinalizedStreamBuilderError { - fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { - match *self {} - } -} - -impl std::error::Error for FinalizedStreamBuilderError {} - -/// Number of chunks between periodic flushes during streaming. -const STREAMING_FLUSH_CHUNK_INTERVAL: usize = 10; - -/// Maximum time between periodic flushes during streaming. -const STREAMING_FLUSH_TIME_INTERVAL: Duration = Duration::from_secs(1); - -/// A tool call in a chat message. -#[derive(Clone, Debug, Default, Serialize)] -#[non_exhaustive] -pub struct ToolCall { - id: String, - #[serde(rename = "type")] - call_type: String, // Always "function" - function: FunctionCall, -} - -impl ToolCall { - /// Create a new ToolCallBuilder. - pub fn builder() -> ToolCallBuilder { - ToolCallBuilder::new() - } - - /// Create a new ToolCall (internal use). - #[allow(dead_code)] - pub(crate) fn new(id: String, call_type: String, function: FunctionCall) -> Self { - Self { - id, - call_type, - function, - } - } - - /// Get the tool call ID. - pub fn id(&self) -> &str { - &self.id - } - - /// Get the tool call type (typically "function"). - pub fn call_type(&self) -> &str { - &self.call_type - } - - /// Get the function details. - pub fn function(&self) -> &FunctionCall { - &self.function - } -} - -/// Builder for ToolCall. -#[derive(Clone, Default)] -pub struct ToolCallBuilder { - id: String, - call_type: String, - function: FunctionCall, -} - -impl ToolCallBuilder { - /// Create a new ToolCallBuilder. - pub fn new() -> Self { - Self { - call_type: "function".to_string(), - ..Default::default() - } - } - - /// Set the tool call ID. - pub fn id(mut self, id: impl Into) -> Self { - self.id = id.into(); - self - } - - /// Set the call type. - pub fn call_type(mut self, call_type: impl Into) -> Self { - self.call_type = call_type.into(); - self - } - - /// Set the function details. - pub fn function(mut self, function: FunctionCall) -> Self { - self.function = function; - self - } - - /// Build the ToolCall. - pub fn build(self) -> std::result::Result { - Ok(ToolCall { - id: self.id, - call_type: self.call_type, - function: self.function, - }) - } -} - -/// Function details in a tool call. -#[derive(Clone, Debug, Default, Serialize)] -#[non_exhaustive] -pub struct FunctionCall { - name: String, - arguments: String, -} - -impl FunctionCall { - /// Create a new FunctionCallBuilder. - pub fn builder() -> FunctionCallBuilder { - FunctionCallBuilder::new() - } - - /// Create a new FunctionCall (internal use). - #[allow(dead_code)] - pub(crate) fn new(name: String, arguments: String) -> Self { - Self { name, arguments } - } - - /// Get the function name. - pub fn name(&self) -> &str { - &self.name - } - - /// Get the function arguments (typically JSON string). - pub fn arguments(&self) -> &str { - &self.arguments - } -} - -/// Builder for FunctionCall. -#[derive(Clone, Default)] -pub struct FunctionCallBuilder { - name: String, - arguments: String, -} - -impl FunctionCallBuilder { - /// Create a new FunctionCallBuilder. - pub fn new() -> Self { - Self::default() - } - - /// Set the function name. - pub fn name(mut self, name: impl Into) -> Self { - self.name = name.into(); - self - } - - /// Set the function arguments. - pub fn arguments(mut self, arguments: impl Into) -> Self { - self.arguments = arguments.into(); - self - } - - /// Build the FunctionCall. - pub fn build(self) -> std::result::Result { - Ok(FunctionCall { - name: self.name, - arguments: self.arguments, - }) - } -} - -/// A chat message in the output. -#[derive(Clone, Debug, Default, Serialize)] -#[non_exhaustive] -pub struct ChatMessage { - #[serde(skip_serializing_if = "Option::is_none")] - role: Option, - #[serde(skip_serializing_if = "Option::is_none")] - content: Option, - #[serde(skip_serializing_if = "Option::is_none")] - tool_calls: Option>, -} - -impl ChatMessage { - /// Create a new ChatMessageBuilder. - pub fn builder() -> ChatMessageBuilder { - ChatMessageBuilder::new() - } - - /// Create a new ChatMessage (internal use). - pub(crate) fn new( - role: Option, - content: Option, - tool_calls: Option>, - ) -> Self { - Self { - role, - content, - tool_calls, - } - } - - /// Get the message role. - pub fn role(&self) -> Option<&str> { - self.role.as_deref() - } - - /// Get the message content. - pub fn content(&self) -> Option<&str> { - self.content.as_deref() - } - - /// Get the tool calls. - pub fn tool_calls(&self) -> Option<&[ToolCall]> { - self.tool_calls.as_deref() - } -} - -/// Builder for ChatMessage. -#[derive(Clone, Default)] -pub struct ChatMessageBuilder { - role: Option, - content: Option, - tool_calls: Option>, -} - -impl ChatMessageBuilder { - /// Create a new ChatMessageBuilder. - pub fn new() -> Self { - Self::default() - } - - /// Set the message role. - pub fn role(mut self, role: impl Into) -> Self { - self.role = Some(role.into()); - self - } - - /// Set the message content. - pub fn content(mut self, content: impl Into) -> Self { - self.content = Some(content.into()); - self - } - - /// Set the tool calls. - pub fn tool_calls(mut self, tool_calls: Vec) -> Self { - self.tool_calls = Some(tool_calls); - self - } - - /// Build the ChatMessage. - pub fn build(self) -> std::result::Result { - Ok(ChatMessage { - role: self.role, - content: self.content, - tool_calls: self.tool_calls, - }) - } -} - -/// A choice in the output array (matches OpenAI response format). -#[derive(Clone, Debug, Serialize)] -#[non_exhaustive] -pub struct OutputChoice { - index: usize, - message: ChatMessage, - logprobs: Option<()>, // Always None - #[serde(skip_serializing_if = "Option::is_none")] - finish_reason: Option, -} - -impl OutputChoice { - /// Create a new OutputChoiceBuilder. - pub fn builder() -> OutputChoiceBuilder { - OutputChoiceBuilder::new() - } - - /// Create a new OutputChoice (internal use). - pub(crate) fn new(index: usize, message: ChatMessage, finish_reason: Option) -> Self { - Self { - index, - message, - logprobs: None, - finish_reason, - } - } - - /// Get the choice index. - pub fn index(&self) -> usize { - self.index - } - - /// Get the message. - pub fn message(&self) -> &ChatMessage { - &self.message - } - - /// Get the finish reason. - pub fn finish_reason(&self) -> Option<&str> { - self.finish_reason.as_deref() - } -} - -/// Builder for OutputChoice. -#[derive(Clone, Default)] -pub struct OutputChoiceBuilder { - index: usize, - message: ChatMessage, - finish_reason: Option, -} - -impl OutputChoiceBuilder { - /// Create a new OutputChoiceBuilder. - pub fn new() -> Self { - Self::default() - } - - /// Set the choice index. - pub fn index(mut self, index: usize) -> Self { - self.index = index; - self - } - - /// Set the message. - pub fn message(mut self, message: ChatMessage) -> Self { - self.message = message; - self - } - - /// Set the finish reason. - pub fn finish_reason(mut self, finish_reason: impl Into) -> Self { - self.finish_reason = Some(finish_reason.into()); - self - } - - /// Build the OutputChoice. - pub fn build(self) -> std::result::Result { - Ok(OutputChoice { - index: self.index, - message: self.message, - logprobs: None, - finish_reason: self.finish_reason, - }) - } -} - -/// Stream metadata with typed known fields and passthrough for extras. -#[derive(Clone, Debug, Default, Serialize)] -#[non_exhaustive] -pub struct StreamMetadata { - #[serde(skip_serializing_if = "Option::is_none")] - model: Option, - /// Catch-all for additional fields (passthrough behavior). - #[serde(flatten, skip_serializing_if = "HashMap::is_empty")] - extra: HashMap, -} - -impl StreamMetadata { - /// Create a new StreamMetadataBuilder. - pub fn builder() -> StreamMetadataBuilder { - StreamMetadataBuilder::new() - } - - /// Create a new StreamMetadata (internal use). - pub(crate) fn new(model: Option, extra: HashMap) -> Self { - Self { model, extra } - } - - /// Get the model name. - pub fn model(&self) -> Option<&str> { - self.model.as_deref() - } - - /// Get the extra metadata fields. - pub fn extra(&self) -> &HashMap { - &self.extra - } - - /// Returns true if the metadata has no content. - pub fn is_empty(&self) -> bool { - self.model.is_none() && self.extra.is_empty() - } - - /// Convert to a serde_json Map for logging. - pub fn to_map(&self) -> Option> { - if self.is_empty() { - return None; - } - // Serialize to Value and extract the Map - match serde_json::to_value(self) { - Ok(Value::Object(map)) => Some(map), - _ => None, - } - } -} - -/// Builder for StreamMetadata. -#[derive(Clone, Default)] -pub struct StreamMetadataBuilder { - model: Option, - extra: HashMap, -} - -impl StreamMetadataBuilder { - /// Create a new StreamMetadataBuilder. - pub fn new() -> Self { - Self::default() - } - - /// Set the model name. - pub fn model(mut self, model: impl Into) -> Self { - self.model = Some(model.into()); - self - } - - /// Add an extra field. - pub fn extra_field(mut self, key: impl Into, value: impl Into) -> Self { - self.extra.insert(key.into(), value.into()); - self - } - - /// Set all extra fields. - pub fn extra(mut self, extra: HashMap) -> Self { - self.extra = extra; - self - } - - /// Build the StreamMetadata. - pub fn build(self) -> std::result::Result { - Ok(StreamMetadata { - model: self.model, - extra: self.extra, - }) - } -} - -/// Aggregated result from a streaming response. -#[derive(Clone)] -#[non_exhaustive] -pub struct FinalizedStream { - /// The output choices (matches OpenAI response format) - output: Vec, - /// Usage metrics extracted from the stream - usage: Option, - /// Metadata (model and any extras) - metadata: StreamMetadata, -} - -impl FinalizedStream { - /// Create a new FinalizedStreamBuilder. - pub fn builder() -> FinalizedStreamBuilder { - FinalizedStreamBuilder::new() - } - - /// Create a new FinalizedStream (internal use). - pub(crate) fn new( - output: Vec, - usage: Option, - metadata: StreamMetadata, - ) -> Self { - Self { - output, - usage, - metadata, - } - } - - /// Get the output choices. - pub fn output(&self) -> &[OutputChoice] { - &self.output - } - - /// Get the usage metrics. - pub fn usage(&self) -> Option<&UsageMetrics> { - self.usage.as_ref() - } - - /// Get the metadata. - pub fn metadata(&self) -> &StreamMetadata { - &self.metadata - } -} +pub use types::*; -/// Builder for FinalizedStream. -#[derive(Clone, Default)] -pub struct FinalizedStreamBuilder { - output: Vec, - usage: Option, - metadata: StreamMetadata, -} - -impl FinalizedStreamBuilder { - /// Create a new FinalizedStreamBuilder. - pub fn new() -> Self { - Self::default() - } - - /// Set the output choices. - pub fn output(mut self, output: Vec) -> Self { - self.output = output; - self - } - - /// Add an output choice. - pub fn add_output(mut self, choice: OutputChoice) -> Self { - self.output.push(choice); - self - } - - /// Set the usage metrics. - pub fn usage(mut self, usage: UsageMetrics) -> Self { - self.usage = Some(usage); - self - } - - /// Set the metadata. - pub fn metadata(mut self, metadata: StreamMetadata) -> Self { - self.metadata = metadata; - self - } - - /// Build the FinalizedStream. - pub fn build(self) -> std::result::Result { - Ok(FinalizedStream { - output: self.output, - usage: self.usage, - metadata: self.metadata, - }) - } -} - -/// A stream aggregator that collects streaming chunks and produces a final value. -/// -/// This follows the JS/Python SDK pattern where streaming responses are -/// collected and aggregated lazily when `final_value()` is called. -/// -/// Raw chunks are stored as-is during streaming (non-blocking), and transformation -/// to universal format happens during aggregation (which runs async in a spawned task). +/// Accumulates each chunk once without retaining raw provider responses. +/// Snapshots contain the output received so far; `final_value` caches a snapshot +/// until another chunk is pushed. Keep-alive markers are not counted. #[derive(Clone, Default)] pub struct BraintrustStream { - raw_chunks: Vec, + state: Accumulator, + chunks: usize, finalized: Option, } - -/// OpenAI-style streaming chunk structure for deserialization. -#[derive(Debug, Clone, Deserialize, Serialize)] -struct StreamChunk { - #[serde(default)] - model: Option, - #[serde(default)] - choices: Vec, - #[serde(default)] - usage: Option, -} - -/// Delta from a streaming chunk (typed for role/content/tool_calls). -#[derive(Debug, Clone, Default, Deserialize, Serialize)] -struct StreamDelta { - #[serde(default)] - role: Option, - #[serde(default)] - content: Option, - #[serde(default)] - tool_calls: Option>, -} - -/// A partial tool call from a streaming delta. -#[derive(Debug, Clone, Default, Deserialize, Serialize)] -struct StreamToolCall { - #[serde(default)] - index: Option, - #[serde(default)] - id: Option, - #[serde(default)] - #[serde(rename = "type")] - call_type: Option, - #[serde(default)] - function: Option, -} - -/// A partial function call from a streaming delta. -#[derive(Debug, Clone, Default, Deserialize, Serialize)] -struct StreamFunctionCall { - #[serde(default)] - name: Option, - #[serde(default)] - arguments: Option, -} - -#[derive(Debug, Clone, Deserialize, Serialize)] -struct StreamChoice { - #[serde(default)] - delta: Option, - #[serde(default)] - finish_reason: Option, -} - -#[derive(Debug, Clone, Default, Deserialize, Serialize)] -struct StreamUsage { - #[serde(default, alias = "input_tokens")] - prompt_tokens: Option, - #[serde(default, alias = "output_tokens")] - completion_tokens: Option, - #[serde(default, alias = "completion_reasoning_tokens")] - completion_reasoning_tokens: Option, - #[serde(default, alias = "accepted_prediction_tokens")] - accepted_prediction_tokens: Option, - #[serde(default, alias = "rejected_prediction_tokens")] - rejected_prediction_tokens: Option, - #[serde( - default, - alias = "prompt_tokens_details", - alias = "input_tokens_details" - )] - prompt_tokens_details: Option, - #[serde( - default, - alias = "completion_tokens_details", - alias = "output_tokens_details" - )] - completion_tokens_details: Option, -} - -#[derive(Debug, Clone, Default, Deserialize, Serialize)] -struct PromptTokensDetails { - #[serde(default, alias = "audio_tokens", alias = "prompt_audio_tokens")] - pub(crate) audio_tokens: Option, - #[serde(default, alias = "cached_tokens", alias = "prompt_cached_tokens")] - pub(crate) cached_tokens: Option, - #[serde( - default, - alias = "cache_creation_tokens", - alias = "cache_creation_input_tokens", - alias = "prompt_cache_creation_tokens" - )] - pub(crate) cache_creation_tokens: Option, -} - -#[derive(Debug, Clone, Default, Deserialize, Serialize)] -struct CompletionTokensDetails { - #[serde(default, alias = "audio_tokens", alias = "completion_audio_tokens")] - pub(crate) audio_tokens: Option, - #[serde( - default, - alias = "reasoning_tokens", - alias = "completion_reasoning_tokens" - )] - pub(crate) reasoning_tokens: Option, -} - impl BraintrustStream { - /// Create a new empty stream. pub fn new() -> Self { - Self { - raw_chunks: Vec::new(), - finalized: None, - } + Self::default() } - /// Add a raw JSON value to the stream. - /// - /// Stores the raw chunk as-is for later aggregation. This is non-blocking - /// to avoid adding latency to the streaming hot path. Transformation to - /// universal format happens lazily in `aggregate()`. + /// Parse and apply a raw chunk. Unknown or malformed chunks are ignored. pub fn push(&mut self, value: Value) { - // Skip keep-alive markers if value.get("_keep_alive").is_some() { return; } - self.raw_chunks.push(value); - } - - /// Get the final aggregated value. - /// - /// This aggregates all chunks into a final response. The result is cached, - /// so subsequent calls return the same value. - pub fn final_value(&mut self) -> Result<&FinalizedStream> { - if self.finalized.is_none() { - self.finalized = Some(self.aggregate()?); + self.chunks += 1; + self.finalized = None; + if let Some(chunk) = provider::parse(value) { + self.state.apply(chunk); } - Ok(self.finalized.as_ref().unwrap()) } - /// Check if the stream has any chunks. - pub fn is_empty(&self) -> bool { - self.raw_chunks.is_empty() + pub fn final_value(&mut self) -> Result<&FinalizedStream> { + Ok(self.finalized.get_or_insert_with(|| self.state.snapshot())) } - - /// Get current partial aggregated value without finalizing. - /// Used for streaming periodic updates. pub fn current_value(&self) -> Result { - self.aggregate() - } - - /// Get the number of chunks collected so far. - pub fn chunk_count(&self) -> usize { - self.raw_chunks.len() - } - - fn aggregate(&self) -> Result { - let mut usage: Option = None; - let mut model: Option = None; - let mut finish_reason: Option = None; - - // Aggregate content from all chunks - let mut aggregated_content = String::new(); - let mut role: Option = None; - - // Accumulate tool calls by index across chunks - let mut tool_calls_by_index: HashMap = HashMap::new(); - - for raw in &self.raw_chunks { - // Try to parse as OpenAI-style streaming chunk - let chunk: StreamChunk = match serde_json::from_value(raw.clone()) { - Ok(c) => c, - Err(_) => continue, // Skip unparseable chunks - }; - - // Extract model (take first non-None) - if model.is_none() { - model = chunk.model; - } - - // Extract detailed usage (take last non-None, preserving all details) - if let Some(ref u) = chunk.usage { - let prompt_tokens = u.prompt_tokens.and_then(|v| u32::try_from(v).ok()); - let completion_tokens = u.completion_tokens.and_then(|v| u32::try_from(v).ok()); - - let prompt_details = u.prompt_tokens_details.clone().map(|p| { - crate::types::PromptTokensDetails::new( - p.audio_tokens.and_then(|v| u32::try_from(v).ok()), - p.cached_tokens.and_then(|v| u32::try_from(v).ok()), - p.cache_creation_tokens.and_then(|v| u32::try_from(v).ok()), - ) - }); - - let completion_details = u.completion_tokens_details.clone().map(|c| { - crate::types::CompletionTokensDetails::new( - c.audio_tokens.and_then(|v| u32::try_from(v).ok()), - c.reasoning_tokens.and_then(|v| u32::try_from(v).ok()), - u.accepted_prediction_tokens - .and_then(|v| u32::try_from(v).ok()), - u.rejected_prediction_tokens - .and_then(|v| u32::try_from(v).ok()), - ) - }); - - usage = Some(UsageMetrics { - prompt_tokens, - completion_tokens, - total_tokens: match (prompt_tokens, completion_tokens) { - (Some(p), Some(c)) => Some(p + c), - _ => None, - }, - reasoning_tokens: None, - prompt_cached_tokens: prompt_details - .as_ref() - .and_then(|d| d.cached_tokens()) - .or_else(|| { - u.prompt_tokens_details - .as_ref() - .and_then(|p| p.cached_tokens) - .and_then(|v| u32::try_from(v).ok()) - }), - prompt_cache_creation_tokens: prompt_details - .as_ref() - .and_then(|d| d.cache_creation_tokens()) - .or_else(|| { - u.prompt_tokens_details - .as_ref() - .and_then(|p| p.cache_creation_tokens) - .and_then(|v| u32::try_from(v).ok()) - }), - completion_reasoning_tokens: u - .completion_reasoning_tokens - .and_then(|v| u32::try_from(v).ok()) - .or_else(|| { - completion_details - .as_ref() - .and_then(|d| d.reasoning_tokens()) - }), - prompt_tokens_details: prompt_details, - completion_tokens_details: completion_details, - }); - } - - // Process choices - for choice in &chunk.choices { - // Extract finish_reason (take last non-None) - if let Some(ref reason) = choice.finish_reason { - if !reason.is_empty() { - finish_reason = Some(reason.clone()); - } - } - - // Extract content and tool calls from delta - if let Some(ref delta) = choice.delta { - // Extract role (take first) - if role.is_none() { - role = delta.role.clone(); - } - - // Append content - if let Some(ref content) = delta.content { - aggregated_content.push_str(content); - } - - // Accumulate tool calls by index - if let Some(ref tool_calls) = delta.tool_calls { - for tc in tool_calls { - let idx = tc.index.unwrap_or(0); - let entry = tool_calls_by_index.entry(idx).or_default(); - - // Store id if present and not yet stored - if entry.id.is_empty() { - entry.id = tc.id.clone().unwrap_or_default(); - } - - // Store call_type if present and not yet stored - if entry.call_type.is_empty() { - entry.call_type = tc.call_type.clone().unwrap_or_default(); - } - - // Store function name if present - if let Some(ref func) = tc.function { - if let Some(ref name) = func.name { - entry.function_name = name.clone(); - } - // Accumulate raw argument fragments - if let Some(ref args) = func.arguments { - if !args.is_empty() { - entry.raw_arguments.push_str(args); - } - } - } - } - } - } - } - } - - // Build metadata (finish_reason moved to OutputChoice) - let metadata = StreamMetadata::new(model, HashMap::new()); - - // Build typed output (matches OpenAI response format) - let message = if tool_calls_by_index.is_empty() { - ChatMessage::new( - Some(role.unwrap_or_else(|| "assistant".to_string())), - Some(aggregated_content), - None, - ) - } else { - let mut tool_calls = Vec::new(); - for idx in 0..tool_calls_by_index.len() { - if let Some(accumulator) = tool_calls_by_index.remove(&idx) { - tool_calls.push(accumulator.build()); - } - } - ChatMessage::new( - Some(role.unwrap_or_else(|| "assistant".to_string())), - Some(aggregated_content), - Some(tool_calls), - ) - }; - - let choice = OutputChoice::new(0, message, finish_reason); - - Ok(FinalizedStream::new(vec![choice], usage, metadata)) - } -} - -/// Wrap a stream with span logging. -/// -/// This creates a new stream that yields the same chunks as the original, -/// but also: -/// - Records time-to-first-token on first meaningful content -/// - Accumulates chunks for aggregation -/// - Periodically flushes partial updates for real-time dashboard visibility -/// - On stream completion, logs the aggregated output/usage/metadata via `span.log()` -/// -/// # Type Parameters -/// - `S`: The stream type yielding `Result` -/// - `E`: The error type (allows use with any error type) -pub fn wrap_stream_with_span( - stream: S, - span: SpanHandle, -) -> Pin> + Send>> -where - S: Stream> + Send + Unpin + 'static, - E: Send + 'static, -{ - use futures::StreamExt; - - let start_time = Instant::now(); - let ttft_recorded = Arc::new(AtomicBool::new(false)); - let aggregator = Arc::new(Mutex::new(BraintrustStream::new())); - let span_for_complete = span.clone(); - let aggregator_for_complete = Arc::clone(&aggregator); - - // State for periodic flushing - let chunk_count = Arc::new(AtomicUsize::new(0)); - let last_flush = Arc::new(Mutex::new(Instant::now())); - - let logged_stream = stream.then(move |result| { - let span = span.clone(); - let ttft_recorded = ttft_recorded.clone(); - let aggregator = aggregator.clone(); - let chunk_count = chunk_count.clone(); - let last_flush = last_flush.clone(); - async move { - if let Ok(ref value) = result { - // Skip keep-alive markers - if value.get("_keep_alive").is_none() { - // Record TTFT on first meaningful chunk - if !ttft_recorded.swap(true, Ordering::SeqCst) && value_has_content(value) { - let ttft_secs = start_time.elapsed().as_secs_f64(); - if let Ok(log) = SpanLog::builder() - .metric("time_to_first_token", ttft_secs) - .build() - { - span.log(log); - } - } - - // Accumulate chunk for final aggregation - aggregator.lock().await.push(value.clone()); - - // Check if we should do a periodic flush - let count = chunk_count.fetch_add(1, Ordering::SeqCst) + 1; - let elapsed = { - let guard = last_flush.lock().await; - guard.elapsed() - }; - - // Flush every N chunks OR every T seconds (whichever comes first) - let should_flush = count % STREAMING_FLUSH_CHUNK_INTERVAL == 0 - || elapsed > STREAMING_FLUSH_TIME_INTERVAL; - - if should_flush { - // Log partial output for dashboard visibility - let agg = aggregator.lock().await; - if let Ok(partial) = agg.current_value() { - // Serialize typed output to Value - if let Ok(output) = serde_json::to_value(&partial.output) { - span.log(SpanLog { - output: Some(output), - ..Default::default() - }); - } - } - drop(agg); - - // Non-blocking flush for real-time updates - drop(span.flush()); - *last_flush.lock().await = Instant::now(); - } - } - } - result - } - }); - - // Wrap in a stream that finalizes on completion - Box::pin(SpanCompleteWrapper { - inner: Box::pin(logged_stream), - span: Some(span_for_complete), - aggregator: Some(aggregator_for_complete), - finalize_state: FinalizeState::Idle, - }) -} - -/// Check if a JSON value contains meaningful output (for TTFT detection). -fn value_has_content(value: &Value) -> bool { - // Check for choices array with content - if let Some(choices) = value.get("choices").and_then(|c| c.as_array()) { - if !choices.is_empty() { - return true; - } - } - // Check for usage with tokens - if let Some(usage) = value.get("usage").and_then(|u| u.as_object()) { - let has_tokens = usage - .get("completion_tokens") - .and_then(|v| v.as_i64()) - .map(|t| t > 0) - .unwrap_or(false) - || usage - .get("prompt_tokens") - .and_then(|v| v.as_i64()) - .map(|t| t > 0) - .unwrap_or(false); - if has_tokens { - return true; - } - } - false -} - -/// State for stream finalization. -enum FinalizeState { - /// Not yet finalizing - Idle, - /// Finalizing in progress - Finalizing(Pin + Send>>), - /// Finalization complete - Done, -} - -/// A wrapper stream that logs aggregated output when the stream is exhausted. -struct SpanCompleteWrapper { - inner: S, - span: Option, - aggregator: Option>>, - finalize_state: FinalizeState, -} - -/// Finalize the stream by logging and flushing the span. -async fn finalize_span(span: SpanHandle, aggregator: Arc>) { - let mut agg = aggregator.lock().await; - if !agg.is_empty() { - match agg.final_value() { - Ok(finalized) => { - // Build metrics from usage - let metrics = finalized - .usage - .as_ref() - .map(|u| usage_metrics_to_map(u.clone())); - - // Convert StreamMetadata to Option - let metadata = finalized.metadata.to_map(); - - // Serialize typed output to Value for SpanLog - let output = serde_json::to_value(&finalized.output).ok(); - - let mut builder = SpanLog::builder(); - if let Some(output) = output { - builder = builder.output(output); - } - if let Some(metadata) = metadata { - builder = builder.metadata(metadata); - } - if let Some(metrics) = metrics { - builder = builder.metrics(metrics); - } - if let Ok(log) = builder.build() { - span.log(log); - } - } - Err(e) => { - tracing::warn!("Failed to finalize stream: {}", e); - } - } - } - let end = span.end(); - let flush = span.flush(); - if let Err(error) = end.await { - tracing::warn!("Failed to deliver final span: {}", error); + Ok(self.state.snapshot()) } - // Flush span with aggregated output - if let Err(e) = flush.await { - tracing::warn!("Failed to flush span: {}", e); - } -} - -impl Stream for SpanCompleteWrapper -where - S: Stream> + Unpin, - E: Send + 'static, -{ - type Item = std::result::Result; - - fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - // SAFETY: We never move the inner stream or finalize future after pinning - let this = unsafe { self.get_unchecked_mut() }; - - // First, check if we're in the middle of finalizing - match &mut this.finalize_state { - FinalizeState::Idle => { - // Not finalizing yet, poll the inner stream - } - FinalizeState::Finalizing(fut) => { - // Poll the finalization future - match fut.as_mut().poll(cx) { - Poll::Ready(()) => { - // Finalization complete - this.finalize_state = FinalizeState::Done; - return Poll::Ready(None); - } - Poll::Pending => { - // Still finalizing - return Poll::Pending; - } - } - } - FinalizeState::Done => { - return Poll::Ready(None); - } - } - - // Poll the inner stream - let result = Pin::new(&mut this.inner).poll_next(cx); - - // If stream is done, start finalization - if matches!(result, Poll::Ready(None)) { - if let (Some(span), Some(aggregator)) = (this.span.take(), this.aggregator.take()) { - // Create the finalization future - let fut = Box::pin(finalize_span(span, aggregator)); - this.finalize_state = FinalizeState::Finalizing(fut); - - // Poll it immediately by recursing - // SAFETY: self is still valid and pinned - return unsafe { Pin::new_unchecked(this) }.poll_next(cx); - } - } - - result - } -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - #[test] - fn aggregates_content_from_streaming_values() { - let chunks = vec![ - json!({ - "id": "chunk1", - "model": "gpt-4", - "choices": [{ - "index": 0, - "delta": { "role": "assistant", "content": "Hello" } - }], - "created": 1 - }), - json!({ - "id": "chunk2", - "model": "gpt-4", - "choices": [{ - "index": 0, - "delta": { "content": " world" } - }], - "created": 1 - }), - json!({ - "id": "chunk3", - "model": "gpt-4", - "choices": [{ - "index": 0, - "delta": { "content": "!" }, - "finish_reason": "stop" - }], - "created": 1 - }), - ]; - - let mut stream = BraintrustStream::new(); - for chunk in chunks { - stream.push(chunk); - } - - let finalized = stream.final_value().expect("should finalize"); - - // Check output is array of choices - assert_eq!(finalized.output().len(), 1); - - let choice = &finalized.output()[0]; - assert_eq!(choice.index(), 0); - assert_eq!(choice.message().role(), Some("assistant")); - assert_eq!(choice.message().content(), Some("Hello world!")); - assert_eq!(choice.finish_reason(), Some("stop")); - - // Check metadata - assert_eq!(finalized.metadata().model(), Some("gpt-4")); - } - - #[test] - fn aggregates_usage_from_final_chunk() { - let chunks = vec![ - json!({ - "id": "chunk1", - "model": "gpt-4", - "choices": [{ - "index": 0, - "delta": { "role": "assistant", "content": "Hi" }, - "finish_reason": "stop" - }], - "created": 1 - }), - json!({ - "id": "chunk2", - "model": "gpt-4", - "choices": [], - "created": 1, - "usage": { - "prompt_tokens": 10, - "completion_tokens": 5 - } - }), - ]; - - let mut stream = BraintrustStream::new(); - for chunk in chunks { - stream.push(chunk); - } - - let finalized = stream.final_value().expect("should finalize"); - - let usage = finalized.usage().expect("should have usage"); - assert_eq!(usage.prompt_tokens(), Some(10)); - assert_eq!(usage.completion_tokens(), Some(5)); - assert_eq!(usage.total_tokens(), Some(15)); - } - - #[test] - fn aggregates_usage_with_input_output_token_aliases() { - let chunks = vec![ - json!({ - "id": "chunk1", - "model": "gpt-4.1", - "choices": [{ - "index": 0, - "delta": { "role": "assistant", "content": "Hi" }, - "finish_reason": "stop" - }], - "created": 1 - }), - json!({ - "id": "chunk2", - "model": "gpt-4.1", - "choices": [], - "created": 1, - "usage": { - "input_tokens": 12, - "output_tokens": 7 - } - }), - ]; - - let mut stream = BraintrustStream::new(); - for chunk in chunks { - stream.push(chunk); - } - - let finalized = stream.final_value().expect("should finalize"); - let usage = finalized.usage().expect("should have usage"); - - assert_eq!(usage.prompt_tokens(), Some(12)); - assert_eq!(usage.completion_tokens(), Some(7)); - assert_eq!(usage.total_tokens(), Some(19)); - } - - #[test] - fn aggregates_tool_calls_from_streaming_values() { - let chunks = vec![ - json!({ - "id": "chunk1", - "model": "gpt-4o", - "choices": [{ - "index": 0, - "delta": { - "role": "assistant", - "tool_calls": [{ - "index": 0, - "id": "call_abc", - "type": "function", - "function": { - "name": "get_weather", - "arguments": "{\"location\":\"" - } - }] - } - }], - "created": 1 - }), - json!({ - "id": "chunk2", - "model": "gpt-4o", - "choices": [{ - "index": 0, - "delta": { - "tool_calls": [{ - "index": 0, - "function": { - "arguments": "NYC\",\"unit\":\"" - } - }] - } - }], - "created": 1 - }), - json!({ - "id": "chunk3", - "model": "gpt-4o", - "choices": [{ - "index": 0, - "delta": { - "tool_calls": [{ - "index": 0, - "function": { - "arguments": "celsius\"}" - } - }] - }, - "finish_reason": "tool_calls" - }], - "created": 1 - }), - ]; - - let mut stream = BraintrustStream::new(); - for chunk in chunks { - stream.push(chunk); - } - - let finalized = stream.final_value().expect("should finalize"); - let choice = &finalized.output()[0]; - let tool_calls = choice - .message() - .tool_calls() - .expect("should aggregate tool calls"); - - assert_eq!(tool_calls.len(), 1); - assert_eq!(tool_calls[0].id(), "call_abc"); - assert_eq!(tool_calls[0].call_type(), "function"); - assert_eq!(tool_calls[0].function().name(), "get_weather"); - assert_eq!( - tool_calls[0].function().arguments(), - "{\"location\":\"NYC\",\"unit\":\"celsius\"}" - ); - assert_eq!(choice.finish_reason(), Some("tool_calls")); - } - - #[test] - fn aggregates_detailed_usage_from_final_chunk() { - let chunks = vec![ - json!({ - "id": "chunk1", - "model": "o4-mini", - "choices": [{ - "index": 0, - "delta": { "role": "assistant", "content": "Reasoned answer" }, - "finish_reason": "stop" - }], - "created": 1 - }), - json!({ - "id": "chunk2", - "model": "o4-mini", - "choices": [], - "created": 1, - "usage": { - "input_tokens": 100, - "output_tokens": 50, - "input_tokens_details": { - "audio_tokens": 3, - "cached_tokens": 20, - "cache_creation_tokens": 5 - }, - "output_tokens_details": { - "audio_tokens": 2, - "reasoning_tokens": 17 - }, - "accepted_prediction_tokens": 11, - "rejected_prediction_tokens": 4 - } - }), - ]; - - let mut stream = BraintrustStream::new(); - for chunk in chunks { - stream.push(chunk); - } - - let finalized = stream.final_value().expect("should finalize"); - let usage = finalized.usage().expect("should have usage"); - - assert_eq!(usage.prompt_tokens(), Some(100)); - assert_eq!(usage.completion_tokens(), Some(50)); - assert_eq!(usage.total_tokens(), Some(150)); - assert_eq!(usage.completion_reasoning_tokens(), Some(17)); - assert_eq!(usage.reasoning_tokens(), None); - - let prompt_details = usage - .prompt_tokens_details() - .expect("should have prompt token details"); - assert_eq!(prompt_details.audio_tokens(), Some(3)); - assert_eq!(prompt_details.cached_tokens(), Some(20)); - assert_eq!(prompt_details.cache_creation_tokens(), Some(5)); - - let completion_details = usage - .completion_tokens_details() - .expect("should have completion token details"); - assert_eq!(completion_details.audio_tokens(), Some(2)); - assert_eq!(completion_details.reasoning_tokens(), Some(17)); - assert_eq!(completion_details.accepted_prediction_tokens(), Some(11)); - assert_eq!(completion_details.rejected_prediction_tokens(), Some(4)); - } - - #[test] - fn skips_keep_alive_markers() { - let mut stream = BraintrustStream::new(); - - // Push a keep-alive marker - stream.push(json!({"_keep_alive": true})); - - assert!(stream.is_empty()); + pub fn is_empty(&self) -> bool { + self.chunks == 0 } - - #[test] - fn caches_finalized_result() { - let chunk = json!({ - "id": "chunk1", - "model": "gpt-4", - "choices": [{ - "index": 0, - "delta": { "role": "assistant", "content": "test" } - }], - "created": 1 - }); - - let mut stream = BraintrustStream::new(); - stream.push(chunk); - - // First call computes - extract content and drop borrow - let first_content = { - let first = stream.final_value().expect("should finalize"); - first - .output() - .first() - .and_then(|c| c.message().content().map(String::from)) - }; - - // Second call returns cached - let second_content = { - let second = stream.final_value().expect("should finalize"); - second - .output() - .first() - .and_then(|c| c.message().content().map(String::from)) - }; - - assert_eq!(first_content, second_content); - assert_eq!(first_content, Some("test".to_string())); + pub fn chunk_count(&self) -> usize { + self.chunks } } diff --git a/src/stream/accumulator.rs b/src/stream/accumulator.rs new file mode 100644 index 0000000..013d206 --- /dev/null +++ b/src/stream/accumulator.rs @@ -0,0 +1,157 @@ +//! Provider-independent retained output state; each parsed chunk is applied once. +use super::update::Update; +use super::{ChatMessage, FinalizedStream, OutputChoice, StreamMetadata}; +use crate::UsageMetrics; +use std::collections::{BTreeMap, HashMap}; + +/// Accumulates partial tool call information from streaming chunks. +/// +/// In OpenAI streaming, each chunk may contain partial tool call data. +/// This struct accumulates the id, type, function name, and raw arguments +/// across multiple chunks, then parses the complete arguments string +/// as JSON when aggregation is complete. +#[derive(Clone, Debug, Default)] +struct StreamToolCallAccumulator { + id: String, + call_type: String, + function_name: String, + raw_arguments: String, +} + +impl StreamToolCallAccumulator { + /// Creates a `ToolCall` from the accumulated data. + /// + /// The raw arguments string is parsed as JSON to extract the function + /// name and arguments, falling back to the raw string if parsing fails. + pub fn build(self) -> super::ToolCall { + let args: serde_json::Value = serde_json::from_str(&self.raw_arguments).unwrap_or_default(); + + // Extract function name from JSON object, falling back to stored name + let function_name = if let Some(obj) = args.as_object() { + obj.get("name") + .and_then(|v| v.as_str()) + .unwrap_or(&self.function_name) + .to_string() + } else { + self.function_name.clone() + }; + + // Extract arguments string from JSON object, falling back to raw string + let arguments = if let Some(obj) = args.as_object() { + obj.get("arguments") + .and_then(|v| v.as_str()) + .unwrap_or(&self.raw_arguments) + .to_string() + } else { + self.raw_arguments.clone() + }; + + let function = super::FunctionCall::new(function_name, arguments); + super::ToolCall::new(self.id, self.call_type, function) + } +} + +#[derive(Clone, Default)] +pub(super) struct Accumulator { + model: Option, + usage: Option, + finish_reason: Option, + role: Option, + aggregated_content: String, + tool_calls_by_index: BTreeMap, +} +impl Accumulator { + pub(super) fn apply(&mut self, chunk: Update) { + // Extract model (take first non-None) + if self.model.is_none() { + self.model = chunk.model; + } + + if let Some(usage) = chunk.usage { + self.usage = Some(usage); + } + + // Process choices + for choice in &chunk.choices { + // Extract finish_reason (take last non-None) + if let Some(ref reason) = choice.finish_reason { + if !reason.is_empty() { + self.finish_reason = Some(reason.clone()); + } + } + + // Extract content and tool calls from delta + if let Some(ref delta) = choice.delta { + // Extract role (take first) + if self.role.is_none() { + self.role = delta.role.clone(); + } + + // Append content + if let Some(ref content) = delta.content { + self.aggregated_content.push_str(content); + } + + // Accumulate tool calls by index + if let Some(ref tool_calls) = delta.tool_calls { + for tc in tool_calls { + let idx = tc.index.unwrap_or(0); + let entry = self.tool_calls_by_index.entry(idx).or_default(); + + // Store id if present and not yet stored + if entry.id.is_empty() { + entry.id = tc.id.clone().unwrap_or_default(); + } + + // Store call_type if present and not yet stored + if entry.call_type.is_empty() { + entry.call_type = tc.call_type.clone().unwrap_or_default(); + } + + // Store function name if present + if let Some(ref func) = tc.function { + if let Some(ref name) = func.name { + entry.function_name = name.clone(); + } + // Accumulate raw argument fragments + if let Some(ref args) = func.arguments { + if !args.is_empty() { + entry.raw_arguments.push_str(args); + } + } + } + } + } + } + } + } + pub(super) fn snapshot(&self) -> FinalizedStream { + // Build metadata (finish_reason moved to OutputChoice) + let metadata = StreamMetadata::new(self.model.clone(), HashMap::new()); + + // Build typed output (matches OpenAI response format) + let message = if self.tool_calls_by_index.is_empty() { + ChatMessage::new( + Some(self.role.clone().unwrap_or_else(|| "assistant".to_string())), + Some(self.aggregated_content.clone()), + None, + ) + } else { + let tool_calls = self + .tool_calls_by_index + .values() + .cloned() + .map(StreamToolCallAccumulator::build) + .collect(); + ChatMessage::new( + Some(self.role.clone().unwrap_or_else(|| "assistant".to_string())), + Some(self.aggregated_content.clone()), + Some(tool_calls), + ) + }; + + let choice = OutputChoice::new(0, message, self.finish_reason.clone()); + + FinalizedStream::new(vec![choice], self.usage.clone(), metadata) + } +} diff --git a/src/stream/instrumentation.rs b/src/stream/instrumentation.rs new file mode 100644 index 0000000..4cdcebc --- /dev/null +++ b/src/stream/instrumentation.rs @@ -0,0 +1,162 @@ +//! Stream polling and span lifecycle, independent of provider accumulation. +use super::BraintrustStream; +use crate::{types::usage_metrics_to_map, SpanHandle, SpanLog}; +use futures::{future::BoxFuture, Stream}; +use serde_json::Value; +use std::{ + pin::Pin, + task::{Context, Poll}, + time::{Duration, Instant}, +}; + +const UPDATE_CHUNKS: usize = 10; +const UPDATE_INTERVAL: Duration = Duration::from_secs(1); + +/// Yield the original stream items while recording incremental output and usage. +/// Exhaustion waits for final delivery. Dropping early submits the partial output +/// and final span immediately; drain the client before exiting to observe delivery. +/// Provider errors are yielded unchanged and mark the span as failed. +pub fn wrap_stream_with_span( + stream: S, + span: SpanHandle, +) -> Pin> + Send>> +where + S: Stream> + Send + Unpin + 'static, + E: Send + 'static, +{ + Box::pin(RecordedStream { + inner: stream, + recorder: Recorder { + span: Some(span), + state: BraintrustStream::new(), + started: Instant::now(), + last_update: Instant::now(), + ttft: false, + }, + finalizing: None, + done: false, + }) +} + +struct Recorder { + span: Option, + state: BraintrustStream, + started: Instant, + last_update: Instant, + ttft: bool, +} +impl Recorder { + fn push(&mut self, value: &Value) { + if value.get("_keep_alive").is_some() { + return; + } + let Some(span) = &self.span else { + return; + }; + if !self.ttft && super::provider::has_content(value) { + self.ttft = true; + if let Ok(log) = SpanLog::builder() + .metric("time_to_first_token", self.started.elapsed().as_secs_f64()) + .build() + { + drop(span.log(log)); + } + } + self.state.push(value.clone()); + if self.state.chunk_count() % UPDATE_CHUNKS == 0 + || self.last_update.elapsed() > UPDATE_INTERVAL + { + if let Ok(partial) = self.state.current_value() { + if let Ok(output) = serde_json::to_value(partial.output()) { + drop(span.log(SpanLog { + output: Some(output), + ..Default::default() + })); + } + } + drop(span.flush()); + self.last_update = Instant::now(); + } + } + // Submission is synchronous so Drop can finish accepted partial output too. + fn finish(&mut self) -> Option> { + let span = self.span.take()?; + if !self.state.is_empty() { + if let Ok(finalized) = self.state.final_value() { + let mut builder = SpanLog::builder(); + if let Ok(output) = serde_json::to_value(finalized.output()) { + builder = builder.output(output); + } + if let Some(metadata) = finalized.metadata().to_map() { + builder = builder.metadata(metadata); + } + if let Some(usage) = finalized.usage() { + builder = builder.metrics(usage_metrics_to_map(usage.clone())); + } + if let Ok(log) = builder.build() { + drop(span.log(log)); + } + } + } + let end = span.end(); + let flush = span.flush(); + Some(Box::pin(async move { + if let Err(error) = end.await { + tracing::warn!(%error, "Failed to deliver final stream span"); + } + if let Err(error) = flush.await { + tracing::warn!(%error, "Failed to flush stream span"); + } + })) + } +} +impl Drop for Recorder { + fn drop(&mut self) { + drop(self.finish()); + } +} +struct RecordedStream { + inner: S, + recorder: Recorder, + finalizing: Option>, + done: bool, +} +impl Stream for RecordedStream +where + S: Stream> + Unpin, +{ + type Item = Result; + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if this.done { + return Poll::Ready(None); + } + if this.finalizing.is_none() { + match Pin::new(&mut this.inner).poll_next(cx) { + Poll::Ready(Some(item)) => { + match &item { + Ok(value) => this.recorder.push(value), + Err(_) => { + if let Some(span) = &this.recorder.span { + if let Ok(log) = SpanLog::builder().error("stream failed").build() { + drop(span.log(log)); + } + } + } + } + return Poll::Ready(Some(item)); + } + Poll::Pending => return Poll::Pending, + Poll::Ready(None) => this.finalizing = this.recorder.finish(), + } + } + if let Some(future) = &mut this.finalizing { + if future.as_mut().poll(cx).is_pending() { + return Poll::Pending; + } + } + this.finalizing = None; + this.done = true; + Poll::Ready(None) + } +} diff --git a/src/stream/provider.rs b/src/stream/provider.rs new file mode 100644 index 0000000..8e3e842 --- /dev/null +++ b/src/stream/provider.rs @@ -0,0 +1,227 @@ +//! Parsing of the existing OpenAI-compatible stream format. +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +/// OpenAI-style streaming chunk structure for deserialization. +#[derive(Debug, Clone, Deserialize, Serialize)] +struct StreamChunk { + #[serde(default)] + model: Option, + #[serde(default)] + choices: Vec, + #[serde(default)] + usage: Option, +} + +/// Delta from a streaming chunk (typed for role/content/tool_calls). +#[derive(Debug, Clone, Default, Deserialize, Serialize)] +struct StreamDelta { + #[serde(default)] + role: Option, + #[serde(default)] + content: Option, + #[serde(default)] + tool_calls: Option>, +} + +/// A partial tool call from a streaming delta. +#[derive(Debug, Clone, Default, Deserialize, Serialize)] +struct StreamToolCall { + #[serde(default)] + index: Option, + #[serde(default)] + id: Option, + #[serde(default)] + #[serde(rename = "type")] + call_type: Option, + #[serde(default)] + function: Option, +} + +/// A partial function call from a streaming delta. +#[derive(Debug, Clone, Default, Deserialize, Serialize)] +struct StreamFunctionCall { + #[serde(default)] + name: Option, + #[serde(default)] + arguments: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +struct StreamChoice { + #[serde(default)] + delta: Option, + #[serde(default)] + finish_reason: Option, +} + +#[derive(Debug, Clone, Default, Deserialize, Serialize)] +struct StreamUsage { + #[serde(default, alias = "input_tokens")] + prompt_tokens: Option, + #[serde(default, alias = "output_tokens")] + completion_tokens: Option, + #[serde(default, alias = "completion_reasoning_tokens")] + completion_reasoning_tokens: Option, + #[serde(default, alias = "accepted_prediction_tokens")] + accepted_prediction_tokens: Option, + #[serde(default, alias = "rejected_prediction_tokens")] + rejected_prediction_tokens: Option, + #[serde( + default, + alias = "prompt_tokens_details", + alias = "input_tokens_details" + )] + prompt_tokens_details: Option, + #[serde( + default, + alias = "completion_tokens_details", + alias = "output_tokens_details" + )] + completion_tokens_details: Option, +} + +#[derive(Debug, Clone, Default, Deserialize, Serialize)] +struct PromptTokensDetails { + #[serde(default, alias = "audio_tokens", alias = "prompt_audio_tokens")] + audio_tokens: Option, + #[serde(default, alias = "cached_tokens", alias = "prompt_cached_tokens")] + cached_tokens: Option, + #[serde( + default, + alias = "cache_creation_tokens", + alias = "cache_creation_input_tokens", + alias = "prompt_cache_creation_tokens" + )] + cache_creation_tokens: Option, +} + +#[derive(Debug, Clone, Default, Deserialize, Serialize)] +struct CompletionTokensDetails { + #[serde(default, alias = "audio_tokens", alias = "completion_audio_tokens")] + audio_tokens: Option, + #[serde( + default, + alias = "reasoning_tokens", + alias = "completion_reasoning_tokens" + )] + reasoning_tokens: Option, +} + +pub(super) fn parse(value: Value) -> Option { + use super::update::*; + let chunk: StreamChunk = serde_json::from_value(value).ok()?; + Some(Update { + model: chunk.model, + usage: chunk.usage.as_ref().map(normalize_usage), + choices: chunk + .choices + .into_iter() + .map(|choice| Choice { + finish_reason: choice.finish_reason, + delta: choice.delta.map(|delta| Delta { + role: delta.role, + content: delta.content, + tool_calls: delta.tool_calls.map(|calls| { + calls + .into_iter() + .map(|call| ToolDelta { + index: call.index, + id: call.id, + call_type: call.call_type, + function: call.function.map(|function| FunctionDelta { + name: function.name, + arguments: function.arguments, + }), + }) + .collect() + }), + }), + }) + .collect(), + }) +} + +/// Meaningful output, excluding role-only, empty-choice and usage-only chunks. +pub(super) fn has_content(value: &Value) -> bool { + value + .get("choices") + .and_then(Value::as_array) + .is_some_and(|choices| { + choices.iter().any(|choice| { + let Some(delta) = choice.get("delta") else { + return false; + }; + delta + .get("content") + .and_then(Value::as_str) + .is_some_and(|s| !s.is_empty()) + || delta + .get("tool_calls") + .and_then(Value::as_array) + .is_some_and(|calls| !calls.is_empty()) + }) + }) +} + +fn normalize_usage(u: &StreamUsage) -> crate::UsageMetrics { + let prompt_tokens = u.prompt_tokens.and_then(|v| u32::try_from(v).ok()); + let completion_tokens = u.completion_tokens.and_then(|v| u32::try_from(v).ok()); + + let prompt_details = u.prompt_tokens_details.clone().map(|p| { + crate::types::PromptTokensDetails::new( + p.audio_tokens.and_then(|v| u32::try_from(v).ok()), + p.cached_tokens.and_then(|v| u32::try_from(v).ok()), + p.cache_creation_tokens.and_then(|v| u32::try_from(v).ok()), + ) + }); + + let completion_details = u.completion_tokens_details.clone().map(|c| { + crate::types::CompletionTokensDetails::new( + c.audio_tokens.and_then(|v| u32::try_from(v).ok()), + c.reasoning_tokens.and_then(|v| u32::try_from(v).ok()), + u.accepted_prediction_tokens + .and_then(|v| u32::try_from(v).ok()), + u.rejected_prediction_tokens + .and_then(|v| u32::try_from(v).ok()), + ) + }); + + crate::UsageMetrics { + prompt_tokens, + completion_tokens, + total_tokens: match (prompt_tokens, completion_tokens) { + (Some(p), Some(c)) => p.checked_add(c), + _ => None, + }, + reasoning_tokens: None, + prompt_cached_tokens: prompt_details + .as_ref() + .and_then(|d| d.cached_tokens()) + .or_else(|| { + u.prompt_tokens_details + .as_ref() + .and_then(|p| p.cached_tokens) + .and_then(|v| u32::try_from(v).ok()) + }), + prompt_cache_creation_tokens: prompt_details + .as_ref() + .and_then(|d| d.cache_creation_tokens()) + .or_else(|| { + u.prompt_tokens_details + .as_ref() + .and_then(|p| p.cache_creation_tokens) + .and_then(|v| u32::try_from(v).ok()) + }), + completion_reasoning_tokens: u + .completion_reasoning_tokens + .and_then(|v| u32::try_from(v).ok()) + .or_else(|| { + completion_details + .as_ref() + .and_then(|d| d.reasoning_tokens()) + }), + prompt_tokens_details: prompt_details, + completion_tokens_details: completion_details, + } +} diff --git a/src/stream/tests.rs b/src/stream/tests.rs new file mode 100644 index 0000000..f516887 --- /dev/null +++ b/src/stream/tests.rs @@ -0,0 +1,363 @@ +use super::*; +use serde_json::json; + +#[test] +fn aggregates_content_from_streaming_values() { + let chunks = vec![ + json!({ + "id": "chunk1", + "model": "gpt-4", + "choices": [{ + "index": 0, + "delta": { "role": "assistant", "content": "Hello" } + }], + "created": 1 + }), + json!({ + "id": "chunk2", + "model": "gpt-4", + "choices": [{ + "index": 0, + "delta": { "content": " world" } + }], + "created": 1 + }), + json!({ + "id": "chunk3", + "model": "gpt-4", + "choices": [{ + "index": 0, + "delta": { "content": "!" }, + "finish_reason": "stop" + }], + "created": 1 + }), + ]; + + let mut stream = BraintrustStream::new(); + for chunk in chunks { + stream.push(chunk); + } + + let finalized = stream.final_value().expect("should finalize"); + + // Check output is array of choices + assert_eq!(finalized.output().len(), 1); + + let choice = &finalized.output()[0]; + assert_eq!(choice.index(), 0); + assert_eq!(choice.message().role(), Some("assistant")); + assert_eq!(choice.message().content(), Some("Hello world!")); + assert_eq!(choice.finish_reason(), Some("stop")); + + // Check metadata + assert_eq!(finalized.metadata().model(), Some("gpt-4")); +} + +#[test] +fn aggregates_usage_from_final_chunk() { + let chunks = vec![ + json!({ + "id": "chunk1", + "model": "gpt-4", + "choices": [{ + "index": 0, + "delta": { "role": "assistant", "content": "Hi" }, + "finish_reason": "stop" + }], + "created": 1 + }), + json!({ + "id": "chunk2", + "model": "gpt-4", + "choices": [], + "created": 1, + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5 + } + }), + ]; + + let mut stream = BraintrustStream::new(); + for chunk in chunks { + stream.push(chunk); + } + + let finalized = stream.final_value().expect("should finalize"); + + let usage = finalized.usage().expect("should have usage"); + assert_eq!(usage.prompt_tokens(), Some(10)); + assert_eq!(usage.completion_tokens(), Some(5)); + assert_eq!(usage.total_tokens(), Some(15)); +} + +#[test] +fn aggregates_usage_with_input_output_token_aliases() { + let chunks = vec![ + json!({ + "id": "chunk1", + "model": "gpt-4.1", + "choices": [{ + "index": 0, + "delta": { "role": "assistant", "content": "Hi" }, + "finish_reason": "stop" + }], + "created": 1 + }), + json!({ + "id": "chunk2", + "model": "gpt-4.1", + "choices": [], + "created": 1, + "usage": { + "input_tokens": 12, + "output_tokens": 7 + } + }), + ]; + + let mut stream = BraintrustStream::new(); + for chunk in chunks { + stream.push(chunk); + } + + let finalized = stream.final_value().expect("should finalize"); + let usage = finalized.usage().expect("should have usage"); + + assert_eq!(usage.prompt_tokens(), Some(12)); + assert_eq!(usage.completion_tokens(), Some(7)); + assert_eq!(usage.total_tokens(), Some(19)); +} + +#[test] +fn aggregates_tool_calls_from_streaming_values() { + let chunks = vec![ + json!({ + "id": "chunk1", + "model": "gpt-4o", + "choices": [{ + "index": 0, + "delta": { + "role": "assistant", + "tool_calls": [{ + "index": 0, + "id": "call_abc", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{\"location\":\"" + } + }] + } + }], + "created": 1 + }), + json!({ + "id": "chunk2", + "model": "gpt-4o", + "choices": [{ + "index": 0, + "delta": { + "tool_calls": [{ + "index": 0, + "function": { + "arguments": "NYC\",\"unit\":\"" + } + }] + } + }], + "created": 1 + }), + json!({ + "id": "chunk3", + "model": "gpt-4o", + "choices": [{ + "index": 0, + "delta": { + "tool_calls": [{ + "index": 0, + "function": { + "arguments": "celsius\"}" + } + }] + }, + "finish_reason": "tool_calls" + }], + "created": 1 + }), + ]; + + let mut stream = BraintrustStream::new(); + for chunk in chunks { + stream.push(chunk); + } + + let finalized = stream.final_value().expect("should finalize"); + let choice = &finalized.output()[0]; + let tool_calls = choice + .message() + .tool_calls() + .expect("should aggregate tool calls"); + + assert_eq!(tool_calls.len(), 1); + assert_eq!(tool_calls[0].id(), "call_abc"); + assert_eq!(tool_calls[0].call_type(), "function"); + assert_eq!(tool_calls[0].function().name(), "get_weather"); + assert_eq!( + tool_calls[0].function().arguments(), + "{\"location\":\"NYC\",\"unit\":\"celsius\"}" + ); + assert_eq!(choice.finish_reason(), Some("tool_calls")); +} + +#[test] +fn aggregates_detailed_usage_from_final_chunk() { + let chunks = vec![ + json!({ + "id": "chunk1", + "model": "o4-mini", + "choices": [{ + "index": 0, + "delta": { "role": "assistant", "content": "Reasoned answer" }, + "finish_reason": "stop" + }], + "created": 1 + }), + json!({ + "id": "chunk2", + "model": "o4-mini", + "choices": [], + "created": 1, + "usage": { + "input_tokens": 100, + "output_tokens": 50, + "input_tokens_details": { + "audio_tokens": 3, + "cached_tokens": 20, + "cache_creation_tokens": 5 + }, + "output_tokens_details": { + "audio_tokens": 2, + "reasoning_tokens": 17 + }, + "accepted_prediction_tokens": 11, + "rejected_prediction_tokens": 4 + } + }), + ]; + + let mut stream = BraintrustStream::new(); + for chunk in chunks { + stream.push(chunk); + } + + let finalized = stream.final_value().expect("should finalize"); + let usage = finalized.usage().expect("should have usage"); + + assert_eq!(usage.prompt_tokens(), Some(100)); + assert_eq!(usage.completion_tokens(), Some(50)); + assert_eq!(usage.total_tokens(), Some(150)); + assert_eq!(usage.completion_reasoning_tokens(), Some(17)); + assert_eq!(usage.reasoning_tokens(), None); + + let prompt_details = usage + .prompt_tokens_details() + .expect("should have prompt token details"); + assert_eq!(prompt_details.audio_tokens(), Some(3)); + assert_eq!(prompt_details.cached_tokens(), Some(20)); + assert_eq!(prompt_details.cache_creation_tokens(), Some(5)); + + let completion_details = usage + .completion_tokens_details() + .expect("should have completion token details"); + assert_eq!(completion_details.audio_tokens(), Some(2)); + assert_eq!(completion_details.reasoning_tokens(), Some(17)); + assert_eq!(completion_details.accepted_prediction_tokens(), Some(11)); + assert_eq!(completion_details.rejected_prediction_tokens(), Some(4)); +} + +#[test] +fn skips_keep_alive_markers() { + let mut stream = BraintrustStream::new(); + + // Push a keep-alive marker + stream.push(json!({"_keep_alive": true})); + + assert!(stream.is_empty()); +} + +#[test] +fn caches_finalized_result() { + let chunk = json!({ + "id": "chunk1", + "model": "gpt-4", + "choices": [{ + "index": 0, + "delta": { "role": "assistant", "content": "test" } + }], + "created": 1 + }); + + let mut stream = BraintrustStream::new(); + stream.push(chunk); + + // First call computes - extract content and drop borrow + let first_content = { + let first = stream.final_value().expect("should finalize"); + first + .output() + .first() + .and_then(|c| c.message().content().map(String::from)) + }; + + // Second call returns cached + let second_content = { + let second = stream.final_value().expect("should finalize"); + second + .output() + .first() + .and_then(|c| c.message().content().map(String::from)) + }; + + assert_eq!(first_content, second_content); + assert_eq!(first_content, Some("test".to_string())); +} + +#[test] +fn snapshots_follow_new_chunks_and_preserve_sparse_tool_indices() { + let mut stream = BraintrustStream::new(); + stream.push(json!({"choices":[{"delta":{"content":"first", "tool_calls":[{"index":7,"id":"tool-seven","type":"function","function":{"name":"lookup","arguments":"{\"query\":"}}]}}]})); + assert_eq!( + stream.final_value().unwrap().output()[0] + .message() + .content(), + Some("first") + ); + stream.push(json!({"choices":[{"delta":{"content":" second", "tool_calls":[{"index":7,"function":{"arguments":"\"value\"}"}}]}}]})); + let snapshot = stream.final_value().unwrap(); + let message = snapshot.output()[0].message(); + assert_eq!(message.content(), Some("first second")); + let calls = message.tool_calls().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].id(), "tool-seven"); + assert_eq!(calls[0].function().arguments(), "{\"query\":\"value\"}"); +} + +#[test] +fn ignores_unknown_and_malformed_chunks_without_losing_valid_output() { + let mut stream = BraintrustStream::new(); + for chunk in [ + json!(null), + json!({"unknown": true}), + json!({"choices":"invalid"}), + json!({"choices":[{"delta":{"content":"valid"}}]}), + json!({"usage":{"prompt_tokens":4294967295_u64, "completion_tokens":1}}), + ] { + stream.push(chunk); + } + let snapshot = stream.current_value().unwrap(); + assert_eq!(snapshot.output()[0].message().content(), Some("valid")); + assert_eq!(snapshot.usage().unwrap().total_tokens(), None); + assert_eq!(stream.chunk_count(), 5); +} diff --git a/src/stream/types.rs b/src/stream/types.rs new file mode 100644 index 0000000..94ab915 --- /dev/null +++ b/src/stream/types.rs @@ -0,0 +1,598 @@ +use crate::UsageMetrics; +use serde::Serialize; +use serde_json::Value; +use std::{collections::HashMap, fmt}; + +// ============================================================================= +// Error types for stream output builders (empty for now, extensible later) +// ============================================================================= + +/// Error type for ToolCall builder validation. +#[derive(Debug, Clone, PartialEq)] +#[non_exhaustive] +pub enum ToolCallBuilderError {} + +impl fmt::Display for ToolCallBuilderError { + fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self {} + } +} + +impl std::error::Error for ToolCallBuilderError {} + +/// Error type for FunctionCall builder validation. +#[derive(Debug, Clone, PartialEq)] +#[non_exhaustive] +pub enum FunctionCallBuilderError {} + +impl fmt::Display for FunctionCallBuilderError { + fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self {} + } +} + +impl std::error::Error for FunctionCallBuilderError {} + +/// Error type for ChatMessage builder validation. +#[derive(Debug, Clone, PartialEq)] +#[non_exhaustive] +pub enum ChatMessageBuilderError {} + +impl fmt::Display for ChatMessageBuilderError { + fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self {} + } +} + +impl std::error::Error for ChatMessageBuilderError {} + +/// Error type for OutputChoice builder validation. +#[derive(Debug, Clone, PartialEq)] +#[non_exhaustive] +pub enum OutputChoiceBuilderError {} + +impl fmt::Display for OutputChoiceBuilderError { + fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self {} + } +} + +impl std::error::Error for OutputChoiceBuilderError {} + +/// Error type for StreamMetadata builder validation. +#[derive(Debug, Clone, PartialEq)] +#[non_exhaustive] +pub enum StreamMetadataBuilderError {} + +impl fmt::Display for StreamMetadataBuilderError { + fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self {} + } +} + +impl std::error::Error for StreamMetadataBuilderError {} + +/// Error type for FinalizedStream builder validation. +#[derive(Debug, Clone, PartialEq)] +#[non_exhaustive] +pub enum FinalizedStreamBuilderError {} + +impl fmt::Display for FinalizedStreamBuilderError { + fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self {} + } +} + +impl std::error::Error for FinalizedStreamBuilderError {} + +/// A tool call in a chat message. +#[derive(Clone, Debug, Default, Serialize)] +#[non_exhaustive] +pub struct ToolCall { + id: String, + #[serde(rename = "type")] + call_type: String, // Always "function" + function: FunctionCall, +} + +impl ToolCall { + /// Create a new ToolCallBuilder. + pub fn builder() -> ToolCallBuilder { + ToolCallBuilder::new() + } + + /// Create a new ToolCall (internal use). + #[allow(dead_code)] + pub(crate) fn new(id: String, call_type: String, function: FunctionCall) -> Self { + Self { + id, + call_type, + function, + } + } + + /// Get the tool call ID. + pub fn id(&self) -> &str { + &self.id + } + + /// Get the tool call type (typically "function"). + pub fn call_type(&self) -> &str { + &self.call_type + } + + /// Get the function details. + pub fn function(&self) -> &FunctionCall { + &self.function + } +} + +/// Builder for ToolCall. +#[derive(Clone, Default)] +pub struct ToolCallBuilder { + id: String, + call_type: String, + function: FunctionCall, +} + +impl ToolCallBuilder { + /// Create a new ToolCallBuilder. + pub fn new() -> Self { + Self { + call_type: "function".to_string(), + ..Default::default() + } + } + + /// Set the tool call ID. + pub fn id(mut self, id: impl Into) -> Self { + self.id = id.into(); + self + } + + /// Set the call type. + pub fn call_type(mut self, call_type: impl Into) -> Self { + self.call_type = call_type.into(); + self + } + + /// Set the function details. + pub fn function(mut self, function: FunctionCall) -> Self { + self.function = function; + self + } + + /// Build the ToolCall. + pub fn build(self) -> std::result::Result { + Ok(ToolCall { + id: self.id, + call_type: self.call_type, + function: self.function, + }) + } +} + +/// Function details in a tool call. +#[derive(Clone, Debug, Default, Serialize)] +#[non_exhaustive] +pub struct FunctionCall { + name: String, + arguments: String, +} + +impl FunctionCall { + /// Create a new FunctionCallBuilder. + pub fn builder() -> FunctionCallBuilder { + FunctionCallBuilder::new() + } + + /// Create a new FunctionCall (internal use). + #[allow(dead_code)] + pub(crate) fn new(name: String, arguments: String) -> Self { + Self { name, arguments } + } + + /// Get the function name. + pub fn name(&self) -> &str { + &self.name + } + + /// Get the function arguments (typically JSON string). + pub fn arguments(&self) -> &str { + &self.arguments + } +} + +/// Builder for FunctionCall. +#[derive(Clone, Default)] +pub struct FunctionCallBuilder { + name: String, + arguments: String, +} + +impl FunctionCallBuilder { + /// Create a new FunctionCallBuilder. + pub fn new() -> Self { + Self::default() + } + + /// Set the function name. + pub fn name(mut self, name: impl Into) -> Self { + self.name = name.into(); + self + } + + /// Set the function arguments. + pub fn arguments(mut self, arguments: impl Into) -> Self { + self.arguments = arguments.into(); + self + } + + /// Build the FunctionCall. + pub fn build(self) -> std::result::Result { + Ok(FunctionCall { + name: self.name, + arguments: self.arguments, + }) + } +} + +/// A chat message in the output. +#[derive(Clone, Debug, Default, Serialize)] +#[non_exhaustive] +pub struct ChatMessage { + #[serde(skip_serializing_if = "Option::is_none")] + role: Option, + #[serde(skip_serializing_if = "Option::is_none")] + content: Option, + #[serde(skip_serializing_if = "Option::is_none")] + tool_calls: Option>, +} + +impl ChatMessage { + /// Create a new ChatMessageBuilder. + pub fn builder() -> ChatMessageBuilder { + ChatMessageBuilder::new() + } + + /// Create a new ChatMessage (internal use). + pub(crate) fn new( + role: Option, + content: Option, + tool_calls: Option>, + ) -> Self { + Self { + role, + content, + tool_calls, + } + } + + /// Get the message role. + pub fn role(&self) -> Option<&str> { + self.role.as_deref() + } + + /// Get the message content. + pub fn content(&self) -> Option<&str> { + self.content.as_deref() + } + + /// Get the tool calls. + pub fn tool_calls(&self) -> Option<&[ToolCall]> { + self.tool_calls.as_deref() + } +} + +/// Builder for ChatMessage. +#[derive(Clone, Default)] +pub struct ChatMessageBuilder { + role: Option, + content: Option, + tool_calls: Option>, +} + +impl ChatMessageBuilder { + /// Create a new ChatMessageBuilder. + pub fn new() -> Self { + Self::default() + } + + /// Set the message role. + pub fn role(mut self, role: impl Into) -> Self { + self.role = Some(role.into()); + self + } + + /// Set the message content. + pub fn content(mut self, content: impl Into) -> Self { + self.content = Some(content.into()); + self + } + + /// Set the tool calls. + pub fn tool_calls(mut self, tool_calls: Vec) -> Self { + self.tool_calls = Some(tool_calls); + self + } + + /// Build the ChatMessage. + pub fn build(self) -> std::result::Result { + Ok(ChatMessage { + role: self.role, + content: self.content, + tool_calls: self.tool_calls, + }) + } +} + +/// A choice in the output array (matches OpenAI response format). +#[derive(Clone, Debug, Serialize)] +#[non_exhaustive] +pub struct OutputChoice { + index: usize, + message: ChatMessage, + logprobs: Option<()>, // Always None + #[serde(skip_serializing_if = "Option::is_none")] + finish_reason: Option, +} + +impl OutputChoice { + /// Create a new OutputChoiceBuilder. + pub fn builder() -> OutputChoiceBuilder { + OutputChoiceBuilder::new() + } + + /// Create a new OutputChoice (internal use). + pub(crate) fn new(index: usize, message: ChatMessage, finish_reason: Option) -> Self { + Self { + index, + message, + logprobs: None, + finish_reason, + } + } + + /// Get the choice index. + pub fn index(&self) -> usize { + self.index + } + + /// Get the message. + pub fn message(&self) -> &ChatMessage { + &self.message + } + + /// Get the finish reason. + pub fn finish_reason(&self) -> Option<&str> { + self.finish_reason.as_deref() + } +} + +/// Builder for OutputChoice. +#[derive(Clone, Default)] +pub struct OutputChoiceBuilder { + index: usize, + message: ChatMessage, + finish_reason: Option, +} + +impl OutputChoiceBuilder { + /// Create a new OutputChoiceBuilder. + pub fn new() -> Self { + Self::default() + } + + /// Set the choice index. + pub fn index(mut self, index: usize) -> Self { + self.index = index; + self + } + + /// Set the message. + pub fn message(mut self, message: ChatMessage) -> Self { + self.message = message; + self + } + + /// Set the finish reason. + pub fn finish_reason(mut self, finish_reason: impl Into) -> Self { + self.finish_reason = Some(finish_reason.into()); + self + } + + /// Build the OutputChoice. + pub fn build(self) -> std::result::Result { + Ok(OutputChoice { + index: self.index, + message: self.message, + logprobs: None, + finish_reason: self.finish_reason, + }) + } +} + +/// Stream metadata with typed known fields and passthrough for extras. +#[derive(Clone, Debug, Default, Serialize)] +#[non_exhaustive] +pub struct StreamMetadata { + #[serde(skip_serializing_if = "Option::is_none")] + model: Option, + /// Catch-all for additional fields (passthrough behavior). + #[serde(flatten, skip_serializing_if = "HashMap::is_empty")] + extra: HashMap, +} + +impl StreamMetadata { + /// Create a new StreamMetadataBuilder. + pub fn builder() -> StreamMetadataBuilder { + StreamMetadataBuilder::new() + } + + /// Create a new StreamMetadata (internal use). + pub(crate) fn new(model: Option, extra: HashMap) -> Self { + Self { model, extra } + } + + /// Get the model name. + pub fn model(&self) -> Option<&str> { + self.model.as_deref() + } + + /// Get the extra metadata fields. + pub fn extra(&self) -> &HashMap { + &self.extra + } + + /// Returns true if the metadata has no content. + pub fn is_empty(&self) -> bool { + self.model.is_none() && self.extra.is_empty() + } + + /// Convert to a serde_json Map for logging. + pub fn to_map(&self) -> Option> { + if self.is_empty() { + return None; + } + // Serialize to Value and extract the Map + match serde_json::to_value(self) { + Ok(Value::Object(map)) => Some(map), + _ => None, + } + } +} + +/// Builder for StreamMetadata. +#[derive(Clone, Default)] +pub struct StreamMetadataBuilder { + model: Option, + extra: HashMap, +} + +impl StreamMetadataBuilder { + /// Create a new StreamMetadataBuilder. + pub fn new() -> Self { + Self::default() + } + + /// Set the model name. + pub fn model(mut self, model: impl Into) -> Self { + self.model = Some(model.into()); + self + } + + /// Add an extra field. + pub fn extra_field(mut self, key: impl Into, value: impl Into) -> Self { + self.extra.insert(key.into(), value.into()); + self + } + + /// Set all extra fields. + pub fn extra(mut self, extra: HashMap) -> Self { + self.extra = extra; + self + } + + /// Build the StreamMetadata. + pub fn build(self) -> std::result::Result { + Ok(StreamMetadata { + model: self.model, + extra: self.extra, + }) + } +} + +/// Aggregated result from a streaming response. +#[derive(Clone)] +#[non_exhaustive] +pub struct FinalizedStream { + /// The output choices (matches OpenAI response format) + output: Vec, + /// Usage metrics extracted from the stream + usage: Option, + /// Metadata (model and any extras) + metadata: StreamMetadata, +} + +impl FinalizedStream { + /// Create a new FinalizedStreamBuilder. + pub fn builder() -> FinalizedStreamBuilder { + FinalizedStreamBuilder::new() + } + + /// Create a new FinalizedStream (internal use). + pub(crate) fn new( + output: Vec, + usage: Option, + metadata: StreamMetadata, + ) -> Self { + Self { + output, + usage, + metadata, + } + } + + /// Get the output choices. + pub fn output(&self) -> &[OutputChoice] { + &self.output + } + + /// Get the usage metrics. + pub fn usage(&self) -> Option<&UsageMetrics> { + self.usage.as_ref() + } + + /// Get the metadata. + pub fn metadata(&self) -> &StreamMetadata { + &self.metadata + } +} + +/// Builder for FinalizedStream. +#[derive(Clone, Default)] +pub struct FinalizedStreamBuilder { + output: Vec, + usage: Option, + metadata: StreamMetadata, +} + +impl FinalizedStreamBuilder { + /// Create a new FinalizedStreamBuilder. + pub fn new() -> Self { + Self::default() + } + + /// Set the output choices. + pub fn output(mut self, output: Vec) -> Self { + self.output = output; + self + } + + /// Add an output choice. + pub fn add_output(mut self, choice: OutputChoice) -> Self { + self.output.push(choice); + self + } + + /// Set the usage metrics. + pub fn usage(mut self, usage: UsageMetrics) -> Self { + self.usage = Some(usage); + self + } + + /// Set the metadata. + pub fn metadata(mut self, metadata: StreamMetadata) -> Self { + self.metadata = metadata; + self + } + + /// Build the FinalizedStream. + pub fn build(self) -> std::result::Result { + Ok(FinalizedStream { + output: self.output, + usage: self.usage, + metadata: self.metadata, + }) + } +} diff --git a/src/stream/update.rs b/src/stream/update.rs new file mode 100644 index 0000000..9b27a11 --- /dev/null +++ b/src/stream/update.rs @@ -0,0 +1,27 @@ +//! Normalized changes applied to retained output, without wire serialization. +use crate::UsageMetrics; +#[derive(Default)] +pub(super) struct Update { + pub model: Option, + pub usage: Option, + pub choices: Vec, +} +pub(super) struct Choice { + pub finish_reason: Option, + pub delta: Option, +} +pub(super) struct Delta { + pub role: Option, + pub content: Option, + pub tool_calls: Option>, +} +pub(super) struct ToolDelta { + pub index: Option, + pub id: Option, + pub call_type: Option, + pub function: Option, +} +pub(super) struct FunctionDelta { + pub name: Option, + pub arguments: Option, +} diff --git a/tests/stream_lifecycle.rs b/tests/stream_lifecycle.rs new file mode 100644 index 0000000..c065da4 --- /dev/null +++ b/tests/stream_lifecycle.rs @@ -0,0 +1,135 @@ +use braintrust_sdk_rust::{wrap_stream_with_span, BraintrustClient}; +use futures::{stream, StreamExt}; +use serde_json::{json, Value}; +use std::time::Duration; +use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, +}; + +async fn client(server: &MockServer) -> BraintrustClient { + Mock::given(method("POST")) + .and(path("/api/project/register")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"test-project"}})), + ) + .mount(server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200)) + .mount(server) + .await; + BraintrustClient::builder() + .skip_login(true) + .api_url(server.uri()) + .app_url(server.uri()) + .build() + .await + .unwrap() +} +async fn rows(server: &MockServer) -> Vec { + server + .received_requests() + .await + .unwrap() + .iter() + .filter(|r| r.url.path() == "/logs3") + .flat_map(|r| { + serde_json::from_slice::(&r.body).unwrap()["rows"] + .as_array() + .unwrap() + .clone() + }) + .collect() +} +// Error values need neither Display nor Debug; original provider items survive. +struct ProviderError; + +#[tokio::test] +async fn exhaustion_finalizes_once_after_role_only_and_usage_only_chunks() { + let server = MockServer::start().await; + let client = client(&server).await; + let span = client + .span_builder_with_credentials("token", "org") + .project_name("test-project") + .build(); + let chunks = vec![ + json!({"choices":[{"delta":{"role":"assistant"}}]}), + json!({"usage":{"prompt_tokens":3}}), + json!({"choices":[{"delta":{"content":"hello"}}]}), + json!({"usage":{"prompt_tokens":3,"completion_tokens":2}}), + ]; + let mut wrapped = wrap_stream_with_span( + stream::iter(chunks.clone().into_iter().map(Ok::<_, ProviderError>)), + span.clone(), + ); + for expected in chunks { + assert_eq!(wrapped.next().await.map(|r| r.ok()), Some(Some(expected))); + } + assert!(wrapped.next().await.is_none()); + let end = span.end_time().unwrap(); + assert!(wrapped.next().await.is_none()); + drop(wrapped); + client.shutdown(Duration::from_secs(5)).await.unwrap(); + assert_eq!(span.end_time(), Some(end)); + let rows = rows(&server).await; + assert!(rows + .iter() + .any(|r| r["output"][0]["message"]["content"] == "hello")); + assert!(rows + .iter() + .any(|r| r["metrics"]["completion_tokens"].as_f64() == Some(2.0))); + assert!(rows + .iter() + .any(|r| r["metrics"].get("time_to_first_token").is_some())); +} + +#[tokio::test] +async fn early_drop_submits_partial_output_and_error_without_waiting_for_eof() { + let server = MockServer::start().await; + let client = client(&server).await; + let span = client + .span_builder_with_credentials("token", "org") + .project_name("test-project") + .build(); + let inner = stream::iter(vec![ + Ok(json!({"choices":[{"delta":{"content":"partial"}}]})), + Err(ProviderError), + ]) + .chain(stream::pending()); + let mut wrapped = wrap_stream_with_span(inner, span.clone()); + assert!(matches!(wrapped.next().await, Some(Ok(_)))); + assert!(matches!(wrapped.next().await, Some(Err(ProviderError)))); + drop(wrapped); + assert!(span.end_time().is_some()); + client.shutdown(Duration::from_secs(5)).await.unwrap(); + let rows = rows(&server).await; + assert!(rows + .iter() + .any(|r| r["output"][0]["message"]["content"] == "partial")); + assert!(rows.iter().any(|r| r["error"] == "stream failed")); + assert!(rows.iter().any(|r| r["metrics"].get("end").is_some())); +} + +#[tokio::test] +async fn role_only_stream_does_not_record_time_to_first_token() { + let server = MockServer::start().await; + let client = client(&server).await; + let span = client + .span_builder_with_credentials("token", "org") + .project_name("test-project") + .build(); + let wrapped = wrap_stream_with_span( + stream::iter([Ok::<_, ProviderError>( + json!({"choices":[{"delta":{"role":"assistant"}}]}), + )]), + span, + ); + let _ = wrapped.collect::>().await; + client.shutdown(Duration::from_secs(5)).await.unwrap(); + assert!(rows(&server) + .await + .iter() + .all(|r| r["metrics"].get("time_to_first_token").is_none())); +} From d5094cf8826f18c5050386a64f818567154c36b8 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Fri, 23 Oct 2026 23:26:06 +0800 Subject: [PATCH 2/2] Keep dataset test servers alive and allow startup capability discovery --- tests/dataset_case_source.rs | 37 ++++++++++++++++++++++++++++-------- 1 file changed, 29 insertions(+), 8 deletions(-) diff --git a/tests/dataset_case_source.rs b/tests/dataset_case_source.rs index 883fb35..114049e 100644 --- a/tests/dataset_case_source.rs +++ b/tests/dataset_case_source.rs @@ -165,7 +165,7 @@ async fn typed_pages_pin_version_and_preserve_case_fields_without_clone_or_seria #[tokio::test] async fn recorded_source_admitted_before_shutdown_can_fetch_all_pages_and_finish_recording() { - let (client, dataset, _, data) = setup(true).await; + let (client, dataset, _control, data) = setup(true).await; page( ResponseTemplate::new(200) .set_body_json(json!({"data":[{"id":"one","input":1}],"cursor":"next"})), @@ -251,13 +251,20 @@ async fn registration_failure_is_a_typed_source_error_before_cases_are_consumed( let partial = error.partial.unwrap(); assert_eq!(partial.total_cases, 0); assert!(partial.dataset.is_none()); - assert!(data.received_requests().await.unwrap().is_empty()); + // Client startup may discover ingestion capabilities; no dataset reads + // may run outside the source's fixed execution binding. + assert!(data + .received_requests() + .await + .unwrap() + .iter() + .all(|request| request.url.path() == "/version")); client.shutdown(Duration::from_secs(5)).await.unwrap(); } #[tokio::test] async fn later_page_failure_retains_completed_cases_and_snapshot_identity() { - let (client, dataset, _, data) = setup(false).await; + let (client, dataset, _control, data) = setup(false).await; page( ResponseTemplate::new(200) .set_body_json(json!({"data":[{"id":"one","input":{"value":4}}],"cursor":"next"})), @@ -288,7 +295,7 @@ async fn later_page_failure_retains_completed_cases_and_snapshot_identity() { #[tokio::test] async fn malformed_row_keeps_decode_cause_and_preceding_results() { - let (client, dataset, _, data) = setup(false).await; + let (client, dataset, _control, data) = setup(false).await; page(ResponseTemplate::new(200).set_body_json(json!({"data":[{"id":"good","input":{"value":2}},{"id":"bad","input":"invalid input"}],"cursor":null}))).mount(&data).await; let error = run(EvalRequest::local(options( DatasetCaseSource::new(dataset, None, None).unwrap(), @@ -320,11 +327,18 @@ async fn malformed_row_keeps_decode_cause_and_preceding_results() { #[tokio::test] async fn dropping_an_unread_source_releases_accepted_work_and_closed_admission_is_immediate() { - let (client, dataset, _, data) = setup(false).await; + let (client, dataset, _control, data) = setup(false).await; let source = DatasetCaseSource::::new(dataset.clone(), None, None).unwrap(); drop(source); client.shutdown(Duration::from_secs(5)).await.unwrap(); - assert!(data.received_requests().await.unwrap().is_empty()); + // Client startup may discover ingestion capabilities; no dataset reads + // may run outside the source's fixed execution binding. + assert!(data + .received_requests() + .await + .unwrap() + .iter() + .all(|request| request.url.path() == "/version")); let error = DatasetCaseSource::::new(dataset, None, None) .err() .expect("source admission must be closed"); @@ -336,7 +350,7 @@ async fn dropping_an_unread_source_releases_accepted_work_and_closed_admission_i #[tokio::test] async fn empty_live_source_reports_identity_without_claiming_a_snapshot() { - let (client, dataset, _, data) = setup(false).await; + let (client, dataset, _control, data) = setup(false).await; page(ResponseTemplate::new(200).set_body_json(json!({"data":[],"cursor":null}))) .expect(1) .mount(&data) @@ -413,7 +427,14 @@ async fn source_runtime_loss_reports_original_worker_loss_without_running_reads_ )) ))); assert!(control.received_requests().await.unwrap().is_empty()); - assert!(data.received_requests().await.unwrap().is_empty()); + // Client startup may discover ingestion capabilities; no dataset reads + // may run outside the source's fixed execution binding. + assert!(data + .received_requests() + .await + .unwrap() + .iter() + .all(|request| request.url.path() == "/version")); } #[tokio::test]