support generic openai api

src-tauri/src/lib.rs

@@ -3,38 +3,61 @@ use serde::{Deserialize, Serialize};
 use futures_util::StreamExt;
 use reqwest::Client;
 
+#[derive(Serialize, Deserialize, Clone)]
+struct Message {
+    role: String,
+    content: String,
+}
+
 #[derive(Serialize, Deserialize, Clone)]
 struct TranslationPayload {
     model: String,
-    prompt: String,
+    messages: Vec<Message>,
     stream: bool,
 }
 
 #[derive(Deserialize)]
-struct OllamaResponse {
-    response: Option<String>,
-    // done: bool,
+struct OpenAIResponse {
+    choices: Vec<Choice>,
+}
+
+#[derive(Deserialize)]
+struct Choice {
+    message: Option<Message>,
+    delta: Option<Delta>,
+}
+
+#[derive(Deserialize)]
+struct Delta {
+    content: Option<String>,
 }
 
 #[tauri::command]
 async fn translate(
     app: AppHandle,
     api_address: String,
+    api_key: String,
     payload: TranslationPayload,
 ) -> Result<String, String> {
     let client = Client::new();
-    let url = format!("{}/api/generate", api_address);
+    // Ensure URL doesn't have double slashes if api_address ends with /
+    let base_url = api_address.trim_end_matches('/');
+    let url = format!("{}/chat/completions", base_url);
 
-    let res = client
-        .post(&url)
-        .json(&payload)
+    let mut request = client.post(&url).json(&payload);
+
+    if !api_key.is_empty() {
+        request = request.header("Authorization", format!("Bearer {}", api_key));
+    }
+
+    let res = request
         .send()
         .await
         .map_err(|e| e.to_string())?;
 
     if !payload.stream {
-        let data = res.json::<OllamaResponse>().await.map_err(|e| e.to_string())?;
-        return Ok(data.response.unwrap_or_default());
+        let data = res.json::<OpenAIResponse>().await.map_err(|e| e.to_string())?;
+        return Ok(data.choices.get(0).and_then(|c| c.message.as_ref()).map(|m| m.content.clone()).unwrap_or_default());
     }
 
     let mut stream = res.bytes_stream();
@@ -44,13 +67,21 @@ async fn translate(
         let chunk = item.map_err(|e| e.to_string())?;
         let text = String::from_utf8_lossy(&chunk);
         
-        // Handle potential multiple JSON objects in one chunk
         for line in text.lines() {
-            if line.trim().is_empty() { continue; }
-            if let Ok(json) = serde_json::from_str::<OllamaResponse>(line) {
-                if let Some(token) = json.response {
-                    full_response.push_str(&token);
-                    app.emit("translation-chunk", token).map_err(|e| e.to_string())?;
+            let line = line.trim();
+            if line.is_empty() { continue; }
+            if line == "data: [DONE]" { break; }
+            
+            if let Some(data_str) = line.strip_prefix("data: ") {
+                if let Ok(json) = serde_json::from_str::<OpenAIResponse>(data_str) {
+                    if let Some(choice) = json.choices.get(0) {
+                        if let Some(delta) = &choice.delta {
+                            if let Some(content) = &delta.content {
+                                full_response.push_str(content);
+                                app.emit("translation-chunk", content).map_err(|e| e.to_string())?;
+                            }
+                        }
+                    }
                 }
             }
         }