Compare commits
3 Commits
bfd804f719
...
4d093548de
| Author | SHA1 | Date | |
|---|---|---|---|
| 4d093548de | |||
| aaa593ac31 | |||
| 24a8dec151 |
@@ -1159,3 +1159,215 @@ pub async fn execute_image_edit(
|
|||||||
|
|
||||||
Ok(json)
|
Ok(json)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ── Gemini native passthrough ───────────────────────────────────────────
|
||||||
|
//
|
||||||
|
// Gemini's native protocol carries the model in the URL path
|
||||||
|
// (`/v1beta/models/{model}:{op}`) and the API key in the `?key=` query, so the
|
||||||
|
// request body needs no rewrite — we forward it verbatim. This mirrors AQBot's
|
||||||
|
// `handle_native_request`, which deliberately skips `rewrite_body_model` for
|
||||||
|
// Gemini, and BoosAPI's own `/v1/responses` passthrough.
|
||||||
|
|
||||||
|
/// Normalize an upstream base into a Gemini `v1beta` root. Accepts bases that
|
||||||
|
/// already end in `/v1beta` (or contain `/v1beta/`) and leaves them as-is;
|
||||||
|
/// otherwise appends `/v1beta`. Trailing slashes are trimmed.
|
||||||
|
fn gemini_base_url(base_url: &str) -> String {
|
||||||
|
let trimmed = base_url.trim().trim_end_matches('/');
|
||||||
|
if trimmed.ends_with("/v1beta") || trimmed.contains("/v1beta/") {
|
||||||
|
trimmed.to_string()
|
||||||
|
} else {
|
||||||
|
format!("{trimmed}/v1beta")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build `{base}/v1beta/models/{model}:{op}?key={api_key}[&{extra_query}]`.
|
||||||
|
/// `extra_query` carries client query params already filtered (e.g. `alt=sse`),
|
||||||
|
/// without a leading `?`/`&`.
|
||||||
|
fn gemini_operation_url(
|
||||||
|
base_url: &str,
|
||||||
|
model: &str,
|
||||||
|
operation: &str,
|
||||||
|
api_key: &str,
|
||||||
|
extra_query: Option<&str>,
|
||||||
|
) -> String {
|
||||||
|
let base = gemini_base_url(base_url);
|
||||||
|
let model = model.trim().trim_start_matches('/');
|
||||||
|
let model = model.strip_prefix("models/").unwrap_or(model);
|
||||||
|
let mut url = format!("{base}/models/{model}:{operation}");
|
||||||
|
let mut sep = '?';
|
||||||
|
if !api_key.is_empty() {
|
||||||
|
url.push(sep);
|
||||||
|
url.push_str("key=");
|
||||||
|
url.push_str(api_key);
|
||||||
|
sep = '&';
|
||||||
|
}
|
||||||
|
if let Some(query) = extra_query.filter(|q| !q.is_empty()) {
|
||||||
|
url.push(sep);
|
||||||
|
url.push_str(query);
|
||||||
|
}
|
||||||
|
url
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute a non-streaming Gemini operation (`generateContent` / `countTokens`)
|
||||||
|
/// and return the upstream JSON verbatim.
|
||||||
|
pub async fn execute_gemini_operation(
|
||||||
|
state: AppState,
|
||||||
|
route: UpstreamRoute,
|
||||||
|
operation: &str,
|
||||||
|
body: Value,
|
||||||
|
extra_query: Option<&str>,
|
||||||
|
) -> AppResult<Value> {
|
||||||
|
let client = state.http.clone();
|
||||||
|
let upstream_url =
|
||||||
|
gemini_operation_url(&route.base_url, &route.model, operation, &route.api_key, extra_query);
|
||||||
|
|
||||||
|
let resp = client
|
||||||
|
.post(&upstream_url)
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.json(&body)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
|
||||||
|
|
||||||
|
let status = resp.status();
|
||||||
|
let text = resp
|
||||||
|
.text()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::UpstreamError(format!("Failed to read upstream response: {e}")))?;
|
||||||
|
|
||||||
|
if !status.is_success() {
|
||||||
|
return Err(AppError::UpstreamError(format!(
|
||||||
|
"Upstream returned {status}: {text}"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
parse_upstream_json(&text, &upstream_url, status)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute a streaming Gemini operation (`streamGenerateContent`) and return the
|
||||||
|
/// live upstream response for SSE passthrough. `alt=sse` is enforced in the
|
||||||
|
/// query so the upstream emits `data:`-framed events the accumulator can parse.
|
||||||
|
pub async fn execute_gemini_stream_response(
|
||||||
|
state: AppState,
|
||||||
|
route: UpstreamRoute,
|
||||||
|
operation: &str,
|
||||||
|
body: Value,
|
||||||
|
extra_query: Option<&str>,
|
||||||
|
) -> AppResult<reqwest::Response> {
|
||||||
|
let client = state.http.clone();
|
||||||
|
let query = merge_alt_sse(extra_query);
|
||||||
|
let upstream_url = gemini_operation_url(
|
||||||
|
&route.base_url,
|
||||||
|
&route.model,
|
||||||
|
operation,
|
||||||
|
&route.api_key,
|
||||||
|
Some(&query),
|
||||||
|
);
|
||||||
|
|
||||||
|
let resp = client
|
||||||
|
.post(&upstream_url)
|
||||||
|
.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)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Ensure `alt=sse` is present in the extra query (Gemini streams plain JSON
|
||||||
|
/// arrays without it). Preserves any other client-supplied params.
|
||||||
|
fn merge_alt_sse(extra_query: Option<&str>) -> String {
|
||||||
|
match extra_query.filter(|q| !q.is_empty()) {
|
||||||
|
Some(q) if q.split('&').any(|p| p == "alt=sse" || p.starts_with("alt=")) => q.to_string(),
|
||||||
|
Some(q) => format!("{q}&alt=sse"),
|
||||||
|
None => "alt=sse".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod gemini_tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn base_url_appends_v1beta_when_missing() {
|
||||||
|
assert_eq!(
|
||||||
|
gemini_base_url("https://generativelanguage.googleapis.com"),
|
||||||
|
"https://generativelanguage.googleapis.com/v1beta"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn base_url_preserves_existing_v1beta() {
|
||||||
|
assert_eq!(
|
||||||
|
gemini_base_url("https://generativelanguage.googleapis.com/v1beta/"),
|
||||||
|
"https://generativelanguage.googleapis.com/v1beta"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn operation_url_embeds_model_op_and_key() {
|
||||||
|
let url = gemini_operation_url(
|
||||||
|
"https://generativelanguage.googleapis.com",
|
||||||
|
"gemini-2.0-flash",
|
||||||
|
"generateContent",
|
||||||
|
"secret",
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
url,
|
||||||
|
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key=secret"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn operation_url_strips_models_prefix_and_appends_query() {
|
||||||
|
let url = gemini_operation_url(
|
||||||
|
"https://generativelanguage.googleapis.com/v1beta",
|
||||||
|
"models/gemini-1.5-pro",
|
||||||
|
"streamGenerateContent",
|
||||||
|
"k",
|
||||||
|
Some("alt=sse"),
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
url,
|
||||||
|
"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro:streamGenerateContent?key=k&alt=sse"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn operation_url_without_key_uses_question_mark_for_query() {
|
||||||
|
let url = gemini_operation_url(
|
||||||
|
"https://example.test/v1beta",
|
||||||
|
"gemini-x",
|
||||||
|
"countTokens",
|
||||||
|
"",
|
||||||
|
Some("foo=bar"),
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
url,
|
||||||
|
"https://example.test/v1beta/models/gemini-x:countTokens?foo=bar"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn merge_alt_sse_variants() {
|
||||||
|
assert_eq!(merge_alt_sse(None), "alt=sse");
|
||||||
|
assert_eq!(merge_alt_sse(Some("")), "alt=sse");
|
||||||
|
assert_eq!(merge_alt_sse(Some("foo=bar")), "foo=bar&alt=sse");
|
||||||
|
// already present → unchanged
|
||||||
|
assert_eq!(merge_alt_sse(Some("alt=sse")), "alt=sse");
|
||||||
|
assert_eq!(merge_alt_sse(Some("alt=json")), "alt=json");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -2840,6 +2840,8 @@ mod tests {
|
|||||||
template_proxy_url: None,
|
template_proxy_url: None,
|
||||||
},
|
},
|
||||||
comfyui: Default::default(),
|
comfyui: Default::default(),
|
||||||
|
payment: Default::default(),
|
||||||
|
pixelle_video: Default::default(),
|
||||||
},
|
},
|
||||||
pool,
|
pool,
|
||||||
BuiltinKeys {
|
BuiltinKeys {
|
||||||
|
|||||||
@@ -412,6 +412,8 @@ mod tests {
|
|||||||
template_proxy_url: None,
|
template_proxy_url: None,
|
||||||
},
|
},
|
||||||
comfyui: Default::default(),
|
comfyui: Default::default(),
|
||||||
|
payment: Default::default(),
|
||||||
|
pixelle_video: Default::default(),
|
||||||
},
|
},
|
||||||
pool,
|
pool,
|
||||||
BuiltinKeys {
|
BuiltinKeys {
|
||||||
|
|||||||
@@ -19,6 +19,11 @@ use serde_json::{json, Value};
|
|||||||
use std::net::{IpAddr, SocketAddr};
|
use std::net::{IpAddr, SocketAddr};
|
||||||
use tracing::{info, warn};
|
use tracing::{info, warn};
|
||||||
|
|
||||||
|
/// Default spending cap (USD) granted to a newly created user's auto-generated
|
||||||
|
/// API key. Per the billing design (新用户默认 $1 额度), new users start with a
|
||||||
|
/// $1 allowance on their default downstream key.
|
||||||
|
pub const DEFAULT_NEW_USER_MAX_COST_USD: f64 = 1.0;
|
||||||
|
|
||||||
// ── Request / Response types ─────────────────────────────────────────
|
// ── Request / Response types ─────────────────────────────────────────
|
||||||
|
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
@@ -1019,6 +1024,7 @@ mod tests {
|
|||||||
},
|
},
|
||||||
comfyui: Default::default(),
|
comfyui: Default::default(),
|
||||||
payment: Default::default(),
|
payment: Default::default(),
|
||||||
|
pixelle_video: Default::default(),
|
||||||
},
|
},
|
||||||
pool,
|
pool,
|
||||||
BuiltinKeys {
|
BuiltinKeys {
|
||||||
|
|||||||
@@ -127,6 +127,8 @@ mod tests {
|
|||||||
template_proxy_url: None,
|
template_proxy_url: None,
|
||||||
},
|
},
|
||||||
comfyui: ComfyUIConfig::default(),
|
comfyui: ComfyUIConfig::default(),
|
||||||
|
payment: Default::default(),
|
||||||
|
pixelle_video: Default::default(),
|
||||||
},
|
},
|
||||||
pool,
|
pool,
|
||||||
BuiltinKeys {
|
BuiltinKeys {
|
||||||
|
|||||||
@@ -1305,6 +1305,8 @@ mod tests {
|
|||||||
template_proxy_url: None,
|
template_proxy_url: None,
|
||||||
},
|
},
|
||||||
comfyui: Default::default(),
|
comfyui: Default::default(),
|
||||||
|
payment: Default::default(),
|
||||||
|
pixelle_video: Default::default(),
|
||||||
},
|
},
|
||||||
pool,
|
pool,
|
||||||
BuiltinKeys {
|
BuiltinKeys {
|
||||||
|
|||||||
@@ -536,6 +536,8 @@ mod tests {
|
|||||||
template_proxy_url: None,
|
template_proxy_url: None,
|
||||||
},
|
},
|
||||||
comfyui: Default::default(),
|
comfyui: Default::default(),
|
||||||
|
payment: Default::default(),
|
||||||
|
pixelle_video: Default::default(),
|
||||||
},
|
},
|
||||||
pool,
|
pool,
|
||||||
BuiltinKeys {
|
BuiltinKeys {
|
||||||
|
|||||||
@@ -2071,6 +2071,8 @@ mod tests {
|
|||||||
template_proxy_url: None,
|
template_proxy_url: None,
|
||||||
},
|
},
|
||||||
comfyui,
|
comfyui,
|
||||||
|
payment: Default::default(),
|
||||||
|
pixelle_video: Default::default(),
|
||||||
},
|
},
|
||||||
pool,
|
pool,
|
||||||
BuiltinKeys {
|
BuiltinKeys {
|
||||||
|
|||||||
@@ -2550,6 +2550,8 @@ mod tests {
|
|||||||
template_proxy_url: None,
|
template_proxy_url: None,
|
||||||
},
|
},
|
||||||
comfyui: ComfyUIConfig::default(),
|
comfyui: ComfyUIConfig::default(),
|
||||||
|
payment: Default::default(),
|
||||||
|
pixelle_video: Default::default(),
|
||||||
},
|
},
|
||||||
pool,
|
pool,
|
||||||
BuiltinKeys {
|
BuiltinKeys {
|
||||||
|
|||||||
@@ -73,6 +73,16 @@ pub fn routes() -> Router<AppState> {
|
|||||||
.route("/v1/videos", post(videos_handler))
|
.route("/v1/videos", post(videos_handler))
|
||||||
.route("/v1/videos/{id}", get(video_task_handler))
|
.route("/v1/videos/{id}", get(video_task_handler))
|
||||||
.route("/v1/models", get(list_models))
|
.route("/v1/models", get(list_models))
|
||||||
|
// ── Gemini native inbound (v1beta) ────────────────────────────
|
||||||
|
// Lets Gemini CLI / Google GenAI SDKs point their base URL straight
|
||||||
|
// at this gateway. The model lives in the path and the key in the
|
||||||
|
// `?key=` query (or x-goog-api-key header), so requests pass through
|
||||||
|
// to a Gemini-shaped upstream without body rewriting.
|
||||||
|
.route("/v1beta/models", get(gemini_list_models))
|
||||||
|
.route(
|
||||||
|
"/v1beta/models/{model_action}",
|
||||||
|
post(gemini_model_operation),
|
||||||
|
)
|
||||||
// ── ComfyUI shim aliases ──────────────────────────────────────
|
// ── ComfyUI shim aliases ──────────────────────────────────────
|
||||||
// ComfyUI's OpenAI nodes hit /proxy/openai/... paths (see
|
// ComfyUI's OpenAI nodes hit /proxy/openai/... paths (see
|
||||||
// comfy_api_nodes/nodes_openai.py). Mounting them here as aliases
|
// comfy_api_nodes/nodes_openai.py). Mounting them here as aliases
|
||||||
@@ -157,6 +167,209 @@ async fn responses_handler(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ── Gemini native inbound handlers (v1beta) ────────────────────────────
|
||||||
|
//
|
||||||
|
// Gemini clients (Gemini CLI / Google GenAI SDK) carry the model in the URL
|
||||||
|
// path (`/v1beta/models/{model}:{op}`) and the gateway key in `?key=` or the
|
||||||
|
// `x-goog-api-key` header. We authenticate against our downstream keys, route
|
||||||
|
// the Gemini model through the same ProxyConductor as the OpenAI endpoints, and
|
||||||
|
// forward the body verbatim to a Gemini-shaped upstream (no body rewrite).
|
||||||
|
|
||||||
|
/// Gemini operations carried in the `{model}:{op}` path segment.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
enum GeminiOperation {
|
||||||
|
GenerateContent,
|
||||||
|
StreamGenerateContent,
|
||||||
|
CountTokens,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl GeminiOperation {
|
||||||
|
fn as_str(self) -> &'static str {
|
||||||
|
match self {
|
||||||
|
GeminiOperation::GenerateContent => "generateContent",
|
||||||
|
GeminiOperation::StreamGenerateContent => "streamGenerateContent",
|
||||||
|
GeminiOperation::CountTokens => "countTokens",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Parse the `{model}:{operation}` path segment (e.g.
|
||||||
|
/// `gemini-2.5-pro:streamGenerateContent` or `models/gemini-x:countTokens`).
|
||||||
|
/// Splits on the last `:` so model names that contain none are rejected.
|
||||||
|
fn parse_gemini_model_action(model_action: &str) -> Option<(String, GeminiOperation)> {
|
||||||
|
let (model, action) = model_action.rsplit_once(':')?;
|
||||||
|
let operation = match action {
|
||||||
|
"generateContent" => GeminiOperation::GenerateContent,
|
||||||
|
"streamGenerateContent" => GeminiOperation::StreamGenerateContent,
|
||||||
|
"countTokens" => GeminiOperation::CountTokens,
|
||||||
|
_ => return None,
|
||||||
|
};
|
||||||
|
if model.trim().is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
Some((model.to_string(), operation))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Authenticate a Gemini-native request: the key arrives in `?key=` or the
|
||||||
|
/// `x-goog-api-key` header (Gemini convention), not `Authorization: Bearer`.
|
||||||
|
/// Reuses the same downstream-key authorization + limits as `check_auth`.
|
||||||
|
async fn check_gemini_auth(
|
||||||
|
state: &AppState,
|
||||||
|
headers: &HeaderMap,
|
||||||
|
query: &std::collections::HashMap<String, String>,
|
||||||
|
) -> Result<Option<i64>, AppError> {
|
||||||
|
let key = query
|
||||||
|
.get("key")
|
||||||
|
.map(|s| s.as_str())
|
||||||
|
.or_else(|| headers.get("x-goog-api-key").and_then(|v| v.to_str().ok()))
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|s| !s.is_empty())
|
||||||
|
.ok_or_else(|| {
|
||||||
|
AppError::Unauthorized("Missing API key (use ?key= or x-goog-api-key)".into())
|
||||||
|
})?;
|
||||||
|
|
||||||
|
let proxy_token = state.runtime_proxy_token();
|
||||||
|
let result = authorize_downstream_token(key, &state.db, Some(&proxy_token))
|
||||||
|
.await
|
||||||
|
.map_err(AppError::Unauthorized)?;
|
||||||
|
|
||||||
|
match result {
|
||||||
|
DownstreamAuthResult::Managed { key } => {
|
||||||
|
check_downstream_key_limits(&key)?;
|
||||||
|
Ok(Some(key.id))
|
||||||
|
}
|
||||||
|
DownstreamAuthResult::Global { .. } => Ok(None),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Rebuild the upstream query string from client params, dropping the inbound
|
||||||
|
/// gateway `key` (the upstream key is injected by the executor) and `alt` (the
|
||||||
|
/// streaming path enforces its own `alt=sse`). Returns `None` when nothing
|
||||||
|
/// remains. Values are forwarded as-is: real Gemini clients only put `key`/`alt`
|
||||||
|
/// in the query and carry everything else (generationConfig, etc.) in the body.
|
||||||
|
fn gemini_forward_query(
|
||||||
|
query: &std::collections::HashMap<String, String>,
|
||||||
|
drop_alt: bool,
|
||||||
|
) -> Option<String> {
|
||||||
|
let mut parts: Vec<(&String, &String)> = query
|
||||||
|
.iter()
|
||||||
|
.filter(|(k, _)| k.as_str() != "key" && !(drop_alt && k.as_str() == "alt"))
|
||||||
|
.collect();
|
||||||
|
if parts.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
// Stable order so the URL is deterministic (and testable).
|
||||||
|
parts.sort_by(|a, b| a.0.cmp(b.0));
|
||||||
|
let encoded = parts
|
||||||
|
.iter()
|
||||||
|
.map(|(k, v)| format!("{k}={v}"))
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join("&");
|
||||||
|
Some(encoded)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GET /v1beta/models — list models in Gemini's native shape.
|
||||||
|
/// Like `/v1/models`, this is a directory listing from configured routes; no
|
||||||
|
/// upstream call, no billing.
|
||||||
|
async fn gemini_list_models(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Query(query): Query<std::collections::HashMap<String, String>>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
info!("GET /v1beta/models");
|
||||||
|
|
||||||
|
let _ = check_gemini_auth(&state, &headers, &query).await?;
|
||||||
|
|
||||||
|
let rows: Vec<(String, String)> = sqlx::query_as(
|
||||||
|
"SELECT DISTINCT
|
||||||
|
tr.model_pattern,
|
||||||
|
COALESCE(NULLIF(TRIM(tr.model_mapping), ''), tr.model_pattern)
|
||||||
|
FROM token_routes tr
|
||||||
|
JOIN route_channels rc ON rc.route_id = tr.id
|
||||||
|
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 tr.enabled = 1
|
||||||
|
AND rc.enabled = 1
|
||||||
|
AND a.status = 'active'
|
||||||
|
AND s.status = 'active'
|
||||||
|
AND TRIM(COALESCE(s.url, '')) <> ''
|
||||||
|
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 tr.id",
|
||||||
|
)
|
||||||
|
.fetch_all(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let models: Vec<Value> = rows
|
||||||
|
.into_iter()
|
||||||
|
.map(|(_pattern, model)| {
|
||||||
|
let name = if model.starts_with("models/") {
|
||||||
|
model.clone()
|
||||||
|
} else {
|
||||||
|
format!("models/{model}")
|
||||||
|
};
|
||||||
|
json!({
|
||||||
|
"name": name,
|
||||||
|
"supportedGenerationMethods": [
|
||||||
|
"generateContent",
|
||||||
|
"streamGenerateContent",
|
||||||
|
"countTokens",
|
||||||
|
],
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
Ok(Json(json!({ "models": models })))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /v1beta/models/{model}:{op} — dispatch to the matching Gemini op.
|
||||||
|
async fn gemini_model_operation(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
axum::extract::Path(model_action): axum::extract::Path<String>,
|
||||||
|
Query(query): Query<std::collections::HashMap<String, String>>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> Result<Response, AppError> {
|
||||||
|
let Some((model, operation)) = parse_gemini_model_action(&model_action) else {
|
||||||
|
return Err(AppError::BadRequest(format!(
|
||||||
|
"Unsupported Gemini model action '{model_action}' \
|
||||||
|
(expected {{model}}:generateContent|streamGenerateContent|countTokens)"
|
||||||
|
)));
|
||||||
|
};
|
||||||
|
info!("POST /v1beta/models/{model}:{}", operation.as_str());
|
||||||
|
|
||||||
|
let key_id = check_gemini_auth(&state, &headers, &query).await?;
|
||||||
|
|
||||||
|
match operation {
|
||||||
|
GeminiOperation::GenerateContent => {
|
||||||
|
let extra_query = gemini_forward_query(&query, false);
|
||||||
|
let result =
|
||||||
|
execute_gemini_with_conductor(state, key_id, model, body, extra_query).await?;
|
||||||
|
Ok(Json(result).into_response())
|
||||||
|
}
|
||||||
|
GeminiOperation::StreamGenerateContent => {
|
||||||
|
let extra_query = gemini_forward_query(&query, true);
|
||||||
|
execute_gemini_stream_with_conductor(state, key_id, model, body, extra_query).await
|
||||||
|
}
|
||||||
|
GeminiOperation::CountTokens => {
|
||||||
|
// countTokens is a metadata call (no generation) — route for a
|
||||||
|
// channel but skip billing, mirroring the lighter helper endpoints.
|
||||||
|
let extra_query = gemini_forward_query(&query, false);
|
||||||
|
let route = resolve_route(&state, &model).await?;
|
||||||
|
let result = executors::execute_gemini_operation(
|
||||||
|
state,
|
||||||
|
route,
|
||||||
|
"countTokens",
|
||||||
|
body,
|
||||||
|
extra_query.as_deref(),
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
Ok(Json(result).into_response())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// POST /v1/chat/completions
|
/// POST /v1/chat/completions
|
||||||
async fn chat_completions(
|
async fn chat_completions(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
@@ -1334,6 +1547,38 @@ fn responses_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
|
|||||||
(prompt != 0 || completion != 0).then_some((prompt, completion))
|
(prompt != 0 || completion != 0).then_some((prompt, completion))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Extract `(prompt, completion)` tokens from a Gemini response body's
|
||||||
|
/// `usageMetadata` (`promptTokenCount` / `candidatesTokenCount`). Returns
|
||||||
|
/// `(0, 0)` when the field is absent — mirrors `extract_usage` for OpenAI.
|
||||||
|
fn gemini_extract_usage(response: &Value) -> (i64, i64) {
|
||||||
|
let usage = response.get("usageMetadata");
|
||||||
|
let prompt = usage
|
||||||
|
.and_then(|u| u.get("promptTokenCount"))
|
||||||
|
.and_then(|v| v.as_i64())
|
||||||
|
.unwrap_or(0);
|
||||||
|
let completion = usage
|
||||||
|
.and_then(|u| u.get("candidatesTokenCount"))
|
||||||
|
.and_then(|v| v.as_i64())
|
||||||
|
.unwrap_or(0);
|
||||||
|
(prompt, completion)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Per-line usage parser for Gemini `streamGenerateContent?alt=sse` output.
|
||||||
|
/// Gemini emits `data:`-framed JSON chunks; the final chunk carries the
|
||||||
|
/// cumulative `usageMetadata`. Feeds the shared `StreamUsageAccumulator`.
|
||||||
|
fn gemini_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
|
||||||
|
let payload = line
|
||||||
|
.strip_prefix("data: ")
|
||||||
|
.or_else(|| line.strip_prefix("data:"))?
|
||||||
|
.trim();
|
||||||
|
if payload.is_empty() || payload == "[DONE]" {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let value: Value = serde_json::from_str(payload).ok()?;
|
||||||
|
let (prompt, completion) = gemini_extract_usage(&value);
|
||||||
|
(prompt != 0 || completion != 0).then_some((prompt, completion))
|
||||||
|
}
|
||||||
|
|
||||||
fn build_conductor(state: &AppState) -> ProxyConductor {
|
fn build_conductor(state: &AppState) -> ProxyConductor {
|
||||||
let db = state.db.clone();
|
let db = state.db.clone();
|
||||||
let select_db = db.clone();
|
let select_db = db.clone();
|
||||||
@@ -1881,6 +2126,185 @@ async fn execute_responses_stream_with_conductor(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Non-streaming Gemini op (`generateContent`) through the ProxyConductor.
|
||||||
|
/// Mirrors `execute_responses_with_conductor` but builds a Gemini upstream URL
|
||||||
|
/// (model in path, key in query) and forwards the body verbatim — Gemini bodies
|
||||||
|
/// carry no `model` field, so there is no body rewrite.
|
||||||
|
async fn execute_gemini_with_conductor(
|
||||||
|
state: AppState,
|
||||||
|
key_id: Option<i64>,
|
||||||
|
model: String,
|
||||||
|
body: Value,
|
||||||
|
extra_query: Option<String>,
|
||||||
|
) -> AppResult<Value> {
|
||||||
|
let started_at = Instant::now();
|
||||||
|
let request_log_body = body.clone();
|
||||||
|
let conductor = build_conductor(&state);
|
||||||
|
let billing_state = state.clone();
|
||||||
|
let billing_model = model.clone();
|
||||||
|
let result = conductor
|
||||||
|
.execute(ExecuteInput {
|
||||||
|
requested_model: model,
|
||||||
|
downstream_policy: None,
|
||||||
|
attempt_fn: Box::new(move |ctx| {
|
||||||
|
let state = state.clone();
|
||||||
|
let body = body.clone();
|
||||||
|
let extra_query = extra_query.clone();
|
||||||
|
async move {
|
||||||
|
match executors::execute_gemini_operation(
|
||||||
|
state,
|
||||||
|
selected_channel_to_route(ctx.selected),
|
||||||
|
"generateContent",
|
||||||
|
body,
|
||||||
|
extra_query.as_deref(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(response) => AttemptResult::Success {
|
||||||
|
response,
|
||||||
|
latency_ms: None,
|
||||||
|
cost: None,
|
||||||
|
},
|
||||||
|
Err(err) => AttemptResult::Failure {
|
||||||
|
action: AttemptFailureAction::Failover,
|
||||||
|
status: upstream_status_from_error(&err),
|
||||||
|
raw_error_text: Some(err.to_string()),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
}),
|
||||||
|
on_terminal_failure: None,
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let (selected, response, attempts) = match result {
|
||||||
|
ExecuteResult::Success {
|
||||||
|
selected,
|
||||||
|
response,
|
||||||
|
attempts,
|
||||||
|
} => (selected, response, attempts),
|
||||||
|
ExecuteResult::Failure {
|
||||||
|
reason,
|
||||||
|
selected,
|
||||||
|
status,
|
||||||
|
raw_error_text,
|
||||||
|
attempts,
|
||||||
|
} => {
|
||||||
|
record_proxy_log_failure(
|
||||||
|
&billing_state,
|
||||||
|
key_id,
|
||||||
|
selected.as_ref(),
|
||||||
|
&billing_model,
|
||||||
|
status,
|
||||||
|
Some(elapsed_ms_since(started_at)),
|
||||||
|
raw_error_text.as_deref(),
|
||||||
|
attempts,
|
||||||
|
false,
|
||||||
|
proxy_log_content_from_request(&request_log_body),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
return Err(AppError::UpstreamError(format!(
|
||||||
|
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
|
||||||
|
status,
|
||||||
|
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let (prompt, completion) = gemini_extract_usage(&response);
|
||||||
|
let cost = estimate_request_cost(&selected, prompt, completion);
|
||||||
|
record_proxy_log_success(
|
||||||
|
&billing_state,
|
||||||
|
key_id,
|
||||||
|
&selected,
|
||||||
|
&billing_model,
|
||||||
|
Some(elapsed_ms_since(started_at)),
|
||||||
|
prompt,
|
||||||
|
completion,
|
||||||
|
cost,
|
||||||
|
attempts,
|
||||||
|
false,
|
||||||
|
proxy_log_content_from_request_response(&request_log_body, Some(&response)),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
record_billing(
|
||||||
|
&billing_state,
|
||||||
|
key_id,
|
||||||
|
&billing_model,
|
||||||
|
prompt,
|
||||||
|
completion,
|
||||||
|
cost,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
Ok(response)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Streaming Gemini op (`streamGenerateContent`) through the ProxyConductor.
|
||||||
|
/// Mirrors `execute_responses_stream_with_conductor`, opening the upstream with
|
||||||
|
/// failover and SSE-passing it through with the Gemini usage parser so streamed
|
||||||
|
/// `usageMetadata` is billed.
|
||||||
|
async fn execute_gemini_stream_with_conductor(
|
||||||
|
state: AppState,
|
||||||
|
key_id: Option<i64>,
|
||||||
|
model: String,
|
||||||
|
body: Value,
|
||||||
|
extra_query: Option<String>,
|
||||||
|
) -> Result<Response, AppError> {
|
||||||
|
let started_at = Instant::now();
|
||||||
|
let request_log_body = body.clone();
|
||||||
|
let body_for_attempt = body.clone();
|
||||||
|
let opened = open_stream_with_failover(&state, &model, move |state, selected| {
|
||||||
|
let body = body_for_attempt.clone();
|
||||||
|
let extra_query = extra_query.clone();
|
||||||
|
async move {
|
||||||
|
executors::execute_gemini_stream_response(
|
||||||
|
state,
|
||||||
|
selected_channel_to_route(selected),
|
||||||
|
"streamGenerateContent",
|
||||||
|
body,
|
||||||
|
extra_query.as_deref(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
|
||||||
|
match opened {
|
||||||
|
Ok(opened) => sse_passthrough_response(
|
||||||
|
state,
|
||||||
|
opened,
|
||||||
|
key_id,
|
||||||
|
model,
|
||||||
|
started_at,
|
||||||
|
proxy_log_content_from_request(&request_log_body),
|
||||||
|
gemini_stream_usage_from_line,
|
||||||
|
),
|
||||||
|
Err(failure) => {
|
||||||
|
record_proxy_log_failure(
|
||||||
|
&state,
|
||||||
|
key_id,
|
||||||
|
failure.selected.as_ref(),
|
||||||
|
&model,
|
||||||
|
failure.status,
|
||||||
|
Some(elapsed_ms_since(started_at)),
|
||||||
|
failure.raw_error_text.as_deref(),
|
||||||
|
failure.attempts,
|
||||||
|
true,
|
||||||
|
proxy_log_content_from_request(&request_log_body),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
Err(AppError::UpstreamError(format!(
|
||||||
|
"Proxy conductor failed after {} attempts (stream open, status: {:?}): {}",
|
||||||
|
failure.attempts,
|
||||||
|
failure.status,
|
||||||
|
failure
|
||||||
|
.raw_error_text
|
||||||
|
.unwrap_or_else(|| "no upstream error text".to_string())
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -2121,6 +2545,8 @@ mod tests {
|
|||||||
template_proxy_url: None,
|
template_proxy_url: None,
|
||||||
},
|
},
|
||||||
comfyui: Default::default(),
|
comfyui: Default::default(),
|
||||||
|
payment: Default::default(),
|
||||||
|
pixelle_video: Default::default(),
|
||||||
},
|
},
|
||||||
pool,
|
pool,
|
||||||
crate::services::builtin_keys::BuiltinKeys {
|
crate::services::builtin_keys::BuiltinKeys {
|
||||||
@@ -2155,4 +2581,87 @@ mod tests {
|
|||||||
.unwrap();
|
.unwrap();
|
||||||
assert!(fail_count > 0);
|
assert!(fail_count > 0);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ── Gemini native inbound unit tests ───────────────────────────────
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_gemini_model_action_known_ops() {
|
||||||
|
let (m, op) = parse_gemini_model_action("gemini-2.5-pro:generateContent").unwrap();
|
||||||
|
assert_eq!(m, "gemini-2.5-pro");
|
||||||
|
assert_eq!(op, GeminiOperation::GenerateContent);
|
||||||
|
|
||||||
|
let (m, op) = parse_gemini_model_action("gemini-2.5-pro:streamGenerateContent").unwrap();
|
||||||
|
assert_eq!(m, "gemini-2.5-pro");
|
||||||
|
assert_eq!(op, GeminiOperation::StreamGenerateContent);
|
||||||
|
|
||||||
|
// model names may carry the `models/` prefix
|
||||||
|
let (m, op) = parse_gemini_model_action("models/gemini-1.5-pro:countTokens").unwrap();
|
||||||
|
assert_eq!(m, "models/gemini-1.5-pro");
|
||||||
|
assert_eq!(op, GeminiOperation::CountTokens);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parse_gemini_model_action_rejects_bad_input() {
|
||||||
|
// no operation suffix
|
||||||
|
assert!(parse_gemini_model_action("gemini-2.5-pro").is_none());
|
||||||
|
// unknown operation
|
||||||
|
assert!(parse_gemini_model_action("gemini-2.5-pro:embedContent").is_none());
|
||||||
|
// empty model
|
||||||
|
assert!(parse_gemini_model_action(":generateContent").is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn gemini_extract_usage_reads_usage_metadata() {
|
||||||
|
let body = json!({
|
||||||
|
"candidates": [],
|
||||||
|
"usageMetadata": {
|
||||||
|
"promptTokenCount": 11,
|
||||||
|
"candidatesTokenCount": 22,
|
||||||
|
"totalTokenCount": 33
|
||||||
|
}
|
||||||
|
});
|
||||||
|
assert_eq!(gemini_extract_usage(&body), (11, 22));
|
||||||
|
// absent usageMetadata → zeros
|
||||||
|
assert_eq!(gemini_extract_usage(&json!({})), (0, 0));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn gemini_stream_usage_from_line_parses_data_frame() {
|
||||||
|
let line = r#"data: {"candidates":[{"content":{"parts":[]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":6,"totalTokenCount":10}}"#;
|
||||||
|
assert_eq!(gemini_stream_usage_from_line(line), Some((4, 6)));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn gemini_stream_usage_from_line_ignores_non_usage_and_done() {
|
||||||
|
// a chunk with no usageMetadata yields None (don't overwrite accumulator)
|
||||||
|
let no_usage = r#"data: {"candidates":[{"content":{"parts":[{"text":"hi"}]}}]}"#;
|
||||||
|
assert_eq!(gemini_stream_usage_from_line(no_usage), None);
|
||||||
|
// sentinels / blanks / non-data lines
|
||||||
|
assert_eq!(gemini_stream_usage_from_line("data: [DONE]"), None);
|
||||||
|
assert_eq!(gemini_stream_usage_from_line("data:"), None);
|
||||||
|
assert_eq!(gemini_stream_usage_from_line(": comment"), None);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn gemini_forward_query_drops_key_and_alt() {
|
||||||
|
let mut q = std::collections::HashMap::new();
|
||||||
|
q.insert("key".to_string(), "secret".to_string());
|
||||||
|
q.insert("alt".to_string(), "sse".to_string());
|
||||||
|
// stream path drops both key and alt → nothing left
|
||||||
|
assert_eq!(gemini_forward_query(&q, true), None);
|
||||||
|
// non-stream keeps alt but always drops key
|
||||||
|
assert_eq!(gemini_forward_query(&q, false), Some("alt=sse".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn gemini_forward_query_preserves_other_params_sorted() {
|
||||||
|
let mut q = std::collections::HashMap::new();
|
||||||
|
q.insert("key".to_string(), "secret".to_string());
|
||||||
|
q.insert("pageSize".to_string(), "50".to_string());
|
||||||
|
q.insert("filter".to_string(), "x".to_string());
|
||||||
|
assert_eq!(
|
||||||
|
gemini_forward_query(&q, true),
|
||||||
|
Some("filter=x&pageSize=50".to_string())
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -696,6 +696,8 @@ mod tests {
|
|||||||
template_proxy_url: None,
|
template_proxy_url: None,
|
||||||
},
|
},
|
||||||
comfyui: ComfyUIConfig::default(),
|
comfyui: ComfyUIConfig::default(),
|
||||||
|
payment: Default::default(),
|
||||||
|
pixelle_video: Default::default(),
|
||||||
},
|
},
|
||||||
pool,
|
pool,
|
||||||
BuiltinKeys {
|
BuiltinKeys {
|
||||||
|
|||||||
Reference in New Issue
Block a user