rig/providers/cohere/
streaming.rs

1use crate::completion::{CompletionError, CompletionRequest, GetTokenUsage};
2use crate::http_client::HttpClientExt;
3use crate::http_client::sse::{Event, GenericEventSource};
4use crate::providers::cohere::CompletionModel;
5use crate::providers::cohere::completion::{
6    AssistantContent, Message, ToolCall, ToolCallFunction, ToolType, Usage,
7};
8use crate::streaming::RawStreamingChoice;
9use crate::telemetry::SpanCombinator;
10use crate::{json_utils, streaming};
11use async_stream::stream;
12use futures::StreamExt;
13use http::Method;
14use serde::{Deserialize, Serialize};
15use tracing::info_span;
16use tracing_futures::Instrument;
17
18#[derive(Debug, Deserialize)]
19#[serde(rename_all = "kebab-case", tag = "type")]
20enum StreamingEvent {
21    MessageStart,
22    ContentStart,
23    ContentDelta { delta: Option<Delta> },
24    ContentEnd,
25    ToolPlan,
26    ToolCallStart { delta: Option<Delta> },
27    ToolCallDelta { delta: Option<Delta> },
28    ToolCallEnd,
29    MessageEnd { delta: Option<MessageEndDelta> },
30}
31
32#[derive(Debug, Deserialize)]
33struct MessageContentDelta {
34    text: Option<String>,
35}
36
37#[derive(Debug, Deserialize)]
38struct MessageToolFunctionDelta {
39    name: Option<String>,
40    arguments: Option<String>,
41}
42
43#[derive(Debug, Deserialize)]
44struct MessageToolCallDelta {
45    id: Option<String>,
46    function: Option<MessageToolFunctionDelta>,
47}
48
49#[derive(Debug, Deserialize)]
50struct MessageDelta {
51    content: Option<MessageContentDelta>,
52    tool_calls: Option<MessageToolCallDelta>,
53}
54
55#[derive(Debug, Deserialize)]
56struct Delta {
57    message: Option<MessageDelta>,
58}
59
60#[derive(Debug, Deserialize)]
61struct MessageEndDelta {
62    usage: Option<Usage>,
63}
64
65#[derive(Clone, Serialize, Deserialize)]
66pub struct StreamingCompletionResponse {
67    pub usage: Option<Usage>,
68}
69
70impl GetTokenUsage for StreamingCompletionResponse {
71    fn token_usage(&self) -> Option<crate::completion::Usage> {
72        let tokens = self
73            .usage
74            .clone()
75            .and_then(|response| response.tokens)
76            .map(|tokens| {
77                (
78                    tokens.input_tokens.map(|x| x as u64),
79                    tokens.output_tokens.map(|y| y as u64),
80                )
81            });
82        let Some((Some(input), Some(output))) = tokens else {
83            return None;
84        };
85        let mut usage = crate::completion::Usage::new();
86        usage.input_tokens = input;
87        usage.output_tokens = output;
88        usage.total_tokens = input + output;
89
90        Some(usage)
91    }
92}
93
94impl<T> CompletionModel<T>
95where
96    T: HttpClientExt + Clone + 'static,
97{
98    pub(crate) async fn stream(
99        &self,
100        request: CompletionRequest,
101    ) -> Result<streaming::StreamingCompletionResponse<StreamingCompletionResponse>, CompletionError>
102    {
103        let request = self.create_completion_request(request)?;
104        let span = if tracing::Span::current().is_disabled() {
105            info_span!(
106                target: "rig::completions",
107                "chat_streaming",
108                gen_ai.operation.name = "chat_streaming",
109                gen_ai.provider.name = "cohere",
110                gen_ai.request.model = self.model,
111                gen_ai.response.id = tracing::field::Empty,
112                gen_ai.response.model = self.model,
113                gen_ai.usage.output_tokens = tracing::field::Empty,
114                gen_ai.usage.input_tokens = tracing::field::Empty,
115                gen_ai.input.messages = serde_json::to_string(&request.get("messages").unwrap()).unwrap(),
116                gen_ai.output.messages = tracing::field::Empty,
117            )
118        } else {
119            tracing::Span::current()
120        };
121
122        let request = json_utils::merge(request, serde_json::json!({"stream": true}));
123
124        tracing::debug!(
125            "Cohere streaming completion input: {}",
126            serde_json::to_string_pretty(&request)?
127        );
128
129        let body = serde_json::to_vec(&request)?;
130
131        let req = self
132            .client
133            .req(Method::POST, "/v2/chat")?
134            .body(body)
135            .unwrap();
136
137        let mut event_source = GenericEventSource::new(self.client.http_client(), req);
138
139        let stream = stream! {
140            let mut current_tool_call: Option<(String, String, String)> = None;
141            let mut text_response = String::new();
142            let mut tool_calls = Vec::new();
143
144            while let Some(event_result) = event_source.next().await {
145                match event_result {
146                    Ok(Event::Open) => {
147                        tracing::trace!("SSE connection opened");
148                        continue;
149                    }
150
151                    Ok(Event::Message(message)) => {
152                        let data_str = message.data.trim();
153                        if data_str.is_empty() || data_str == "[DONE]" {
154                            continue;
155                        }
156
157                        let event: StreamingEvent = match serde_json::from_str(data_str) {
158                            Ok(ev) => ev,
159                            Err(_) => {
160                                tracing::debug!("Couldn't parse SSE payload as StreamingEvent");
161                                continue;
162                            }
163                        };
164
165                        match event {
166                            StreamingEvent::ContentDelta { delta: Some(delta) } => {
167                                let Some(message) = &delta.message else { continue; };
168                                let Some(content) = &message.content else { continue; };
169                                let Some(text) = &content.text else { continue; };
170
171                                text_response += text;
172
173                                yield Ok(RawStreamingChoice::Message(text.clone()));
174                            },
175
176                            StreamingEvent::MessageEnd { delta: Some(delta) } => {
177                                let message = Message::Assistant {
178                                    tool_calls: tool_calls.clone(),
179                                    content: vec![AssistantContent::Text { text: text_response.clone() }],
180                                    tool_plan: None,
181                                    citations: vec![]
182                                };
183
184                                let span = tracing::Span::current();
185                                span.record_token_usage(&delta.usage);
186                                span.record_model_output(&vec![message]);
187
188                                yield Ok(RawStreamingChoice::FinalResponse(StreamingCompletionResponse {
189                                    usage: delta.usage.clone()
190                                }));
191                            },
192
193                            StreamingEvent::ToolCallStart { delta: Some(delta) } => {
194                                let Some(message) = &delta.message else { continue; };
195                                let Some(tool_calls) = &message.tool_calls else { continue; };
196                                let Some(id) = tool_calls.id.clone() else { continue; };
197                                let Some(function) = &tool_calls.function else { continue; };
198                                let Some(name) = function.name.clone() else { continue; };
199                                let Some(arguments) = function.arguments.clone() else { continue; };
200
201                                current_tool_call = Some((id, name, arguments));
202                            },
203
204                            StreamingEvent::ToolCallDelta { delta: Some(delta) } => {
205                                let Some(message) = &delta.message else { continue; };
206                                let Some(tool_calls) = &message.tool_calls else { continue; };
207                                let Some(function) = &tool_calls.function else { continue; };
208                                let Some(arguments) = function.arguments.clone() else { continue; };
209
210                                let Some(tc) = current_tool_call.clone() else { continue; };
211                                current_tool_call = Some((tc.0, tc.1, format!("{}{}", tc.2, arguments)));
212                            },
213
214                            StreamingEvent::ToolCallEnd => {
215                                let Some(tc) = current_tool_call.clone() else { continue; };
216                                let Ok(args) = serde_json::from_str::<serde_json::Value>(&tc.2) else { continue; };
217
218                                tool_calls.push(ToolCall {
219                                    id: Some(tc.0.clone()),
220                                    r#type: Some(ToolType::Function),
221                                    function: Some(ToolCallFunction {
222                                        name: tc.1.clone(),
223                                        arguments: args.clone()
224                                    })
225                                });
226
227                                yield Ok(RawStreamingChoice::ToolCall {
228                                    id: tc.0,
229                                    name: tc.1,
230                                    arguments: args,
231                                    call_id: None
232                                });
233
234                                current_tool_call = None;
235                            },
236
237                            _ => {}
238                        }
239                    },
240                    Err(crate::http_client::Error::StreamEnded) => {
241                        break;
242                    }
243                    Err(err) => {
244                        tracing::error!(?err, "SSE error");
245                        yield Err(CompletionError::ResponseError(err.to_string()));
246                        break;
247                    }
248                }
249            }
250
251            event_source.close();
252        }.instrument(span);
253
254        Ok(streaming::StreamingCompletionResponse::stream(Box::pin(
255            stream,
256        )))
257    }
258}
rig/providers/cohere/streaming.rs

rig/providers/cohere/
streaming.rs
