use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use crate::domain::errors::DomainError;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum MessageRole {
System,
User,
Assistant,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
pub role: MessageRole,
pub content: String,
}
impl ChatMessage {
pub fn system(content: impl Into<String>) -> Self {
Self {
role: MessageRole::System,
content: content.into(),
}
}
pub fn user(content: impl Into<String>) -> Self {
Self {
role: MessageRole::User,
content: content.into(),
}
}
pub fn assistant(content: impl Into<String>) -> Self {
Self {
role: MessageRole::Assistant,
content: content.into(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompletionOptions {
pub max_tokens: Option<u32>,
pub temperature: Option<f32>,
pub top_p: Option<f32>,
pub stop_sequences: Option<Vec<String>>,
}
impl Default for CompletionOptions {
fn default() -> Self {
Self {
max_tokens: Some(4096),
temperature: Some(0.7),
top_p: None,
stop_sequences: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompletionResponse {
pub content: String,
pub model: String,
pub usage: TokenUsage,
pub finish_reason: Option<String>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TokenUsage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub total_tokens: u32,
}
#[async_trait]
pub trait TeiLlmProvider: Send + Sync {
async fn complete(
&self,
messages: &[ChatMessage],
options: &CompletionOptions,
) -> Result<CompletionResponse, DomainError>;
async fn complete_simple(&self, prompt: &str) -> Result<String, DomainError> {
let messages = vec![ChatMessage::user(prompt)];
let response = self
.complete(&messages, &CompletionOptions::default())
.await?;
Ok(response.content)
}
fn provider_name(&self) -> &str;
fn model_id(&self) -> &str;
async fn health_check(&self) -> Result<bool, DomainError> {
Ok(true)
}
fn estimate_tokens(&self, text: &str) -> u32 {
(text.len() / 4) as u32
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamChunk {
pub content: String,
pub is_final: bool,
}