feat(proxy): update Gemini streaming and transformation logic

This commit is contained in:
YoVinchen
2026-04-06 12:23:19 +08:00
parent 8f300f4591
commit 7cf554c011
2 changed files with 112 additions and 1 deletions
@@ -665,4 +665,70 @@ mod tests {
assert!(output.contains("\"partial_json\":\"{\\\"command\\\":\\\"git status\\\"}\""));
}
#[test]
fn rectifies_streamed_skill_args_from_nested_parameters() {
let payload = json!({
"responseId": "resp_6",
"modelVersion": "gemini-2.5-pro",
"candidates": [{
"finishReason": "STOP",
"content": {
"parts": [{
"functionCall": {
"id": "call_1",
"name": "Skill",
"args": {
"name": "git-commit",
"parameters": {
"args": ["详细分析内容 编写提交信息 分多次提交代码"]
}
}
}
}]
}
}],
"usageMetadata": {
"promptTokenCount": 5,
"totalTokenCount": 8
}
});
let owned_chunks = vec![format!(
"data: {}\n\n",
serde_json::to_string(&payload).unwrap()
)];
let stream = futures::stream::iter(
owned_chunks
.into_iter()
.map(|chunk| Ok::<Bytes, std::io::Error>(Bytes::from(chunk))),
);
let hints = super::super::transform_gemini::extract_anthropic_tool_schema_hints(&json!({
"tools": [{
"name": "Skill",
"input_schema": {
"type": "object",
"properties": {
"skill": { "type": "string" },
"args": { "type": "string" }
},
"required": ["skill"]
}
}]
}));
let converted =
create_anthropic_sse_stream_from_gemini(stream, None, None, None, Some(hints));
let output = futures::executor::block_on(async move {
converted
.collect::<Vec<_>>()
.await
.into_iter()
.map(|item| String::from_utf8(item.unwrap().to_vec()).unwrap())
.collect::<Vec<_>>()
.join("")
});
assert!(output.contains("git-commit"));
assert!(output.contains("详细分析内容 编写提交信息 分多次提交代码"));
assert!(!output.contains("\\\"parameters\\\""));
}
}
@@ -576,6 +576,41 @@ pub fn rectify_tool_call_args(
if args_object.is_empty() || hint.expected_keys.is_empty() {
return false;
}
let mut changed = false;
if hint.expected_keys.iter().any(|key| key == "skill") && !args_object.contains_key("skill") {
if let Some(value) = args_object.remove("name") {
args_object.insert("skill".to_string(), value);
changed = true;
}
}
if let Some(parameters_value) = args_object.remove("parameters") {
if let Some(parameters_object) = parameters_value.as_object() {
for expected_key in &hint.expected_keys {
if args_object.contains_key(expected_key) {
continue;
}
let Some(value) = parameters_object.get(expected_key) else {
continue;
};
let normalized_value = match value {
Value::Array(values) if values.len() == 1 => values[0].clone(),
_ => value.clone(),
};
args_object.insert(expected_key.clone(), normalized_value);
changed = true;
}
}
}
if hint
.required_keys
.iter()
.all(|key| args_object.contains_key(key.as_str()))
{
return changed;
}
let expected_key_set = hint
.expected_keys
@@ -1031,7 +1066,12 @@ mod tests {
"functionCall": {
"id": "call_1",
"name": "Skill",
"args": { "name": "git-commit" }
"args": {
"name": "git-commit",
"parameters": {
"args": ["详细分析内容 编写提交信息 分多次提交代码"]
}
}
}
}]
}
@@ -1056,7 +1096,12 @@ mod tests {
.unwrap();
assert_eq!(result["content"][0]["input"]["skill"], "git-commit");
assert_eq!(
result["content"][0]["input"]["args"],
"详细分析内容 编写提交信息 分多次提交代码"
);
assert!(result["content"][0]["input"].get("name").is_none());
assert!(result["content"][0]["input"].get("parameters").is_none());
}
#[test]