1use std::sync::Arc;
2
3use async_trait::async_trait;
4use serde::{Deserialize, Serialize};
5use tracing::instrument;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct ChatMessage {
9 pub role: ChatRole,
10 pub content: String,
11}
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14#[serde(rename_all = "lowercase")]
15pub enum ChatRole {
16 System,
17 User,
18 Assistant,
19}
20
21#[derive(Debug, thiserror::Error)]
22pub enum AiError {
23 #[error("HTTP request failed: {0}")]
24 Http(#[from] reqwest::Error),
25
26 #[error("serialization error: {0}")]
27 Serde(#[from] serde_json::Error),
28
29 #[error("API returned error: {status} {body}")]
30 Api { status: u16, body: String },
31
32 #[error("not configured: {0}")]
33 NotConfigured(String),
34}
35
36#[async_trait]
37pub trait AiProvider: Send + Sync {
38 async fn chat(&self, system: &str, messages: &[ChatMessage]) -> Result<String, AiError>;
39 async fn embed(&self, text: &str) -> Result<Vec<f32>, AiError>;
40 async fn transcribe(&self, audio_data: &[u8], mime_type: &str) -> Result<String, AiError>;
41}
42
43pub fn create_provider() -> Result<Arc<dyn AiProvider>, AiError> {
44 let provider = std::env::var("AI_PROVIDER").unwrap_or_else(|_| "openrouter".into());
45 match provider.as_str() {
46 "openai" | "openrouter" => {
47 let base_url =
48 std::env::var("AI_BASE_URL").unwrap_or_else(|_| "https://openrouter.ai/api/v1".into());
49 let api_key = std::env::var("AI_API_KEY")
50 .or_else(|_| std::env::var("OPENAI_API_KEY"))
51 .or_else(|_| std::env::var("OPENROUTER_API_KEY"))
52 .map_err(|_| {
53 AiError::NotConfigured(
54 "set AI_API_KEY, OPENAI_API_KEY, or OPENROUTER_API_KEY".into(),
55 )
56 })?;
57 let model = std::env::var("AI_MODEL").unwrap_or_else(|_| "deepseek/deepseek-v4-flash".into());
58 let embedding_model = std::env::var("AI_EMBEDDING_MODEL")
59 .unwrap_or_else(|_| "baai/bge-m3".into());
60 let whisper_model =
61 std::env::var("AI_WHISPER_MODEL").unwrap_or_else(|_| "qwen/qwen3-asr-flash-2026-02-10".into());
62
63 let client =
64 OpenAiProvider::new(base_url, api_key, model, embedding_model, whisper_model);
65 tracing::info!(provider = %provider, model = %client.model, "AI provider configured");
66 Ok(Arc::new(client))
67 }
68 other => Err(AiError::NotConfigured(format!(
69 "unknown AI provider: {other}"
70 ))),
71 }
72}
73
74pub struct OpenAiProvider {
75 http: reqwest::Client,
76 base_url: String,
77 api_key: String,
78 model: String,
79 embedding_model: String,
80 whisper_model: String,
81}
82
83impl OpenAiProvider {
84 pub fn new(
85 base_url: String,
86 api_key: String,
87 model: String,
88 embedding_model: String,
89 whisper_model: String,
90 ) -> Self {
91 Self {
92 http: reqwest::Client::builder()
93 .timeout(std::time::Duration::from_mins(2))
94 .build()
95 .expect("failed to create HTTP client"),
96 base_url,
97 api_key,
98 model,
99 embedding_model,
100 whisper_model,
101 }
102 }
103}
104
105#[async_trait]
106impl AiProvider for OpenAiProvider {
107 #[instrument(skip(self, system, messages), fields(model = %self.model))]
108 async fn chat(&self, system: &str, messages: &[ChatMessage]) -> Result<String, AiError> {
109 let mut api_messages: Vec<serde_json::Value> =
110 vec![serde_json::json!({"role": "system", "content": system})];
111 for msg in messages {
112 api_messages.push(serde_json::json!({"role": msg.role, "content": msg.content}));
113 }
114
115 let body = serde_json::json!({
116 "model": self.model,
117 "messages": api_messages,
118 });
119
120 let resp = self
121 .http
122 .post(format!("{}/chat/completions", self.base_url))
123 .header("Authorization", format!("Bearer {}", self.api_key))
124 .json(&body)
125 .send()
126 .await?;
127
128 let status = resp.status();
129 let text = resp.text().await?;
130
131 if !status.is_success() {
132 return Err(AiError::Api {
133 status: status.as_u16(),
134 body: text,
135 });
136 }
137
138 let json: serde_json::Value = serde_json::from_str(&text)?;
139 let content = json["choices"][0]["message"]["content"]
140 .as_str()
141 .ok_or_else(|| AiError::Api {
142 status: status.as_u16(),
143 body: "unexpected response format".into(),
144 })?;
145
146 Ok(content.to_string())
147 }
148
149 #[instrument(skip(self), fields(model = %self.embedding_model))]
150 async fn embed(&self, text: &str) -> Result<Vec<f32>, AiError> {
151 let body = serde_json::json!({
152 "model": self.embedding_model,
153 "input": text,
154 });
155
156 let resp = self
157 .http
158 .post(format!("{}/embeddings", self.base_url))
159 .header("Authorization", format!("Bearer {}", self.api_key))
160 .json(&body)
161 .send()
162 .await?;
163
164 let status = resp.status();
165 let resp_text = resp.text().await?;
166
167 if !status.is_success() {
168 return Err(AiError::Api {
169 status: status.as_u16(),
170 body: resp_text,
171 });
172 }
173
174 let json: serde_json::Value = serde_json::from_str(&resp_text)?;
175 #[expect(clippy::cast_possible_truncation)]
176 let embedding: Vec<f32> = json["data"][0]["embedding"]
177 .as_array()
178 .ok_or_else(|| AiError::Api {
179 status: status.as_u16(),
180 body: "no embedding in response".into(),
181 })?
182 .iter()
183 .map(|v| v.as_f64().unwrap_or(0.0) as f32)
184 .collect();
185
186 Ok(embedding)
187 }
188
189 #[instrument(skip(self, audio_data), fields(model = %self.whisper_model))]
190 async fn transcribe(&self, audio_data: &[u8], mime_type: &str) -> Result<String, AiError> {
191 let ext = mime_to_ext(mime_type);
192 let part = reqwest::multipart::Part::bytes(audio_data.to_vec())
193 .file_name(format!("audio.{ext}"))
194 .mime_str(mime_type)
195 .map_err(|e| AiError::Api {
196 status: 0,
197 body: e.to_string(),
198 })?;
199
200 let form = reqwest::multipart::Form::new()
201 .part("file", part)
202 .text("model", self.whisper_model.clone());
203
204 let resp = self
205 .http
206 .post(format!("{}/audio/transcriptions", self.base_url))
207 .header("Authorization", format!("Bearer {}", self.api_key))
208 .multipart(form)
209 .send()
210 .await?;
211
212 let status = resp.status();
213 let text = resp.text().await?;
214
215 if !status.is_success() {
216 return Err(AiError::Api {
217 status: status.as_u16(),
218 body: text,
219 });
220 }
221
222 let json: serde_json::Value = serde_json::from_str(&text)?;
223 let transcription = json["text"].as_str().unwrap_or_default().to_string();
224 Ok(transcription)
225 }
226}
227
228fn mime_to_ext(mime: &str) -> &'static str {
229 match mime {
230 "audio/wav" | "audio/wave" => "wav",
231 "audio/mp3" | "audio/mpeg" => "mp3",
232 "audio/mp4" => "mp4",
233 "audio/x-m4a" => "m4a",
234 "audio/flac" => "flac",
235 _ => "webm",
236 }
237}