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

llm/providers/openai/
streaming.rs
