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 = std::env::var("AI_BASE_URL")
48 .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 =
58 std::env::var("AI_MODEL").unwrap_or_else(|_| "deepseek/deepseek-v4-flash".into());
59 let embedding_model =
60 std::env::var("AI_EMBEDDING_MODEL").unwrap_or_else(|_| "baai/bge-m3".into());
61 let whisper_model = std::env::var("AI_WHISPER_MODEL")
62 .unwrap_or_else(|_| "qwen/qwen3-asr-flash-2026-02-10".into());
63
64 let client =
65 OpenAiProvider::new(base_url, api_key, model, embedding_model, whisper_model);
66 tracing::info!(provider = %provider, model = %client.model, "AI provider configured");
67 Ok(Arc::new(client))
68 }
69 other => Err(AiError::NotConfigured(format!(
70 "unknown AI provider: {other}"
71 ))),
72 }
73}
74
75pub struct OpenAiProvider {
76 http: reqwest::Client,
77 base_url: String,
78 api_key: String,
79 model: String,
80 embedding_model: String,
81 whisper_model: String,
82}
83
84impl OpenAiProvider {
85 pub fn new(
86 base_url: String,
87 api_key: String,
88 model: String,
89 embedding_model: String,
90 whisper_model: String,
91 ) -> Self {
92 Self {
93 http: reqwest::Client::builder()
94 .timeout(std::time::Duration::from_mins(2))
95 .build()
96 .expect("failed to create HTTP client"),
97 base_url,
98 api_key,
99 model,
100 embedding_model,
101 whisper_model,
102 }
103 }
104}
105
106#[async_trait]
107impl AiProvider for OpenAiProvider {
108 #[instrument(skip(self, system, messages), fields(model = %self.model))]
109 async fn chat(&self, system: &str, messages: &[ChatMessage]) -> Result<String, AiError> {
110 let mut api_messages: Vec<serde_json::Value> =
111 vec![serde_json::json!({"role": "system", "content": system})];
112 for msg in messages {
113 api_messages.push(serde_json::json!({"role": msg.role, "content": msg.content}));
114 }
115
116 let body = serde_json::json!({
117 "model": self.model,
118 "messages": api_messages,
119 });
120
121 let resp = self
122 .http
123 .post(format!("{}/chat/completions", self.base_url))
124 .header("Authorization", format!("Bearer {}", self.api_key))
125 .json(&body)
126 .send()
127 .await?;
128
129 let status = resp.status();
130 let text = resp.text().await?;
131
132 if !status.is_success() {
133 return Err(AiError::Api {
134 status: status.as_u16(),
135 body: text,
136 });
137 }
138
139 let json: serde_json::Value = serde_json::from_str(&text)?;
140 let content = json["choices"][0]["message"]["content"]
141 .as_str()
142 .ok_or_else(|| AiError::Api {
143 status: status.as_u16(),
144 body: "unexpected response format".into(),
145 })?;
146
147 Ok(content.to_string())
148 }
149
150 #[instrument(skip(self), fields(model = %self.embedding_model))]
151 async fn embed(&self, text: &str) -> Result<Vec<f32>, AiError> {
152 let body = serde_json::json!({
153 "model": self.embedding_model,
154 "input": text,
155 });
156
157 let resp = self
158 .http
159 .post(format!("{}/embeddings", self.base_url))
160 .header("Authorization", format!("Bearer {}", self.api_key))
161 .json(&body)
162 .send()
163 .await?;
164
165 let status = resp.status();
166 let resp_text = resp.text().await?;
167
168 if !status.is_success() {
169 return Err(AiError::Api {
170 status: status.as_u16(),
171 body: resp_text,
172 });
173 }
174
175 let json: serde_json::Value = serde_json::from_str(&resp_text)?;
176 #[expect(clippy::cast_possible_truncation)]
177 let embedding: Vec<f32> = json["data"][0]["embedding"]
178 .as_array()
179 .ok_or_else(|| AiError::Api {
180 status: status.as_u16(),
181 body: "no embedding in response".into(),
182 })?
183 .iter()
184 .map(|v| v.as_f64().unwrap_or(0.0) as f32)
185 .collect();
186
187 Ok(embedding)
188 }
189
190 #[instrument(skip(self, audio_data), fields(model = %self.whisper_model))]
191 async fn transcribe(&self, audio_data: &[u8], mime_type: &str) -> Result<String, AiError> {
192 let ext = mime_to_ext(mime_type);
193 let part = reqwest::multipart::Part::bytes(audio_data.to_vec())
194 .file_name(format!("audio.{ext}"))
195 .mime_str(mime_type)
196 .map_err(|e| AiError::Api {
197 status: 0,
198 body: e.to_string(),
199 })?;
200
201 let form = reqwest::multipart::Form::new()
202 .part("file", part)
203 .text("model", self.whisper_model.clone());
204
205 let resp = self
206 .http
207 .post(format!("{}/audio/transcriptions", self.base_url))
208 .header("Authorization", format!("Bearer {}", self.api_key))
209 .multipart(form)
210 .send()
211 .await?;
212
213 let status = resp.status();
214 let text = resp.text().await?;
215
216 if !status.is_success() {
217 return Err(AiError::Api {
218 status: status.as_u16(),
219 body: text,
220 });
221 }
222
223 let json: serde_json::Value = serde_json::from_str(&text)?;
224 let transcription = json["text"].as_str().unwrap_or_default().to_string();
225 Ok(transcription)
226 }
227}
228
229fn mime_to_ext(mime: &str) -> &'static str {
230 match mime {
231 "audio/wav" | "audio/wave" => "wav",
232 "audio/mp3" | "audio/mpeg" => "mp3",
233 "audio/mp4" => "mp4",
234 "audio/x-m4a" => "m4a",
235 "audio/flac" => "flac",
236 _ => "webm",
237 }
238}