rig/providers/openai/completion/
streaming.rs

1use crate::completion::{CompletionError, CompletionRequest, GetTokenUsage};
2use crate::http_client::HttpClientExt;
3use crate::http_client::sse::{Event, GenericEventSource};
4use crate::json_utils;
5use crate::json_utils::merge;
6use crate::providers::openai::completion::{CompletionModel, Usage};
7use crate::streaming;
8use crate::streaming::RawStreamingChoice;
9use async_stream::stream;
10use futures::StreamExt;
11use http::Request;
12use serde::{Deserialize, Serialize};
13use serde_json::json;
14use std::collections::HashMap;
15use tracing::{debug, info_span};
16use tracing_futures::Instrument;
17
18// ================================================================
19// OpenAI Completion Streaming API
20// ================================================================
21#[derive(Debug, Serialize, Deserialize, Clone)]
22pub struct StreamingFunction {
23    #[serde(default)]
24    pub name: Option<String>,
25    #[serde(default)]
26    pub arguments: String,
27}
28
29#[derive(Debug, Serialize, Deserialize, Clone)]
30pub struct StreamingToolCall {
31    pub index: usize,
32    pub id: Option<String>,
33    pub function: StreamingFunction,
34}
35
36#[derive(Deserialize, Debug)]
37struct StreamingDelta {
38    #[serde(default)]
39    content: Option<String>,
40    #[serde(default, deserialize_with = "json_utils::null_or_vec")]
41    tool_calls: Vec<StreamingToolCall>,
42}
43
44#[derive(Deserialize, Debug)]
45struct StreamingChoice {
46    delta: StreamingDelta,
47}
48
49#[derive(Deserialize, Debug)]
50struct StreamingCompletionChunk {
51    choices: Vec<StreamingChoice>,
52    usage: Option<Usage>,
53}
54
55#[derive(Clone, Serialize, Deserialize)]
56pub struct StreamingCompletionResponse {
57    pub usage: Usage,
58}
59
60impl GetTokenUsage for StreamingCompletionResponse {
61    fn token_usage(&self) -> Option<crate::completion::Usage> {
62        let mut usage = crate::completion::Usage::new();
63        usage.input_tokens = self.usage.prompt_tokens as u64;
64        usage.output_tokens = self.usage.total_tokens as u64 - self.usage.prompt_tokens as u64;
65        usage.total_tokens = self.usage.total_tokens as u64;
66        Some(usage)
67    }
68}
69
70impl CompletionModel<reqwest::Client> {
71    pub(crate) async fn stream(
72        &self,
73        completion_request: CompletionRequest,
74    ) -> Result<streaming::StreamingCompletionResponse<StreamingCompletionResponse>, CompletionError>
75    {
76        let request = super::CompletionRequest::try_from((self.model.clone(), completion_request))?;
77        let request_messages = serde_json::to_string(&request.messages)
78            .expect("Converting to JSON from a Rust struct shouldn't fail");
79        let mut request_as_json = serde_json::to_value(request).expect("this should never fail");
80
81        request_as_json = merge(
82            request_as_json,
83            json!({"stream": true, "stream_options": {"include_usage": true}}),
84        );
85
86        let req_body = serde_json::to_vec(&request_as_json)?;
87
88        let req = self
89            .client
90            .post("/chat/completions")?
91            .header("Content-Type", "application/json")
92            .body(req_body)
93            .map_err(|e| CompletionError::HttpError(e.into()))?;
94
95        let span = if tracing::Span::current().is_disabled() {
96            info_span!(
97                target: "rig::completions",
98                "chat",
99                gen_ai.operation.name = "chat",
100                gen_ai.provider.name = "openai",
101                gen_ai.request.model = self.model,
102                gen_ai.response.id = tracing::field::Empty,
103                gen_ai.response.model = self.model,
104                gen_ai.usage.output_tokens = tracing::field::Empty,
105                gen_ai.usage.input_tokens = tracing::field::Empty,
106                gen_ai.input.messages = request_messages,
107                gen_ai.output.messages = tracing::field::Empty,
108            )
109        } else {
110            tracing::Span::current()
111        };
112
113        tracing::Instrument::instrument(
114            send_compatible_streaming_request(self.client.http_client.clone(), req),
115            span,
116        )
117        .await
118    }
119}
120
121pub async fn send_compatible_streaming_request<T>(
122    http_client: T,
123    req: Request<Vec<u8>>,
124) -> Result<streaming::StreamingCompletionResponse<StreamingCompletionResponse>, CompletionError>
125where
126    T: HttpClientExt + Clone + 'static,
127{
128    let span = tracing::Span::current();
129    // Build the request with proper headers for SSE
130    let mut event_source = GenericEventSource::new(http_client, req);
131
132    let stream = stream! {
133        let span = tracing::Span::current();
134        let mut final_usage = Usage::new();
135
136        // Track in-progress tool calls
137        let mut tool_calls: HashMap<usize, (String, String, String)> = HashMap::new();
138
139        let mut text_content = String::new();
140
141        while let Some(event_result) = event_source.next().await {
142            match event_result {
143                Ok(Event::Open) => {
144                    tracing::trace!("SSE connection opened");
145                    continue;
146                }
147                Ok(Event::Message(message)) => {
148                    if message.data.trim().is_empty() || message.data == "[DONE]" {
149                        continue;
150                    }
151
152                    let data = serde_json::from_str::<StreamingCompletionChunk>(&message.data);
153                    let Ok(data) = data else {
154                        let err = data.unwrap_err();
155                        debug!("Couldn't serialize data as StreamingCompletionChunk: {:?}", err);
156                        continue;
157                    };
158
159                    if let Some(choice) = data.choices.first() {
160                        let delta = &choice.delta;
161
162                        // Tool calls
163                        if !delta.tool_calls.is_empty() {
164                            for tool_call in &delta.tool_calls {
165                                let function = tool_call.function.clone();
166
167                                // Start of tool call
168                                if function.name.is_some() && function.arguments.is_empty() {
169                                    let id = tool_call.id.clone().unwrap_or_default();
170                                    tool_calls.insert(
171                                        tool_call.index,
172                                        (id, function.name.clone().unwrap(), "".to_string()),
173                                    );
174                                }
175                                // tool call partial (ie, a continuation of a previously received tool call)
176                                // name: None or Empty String
177                                // arguments: Some(String)
178                                else if function.name.clone().is_none_or(|s| s.is_empty())
179                                    && !function.arguments.is_empty()
180                                {
181                                    if let Some((id, name, arguments)) =
182                                        tool_calls.get(&tool_call.index)
183                                    {
184                                        let new_arguments = &tool_call.function.arguments;
185                                        let arguments = format!("{arguments}{new_arguments}");
186                                        tool_calls.insert(
187                                            tool_call.index,
188                                            (id.clone(), name.clone(), arguments),
189                                        );
190                                    } else {
191                                        debug!("Partial tool call received but tool call was never started.");
192                                    }
193                                }
194                                // Complete tool call
195                                else {
196                                    let id = tool_call.id.clone().unwrap_or_default();
197                                    let name = function.name.expect("tool call should have a name");
198                                    let arguments = function.arguments;
199                                    let Ok(arguments) = serde_json::from_str(&arguments) else {
200                                        debug!("Couldn't serialize '{arguments}' as JSON");
201                                        continue;
202                                    };
203
204                                    yield Ok(streaming::RawStreamingChoice::ToolCall {
205                                        id,
206                                        name,
207                                        arguments,
208                                        call_id: None,
209                                    });
210                                }
211                            }
212                        }
213
214                        // Message content
215                        if let Some(content) = &choice.delta.content {
216                            text_content += content;
217                            yield Ok(streaming::RawStreamingChoice::Message(content.clone()))
218                        }
219                    }
220
221                    // Usage updates
222                    if let Some(usage) = data.usage {
223                        final_usage = usage.clone();
224                    }
225                }
226                Err(crate::http_client::Error::StreamEnded) => {
227                    break;
228                }
229                Err(error) => {
230                    tracing::error!(?error, "SSE error");
231                    yield Err(CompletionError::ResponseError(error.to_string()));
232                    break;
233                }
234            }
235        }
236
237        // Ensure event source is closed when stream ends
238        event_source.close();
239
240        let mut vec_toolcalls = vec![];
241
242        // Flush any tool calls that weren’t fully yielded
243        for (_, (id, name, arguments)) in tool_calls {
244            let Ok(arguments) = serde_json::from_str::<serde_json::Value>(&arguments) else {
245                continue;
246            };
247
248            vec_toolcalls.push(super::ToolCall {
249                r#type: super::ToolType::Function,
250                id: id.clone(),
251                function: super::Function {
252                    name: name.clone(), arguments: arguments.clone()
253                },
254            });
255
256            yield Ok(RawStreamingChoice::ToolCall {
257                id,
258                name,
259                arguments,
260                call_id: None,
261            });
262        }
263
264        let message_output = super::Message::Assistant {
265            content: vec![super::AssistantContent::Text { text: text_content }],
266            refusal: None,
267            audio: None,
268            name: None,
269            tool_calls: vec_toolcalls
270        };
271
272        span.record("gen_ai.usage.input_tokens", final_usage.prompt_tokens);
273        span.record("gen_ai.usage.output_tokens", final_usage.total_tokens - final_usage.prompt_tokens);
274        span.record("gen_ai.output.messages", serde_json::to_string(&vec![message_output]).expect("Converting from a Rust struct should always convert to JSON without failing"));
275
276        yield Ok(RawStreamingChoice::FinalResponse(StreamingCompletionResponse {
277            usage: final_usage.clone()
278        }));
279    }.instrument(span);
280
281    Ok(streaming::StreamingCompletionResponse::stream(Box::pin(
282        stream,
283    )))
284}
rig/providers/openai/completion/streaming.rs

rig/providers/openai/completion/
streaming.rs
