xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(ai): support proxied API channels

This commit is contained in:
t8y2 2026-06-08 20:32:54 +08:00
parent 2db46f10ff
commit 87b8f1f462
1 changed files with 132 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

142
crates/dbx-core/src/ai.rs
View File

@ -3,6 +3,7 @@ use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE};
use serde::{Deserialize, Serialize};
use serde_json::json;
use std::collections::{HashMap, HashSet};
use std::net::IpAddr;
use std::path::Path;
use std::sync::{Arc, LazyLock};
use tokio::sync::{Notify, RwLock};
@ -225,12 +226,53 @@ pub fn claude_stream_text(event: &serde_json::Value) -> Option<&str> {
None
}
pub fn openai_stream_text(event: &serde_json::Value) -> Option<&str> {
fn text_from_content_value(value: &serde_json::Value) -> Option<String> {
if let Some(text) = value.as_str().filter(|text| !text.is_empty()) {
return Some(text.to_string());
}
value.as_array().and_then(|parts| {
let text = parts
.iter()
.filter_map(|part| {
part["text"]
.as_str()
.or_else(|| part["content"].as_str())
.or_else(|| part["input_text"].as_str())
.or_else(|| part["output_text"].as_str())
})
.collect::<Vec<_>>()
.join("");
(!text.is_empty()).then_some(text)
})
}
pub fn openai_response_text(data: &serde_json::Value) -> String {
data["choices"]
.get(0)
.and_then(|choice| {
text_from_content_value(&choice["message"]["content"])
.or_else(|| text_from_content_value(&choice["text"]))
.or_else(|| text_from_content_value(&choice["delta"]["content"]))
})
.or_else(|| text_from_content_value(&data["content"]))
.or_else(|| {
let text = responses_text(data);
(!text.is_empty()).then_some(text)
})
.unwrap_or_default()
}
pub fn openai_stream_text(event: &serde_json::Value) -> Option<String> {
event["choices"]
.get(0)
.and_then(|choice| choice["delta"]["content"].as_str().or_else(|| choice["message"]["content"].as_str()))
.or_else(|| event["content"].as_str())
.filter(|text| !text.is_empty())
.and_then(|choice| {
text_from_content_value(&choice["delta"]["content"])
.or_else(|| text_from_content_value(&choice["message"]["content"]))
.or_else(|| text_from_content_value(&choice["text"]))
})
.or_else(|| text_from_content_value(&event["content"]))
.or_else(|| event["delta"].as_str().filter(|text| !text.is_empty()).map(ToString::to_string))
}
pub fn openai_stream_reasoning(event: &serde_json::Value) -> Option<&str> {
@ -337,10 +379,31 @@ fn claude_headers(config: &AiConfig) -> Result<HeaderMap, String> {
Ok(headers)
}
fn normalize_ai_proxy_url(proxy_url: &str) -> String {
let proxy_url = proxy_url.trim();
if proxy_url.contains("://") || proxy_url.is_empty() {
proxy_url.to_string()
} else {
format!("http://{proxy_url}")
}
}
fn ai_endpoint_is_loopback(config: &AiConfig) -> bool {
let endpoint = resolve_endpoint(config);
let Ok(url) = reqwest::Url::parse(&endpoint) else {
return false;
};
let Some(host) = url.host_str() else {
return false;
};
host.eq_ignore_ascii_case("localhost") || host.parse::<IpAddr>().map(|addr| addr.is_loopback()).unwrap_or(false)
}
pub fn build_ai_http_client(config: &AiConfig, timeout_secs: u64) -> Result<reqwest::Client, String> {
let mut builder = reqwest::Client::builder().timeout(std::time::Duration::from_secs(timeout_secs));
if config.proxy_enabled && !config.proxy_url.trim().is_empty() {
let proxy = reqwest::Proxy::all(config.proxy_url.trim()).map_err(|e| format!("Invalid AI proxy URL: {e}"))?;
if config.proxy_enabled && !config.proxy_url.trim().is_empty() && !ai_endpoint_is_loopback(config) {
let proxy_url = normalize_ai_proxy_url(&config.proxy_url);
let proxy = reqwest::Proxy::all(&proxy_url).map_err(|e| format!("Invalid AI proxy URL: {e}"))?;
builder = builder.proxy(proxy);
}
builder.build().map_err(|e| e.to_string())
@ -497,7 +560,7 @@ pub async fn call_openai_compatible(client: &reqwest::Client, request: AiComplet
return Err(extract_error(&data).unwrap_or_else(|| format!("API error: {status}")));
}
Ok(data["choices"][0]["message"]["content"].as_str().unwrap_or_default().to_string())
Ok(openai_response_text(&data))
}
pub async fn call_responses_api(client: &reqwest::Client, request: AiCompletionRequest) -> Result<String, String> {
@ -807,7 +870,7 @@ async fn stream_openai(
if let Some(text) = openai_stream_text(&event) {
on_chunk(AiStreamChunk {
session_id: session_id.to_string(),
delta: text.to_string(),
delta: text,
reasoning_delta: None,
done: false,
});
@ -1048,8 +1111,9 @@ pub fn load_config(path: &Path) -> Result<Option<AiConfig>, String> {
#[cfg(test)]
mod tests {
use super::{
build_ai_http_client, gemini_text, parse_model_list_response, resolve_endpoint, resolve_model_list_endpoint,
responses_max_output_tokens, responses_text, validate_config, AiApiStyle, AiConfig, AiModelInfo, AiProvider,
build_ai_http_client, gemini_text, openai_response_text, openai_stream_text, parse_model_list_response,
resolve_endpoint, resolve_model_list_endpoint, responses_max_output_tokens, responses_text, validate_config,
AiApiStyle, AiConfig, AiModelInfo, AiProvider,
};
#[test]
@ -1086,6 +1150,38 @@ mod tests {
assert!(err.contains("Invalid AI proxy URL"));
}
#[test]
fn ai_http_client_accepts_proxy_host_port_without_scheme() {
let config = AiConfig {
provider: AiProvider::Openai,
api_key: "key".to_string(),
endpoint: "https://api.openai.com/v1/chat/completions".to_string(),
model: "gpt-4o".to_string(),
api_style: AiApiStyle::Completions,
proxy_enabled: true,
proxy_url: "127.0.0.1:7890".to_string(),
enable_thinking: true,
};
build_ai_http_client(&config, 1).unwrap();
}
#[test]
fn ai_http_client_bypasses_proxy_for_loopback_endpoint() {
let config = AiConfig {
provider: AiProvider::OpenaiCompatible,
api_key: "key".to_string(),
endpoint: "http://127.0.0.1:3456/v1".to_string(),
model: "gpt-4o".to_string(),
api_style: AiApiStyle::Completions,
proxy_enabled: true,
proxy_url: "not a proxy url".to_string(),
enable_thinking: true,
};
build_ai_http_client(&config, 1).unwrap();
}
#[test]
fn resolves_gemini_and_ollama_endpoints() {
let gemini = AiConfig {
@ -1196,6 +1292,32 @@ mod tests {
);
}
#[test]
fn parses_openai_compatible_proxy_response_shapes() {
assert_eq!(
openai_response_text(&serde_json::json!({
"choices": [{
"message": {
"content": [
{ "type": "text", "text": "SELECT " },
{ "type": "text", "text": "1;" }
]
}
}]
})),
"SELECT 1;"
);
assert_eq!(
openai_stream_text(&serde_json::json!({
"type": "response.output_text.delta",
"delta": "SELECT 2;"
}))
.as_deref(),
Some("SELECT 2;")
);
}
#[test]
fn parses_gemini_text_and_provider_aliases() {
let data = serde_json::json!({