use serde::{Deserialize, Serialize};
use crate::ModelIden;
use crate::chat::{ChatStream, MessageContent, ToolCall, Usage};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum StopReason {
Completed(String),
MaxTokens(String),
ToolCall(String),
ContentFilter(String),
StopSequence(String),
Other(String),
}
impl From<String> for StopReason {
fn from(reason: String) -> Self {
match reason.as_str() {
"stop" | "end_turn" | "STOP" | "COMPLETE" | "completed" => Self::Completed(reason),
"length" | "max_tokens" | "MAX_TOKENS" | "incomplete" => Self::MaxTokens(reason),
"tool_calls" | "tool_use" | "function_call" => Self::ToolCall(reason),
"content_filter" | "SAFETY" | "RECITATION" | "BLOCKLIST" | "PROHIBITED_CONTENT" | "SPII"
| "IMAGE_SAFETY" | "ERROR_TOXIC" => Self::ContentFilter(reason),
"stop_sequence" | "STOP_SEQUENCE" => Self::StopSequence(reason),
_ => Self::Other(reason),
}
}
}
impl StopReason {
pub fn raw(&self) -> &str {
match self {
Self::Completed(s)
| Self::MaxTokens(s)
| Self::ToolCall(s)
| Self::ContentFilter(s)
| Self::StopSequence(s)
| Self::Other(s) => s,
}
}
pub fn is_max_tokens(&self) -> bool {
matches!(self, Self::MaxTokens(_))
}
}
impl PartialEq for StopReason {
fn eq(&self, other: &Self) -> bool {
core::mem::discriminant(self) == core::mem::discriminant(other)
}
}
impl Eq for StopReason {}
impl std::fmt::Display for StopReason {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.raw())
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatResponse {
pub content: MessageContent,
pub reasoning_content: Option<String>,
pub model_iden: ModelIden,
pub provider_model_iden: ModelIden,
pub stop_reason: Option<StopReason>,
pub usage: Usage,
pub captured_raw_body: Option<serde_json::Value>,
}
impl ChatResponse {
pub fn first_text(&self) -> Option<&str> {
self.content.first_text()
}
pub fn into_first_text(self) -> Option<String> {
self.content.into_first_text()
}
pub fn texts(&self) -> Vec<&str> {
self.content.texts()
}
pub fn into_texts(self) -> Vec<String> {
self.content.into_texts()
}
pub fn tool_calls(&self) -> Vec<&ToolCall> {
self.content.tool_calls()
}
pub fn into_tool_calls(self) -> Vec<ToolCall> {
self.content.into_tool_calls()
}
}
impl ChatResponse {
#[deprecated(note = "Use '.first_text()' or '.texts()'")]
pub fn content_text_as_str(&self) -> Option<&str> {
self.first_text()
}
#[deprecated(note = "Use '.into_first_text()' or '.into_texts()")]
pub fn content_text_into_string(self) -> Option<String> {
self.into_first_text()
}
}
pub struct ChatStreamResponse {
pub stream: ChatStream,
pub model_iden: ModelIden,
}