llm/providers/openai/
streaming.rs

1use async_openai::types::chat::{
2    CompletionUsage, CreateChatCompletionStreamResponse, FinishReason as OpenAiFinishReason,
3};
4use async_stream;
5use tokio_stream::{Stream, StreamExt};
6use tracing::debug;
7
8use crate::providers::tool_call_collector::ToolCallCollector;
9use crate::{LlmError, LlmResponse, Result, StopReason, TokenUsage};
10
11impl From<CompletionUsage> for TokenUsage {
12    fn from(usage: CompletionUsage) -> Self {
13        let prompt = usage.prompt_tokens_details.unwrap_or_default();
14        let completion = usage.completion_tokens_details.unwrap_or_default();
15        TokenUsage {
16            input_tokens: usage.prompt_tokens,
17            output_tokens: usage.completion_tokens,
18            cache_read_tokens: prompt.cached_tokens,
19            cache_reporting_exclusive: Some(false),
20            input_audio_tokens: prompt.audio_tokens,
21            reasoning_tokens: completion.reasoning_tokens,
22            output_audio_tokens: completion.audio_tokens,
23            accepted_prediction_tokens: completion.accepted_prediction_tokens,
24            rejected_prediction_tokens: completion.rejected_prediction_tokens,
25            ..TokenUsage::default()
26        }
27    }
28}
29
30/// Common stream processing logic that handles tool call state tracking and event emission.
31/// Works with standard `async_openai` `CreateChatCompletionStreamResponse` types.
32pub fn process_completion_stream<E: Into<LlmError> + Send>(
33    mut stream: impl Stream<Item = std::result::Result<CreateChatCompletionStreamResponse, E>> + Send + Unpin,
34) -> impl Stream<Item = Result<LlmResponse>> + Send {
35    async_stream::stream! {
36        let message_id = uuid::Uuid::new_v4().to_string();
37        yield Ok(LlmResponse::Start { message_id });
38
39        let mut collector = ToolCallCollector::<u32>::new();
40        let mut last_stop_reason: Option<StopReason> = None;
41
42        while let Some(result) = stream.next().await {
43            match result {
44                Ok(mut response) => {
45                    // Emit usage information if available
46                    // This must be checked on every chunk since usage may come
47                    // in a separate final chunk after finish_reason
48                    if let Some(usage) = response.usage {
49                        yield Ok(LlmResponse::Usage { tokens: usage.into() });
50                    }
51
52                    if let Some(choice) = response.choices.pop() {
53                        let delta = choice.delta;
54
55                        if let Some(content) = delta.content
56                            && !content.is_empty() {
57                                // If we have pending tool calls and now we're getting content,
58                                // complete all tool calls first
59                                for tool_call in collector.complete_all() {
60                                    yield Ok(LlmResponse::ToolRequestComplete { tool_call });
61                                }
62                                yield Ok(LlmResponse::Text { chunk: content });
63                            }
64
65                        if let Some(tool_calls) = delta.tool_calls {
66                            for tc in tool_calls {
67                                let (id, name, args) = match tc.function {
68                                    Some(f) => (tc.id, f.name, f.arguments),
69                                    None => (tc.id, None, None),
70                                };
71                                for response in collector.handle_delta(tc.index, id, name, args) {
72                                    yield Ok(response);
73                                }
74                            }
75                        }
76
77                        if let Some(finish_reason) = choice.finish_reason {
78                            let finish_reason_str = format!("{finish_reason:?}");
79                            debug!("Received finish reason: {finish_reason_str}");
80                            last_stop_reason = Some(map_openai_finish_reason(finish_reason));
81
82                            for tool_call in collector.complete_all() {
83                                yield Ok(LlmResponse::ToolRequestComplete { tool_call });
84                            }
85                            // Don't break yet - continue to capture usage from subsequent chunks
86                            // OpenRouter sends usage in the last SSE message after finish_reason
87                            // See: https://openrouter.ai/docs/guides/usage-accounting
88                        }
89                    } else {
90                        // No choices in this chunk - could be:
91                        // 1. Final usage-only chunk after finish_reason (OpenRouter)
92                        // 2. Stream is done (some providers)
93                        // We already extracted usage above if present
94                        debug!("No choices in response, ending stream");
95                        for tool_call in collector.complete_all() {
96                            yield Ok(LlmResponse::ToolRequestComplete { tool_call });
97                        }
98                        break;
99                    }
100                }
101                Err(e) => {
102                    yield Err(e.into());
103                    break;
104                }
105            }
106        }
107
108        yield Ok(LlmResponse::Done {
109            stop_reason: last_stop_reason,
110        });
111    }
112}
113
114fn map_openai_finish_reason(reason: OpenAiFinishReason) -> StopReason {
115    match reason {
116        OpenAiFinishReason::Stop => StopReason::EndTurn,
117        OpenAiFinishReason::Length => StopReason::Length,
118        OpenAiFinishReason::ToolCalls => StopReason::ToolCalls,
119        OpenAiFinishReason::ContentFilter => StopReason::ContentFilter,
120        OpenAiFinishReason::FunctionCall => StopReason::FunctionCall,
121    }
122}
llm/providers/openai/streaming.rs

llm/providers/openai/
streaming.rs
