use crate::chat::message::{ContentBlock, Message, StopReason};
use crate::cost::usage::Usage;
use crate::model::ModelId;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ChatResponse {
pub message: Message,
pub stop_reason: StopReason,
pub stop_details: Option<String>,
pub usage: Usage,
pub model: ModelId,
}
impl ChatResponse {
pub fn new(message: Message, stop_reason: StopReason, usage: Usage, model: ModelId) -> Self {
Self {
message,
stop_reason,
stop_details: None,
usage,
model,
}
}
#[must_use]
pub fn with_stop_details(mut self, details: impl Into<String>) -> Self {
self.stop_details = Some(details.into());
self
}
pub fn text(&self) -> String {
self.message.text()
}
pub fn tool_calls(&self) -> Vec<&ContentBlock> {
self.message.tool_calls()
}
pub fn is_complete(&self) -> bool {
self.stop_reason.is_complete()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::chat::message::Role;
fn reply(stop: StopReason) -> ChatResponse {
ChatResponse {
message: Message {
role: Role::Assistant,
content: vec![ContentBlock::Text("half an answ".into())],
},
stop_reason: stop,
stop_details: None,
usage: Usage::absent(),
model: "m".into(),
}
}
#[test]
fn a_truncated_reply_says_so_even_though_the_call_succeeded() {
assert!(!reply(StopReason::MaxTokens).is_complete());
assert!(reply(StopReason::EndTurn).is_complete());
}
}