Skip to main content

rig_core/providers/cohere/
completion.rs

1use crate::{
2    OneOrMany,
3    completion::{self, CompletionError, GetTokenUsage},
4    http_client::{self, HttpClientExt},
5    json_utils,
6    message::{self, Reasoning, ToolChoice},
7    telemetry::SpanCombinator,
8};
9use std::collections::HashMap;
10
11use super::client::Client;
12use crate::completion::CompletionRequest;
13use crate::providers::cohere::streaming::StreamingCompletionResponse;
14use serde::{Deserialize, Serialize};
15use tracing::{Instrument, Level, enabled, info_span};
16
17#[derive(Debug, Deserialize, Serialize)]
18pub struct CompletionResponse {
19    pub id: String,
20    pub finish_reason: FinishReason,
21    message: Message,
22    #[serde(default)]
23    pub usage: Option<Usage>,
24}
25
26type AssistantMessageParts = (Vec<AssistantContent>, Vec<Citation>, Vec<ToolCall>);
27
28impl CompletionResponse {
29    /// Return that parts of the response for assistant messages w/o dealing with the other variants
30    pub fn message(&self) -> Result<AssistantMessageParts, CompletionError> {
31        let Message::Assistant {
32            content,
33            citations,
34            tool_calls,
35            ..
36        } = self.message.clone()
37        else {
38            return Err(CompletionError::ResponseError(
39                "completion response did not contain an assistant message".into(),
40            ));
41        };
42
43        Ok((content, citations, tool_calls))
44    }
45}
46
47impl crate::telemetry::ProviderResponseExt for CompletionResponse {
48    type OutputMessage = Message;
49    type Usage = Usage;
50
51    fn get_response_id(&self) -> Option<String> {
52        Some(self.id.clone())
53    }
54
55    fn get_response_model_name(&self) -> Option<String> {
56        None
57    }
58
59    fn get_output_messages(&self) -> Vec<Self::OutputMessage> {
60        vec![self.message.clone()]
61    }
62
63    fn get_text_response(&self) -> Option<String> {
64        let Message::Assistant { ref content, .. } = self.message else {
65            return None;
66        };
67
68        let res = content
69            .iter()
70            .filter_map(|x| {
71                if let AssistantContent::Text { text } = x {
72                    Some(text.to_string())
73                } else {
74                    None
75                }
76            })
77            .collect::<Vec<String>>()
78            .join("\n");
79
80        if res.is_empty() { None } else { Some(res) }
81    }
82
83    fn get_usage(&self) -> Option<Self::Usage> {
84        self.usage.clone()
85    }
86}
87
88#[derive(Debug, Deserialize, PartialEq, Eq, Clone, Serialize)]
89#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
90pub enum FinishReason {
91    MaxTokens,
92    StopSequence,
93    Complete,
94    Error,
95    ToolCall,
96}
97
98#[derive(Debug, Deserialize, Clone, Serialize)]
99pub struct Usage {
100    #[serde(default)]
101    pub billed_units: Option<BilledUnits>,
102    #[serde(default)]
103    pub tokens: Option<Tokens>,
104}
105
106impl GetTokenUsage for Usage {
107    fn token_usage(&self) -> Option<crate::completion::Usage> {
108        let mut usage = crate::completion::Usage::new();
109
110        if let Some(ref billed_units) = self.billed_units {
111            usage.input_tokens = billed_units.input_tokens.unwrap_or_default() as u64;
112            usage.output_tokens = billed_units.output_tokens.unwrap_or_default() as u64;
113            usage.total_tokens = usage.input_tokens + usage.output_tokens;
114        }
115
116        Some(usage)
117    }
118}
119
120#[derive(Debug, Deserialize, Clone, Serialize)]
121pub struct BilledUnits {
122    #[serde(default)]
123    pub output_tokens: Option<f64>,
124    #[serde(default)]
125    pub classifications: Option<f64>,
126    #[serde(default)]
127    pub search_units: Option<f64>,
128    #[serde(default)]
129    pub input_tokens: Option<f64>,
130}
131
132#[derive(Debug, Deserialize, Clone, Serialize)]
133pub struct Tokens {
134    #[serde(default)]
135    pub input_tokens: Option<f64>,
136    #[serde(default)]
137    pub output_tokens: Option<f64>,
138}
139
140impl TryFrom<CompletionResponse> for completion::CompletionResponse<CompletionResponse> {
141    type Error = CompletionError;
142
143    fn try_from(response: CompletionResponse) -> Result<Self, Self::Error> {
144        let (content, _, tool_calls) = response.message()?;
145
146        let model_response = if !tool_calls.is_empty() {
147            OneOrMany::many(
148                tool_calls
149                    .into_iter()
150                    .filter_map(|tool_call| {
151                        let ToolCallFunction { name, arguments } = tool_call.function?;
152                        let id = tool_call.id.unwrap_or_else(|| name.clone());
153
154                        Some(completion::AssistantContent::tool_call(id, name, arguments))
155                    })
156                    .collect::<Vec<_>>(),
157            )
158            .map_err(|_| {
159                CompletionError::ResponseError(
160                    "response contained tool call metadata without any callable tool content"
161                        .to_owned(),
162                )
163            })?
164        } else {
165            OneOrMany::many(content.into_iter().map(|content| match content {
166                AssistantContent::Text { text } => completion::AssistantContent::text(text),
167                AssistantContent::Thinking { thinking } => {
168                    completion::AssistantContent::Reasoning(Reasoning::new(&thinking))
169                }
170            }))
171            .map_err(|_| {
172                CompletionError::ResponseError(
173                    "Response contained no message or tool call (empty)".to_owned(),
174                )
175            })?
176        };
177
178        let usage = response
179            .usage
180            .as_ref()
181            .and_then(|usage| usage.tokens.as_ref())
182            .map(|tokens| {
183                let input_tokens = tokens.input_tokens.unwrap_or(0.0);
184                let output_tokens = tokens.output_tokens.unwrap_or(0.0);
185
186                completion::Usage {
187                    input_tokens: input_tokens as u64,
188                    output_tokens: output_tokens as u64,
189                    total_tokens: (input_tokens + output_tokens) as u64,
190                    cached_input_tokens: 0,
191                    cache_creation_input_tokens: 0,
192                    reasoning_tokens: 0,
193                }
194            })
195            .unwrap_or_default();
196
197        Ok(completion::CompletionResponse {
198            choice: model_response,
199            usage,
200            raw_response: response,
201            message_id: None,
202        })
203    }
204}
205
206#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
207pub struct Document {
208    pub id: String,
209    pub data: HashMap<String, serde_json::Value>,
210}
211
212impl From<completion::Document> for Document {
213    fn from(document: completion::Document) -> Self {
214        let mut data: HashMap<String, serde_json::Value> = HashMap::new();
215
216        // We use `.into()` here explicitly since the `document.additional_props` type will likely
217        //  evolve into `serde_json::Value` in the future.
218        document
219            .additional_props
220            .into_iter()
221            .for_each(|(key, value)| {
222                data.insert(key, value.into());
223            });
224
225        data.insert("text".to_string(), document.text.into());
226
227        Self {
228            id: document.id,
229            data,
230        }
231    }
232}
233
234#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
235pub struct ToolCall {
236    #[serde(default)]
237    pub id: Option<String>,
238    #[serde(default)]
239    pub r#type: Option<ToolType>,
240    #[serde(default)]
241    pub function: Option<ToolCallFunction>,
242}
243
244#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
245pub struct ToolCallFunction {
246    pub name: String,
247    #[serde(with = "json_utils::stringified_json")]
248    pub arguments: serde_json::Value,
249}
250
251#[derive(Clone, Default, Debug, Deserialize, Serialize, PartialEq, Eq)]
252#[serde(rename_all = "lowercase")]
253pub enum ToolType {
254    #[default]
255    Function,
256}
257
258#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
259pub struct Tool {
260    pub r#type: ToolType,
261    pub function: Function,
262}
263
264#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
265pub struct Function {
266    pub name: String,
267    #[serde(default)]
268    pub description: Option<String>,
269    pub parameters: serde_json::Value,
270}
271
272impl From<completion::ToolDefinition> for Tool {
273    fn from(tool: completion::ToolDefinition) -> Self {
274        Self {
275            r#type: ToolType::default(),
276            function: Function {
277                name: tool.name,
278                description: Some(tool.description),
279                parameters: tool.parameters,
280            },
281        }
282    }
283}
284
285#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
286#[serde(tag = "role", rename_all = "lowercase")]
287pub enum Message {
288    User {
289        content: OneOrMany<UserContent>,
290    },
291
292    Assistant {
293        #[serde(default)]
294        content: Vec<AssistantContent>,
295        #[serde(default)]
296        citations: Vec<Citation>,
297        #[serde(default)]
298        tool_calls: Vec<ToolCall>,
299        #[serde(default)]
300        tool_plan: Option<String>,
301    },
302
303    Tool {
304        content: OneOrMany<ToolResultContent>,
305        tool_call_id: String,
306    },
307
308    System {
309        content: String,
310    },
311}
312
313#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
314#[serde(tag = "type", rename_all = "lowercase")]
315pub enum UserContent {
316    Text { text: String },
317    ImageUrl { image_url: ImageUrl },
318}
319
320#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
321#[serde(tag = "type", rename_all = "lowercase")]
322pub enum AssistantContent {
323    Text { text: String },
324    Thinking { thinking: String },
325}
326
327#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
328pub struct ImageUrl {
329    pub url: String,
330}
331
332#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
333pub enum ToolResultContent {
334    Text { text: String },
335    Document { document: Document },
336}
337
338#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
339pub struct Citation {
340    #[serde(default)]
341    pub start: Option<u32>,
342    #[serde(default)]
343    pub end: Option<u32>,
344    #[serde(default)]
345    pub text: Option<String>,
346    #[serde(rename = "type")]
347    pub citation_type: Option<CitationType>,
348    #[serde(default)]
349    pub sources: Vec<Source>,
350}
351
352#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
353#[serde(tag = "type", rename_all = "lowercase")]
354pub enum Source {
355    Document {
356        id: Option<String>,
357        document: Option<serde_json::Map<String, serde_json::Value>>,
358    },
359    Tool {
360        id: Option<String>,
361        tool_output: Option<serde_json::Map<String, serde_json::Value>>,
362    },
363}
364
365#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
366#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
367pub enum CitationType {
368    TextContent,
369    Plan,
370}
371
372impl TryFrom<message::Message> for Vec<Message> {
373    type Error = message::MessageError;
374
375    fn try_from(message: message::Message) -> Result<Self, Self::Error> {
376        Ok(match message {
377            message::Message::User { content } => content
378                .into_iter()
379                .map(|content| match content {
380                    message::UserContent::Text(message::Text { text }) => Ok(Message::User {
381                        content: OneOrMany::one(UserContent::Text { text }),
382                    }),
383                    message::UserContent::ToolResult(message::ToolResult {
384                        id, content, ..
385                    }) => Ok(Message::Tool {
386                        tool_call_id: id,
387                        content: content.try_map(|content| match content {
388                            message::ToolResultContent::Text(text) => {
389                                Ok(ToolResultContent::Text { text: text.text })
390                            }
391                            _ => Err(message::MessageError::ConversionError(
392                                "Only text tool result content is supported by Cohere".to_owned(),
393                            )),
394                        })?,
395                    }),
396                    _ => Err(message::MessageError::ConversionError(
397                        "Only text content is supported by Cohere".to_owned(),
398                    )),
399                })
400                .collect::<Result<Vec<_>, _>>()?,
401            message::Message::System { content } => {
402                vec![Message::System { content }]
403            }
404            message::Message::Assistant { content, .. } => {
405                let mut text_content = vec![];
406                let mut tool_calls = vec![];
407
408                for content in content.into_iter() {
409                    match content {
410                        message::AssistantContent::Text(message::Text { text }) => {
411                            text_content.push(AssistantContent::Text { text });
412                        }
413                        message::AssistantContent::ToolCall(message::ToolCall {
414                            id,
415                            function:
416                                message::ToolFunction {
417                                    name, arguments, ..
418                                },
419                            ..
420                        }) => {
421                            tool_calls.push(ToolCall {
422                                id: Some(id),
423                                r#type: Some(ToolType::Function),
424                                function: Some(ToolCallFunction {
425                                    name,
426                                    arguments: serde_json::to_value(arguments).unwrap_or_default(),
427                                }),
428                            });
429                        }
430                        message::AssistantContent::Reasoning(reasoning) => {
431                            let thinking = reasoning.display_text();
432                            text_content.push(AssistantContent::Thinking { thinking });
433                        }
434                        message::AssistantContent::Image(_) => {
435                            return Err(message::MessageError::ConversionError(
436                                "Cohere currently doesn't support images.".to_owned(),
437                            ));
438                        }
439                    }
440                }
441
442                vec![Message::Assistant {
443                    content: text_content,
444                    citations: vec![],
445                    tool_calls,
446                    tool_plan: None,
447                }]
448            }
449        })
450    }
451}
452
453impl TryFrom<Message> for message::Message {
454    type Error = message::MessageError;
455
456    fn try_from(message: Message) -> Result<Self, Self::Error> {
457        match message {
458            Message::User { content } => Ok(message::Message::User {
459                content: content.map(|content| match content {
460                    UserContent::Text { text } => {
461                        message::UserContent::Text(message::Text { text })
462                    }
463                    UserContent::ImageUrl { image_url } => {
464                        message::UserContent::image_url(image_url.url, None, None)
465                    }
466                }),
467            }),
468            Message::Assistant {
469                content,
470                tool_calls,
471                ..
472            } => {
473                let mut content = content
474                    .into_iter()
475                    .map(|content| match content {
476                        AssistantContent::Text { text } => message::AssistantContent::text(text),
477                        AssistantContent::Thinking { thinking } => {
478                            message::AssistantContent::Reasoning(Reasoning::new(&thinking))
479                        }
480                    })
481                    .collect::<Vec<_>>();
482
483                content.extend(tool_calls.into_iter().filter_map(|tool_call| {
484                    let ToolCallFunction { name, arguments } = tool_call.function?;
485
486                    Some(message::AssistantContent::tool_call(
487                        tool_call.id.unwrap_or_else(|| name.clone()),
488                        name,
489                        arguments,
490                    ))
491                }));
492
493                let content = OneOrMany::many(content).map_err(|_| {
494                    message::MessageError::ConversionError(
495                        "Expected either text content or tool calls".to_string(),
496                    )
497                })?;
498
499                Ok(message::Message::Assistant { id: None, content })
500            }
501            Message::Tool {
502                content,
503                tool_call_id,
504            } => {
505                let content = content.try_map(|content| {
506                    Ok(match content {
507                        ToolResultContent::Text { text } => message::ToolResultContent::text(text),
508                        ToolResultContent::Document { document } => {
509                            message::ToolResultContent::text(
510                                serde_json::to_string(&document.data).map_err(|e| {
511                                    message::MessageError::ConversionError(
512                                        format!("Failed to convert tool result document content into text: {e}"),
513                                    )
514                                })?,
515                            )
516                        }
517                    })
518                })?;
519
520                Ok(message::Message::User {
521                    content: OneOrMany::one(message::UserContent::tool_result(
522                        tool_call_id,
523                        content,
524                    )),
525                })
526            }
527            Message::System { content } => Ok(message::Message::user(content)),
528        }
529    }
530}
531
532#[derive(Clone)]
533pub struct CompletionModel<T = reqwest::Client> {
534    pub(crate) client: Client<T>,
535    pub model: String,
536}
537
538#[derive(Debug, Serialize, Deserialize)]
539pub(super) struct CohereCompletionRequest {
540    model: String,
541    pub messages: Vec<Message>,
542    documents: Vec<crate::completion::Document>,
543    #[serde(skip_serializing_if = "Option::is_none")]
544    temperature: Option<f64>,
545    #[serde(skip_serializing_if = "Vec::is_empty")]
546    tools: Vec<Tool>,
547    #[serde(skip_serializing_if = "Option::is_none")]
548    tool_choice: Option<ToolChoice>,
549    #[serde(flatten, skip_serializing_if = "Option::is_none")]
550    pub additional_params: Option<serde_json::Value>,
551}
552
553impl TryFrom<(&str, CompletionRequest)> for CohereCompletionRequest {
554    type Error = CompletionError;
555
556    fn try_from((model, req): (&str, CompletionRequest)) -> Result<Self, Self::Error> {
557        if req.output_schema.is_some() {
558            tracing::warn!("Structured outputs currently not supported for Cohere");
559        }
560
561        let model = req.model.clone().unwrap_or_else(|| model.to_string());
562        let mut partial_history = vec![];
563        if let Some(docs) = req.normalized_documents() {
564            partial_history.push(docs);
565        }
566        partial_history.extend(req.chat_history);
567
568        let mut full_history: Vec<Message> = req.preamble.map_or_else(Vec::new, |preamble| {
569            vec![Message::System { content: preamble }]
570        });
571
572        full_history.extend(
573            partial_history
574                .into_iter()
575                .map(message::Message::try_into)
576                .collect::<Result<Vec<Vec<Message>>, _>>()?
577                .into_iter()
578                .flatten()
579                .collect::<Vec<_>>(),
580        );
581
582        let tool_choice = if let Some(tool_choice) = req.tool_choice {
583            if !matches!(tool_choice, ToolChoice::Auto) {
584                Some(tool_choice)
585            } else {
586                return Err(CompletionError::RequestError(
587                    "\"auto\" is not an allowed tool_choice value in the Cohere API".into(),
588                ));
589            }
590        } else {
591            None
592        };
593
594        Ok(Self {
595            model: model.to_string(),
596            messages: full_history,
597            documents: req.documents,
598            temperature: req.temperature,
599            tools: req.tools.into_iter().map(Tool::from).collect::<Vec<_>>(),
600            tool_choice,
601            additional_params: req.additional_params,
602        })
603    }
604}
605
606impl<T> CompletionModel<T>
607where
608    T: HttpClientExt,
609{
610    pub fn new(client: Client<T>, model: impl Into<String>) -> Self {
611        Self {
612            client,
613            model: model.into(),
614        }
615    }
616}
617
618impl<T> completion::CompletionModel for CompletionModel<T>
619where
620    T: HttpClientExt + Clone + 'static,
621{
622    type Response = CompletionResponse;
623    type StreamingResponse = StreamingCompletionResponse;
624    type Client = Client<T>;
625
626    fn make(client: &Self::Client, model: impl Into<String>) -> Self {
627        Self::new(client.clone(), model.into())
628    }
629
630    async fn completion(
631        &self,
632        completion_request: completion::CompletionRequest,
633    ) -> Result<completion::CompletionResponse<CompletionResponse>, CompletionError> {
634        let request = CohereCompletionRequest::try_from((self.model.as_ref(), completion_request))?;
635
636        let llm_span = if tracing::Span::current().is_disabled() {
637            info_span!(
638            target: "rig::completions",
639            "chat",
640            gen_ai.operation.name = "chat",
641            gen_ai.provider.name = "cohere",
642            gen_ai.request.model = self.model,
643            gen_ai.response.id = tracing::field::Empty,
644            gen_ai.response.model = self.model,
645            gen_ai.usage.output_tokens = tracing::field::Empty,
646            gen_ai.usage.input_tokens = tracing::field::Empty,
647            gen_ai.usage.cache_read.input_tokens = tracing::field::Empty,
648            )
649        } else {
650            tracing::Span::current()
651        };
652
653        if enabled!(Level::TRACE) {
654            tracing::trace!(
655                "Cohere completion request: {}",
656                serde_json::to_string_pretty(&request)?
657            );
658        }
659
660        let req_body = serde_json::to_vec(&request)?;
661
662        let req = self
663            .client
664            .post("/v2/chat")?
665            .body(req_body)
666            .map_err(|e| CompletionError::HttpError(e.into()))?;
667
668        async {
669            let response = self
670                .client
671                .send::<_, bytes::Bytes>(req)
672                .await
673                .map_err(|e| http_client::Error::Instance(e.into()))?;
674
675            let status = response.status();
676            let body = response.into_body().into_future().await?.to_owned();
677
678            if status.is_success() {
679                let json_response: CompletionResponse = serde_json::from_slice(&body)?;
680                let span = tracing::Span::current();
681                span.record_token_usage(&json_response.usage);
682                span.record_response_metadata(&json_response);
683
684                if enabled!(Level::TRACE) {
685                    tracing::trace!(
686                        target: "rig::completions",
687                        "Cohere completion response: {}",
688                        serde_json::to_string_pretty(&json_response)?
689                    );
690                }
691
692                let completion: completion::CompletionResponse<CompletionResponse> =
693                    json_response.try_into()?;
694                Ok(completion)
695            } else {
696                Err(CompletionError::ProviderError(
697                    String::from_utf8_lossy(&body).to_string(),
698                ))
699            }
700        }
701        .instrument(llm_span)
702        .await
703    }
704
705    async fn stream(
706        &self,
707        request: CompletionRequest,
708    ) -> Result<
709        crate::streaming::StreamingCompletionResponse<Self::StreamingResponse>,
710        CompletionError,
711    > {
712        CompletionModel::stream(self, request).await
713    }
714}
715#[cfg(test)]
716mod tests {
717    use super::*;
718    use serde_path_to_error::deserialize;
719
720    #[test]
721    fn test_deserialize_completion_response() {
722        let json_data = r#"
723        {
724            "id": "abc123",
725            "message": {
726                "role": "assistant",
727                "tool_plan": "I will use the subtract tool to find the difference between 2 and 5.",
728                "tool_calls": [
729                        {
730                            "id": "subtract_sm6ps6fb6y9f",
731                            "type": "function",
732                            "function": {
733                                "name": "subtract",
734                                "arguments": "{\"x\":5,\"y\":2}"
735                            }
736                        }
737                    ]
738                },
739                "finish_reason": "TOOL_CALL",
740                "usage": {
741                "billed_units": {
742                    "input_tokens": 78,
743                    "output_tokens": 27
744                },
745                "tokens": {
746                    "input_tokens": 1028,
747                    "output_tokens": 63
748                }
749            }
750        }
751        "#;
752
753        let mut deserializer = serde_json::Deserializer::from_str(json_data);
754        let result: Result<CompletionResponse, _> = deserialize(&mut deserializer);
755
756        let response = result.unwrap();
757        let (_, citations, tool_calls) = response.message().expect("assistant message");
758        let CompletionResponse {
759            id,
760            finish_reason,
761            usage,
762            ..
763        } = response;
764
765        assert_eq!(id, "abc123");
766        assert_eq!(finish_reason, FinishReason::ToolCall);
767
768        let Usage {
769            billed_units,
770            tokens,
771        } = usage.unwrap();
772        let BilledUnits {
773            input_tokens: billed_input_tokens,
774            output_tokens: billed_output_tokens,
775            ..
776        } = billed_units.unwrap();
777        let Tokens {
778            input_tokens,
779            output_tokens,
780        } = tokens.unwrap();
781
782        assert_eq!(billed_input_tokens.unwrap(), 78.0);
783        assert_eq!(billed_output_tokens.unwrap(), 27.0);
784        assert_eq!(input_tokens.unwrap(), 1028.0);
785        assert_eq!(output_tokens.unwrap(), 63.0);
786
787        assert!(citations.is_empty());
788        assert_eq!(tool_calls.len(), 1);
789
790        let ToolCallFunction { name, arguments } = tool_calls[0].function.clone().unwrap();
791
792        assert_eq!(name, "subtract");
793        assert_eq!(arguments, serde_json::json!({"x": 5, "y": 2}));
794    }
795
796    #[test]
797    fn test_convert_completion_message_to_message_and_back() {
798        let completion_message = completion::Message::User {
799            content: OneOrMany::one(completion::message::UserContent::Text(
800                completion::message::Text {
801                    text: "Hello, world!".to_string(),
802                },
803            )),
804        };
805
806        let messages: Vec<Message> = completion_message.clone().try_into().unwrap();
807        let _converted_back: Vec<completion::Message> = messages
808            .into_iter()
809            .map(|msg| msg.try_into().unwrap())
810            .collect::<Vec<_>>();
811    }
812
813    #[test]
814    fn test_convert_message_to_completion_message_and_back() {
815        let message = Message::User {
816            content: OneOrMany::one(UserContent::Text {
817                text: "Hello, world!".to_string(),
818            }),
819        };
820
821        let completion_message: completion::Message = message.clone().try_into().unwrap();
822        let _converted_back: Vec<Message> = completion_message.try_into().unwrap();
823    }
824}