mod anthropic;
mod http;
mod openai;
use crate::config::Config;
use crate::providers;
use anthropic::Anthropic;
use anyhow::Result;
use openai::OpenAiCompatible;
use serde::Serialize;
#[derive(Debug, Clone, Copy, PartialEq, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
System,
User,
Assistant,
}
#[derive(Debug, Clone, Serialize)]
pub struct Message {
pub role: Role,
pub content: String,
}
impl Message {
pub fn system(content: impl Into<String>) -> Self {
Self {
role: Role::System,
content: content.into(),
}
}
pub fn user(content: impl Into<String>) -> Self {
Self {
role: Role::User,
content: content.into(),
}
}
pub fn assistant(content: impl Into<String>) -> Self {
Self {
role: Role::Assistant,
content: content.into(),
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct ModelInfo {
pub id: String,
pub loaded: Option<bool>,
pub context_length: Option<usize>,
}
impl ModelInfo {
pub fn named(id: &str) -> Self {
Self {
id: id.to_owned(),
loaded: None,
context_length: None,
}
}
}
pub trait ModelServer {
async fn models(&self) -> Result<Vec<ModelInfo>>;
async fn complete(&self, model: &str, messages: &[Message]) -> Result<String>;
}
pub enum Server {
OpenAi(OpenAiCompatible),
Anthropic(Anthropic),
}
impl Server {
pub fn connect(config: &Config) -> Result<Self> {
Ok(if providers::is_anthropic(&config.base_url) {
Self::Anthropic(Anthropic::new(config)?)
} else {
Self::OpenAi(OpenAiCompatible::new(config)?)
})
}
pub fn takes_temperature(&self) -> bool {
matches!(self, Self::OpenAi(_))
}
}
impl ModelServer for Server {
async fn models(&self) -> Result<Vec<ModelInfo>> {
match self {
Self::OpenAi(server) => server.models().await,
Self::Anthropic(server) => server.models().await,
}
}
async fn complete(&self, model: &str, messages: &[Message]) -> Result<String> {
match self {
Self::OpenAi(server) => server.complete(model, messages).await,
Self::Anthropic(server) => server.complete(model, messages).await,
}
}
}