use crate::error::Result;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
pub role: String,
pub content: String,
}
impl ChatMessage {
pub fn new(role: String, content: String) -> Self {
Self { role, content }
}
pub fn user(content: String) -> Self {
Self {
role: "user".to_string(),
content,
}
}
pub fn assistant(content: String) -> Self {
Self {
role: "assistant".to_string(),
content,
}
}
pub fn system(content: String) -> Self {
Self {
role: "system".to_string(),
content,
}
}
}
#[async_trait]
pub trait ChatModel: Send + Sync {
fn model(&self) -> &str;
fn provider(&self) -> &str;
fn name(&self) -> &str {
self.model()
}
async fn chat(&self, messages: &[ChatMessage]) -> Result<ChatInvokeCompletion<String>>;
async fn chat_stream(
&self,
messages: &[ChatMessage],
) -> Result<Box<dyn futures_util::stream::Stream<Item = Result<String>> + Send + Unpin>>;
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatInvokeUsage {
pub prompt_tokens: u32,
pub prompt_cached_tokens: Option<u32>,
pub prompt_cache_creation_tokens: Option<u32>,
pub prompt_image_tokens: Option<u32>,
pub completion_tokens: u32,
pub total_tokens: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatInvokeCompletion<T> {
pub completion: T,
pub thinking: Option<String>,
pub redacted_thinking: Option<String>,
pub usage: Option<ChatInvokeUsage>,
pub stop_reason: Option<String>,
}
impl<T> ChatInvokeCompletion<T> {
pub fn new(completion: T) -> Self {
Self {
completion,
thinking: None,
redacted_thinking: None,
usage: None,
stop_reason: None,
}
}
pub fn with_usage(mut self, usage: ChatInvokeUsage) -> Self {
self.usage = Some(usage);
self
}
}