mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
Improve request trace upstream diagnostics
This commit is contained in:
@@ -69,9 +69,10 @@ pub fn build_gemini_content_url(
|
||||
} else {
|
||||
"generateContent"
|
||||
};
|
||||
let mut url = if trimmed_base_url.ends_with("/v1beta") {
|
||||
let mut url = if trimmed_base_url.ends_with("/v1") || trimmed_base_url.ends_with("/v1beta") {
|
||||
format!("{trimmed_base_url}/models/{trimmed_model}:{operation}")
|
||||
} else if trimmed_base_url.contains("/v1beta/models/") {
|
||||
} else if gemini_content_base_url_contains_model_path(trimmed_base_url) {
|
||||
let trimmed_base_url = strip_gemini_content_action(trimmed_base_url);
|
||||
format!("{trimmed_base_url}:{operation}")
|
||||
} else {
|
||||
format!("{trimmed_base_url}/v1beta/models/{trimmed_model}:{operation}")
|
||||
@@ -80,6 +81,37 @@ pub fn build_gemini_content_url(
|
||||
Some(url)
|
||||
}
|
||||
|
||||
pub fn normalize_gemini_content_action_path(path: &str, stream: bool) -> String {
|
||||
let trimmed = path.trim();
|
||||
let (path, query) = split_path_query(trimmed);
|
||||
let action = if stream {
|
||||
"streamGenerateContent"
|
||||
} else {
|
||||
"generateContent"
|
||||
};
|
||||
let normalized = strip_gemini_content_action(path);
|
||||
let normalized = if normalized.len() == path.len() {
|
||||
path.to_string()
|
||||
} else {
|
||||
format!("{normalized}:{action}")
|
||||
};
|
||||
match query {
|
||||
Some(query) => format!("{normalized}?{query}"),
|
||||
None => normalized,
|
||||
}
|
||||
}
|
||||
|
||||
fn strip_gemini_content_action(value: &str) -> &str {
|
||||
value
|
||||
.strip_suffix(":streamGenerateContent")
|
||||
.or_else(|| value.strip_suffix(":generateContent"))
|
||||
.unwrap_or(value)
|
||||
}
|
||||
|
||||
fn gemini_content_base_url_contains_model_path(value: &str) -> bool {
|
||||
value.contains("/v1/models/") || value.contains("/v1beta/models/")
|
||||
}
|
||||
|
||||
pub fn build_gemini_video_predict_long_running_url(
|
||||
upstream_base_url: &str,
|
||||
model: &str,
|
||||
@@ -92,9 +124,9 @@ pub fn build_gemini_video_predict_long_running_url(
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut url = if trimmed_base_url.ends_with("/v1beta") {
|
||||
let mut url = if trimmed_base_url.ends_with("/v1") || trimmed_base_url.ends_with("/v1beta") {
|
||||
format!("{trimmed_base_url}/models/{trimmed_model}:predictLongRunning")
|
||||
} else if trimmed_base_url.contains("/v1beta/models/") {
|
||||
} else if gemini_content_base_url_contains_model_path(trimmed_base_url) {
|
||||
format!("{trimmed_base_url}:predictLongRunning")
|
||||
} else {
|
||||
format!("{trimmed_base_url}/v1beta/models/{trimmed_model}:predictLongRunning")
|
||||
@@ -248,6 +280,7 @@ mod tests {
|
||||
build_gemini_content_url, build_gemini_files_passthrough_url,
|
||||
build_gemini_video_predict_long_running_url, build_openai_chat_url,
|
||||
build_openai_responses_url, build_passthrough_path_url,
|
||||
normalize_gemini_content_action_path,
|
||||
};
|
||||
|
||||
#[test]
|
||||
@@ -289,6 +322,64 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_content_urls_rewrite_existing_base_action_for_stream_mode() {
|
||||
assert_eq!(
|
||||
build_gemini_content_url(
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent",
|
||||
"ignored-model",
|
||||
true,
|
||||
Some("foo=bar")
|
||||
)
|
||||
.as_deref(),
|
||||
Some(
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:streamGenerateContent?foo=bar"
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
build_gemini_content_url(
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:streamGenerateContent",
|
||||
"ignored-model",
|
||||
false,
|
||||
Some("foo=bar")
|
||||
)
|
||||
.as_deref(),
|
||||
Some(
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent?foo=bar"
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
build_gemini_content_url(
|
||||
"https://generativelanguage.googleapis.com/v1/models/gemini-2.5-pro:generateContent",
|
||||
"ignored-model",
|
||||
true,
|
||||
Some("foo=bar")
|
||||
)
|
||||
.as_deref(),
|
||||
Some(
|
||||
"https://generativelanguage.googleapis.com/v1/models/gemini-2.5-pro:streamGenerateContent?foo=bar"
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalizes_gemini_content_action_in_custom_paths() {
|
||||
assert_eq!(
|
||||
normalize_gemini_content_action_path(
|
||||
"/v1beta/models/gemini-2.5-pro:generateContent?alt=sse",
|
||||
true
|
||||
),
|
||||
"/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse"
|
||||
);
|
||||
assert_eq!(
|
||||
normalize_gemini_content_action_path(
|
||||
"/v1beta/models/gemini-2.5-pro:streamGenerateContent",
|
||||
false
|
||||
),
|
||||
"/v1beta/models/gemini-2.5-pro:generateContent"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merges_base_path_and_request_query_for_passthrough_paths() {
|
||||
assert_eq!(
|
||||
|
||||
Reference in New Issue
Block a user