feat: add AiApiStyle support and integrate into AI completion and streaming functions

This commit is contained in:
t8y2 2026-05-01 12:10:32 +08:00
parent 78ea289f0f
commit 65cd285567
6 changed files with 286 additions and 17 deletions

View File

@ -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<std::path::PathBuf, String> {
Ok(dir.join("ai_config.json"))
}
#[tauri::command]
pub async fn ai_test_connection(config: AiConfig) -> Result<String, String> {
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<String, String>
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<String> {
.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<String, 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),
});
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::<serde_json::Value>(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())
}

View File

@ -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,

View File

@ -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<AiProvider>(settings.aiConfig.provider);
const tempApiKey = ref(settings.aiConfig.apiKey);
const tempEndpoint = ref(settings.aiConfig.endpoint);
const tempModel = ref(settings.aiConfig.model);
const tempApiStyle = ref<AiApiStyle>(settings.aiConfig.apiStyle || "completions");
const providerDefaults: Record<AiProvider, { endpoint: string; model: string }> = {
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 {
</div>
</div>
<div v-else class="flex">
<div v-else-if="msg.content" class="flex">
<div class="max-w-[95%] rounded-lg bg-muted px-3 py-2 text-xs leading-relaxed">
<template v-for="(seg, j) in parseMessage(msg.content)" :key="j">
<div v-if="seg.type === 'text'" class="whitespace-normal">
@ -350,12 +378,12 @@ function formatInlineText(text: string): string {
rows="4"
class="flex-1 resize-none bg-transparent text-xs outline-none placeholder:text-muted-foreground"
:placeholder="t('ai.placeholder')"
:disabled="isGenerating"
:disabled="isGenerating || !props.tab?.database"
@keydown.enter.exact="send"
/>
<button
class="h-7 w-7 shrink-0 rounded-full bg-foreground text-background flex items-center justify-center disabled:opacity-30"
:disabled="isGenerating || !prompt.trim()"
:disabled="isGenerating || !prompt.trim() || !props.tab?.database"
@click="send"
>
<ArrowUp class="h-4 w-4" />
@ -388,14 +416,29 @@ function formatInlineText(text: string): string {
</div>
<div class="grid grid-cols-3 items-center gap-3">
<Label class="text-right text-xs">Endpoint</Label>
<Input v-model="tempEndpoint" class="col-span-2 h-8 text-xs" />
<Input v-model="tempEndpoint" placeholder="https://api.openai.com/v1" class="col-span-2 h-8 text-xs" />
</div>
<div class="grid grid-cols-3 items-center gap-3">
<Label class="text-right text-xs">Model</Label>
<Input v-model="tempModel" class="col-span-2 h-8 text-xs" />
</div>
<div v-if="tempProvider !== 'claude'" class="grid grid-cols-3 items-center gap-3">
<Label class="text-right text-xs">API</Label>
<div class="col-span-2 flex gap-2">
<Button size="sm" variant="outline" class="h-8 flex-1 text-xs" :class="{ 'bg-accent': tempApiStyle === 'completions' }" @click="tempApiStyle = 'completions'">/chat/completions</Button>
<Button size="sm" variant="outline" class="h-8 flex-1 text-xs" :class="{ 'bg-accent': tempApiStyle === 'responses' }" @click="tempApiStyle = 'responses'">/responses</Button>
</div>
</div>
</div>
<DialogFooter>
<DialogFooter class="flex items-center gap-2">
<div class="flex-1 flex items-center gap-2">
<Button size="sm" variant="outline" :disabled="testingAi || !tempApiKey.trim() || !tempEndpoint.trim() || !tempModel.trim()" @click="testAiConnection">
<Loader2 v-if="testingAi" class="h-3 w-3 animate-spin mr-1" />
{{ t('connection.test') }}
</Button>
<span v-if="testResult === 'success'" class="text-xs text-green-500">{{ t('connection.testSuccess') }}</span>
<span v-else-if="testResult === 'error'" class="text-xs text-destructive truncate max-w-[200px]" :title="testError">{{ testError }}</span>
</div>
<Button size="sm" @click="saveSettings">{{ t('grid.save') }}</Button>
</DialogFooter>
</DialogContent>

View File

@ -22,6 +22,8 @@ const modelValue = useVModel(props, 'modelValue', emits, {
<template>
<input
v-model="modelValue"
autocapitalize="off"
autocomplete="off"
data-slot="input"
:class="cn(
'dark:bg-input/30 border-input focus-visible:border-ring focus-visible:ring-ring/50 aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive dark:aria-invalid:border-destructive/50 disabled:bg-input/50 dark:disabled:bg-input/80 h-8 rounded-lg border bg-transparent px-2.5 py-1 text-base transition-colors file:h-6 file:text-sm file:font-medium focus-visible:ring-3 aria-invalid:ring-3 md:text-sm w-full min-w-0 outline-none file:inline-flex file:border-0 file:bg-transparent file:text-foreground placeholder:text-muted-foreground disabled:pointer-events-none disabled:cursor-not-allowed disabled:opacity-50',

View File

@ -58,6 +58,10 @@ export async function saveAiConfig(config: AiConfig): Promise<void> {
return invoke("save_ai_config", { config });
}
export async function aiTestConnection(config: AiConfig): Promise<string> {
return invoke("ai_test_connection", { config });
}
export async function loadAiConfig(): Promise<AiConfig | null> {
return invoke("load_ai_config");
}

View File

@ -3,22 +3,24 @@ import { ref } from "vue";
import * as api from "@/lib/tauri";
export type AiProvider = "claude" | "openai" | "custom";
export type AiApiStyle = "completions" | "responses";
export interface AiConfig {
provider: AiProvider;
apiKey: string;
endpoint: string;
model: string;
apiStyle: AiApiStyle;
}
const defaultConfigs: Record<AiProvider, Omit<AiConfig, "apiKey">> = {
claude: { provider: "claude", endpoint: "https://api.anthropic.com/v1/messages", model: "claude-sonnet-4-20250514" },
openai: { provider: "openai", endpoint: "https://api.openai.com/v1/chat/completions", model: "gpt-4o" },
custom: { provider: "custom", endpoint: "", model: "" },
claude: { provider: "claude", endpoint: "https://api.anthropic.com/v1/messages", model: "claude-sonnet-4-20250514", apiStyle: "completions" },
openai: { provider: "openai", endpoint: "https://api.openai.com/v1/chat/completions", model: "gpt-4o", apiStyle: "completions" },
custom: { provider: "custom", endpoint: "", model: "", apiStyle: "completions" },
};
export const useSettingsStore = defineStore("settings", () => {
const aiConfig = ref<AiConfig>({ ...defaultConfigs.claude, apiKey: "" });
const aiConfig = ref<AiConfig>({ ...defaultConfigs.claude, apiKey: "", apiStyle: "completions" });
const isAiConfigLoaded = ref(false);
async function initAiConfig() {