use crate::{chat::*, common::Usage};
#[must_use = "Builder does nothing until .build() is called"]
#[derive(Clone, Debug)]
pub struct ChatCompletionResponseBuilder {
id: String,
object: String,
created: u64,
model: String,
choices: Vec<ChatChoice>,
usage: Option<Usage>,
system_fingerprint: Option<String>,
}
impl ChatCompletionResponseBuilder {
pub fn new(id: impl Into<String>, model: impl Into<String>) -> Self {
Self {
id: id.into(),
object: "chat.completion".to_string(),
created: chrono::Utc::now().timestamp() as u64,
model: model.into(),
choices: Vec::new(),
usage: None,
system_fingerprint: None,
}
}
pub fn copy_from_request(mut self, request: &ChatCompletionRequest) -> Self {
self.model.clone_from(&request.model);
self
}
pub fn object(mut self, object: impl Into<String>) -> Self {
self.object = object.into();
self
}
pub fn created(mut self, timestamp: u64) -> Self {
self.created = timestamp;
self
}
pub fn choices(mut self, choices: Vec<ChatChoice>) -> Self {
self.choices = choices;
self
}
pub fn add_choice(mut self, choice: ChatChoice) -> Self {
self.choices.push(choice);
self
}
pub fn usage(mut self, usage: Usage) -> Self {
self.usage = Some(usage);
self
}
pub fn maybe_usage(mut self, usage: Option<Usage>) -> Self {
if let Some(u) = usage {
self.usage = Some(u);
}
self
}
pub fn maybe_system_fingerprint(mut self, fingerprint: Option<impl Into<String>>) -> Self {
if let Some(fp) = fingerprint {
self.system_fingerprint = Some(fp.into());
}
self
}
pub fn build(self) -> ChatCompletionResponse {
ChatCompletionResponse {
id: self.id,
object: self.object,
created: self.created,
model: self.model,
choices: self.choices,
usage: self.usage,
system_fingerprint: self.system_fingerprint,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_build_minimal() {
let response = ChatCompletionResponse::builder("chatcmpl_123", "gpt-4").build();
assert_eq!(response.id, "chatcmpl_123");
assert_eq!(response.model, "gpt-4");
assert_eq!(response.object, "chat.completion");
assert!(response.choices.is_empty());
assert!(response.usage.is_none());
assert!(response.system_fingerprint.is_none());
}
#[test]
fn test_build_complete() {
let choice = ChatChoice {
index: 0,
message: ChatCompletionMessage {
role: "assistant".to_string(),
content: Some("Hello!".to_string()),
tool_calls: None,
reasoning_content: None,
},
logprobs: None,
finish_reason: Some("stop".to_string()),
matched_stop: None,
hidden_states: None,
};
let usage = Usage {
prompt_tokens: 10,
completion_tokens: 20,
total_tokens: 30,
prompt_tokens_details: None,
completion_tokens_details: None,
};
let response = ChatCompletionResponse::builder("chatcmpl_456", "gpt-4")
.choices(vec![choice.clone()])
.maybe_usage(Some(usage))
.maybe_system_fingerprint(Some("fp_123abc"))
.build();
assert_eq!(response.id, "chatcmpl_456");
assert_eq!(response.choices.len(), 1);
assert_eq!(response.choices[0].index, 0);
assert!(response.usage.is_some());
assert_eq!(response.system_fingerprint.as_ref().unwrap(), "fp_123abc");
}
#[test]
fn test_add_multiple_choices() {
let choice1 = ChatChoice {
index: 0,
message: ChatCompletionMessage {
role: "assistant".to_string(),
content: Some("Option 1".to_string()),
tool_calls: None,
reasoning_content: None,
},
logprobs: None,
finish_reason: Some("stop".to_string()),
matched_stop: None,
hidden_states: None,
};
let choice2 = ChatChoice {
index: 1,
message: ChatCompletionMessage {
role: "assistant".to_string(),
content: Some("Option 2".to_string()),
tool_calls: None,
reasoning_content: None,
},
logprobs: None,
finish_reason: Some("stop".to_string()),
matched_stop: None,
hidden_states: None,
};
let response = ChatCompletionResponse::builder("chatcmpl_789", "gpt-4")
.add_choice(choice1)
.add_choice(choice2)
.build();
assert_eq!(response.choices.len(), 2);
assert_eq!(response.choices[0].index, 0);
assert_eq!(response.choices[1].index, 1);
}
#[test]
fn test_copy_from_request() {
let request = ChatCompletionRequest {
messages: vec![],
model: "gpt-3.5-turbo".to_string(),
..Default::default()
};
let response = ChatCompletionResponse::builder("chatcmpl_101", "gpt-4")
.copy_from_request(&request)
.build();
assert_eq!(response.model, "gpt-3.5-turbo"); }
}