feat(proxy): add Gemini native upstream executors

Add Gemini-shaped upstream calls that mirror AQBot's native passthrough
and BoosAPI's /v1/responses forwarding: the model lives in the URL path
(/v1beta/models/{model}:{op}) and the key in ?key=, so the body is sent
verbatim with no model rewrite.

- gemini_base_url / gemini_operation_url normalize the v1beta root and
  build {base}/v1beta/models/{model}:{op}?key=...&{extra}, stripping a
  leading models/ prefix and handling the no-key case.
- execute_gemini_operation runs non-streaming ops (generateContent /
  countTokens) and returns upstream JSON verbatim.
- execute_gemini_stream_response opens streamGenerateContent for SSE
  passthrough, enforcing alt=sse via merge_alt_sse.

Unit tests cover URL building and alt=sse merging.
This commit is contained in:
boos
2026-06-22 08:39:26 +08:00
parent 24a8dec151
commit aaa593ac31
+212
View File
@@ -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");
}
}