aibo/BoosAPI
1
0
Fork 0
forked from boos/BoosAPI
Code Pull Requests Activity

perf: optimize proxy streaming and database access

Browse Source
This commit is contained in:
boos
2026-06-07 00:00:24 +08:00
parent f3d2d50f9a
commit aaba6f152e
13 changed files with 960 additions and 492 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/proxy/executors/mod.rs
+83 -49
View File
@@ -14,9 +14,7 @@ use crate::transformers::{
canonical::CanonicalRequest,
};
use crate::AppState;
use reqwest::Client;
use serde_json::{json, Value};
use std::time::Duration;
/// Resolved upstream info for a given model request.
#[derive(Debug, Clone)]
@@ -26,15 +24,6 @@ pub struct UpstreamRoute {
pub model: String,
}
/// Create a shared reqwest client with sensible defaults.
fn http_client() -> Client {
Client::builder()
.timeout(Duration::from_secs(300))
.pool_max_idle_per_host(32)
.build()
.expect("Failed to build HTTP client")
}
fn parse_upstream_json(
text: &str,
upstream_url: &str,
@@ -174,11 +163,11 @@ fn aggregate_chat_sse(text: &str, fallback_model: &str) -> Option<Value> {
/// Execute a non-streaming chat completions request.
pub async fn execute_chat_completions(
_state: AppState,
state: AppState,
route: UpstreamRoute,
body: Value,
) -> AppResult<Value> {
let client = http_client();
let client = state.http.clone();
let upstream_url = upstream_url(&route.base_url, "/chat/completions");
let resp = client
@@ -210,19 +199,29 @@ pub async fn execute_chat_completions(
/// Execute a streaming chat completions request.
/// Returns a vector of SSE data chunks as strings.
pub async fn execute_chat_completions_stream(
_state: AppState,
state: AppState,
route: UpstreamRoute,
body: Value,
) -> AppResult<Vec<String>> {
execute_streaming_request(route, "/chat/completions", body).await
execute_streaming_request(state, route, "/chat/completions", body).await
}
/// Execute a streaming chat completions request and return the live upstream response.
pub async fn execute_chat_completions_stream_response(
state: AppState,
route: UpstreamRoute,
body: Value,
) -> AppResult<reqwest::Response> {
execute_streaming_response(state, route, "/chat/completions", body).await
}
/// Execute an Anthropic Messages API request against an Anthropic-shaped upstream.
pub async fn execute_anthropic_messages(
state: AppState,
route: UpstreamRoute,
request: CanonicalRequest,
) -> AppResult<Value> {
let client = http_client();
let client = state.http.clone();
let upstream_url = format!("{}/messages", route.base_url.trim_end_matches('/'));
let compatibility = AnthropicMessagesCompatibility::default();
let body = format_messages_request_with_compat(&request, &compatibility);
@@ -260,10 +259,21 @@ pub async fn execute_anthropic_messages(
/// Execute a streaming Anthropic Messages request and return Anthropic SSE frames.
pub async fn execute_anthropic_messages_stream(
state: AppState,
route: UpstreamRoute,
request: CanonicalRequest,
) -> AppResult<Vec<String>> {
let client = http_client();
let response = execute_anthropic_messages_stream_response(state, route, request).await?;
collect_anthropic_stream_frames(response).await
}
/// Execute a streaming Anthropic Messages request and return the live upstream response.
pub async fn execute_anthropic_messages_stream_response(
state: AppState,
route: UpstreamRoute,
request: CanonicalRequest,
) -> AppResult<reqwest::Response> {
let client = state.http.clone();
let upstream_url = format!("{}/messages", route.base_url.trim_end_matches('/'));
let compatibility = AnthropicMessagesCompatibility::default();
let body = format_messages_request_with_compat(&request, &compatibility);
@@ -292,6 +302,10 @@ pub async fn execute_anthropic_messages_stream(
)));
}
Ok(resp)
}
async fn collect_anthropic_stream_frames(resp: reqwest::Response) -> AppResult<Vec<String>> {
let mut frames = Vec::new();
let mut stream = resp.bytes_stream();
use futures::StreamExt;
@@ -406,32 +420,12 @@ fn denormalize_anthropic_stop_reason(reason: &str) -> &str {
/// Execute a streaming request against any upstream path (e.g. /responses).
/// Returns SSE `data: ...` lines verbatim.
pub async fn execute_streaming_request(
state: AppState,
route: UpstreamRoute,
path_suffix: &str,
body: Value,
) -> AppResult<Vec<String>> {
let client = http_client();
let upstream_url = upstream_url(&route.base_url, path_suffix);
let resp = client
.post(&upstream_url)
.header("Authorization", format!("Bearer {}", route.api_key))
.header("Content-Type", "application/json")
.json(&body)
.send()
.await
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
let status = resp.status();
if !status.is_success() {
let text = resp
.text()
.await
.unwrap_or_else(|_| "unknown error".to_string());
return Err(AppError::UpstreamError(format!(
"Upstream returned {status}: {text}"
)));
}
let resp = execute_streaming_response(state, route, path_suffix, body).await?;
let mut chunks = Vec::new();
let mut stream = resp.bytes_stream();
@@ -455,14 +449,50 @@ pub async fn execute_streaming_request(
Ok(chunks)
}
/// Execute a streaming request against any upstream path and return the live response.
///
/// The caller owns draining `bytes_stream()` so bytes can be flushed to the
/// downstream client as soon as the upstream sends them.
pub async fn execute_streaming_response(
state: AppState,
route: UpstreamRoute,
path_suffix: &str,
body: Value,
) -> AppResult<reqwest::Response> {
let client = state.http.clone();
let upstream_url = upstream_url(&route.base_url, path_suffix);
let resp = client
.post(&upstream_url)
.header("Authorization", format!("Bearer {}", route.api_key))
.header("Content-Type", "application/json")
.json(&body)
.send()
.await
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
let status = resp.status();
if !status.is_success() {
let text = resp
.text()
.await
.unwrap_or_else(|_| "unknown error".to_string());
return Err(AppError::UpstreamError(format!(
"Upstream returned {status}: {text}"
)));
}
Ok(resp)
}
/// Generic non-streaming proxy request to any upstream endpoint path.
pub async fn execute_proxy_request(
_state: AppState,
state: AppState,
route: UpstreamRoute,
path_suffix: &str,
body: Value,
) -> AppResult<Value> {
let client = http_client();
let client = state.http.clone();
let upstream_url = upstream_url(&route.base_url, path_suffix);
let resp = client
@@ -492,8 +522,12 @@ pub async fn execute_proxy_request(
}
/// Generic GET proxy request to an upstream endpoint path.
pub async fn execute_proxy_get(route: UpstreamRoute, path_suffix: &str) -> AppResult<Value> {
let client = http_client();
pub async fn execute_proxy_get(
state: AppState,
route: UpstreamRoute,
path_suffix: &str,
) -> AppResult<Value> {
let client = state.http.clone();
let upstream_url = upstream_url(&route.base_url, path_suffix);
let resp = client
@@ -532,7 +566,7 @@ pub async fn execute_video_task_status(
} else {
format!("/videos/{id}")
};
execute_proxy_get(route, &suffix).await
execute_proxy_get(state, route, &suffix).await
}
/// Execute an OpenAI-shaped video task request, with Volcengine Ark native
@@ -929,7 +963,7 @@ async fn comfy_json(
path: &str,
body: Option<Value>,
) -> AppResult<Value> {
let client = http_client();
let client = state.http.clone();
let mut request = client.request(method, comfy_native_url(state, path));
if let Some(api_key) = state
.config
@@ -1016,11 +1050,11 @@ fn comfy_output_objects(value: &Value) -> bool {
/// Execute an image generation request.
pub async fn execute_image_generation(
_state: AppState,
state: AppState,
route: UpstreamRoute,
body: Value,
) -> AppResult<Value> {
let client = http_client();
let client = state.http.clone();
let upstream_url = upstream_url(&route.base_url, "/images/generations");
let resp = client
@@ -1073,13 +1107,13 @@ pub struct MultipartField {
/// the route-resolved upstream model and skip any client-side fields
/// that don't belong on the wire.
pub async fn execute_image_edit(
_state: AppState,
state: AppState,
route: UpstreamRoute,
fields: Vec<MultipartField>,
) -> AppResult<Value> {
use reqwest::multipart::{Form, Part};
let client = http_client();
let client = state.http.clone();
let upstream_url = upstream_url(&route.base_url, "/images/edits");
let mut form = Form::new();
src/proxy/token_router/router.rs
+126 -53
View File
@@ -2,16 +2,52 @@ use super::types::{RouteDecision, RouteMatch};
use crate::error::{AppError, AppResult};
use crate::proxy::conductor::{FailureInfo, SelectedChannel, SuccessMetrics};
use crate::proxy::executors::UpstreamRoute;
use once_cell::sync::Lazy;
use parking_lot::RwLock;
use regex::Regex;
use sqlx::{FromRow, SqlitePool};
use std::sync::Arc;
use std::{
collections::HashMap,
sync::Arc,
time::{Duration, Instant},
};
const ROUTE_CACHE_TTL: Duration = Duration::from_secs(2);
static ROUTE_ROWS_CACHE: Lazy<RwLock<Option<TimedCache<Vec<RouteRow>>>>> =
Lazy::new(|| RwLock::new(None));
static CHANNEL_ROWS_CACHE: Lazy<RwLock<HashMap<i64, TimedCache<Vec<ChannelRow>>>>> =
Lazy::new(|| RwLock::new(HashMap::new()));
#[derive(Debug, Clone)]
struct TimedCache<T> {
loaded_at: Instant,
value: T,
}
impl<T> TimedCache<T> {
fn fresh(&self) -> bool {
self.loaded_at.elapsed() <= ROUTE_CACHE_TTL
}
}
pub fn invalidate_route_cache() {
*ROUTE_ROWS_CACHE.write() = None;
CHANNEL_ROWS_CACHE.write().clear();
}
fn invalidate_channel_cache_for_channel(channel_id: i64) {
CHANNEL_ROWS_CACHE
.write()
.retain(|_, cache| !cache.value.iter().any(|row| row.channel_id == channel_id));
}
#[derive(Clone)]
pub struct TokenRouter {
db: Arc<SqlitePool>,
}
#[derive(Debug, FromRow)]
#[derive(Debug, Clone, FromRow)]
struct RouteRow {
id: i64,
model_pattern: String,
@@ -19,7 +55,7 @@ struct RouteRow {
route_mode: Option<String>,
}
#[derive(Debug, FromRow)]
#[derive(Debug, Clone, FromRow)]
struct ChannelRow {
channel_id: i64,
route_id: i64,
@@ -91,14 +127,7 @@ impl TokenRouter {
}
async fn match_route(&self, model: &str) -> AppResult<Option<(RouteRow, RouteMatch)>> {
let rows: Vec<RouteRow> = sqlx::query_as(
"SELECT id, model_pattern, model_mapping, route_mode
FROM token_routes
WHERE enabled = 1
ORDER BY id",
)
.fetch_all(&*self.db)
.await?;
let rows = self.route_rows().await?;
let mut best: Option<(RouteRow, RouteMatch, usize, usize)> = None;
for row in rows {
@@ -131,48 +160,7 @@ impl TokenRouter {
requested_model: &str,
exclude_channel_ids: &[i64],
) -> AppResult<Option<SelectedChannel>> {
let rows: Vec<ChannelRow> = sqlx::query_as(
"SELECT
rc.id AS channel_id,
rc.route_id AS route_id,
rc.account_id AS account_id,
a.site_id AS site_id,
COALESCE((
SELECT sae.url
FROM site_api_endpoints sae
WHERE sae.site_id = s.id
AND COALESCE(sae.enabled, 1) = 1
AND (sae.cooldown_until IS NULL OR sae.cooldown_until <= datetime('now'))
ORDER BY COALESCE(sae.sort_order, 0), sae.id
LIMIT 1
), s.url) AS base_url,
s.api_key AS site_api_key,
a.access_token AS account_access_token,
a.api_token AS account_api_token,
a.unit_cost AS account_unit_cost,
at.name AS token_name,
at.token AS token_value,
rc.source_model AS source_model
FROM route_channels rc
JOIN accounts a ON a.id = rc.account_id
JOIN sites s ON s.id = a.site_id
LEFT JOIN account_tokens at ON at.id = rc.token_id AND at.enabled = 1
WHERE rc.route_id = ?
AND rc.enabled = 1
AND a.status = 'active'
AND s.status = 'active'
AND (rc.cooldown_until IS NULL OR rc.cooldown_until <= datetime('now'))
AND (rc.token_id IS NULL OR at.id IS NOT NULL)
ORDER BY
COALESCE(rc.priority, 0) DESC,
rc.consecutive_fail_count ASC,
COALESCE(rc.weight, 10) DESC,
rc.last_selected_at IS NOT NULL,
rc.id ASC",
)
.bind(route.id)
.fetch_all(&*self.db)
.await?;
let rows = self.channel_rows(route.id).await?;
let row = rows
.into_iter()
@@ -214,6 +202,89 @@ impl TokenRouter {
}))
}
async fn route_rows(&self) -> AppResult<Vec<RouteRow>> {
if let Some(cache) = ROUTE_ROWS_CACHE.read().as_ref() {
if cache.fresh() {
return Ok(cache.value.clone());
}
}
let rows: Vec<RouteRow> = sqlx::query_as(
"SELECT id, model_pattern, model_mapping, route_mode
FROM token_routes
WHERE enabled = 1
ORDER BY id",
)
.fetch_all(&*self.db)
.await?;
*ROUTE_ROWS_CACHE.write() = Some(TimedCache {
loaded_at: Instant::now(),
value: rows.clone(),
});
Ok(rows)
}
async fn channel_rows(&self, route_id: i64) -> AppResult<Vec<ChannelRow>> {
if let Some(cache) = CHANNEL_ROWS_CACHE.read().get(&route_id) {
if cache.fresh() {
return Ok(cache.value.clone());
}
}
let rows: Vec<ChannelRow> = sqlx::query_as(
"SELECT
rc.id AS channel_id,
rc.route_id AS route_id,
rc.account_id AS account_id,
a.site_id AS site_id,
COALESCE((
SELECT sae.url
FROM site_api_endpoints sae
WHERE sae.site_id = s.id
AND COALESCE(sae.enabled, 1) = 1
AND (sae.cooldown_until IS NULL OR sae.cooldown_until <= datetime('now'))
ORDER BY COALESCE(sae.sort_order, 0), sae.id
LIMIT 1
), s.url) AS base_url,
s.api_key AS site_api_key,
a.access_token AS account_access_token,
a.api_token AS account_api_token,
a.unit_cost AS account_unit_cost,
at.name AS token_name,
at.token AS token_value,
rc.source_model AS source_model
FROM route_channels rc
JOIN accounts a ON a.id = rc.account_id
JOIN sites s ON s.id = a.site_id
LEFT JOIN account_tokens at ON at.id = rc.token_id AND at.enabled = 1
WHERE rc.route_id = ?
AND rc.enabled = 1
AND a.status = 'active'
AND s.status = 'active'
AND (rc.cooldown_until IS NULL OR rc.cooldown_until <= datetime('now'))
AND (rc.token_id IS NULL OR at.id IS NOT NULL)
ORDER BY
COALESCE(rc.priority, 0) DESC,
rc.consecutive_fail_count ASC,
COALESCE(rc.weight, 10) DESC,
rc.last_selected_at IS NOT NULL,
rc.id ASC",
)
.bind(route_id)
.fetch_all(&*self.db)
.await?;
CHANNEL_ROWS_CACHE.write().insert(
route_id,
TimedCache {
loaded_at: Instant::now(),
value: rows.clone(),
},
);
Ok(rows)
}
pub async fn record_success(&self, channel_id: i64, metrics: SuccessMetrics) -> AppResult<()> {
sqlx::query(
"UPDATE route_channels
@@ -232,6 +303,7 @@ impl TokenRouter {
.bind(channel_id)
.execute(&*self.db)
.await?;
invalidate_channel_cache_for_channel(channel_id);
Ok(())
}
@@ -254,6 +326,7 @@ impl TokenRouter {
.bind(channel_id)
.execute(&*self.db)
.await?;
invalidate_channel_cache_for_channel(channel_id);
Ok(())
}
}