Skip to main content

oxinbox_backend/ai/
mod.rs

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}