diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index adfbb78d2b..baa77ad145 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -1116,6 +1116,7 @@ dependencies = [ "anyhow", "chrono", "duckdb", + "futures", "log", "mongodb", "oracle-rs", @@ -1688,6 +1689,21 @@ dependencies = [ "new_debug_unreachable", ] +[[package]] +name = "futures" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.32" @@ -1761,6 +1777,7 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ + "futures-channel", "futures-core", "futures-io", "futures-macro", @@ -4506,12 +4523,14 @@ dependencies = [ "tokio", "tokio-native-tls", "tokio-rustls", + "tokio-util", "tower", "tower-http", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams 0.4.2", "web-sys", "webpki-roots 1.0.7", ] @@ -4546,7 +4565,7 @@ dependencies = [ "url", "wasm-bindgen", "wasm-bindgen-futures", - "wasm-streams", + "wasm-streams 0.5.0", "web-sys", ] @@ -6949,6 +6968,19 @@ dependencies = [ "wasmparser", ] +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "wasm-streams" version = "0.5.0" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 5e9b37eb05..1b04f3c061 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -41,7 +41,8 @@ portpicker = "0.1.1" duckdb = { version = "1.3.2", features = ["bundled"] } tiberius = { version = "0.12.3", features = ["tds73", "chrono"] } tokio-util = { version = "0.7", features = ["compat"] } -reqwest = { version = "0.12", features = ["json"] } +reqwest = { version = "0.12", features = ["json", "stream"] } +futures = "0.3" mongodb = "3.2.5" oracle-rs = "0.1" percent-encoding = "2" diff --git a/src-tauri/src/commands/ai.rs b/src-tauri/src/commands/ai.rs index e12e20e796..6db7cef400 100644 --- a/src-tauri/src/commands/ai.rs +++ b/src-tauri/src/commands/ai.rs @@ -1,7 +1,8 @@ +use futures::StreamExt; use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE}; use serde::{Deserialize, Serialize}; use serde_json::json; -use tauri::{AppHandle, Manager}; +use tauri::{AppHandle, Emitter, Manager}; #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] @@ -82,6 +83,185 @@ pub async fn ai_complete(request: AiCompletionRequest) -> Result } } +#[derive(Debug, Clone, Serialize)] +struct AiStreamChunk { + session_id: String, + delta: String, + done: bool, +} + +#[tauri::command] +pub async fn ai_stream(app: AppHandle, session_id: String, request: AiCompletionRequest) -> Result<(), String> { + if request.config.api_key.trim().is_empty() { + return Err("API key is required".to_string()); + } + if request.config.endpoint.trim().is_empty() { + return Err("Endpoint is required".to_string()); + } + if request.config.model.trim().is_empty() { + return Err("Model is required".to_string()); + } + + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(120)) + .build() + .map_err(|e| e.to_string())?; + + match request.config.provider { + AiProvider::Claude => stream_claude(&app, &client, &session_id, request).await, + AiProvider::Openai | AiProvider::Custom => stream_openai(&app, &client, &session_id, request).await, + } +} + +async fn stream_claude(app: &AppHandle, client: &reqwest::Client, session_id: &str, request: AiCompletionRequest) -> Result<(), String> { + let mut headers = HeaderMap::new(); + headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + headers.insert( + "x-api-key", + HeaderValue::from_str(&request.config.api_key).map_err(|e| e.to_string())?, + ); + headers.insert("anthropic-version", HeaderValue::from_static("2023-06-01")); + + let body = json!({ + "model": request.config.model, + "max_tokens": request.max_tokens.unwrap_or(2048), + "temperature": request.temperature.unwrap_or(0.2), + "system": request.system_prompt, + "messages": request.messages, + "stream": true, + }); + + let res = client + .post(&request.config.endpoint) + .headers(headers) + .json(&body) + .send() + .await + .map_err(|e| format!("Claude request failed: {e}"))?; + + if !res.status().is_success() { + let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?; + return Err(extract_error(&data).unwrap_or_else(|| "Claude API error".to_string())); + } + + let mut stream = res.bytes_stream(); + let mut buf = String::new(); + + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|e| e.to_string())?; + buf.push_str(&String::from_utf8_lossy(&chunk)); + + while let Some(pos) = buf.find('\n') { + let line = buf[..pos].to_string(); + buf = buf[pos + 1..].to_string(); + + let line = line.trim(); + if !line.starts_with("data: ") { + continue; + } + let data = &line[6..]; + if data == "[DONE]" { + break; + } + + if let Ok(event) = serde_json::from_str::(data) { + if event["type"] == "content_block_delta" { + if let Some(text) = event["delta"]["text"].as_str() { + let _ = app.emit("ai-stream-chunk", AiStreamChunk { + session_id: session_id.to_string(), + delta: text.to_string(), + done: false, + }); + } + } + } + } + } + + let _ = app.emit("ai-stream-chunk", AiStreamChunk { + session_id: session_id.to_string(), + delta: String::new(), + done: true, + }); + + Ok(()) +} + +async fn stream_openai(app: &AppHandle, client: &reqwest::Client, session_id: &str, request: AiCompletionRequest) -> Result<(), String> { + let mut headers = HeaderMap::new(); + headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + headers.insert( + AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer {}", request.config.api_key)).map_err(|e| e.to_string())?, + ); + + let mut messages = vec![json!({ "role": "system", "content": request.system_prompt })]; + messages.extend( + request.messages.iter().map(|m| json!({ "role": m.role, "content": m.content })), + ); + + let body = json!({ + "model": request.config.model, + "messages": messages, + "max_tokens": request.max_tokens.unwrap_or(2048), + "temperature": request.temperature.unwrap_or(0.2), + "stream": true, + }); + + let res = client + .post(&request.config.endpoint) + .headers(headers) + .json(&body) + .send() + .await + .map_err(|e| format!("AI request failed: {e}"))?; + + if !res.status().is_success() { + let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?; + return Err(extract_error(&data).unwrap_or_else(|| "API error".to_string())); + } + + let mut stream = res.bytes_stream(); + let mut buf = String::new(); + + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|e| e.to_string())?; + buf.push_str(&String::from_utf8_lossy(&chunk)); + + while let Some(pos) = buf.find('\n') { + let line = buf[..pos].to_string(); + buf = buf[pos + 1..].to_string(); + + let line = line.trim(); + if !line.starts_with("data: ") { + continue; + } + let data = &line[6..]; + if data == "[DONE]" { + break; + } + + if let Ok(event) = serde_json::from_str::(data) { + if let Some(text) = event["choices"][0]["delta"]["content"].as_str() { + let _ = app.emit("ai-stream-chunk", AiStreamChunk { + session_id: session_id.to_string(), + delta: text.to_string(), + done: false, + }); + } + } + } + } + + let _ = app.emit("ai-stream-chunk", AiStreamChunk { + session_id: session_id.to_string(), + delta: String::new(), + done: true, + }); + + Ok(()) +} + async fn call_claude(client: &reqwest::Client, request: AiCompletionRequest) -> Result { let mut headers = HeaderMap::new(); headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index ca72984220..4bd9bbaf84 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -37,6 +37,7 @@ pub fn run() { }) .invoke_handler(tauri::generate_handler![ commands::ai::ai_complete, + commands::ai::ai_stream, commands::ai::save_ai_config, commands::ai::load_ai_config, commands::connection::test_connection, diff --git a/src/App.vue b/src/App.vue index 5a1387a4a7..bc5cf53a9b 100644 --- a/src/App.vue +++ b/src/App.vue @@ -1,7 +1,7 @@