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:
@@ -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");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user