From 65cd28556779c49856f0f6526bbf0654d9e059b9 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Fri, 1 May 2026 12:10:32 +0800 Subject: [PATCH] feat: add AiApiStyle support and integrate into AI completion and streaming functions --- src-tauri/src/commands/ai.rs | 229 +++++++++++++++++++++++++- src-tauri/src/lib.rs | 1 + src/components/editor/AiAssistant.vue | 57 ++++++- src/components/ui/input/Input.vue | 2 + src/lib/tauri.ts | 4 + src/stores/settingsStore.ts | 10 +- 6 files changed, 286 insertions(+), 17 deletions(-) diff --git a/src-tauri/src/commands/ai.rs b/src-tauri/src/commands/ai.rs index 89e6cd69a0..97badba55d 100644 --- a/src-tauri/src/commands/ai.rs +++ b/src-tauri/src/commands/ai.rs @@ -12,6 +12,19 @@ pub enum AiProvider { Custom, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "lowercase")] +pub enum AiApiStyle { + Completions, + Responses, +} + +impl Default for AiApiStyle { + fn default() -> Self { + Self::Completions + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct AiConfig { @@ -19,6 +32,8 @@ pub struct AiConfig { pub api_key: String, pub endpoint: String, pub model: String, + #[serde(default)] + pub api_style: AiApiStyle, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -44,6 +59,43 @@ fn ai_config_file(app: &AppHandle) -> Result { Ok(dir.join("ai_config.json")) } +#[tauri::command] +pub async fn ai_test_connection(config: AiConfig) -> Result { + if config.api_key.trim().is_empty() { + return Err("API key is required".to_string()); + } + if config.endpoint.trim().is_empty() { + return Err("Endpoint is required".to_string()); + } + if config.model.trim().is_empty() { + return Err("Model is required".to_string()); + } + + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(15)) + .build() + .map_err(|e| e.to_string())?; + + let request = AiCompletionRequest { + config, + system_prompt: String::new(), + messages: vec![AiMessage { role: "user".into(), content: "hi".into() }], + max_tokens: Some(1), + temperature: Some(0.0), + }; + + match request.config.provider { + AiProvider::Claude => call_claude(&client, request).await, + AiProvider::Openai | AiProvider::Custom => { + if request.config.api_style == AiApiStyle::Responses { + call_responses_api(&client, request).await + } else { + call_openai_compatible(&client, request).await + } + } + }.map(|_| "OK".to_string()) +} + #[tauri::command] pub async fn save_ai_config(app: AppHandle, config: AiConfig) -> Result<(), String> { let json = serde_json::to_string_pretty(&config).map_err(|e| e.to_string())?; @@ -79,7 +131,13 @@ pub async fn ai_complete(request: AiCompletionRequest) -> Result match request.config.provider { AiProvider::Claude => call_claude(&client, request).await, - AiProvider::Openai | AiProvider::Custom => call_openai_compatible(&client, request).await, + AiProvider::Openai | AiProvider::Custom => { + if request.config.api_style == AiApiStyle::Responses { + call_responses_api(&client, request).await + } else { + call_openai_compatible(&client, request).await + } + } } } @@ -109,7 +167,13 @@ pub async fn ai_stream(app: AppHandle, session_id: String, request: AiCompletion 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, + AiProvider::Openai | AiProvider::Custom => { + if request.config.api_style == AiApiStyle::Responses { + stream_responses_api(&app, &client, &session_id, request).await + } else { + stream_openai(&app, &client, &session_id, request).await + } + } } } @@ -132,7 +196,7 @@ async fn stream_claude(app: &AppHandle, client: &reqwest::Client, session_id: &s }); let res = client - .post(&request.config.endpoint) + .post(&resolve_endpoint(&request.config)) .headers(headers) .json(&body) .send() @@ -207,7 +271,7 @@ async fn stream_openai(app: &AppHandle, client: &reqwest::Client, session_id: &s }); let res = client - .post(&request.config.endpoint) + .post(&resolve_endpoint(&request.config)) .headers(headers) .json(&body) .send() @@ -278,7 +342,7 @@ async fn call_claude(client: &reqwest::Client, request: AiCompletionRequest) -> }); let res = client - .post(&request.config.endpoint) + .post(&resolve_endpoint(&request.config)) .headers(headers) .json(&body) .send() @@ -325,7 +389,7 @@ async fn call_openai_compatible( }); let res = client - .post(&request.config.endpoint) + .post(&resolve_endpoint(&request.config)) .headers(headers) .json(&body) .send() @@ -351,6 +415,23 @@ fn extract_error(data: &serde_json::Value) -> Option { .map(ToString::to_string) } +fn resolve_endpoint(config: &AiConfig) -> String { + let ep = config.endpoint.trim().trim_end_matches('/'); + if ep.ends_with("/chat/completions") || ep.ends_with("/responses") || ep.ends_with("/messages") { + return ep.to_string(); + } + match config.provider { + AiProvider::Claude => format!("{ep}/messages"), + AiProvider::Openai | AiProvider::Custom => { + if config.api_style == AiApiStyle::Responses { + format!("{ep}/responses") + } else { + format!("{ep}/chat/completions") + } + } + } +} + fn stream_data_payload(line: &str) -> Option<&str> { let line = line.trim(); if line.is_empty() || line.starts_with(':') || line.starts_with("event:") || line.starts_with("id:") { @@ -392,3 +473,139 @@ fn emit_stream_delta(app: &AppHandle, session_id: &str, delta: &str) { done: false, }); } + +fn build_responses_input(system_prompt: &str, messages: &[AiMessage]) -> serde_json::Value { + let mut input = Vec::new(); + if !system_prompt.is_empty() { + input.push(json!({ + "role": "developer", + "content": system_prompt, + })); + } + for m in messages { + input.push(json!({ + "role": m.role, + "content": m.content, + })); + } + json!(input) +} + +async fn call_responses_api( + client: &reqwest::Client, + request: AiCompletionRequest, +) -> Result { + 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 body = json!({ + "model": request.config.model, + "input": build_responses_input(&request.system_prompt, &request.messages), + "max_output_tokens": request.max_tokens.unwrap_or(2048), + "temperature": request.temperature.unwrap_or(0.2), + }); + + let res = client + .post(&resolve_endpoint(&request.config)) + .headers(headers) + .json(&body) + .send() + .await + .map_err(|e| format!("AI request failed: {e}"))?; + + let status = res.status(); + let data: serde_json::Value = res.json().await.map_err(|e| e.to_string())?; + if !status.is_success() { + return Err(extract_error(&data).unwrap_or_else(|| format!("API error: {status}"))); + } + + Ok(data["output"] + .as_array() + .and_then(|items| { + items.iter().find_map(|item| { + item["content"] + .as_array() + .and_then(|parts| parts.iter().find_map(|p| p["text"].as_str())) + }) + }) + .unwrap_or_default() + .to_string()) +} + +async fn stream_responses_api(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 body = json!({ + "model": request.config.model, + "input": build_responses_input(&request.system_prompt, &request.messages), + "max_output_tokens": request.max_tokens.unwrap_or(2048), + "temperature": request.temperature.unwrap_or(0.2), + "stream": true, + }); + + let res = client + .post(&resolve_endpoint(&request.config)) + .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(); + + let mut finished = false; + 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 Some(data) = stream_data_payload(&line) else { + continue; + }; + if data == "[DONE]" { + finished = true; + break; + } + + if let Ok(event) = serde_json::from_str::(data) { + if let Some(text) = responses_stream_text(&event) { + emit_stream_delta(app, session_id, text); + } + } + } + + if finished { + break; + } + } + + let _ = app.emit("ai-stream-chunk", AiStreamChunk { + session_id: session_id.to_string(), + delta: String::new(), + done: true, + }); + + Ok(()) +} + +fn responses_stream_text(event: &serde_json::Value) -> Option<&str> { + event["delta"].as_str().filter(|s| !s.is_empty()) +} diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index ffee5939da..5ea6a41cc8 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -38,6 +38,7 @@ pub fn run() { .invoke_handler(tauri::generate_handler![ commands::ai::ai_complete, commands::ai::ai_stream, + commands::ai::ai_test_connection, commands::ai::save_ai_config, commands::ai::load_ai_config, commands::connection::test_connection, diff --git a/src/components/editor/AiAssistant.vue b/src/components/editor/AiAssistant.vue index 07d710a14e..b12b649e7c 100644 --- a/src/components/editor/AiAssistant.vue +++ b/src/components/editor/AiAssistant.vue @@ -15,11 +15,11 @@ import { import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { ScrollArea } from "@/components/ui/scroll-area"; -import { useSettingsStore, type AiProvider } from "@/stores/settingsStore"; +import { useSettingsStore, type AiProvider, type AiApiStyle } from "@/stores/settingsStore"; import { useConnectionStore } from "@/stores/connectionStore"; import { useQueryStore } from "@/stores/queryStore"; import { buildAiContext, runAiStream } from "@/lib/ai"; -import { listDatabases, redisListDatabases, mongoListDatabases } from "@/lib/tauri"; +import { listDatabases, redisListDatabases, mongoListDatabases, aiTestConnection } from "@/lib/tauri"; import type { AiMessage } from "@/lib/tauri"; import type { ConnectionConfig, QueryTab } from "@/types/database"; @@ -102,6 +102,7 @@ const tempProvider = ref(settings.aiConfig.provider); const tempApiKey = ref(settings.aiConfig.apiKey); const tempEndpoint = ref(settings.aiConfig.endpoint); const tempModel = ref(settings.aiConfig.model); +const tempApiStyle = ref(settings.aiConfig.apiStyle || "completions"); const providerDefaults: Record = { claude: { endpoint: "https://api.anthropic.com/v1/messages", model: "claude-sonnet-4-20250514" }, @@ -119,6 +120,7 @@ function openSettings() { tempApiKey.value = settings.aiConfig.apiKey; tempEndpoint.value = settings.aiConfig.endpoint; tempModel.value = settings.aiConfig.model; + tempApiStyle.value = settings.aiConfig.apiStyle || "completions"; showSettings.value = true; } @@ -128,10 +130,36 @@ function saveSettings() { apiKey: tempApiKey.value, endpoint: tempEndpoint.value, model: tempModel.value, + apiStyle: tempApiStyle.value, }); showSettings.value = false; } +const testingAi = ref(false); +const testResult = ref<"" | "success" | "error">(""); +const testError = ref(""); + +async function testAiConnection() { + testingAi.value = true; + testResult.value = ""; + testError.value = ""; + try { + await aiTestConnection({ + provider: tempProvider.value, + apiKey: tempApiKey.value, + endpoint: tempEndpoint.value, + model: tempModel.value, + apiStyle: tempApiStyle.value, + }); + testResult.value = "success"; + } catch (e: any) { + testResult.value = "error"; + testError.value = e?.message || String(e); + } finally { + testingAi.value = false; + } +} + function selectProvider(provider: AiProvider) { tempProvider.value = provider; tempEndpoint.value = providerDefaults[provider].endpoint; @@ -279,7 +307,7 @@ function formatInlineText(text: string): string { -
+