feat: add AiApiStyle support and integrate into AI completion and streaming functions
This commit is contained in:
parent
78ea289f0f
commit
65cd285567
|
|
@ -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())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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',
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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() {
|
||||
|
|
|
|||
Loading…
Reference in New Issue