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 = 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}