Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: normalize each Gemini response call before Tool dispatch

Browse Source
This commit is contained in:
Hare
2026-10-30 09:47:54 +09:00
parent 5a317e6d2a
commit 0f82d6fa8f
4 changed files with 804 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/src/llm_client/scheme/gemini/events.rs
+142 -29
View File
@@ -66,11 +66,15 @@ pub(crate) struct CandidatePart {
/// Function call.
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct FunctionCall {
/// Function name.
pub name: String,
/// Arguments.
pub args: Option<serde_json::Value>,
/// Partial-argument protocol is not a complete args object.
pub partial_args: Option<serde_json::Value>,
pub will_continue: Option<bool>,
}
/// Usage metadata.
@@ -110,6 +114,14 @@ pub(crate) struct SafetyRating {
// Event conversion
// ============================================================================
/// Response-local Gemini block allocation. Parts in separate SSE frames may
/// restart their array index, and function names are not call identities.
#[derive(Debug, Default)]
pub struct GeminiStreamState {
next_index: usize,
last_text_index: Option<usize>,
}
impl GeminiScheme {
/// Convert SSE data to Event.
///
@@ -120,7 +132,11 @@ impl GeminiScheme {
/// * `Ok(Some(Vec<Event>))` - Conversion succeeded.
/// * `Ok(None)` - Event ignored.
/// * `Err(ClientError)` - Parse error.
pub(crate) fn parse_event(&self, data: &str) -> Result<Option<Vec<Event>>, ClientError> {
pub(crate) fn parse_event(
&self,
data: &str,
state: &mut GeminiStreamState,
) -> Result<Option<Vec<Event>>, ClientError> {
// Skip empty or invalid data.
if data.is_empty() || data == "[DONE]" {
return Ok(None);
@@ -144,31 +160,39 @@ impl GeminiScheme {
// Process candidates.
if let Some(candidates) = response.candidates {
for candidate in candidates {
let candidate_index = candidate.index.unwrap_or(0);
if let Some(content) = candidate.content {
if let Some(parts) = content.parts {
for (part_index, part) in parts.iter().enumerate() {
for part in &parts {
let index = state.next_index;
state.next_index += 1;
// Text delta.
if let Some(text) = &part.text {
if !text.is_empty() {
// Gemini sends no explicit BlockStart, so emit TextDelta
// directly (Timeline handles the implicit start).
events.push(Event::text_delta(part_index, text.clone()));
state.last_text_index = Some(index);
events.push(Event::text_delta(index, text.clone()));
}
}
// Function call.
if let Some(function_call) = &part.function_call {
// Function call start.
// Gemini usually sends function calls all at once, but may
// send partial calls when streaming arguments are enabled.
// Gemini has no function call IDs, so use the name as the ID.
let function_id = format!("call_{}", function_call.name);
if function_call.partial_args.is_some()
|| function_call.will_continue.is_some()
{
return Err(ClientError::Api {
status: None,
code: Some("unsupported_partial_function_call".into()),
message: "Gemini partial function arguments require an unsupported streaming protocol".into(),
retry_after: None,
});
}
// These synthetic IDs correlate calls within a response;
// durable identity still comes from Host history entries.
let function_id = format!("call_{index}");
events.push(Event::BlockStart(BlockStart {
index: candidate_index * 10 + part_index, // Composite index.
index,
block_type: BlockType::ToolUse,
metadata: BlockMetadata::ToolUse {
id: function_id,
@@ -180,12 +204,13 @@ impl GeminiScheme {
if let Some(args) = &function_call.args {
let args_str = serde_json::to_string(args).unwrap_or_default();
if !args_str.is_empty() && args_str != "null" {
events.push(Event::tool_input_delta(
candidate_index * 10 + part_index,
args_str,
));
events.push(Event::tool_input_delta(index, args_str));
}
}
// Each args object is complete in this wire part. Close
// its collector before the next function call starts;
// dispatch still waits for the response-complete gate.
events.push(Event::tool_use_stop(index));
}
}
}
@@ -200,13 +225,14 @@ impl GeminiScheme {
_ => None,
};
// Text block stop.
events.push(Event::BlockStop(BlockStop {
index: candidate_index,
block_type: BlockType::Text,
stop_reason,
reasoning: None,
}));
if let Some(index) = state.last_text_index.take() {
events.push(Event::BlockStop(BlockStop {
index,
block_type: BlockType::Text,
stop_reason,
reasoning: None,
}));
}
}
}
}
@@ -240,7 +266,10 @@ mod tests {
let data =
r#"{"candidates":[{"content":{"parts":[{"text":"Hello"}],"role":"model"},"index":0}]}"#;
let events = scheme.parse_event(data).unwrap().unwrap();
let events = scheme
.parse_event(data, &mut GeminiStreamState::default())
.unwrap()
.unwrap();
assert_eq!(events.len(), 1);
if let Event::BlockDelta(delta) = &events[0] {
@@ -260,7 +289,10 @@ mod tests {
let scheme = GeminiScheme::new();
let data = r#"{"candidates":[{"content":{"parts":[{"text":"Hi"}],"role":"model"},"index":0}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":5,"totalTokenCount":15}}"#;
let events = scheme.parse_event(data).unwrap().unwrap();
let events = scheme
.parse_event(data, &mut GeminiStreamState::default())
.unwrap()
.unwrap();
// Should contain a Usage event.
let usage_event = events.iter().find(|e| matches!(e, Event::Usage(_)));
@@ -278,7 +310,10 @@ mod tests {
let scheme = GeminiScheme::new();
let data = r#"{"candidates":[{"content":{"parts":[{"functionCall":{"name":"get_weather","args":{"location":"Tokyo"}}}],"role":"model"},"index":0}]}"#;
let events = scheme.parse_event(data).unwrap().unwrap();
let events = scheme
.parse_event(data, &mut GeminiStreamState::default())
.unwrap()
.unwrap();
// Should contain a BlockStart event.
let start_event = events.iter().find(|e| matches!(e, Event::BlockStart(_)));
@@ -309,7 +344,10 @@ mod tests {
let scheme = GeminiScheme::new();
let data = r#"{"candidates":[{"content":{"parts":[{"text":"Done"}],"role":"model"},"finishReason":"STOP","index":0}]}"#;
let events = scheme.parse_event(data).unwrap().unwrap();
let events = scheme
.parse_event(data, &mut GeminiStreamState::default())
.unwrap()
.unwrap();
// Should contain a BlockStop event.
let stop_event = events.iter().find(|e| matches!(e, Event::BlockStop(_)));
@@ -320,10 +358,85 @@ mod tests {
}
}
#[test]
fn same_name_calls_have_distinct_ids_and_block_indexes_across_parts_and_frames() {
use crate::llm_client::scheme::Scheme;
let scheme = GeminiScheme::new();
for frames in [
vec![
r#"{"candidates":[{"content":{"parts":[{"functionCall":{"name":"echo","args":{"text":"first"}}},{"functionCall":{"name":"echo","args":{"text":"second"}}}]}}]}"#,
],
vec![
r#"{"candidates":[{"content":{"parts":[{"functionCall":{"name":"echo","args":{"text":"first"}}}]}}]}"#,
r#"{"candidates":[{"content":{"parts":[{"functionCall":{"name":"echo","args":{"text":"second"}}}]}}]}"#,
],
] {
let mut state = Default::default();
let mut events = Vec::new();
for frame in frames {
events.extend(scheme.parse_sse("", frame, &mut state).unwrap());
}
let calls: Vec<_> = events
.iter()
.filter_map(|event| match event {
Event::BlockStart(BlockStart {
index,
metadata: BlockMetadata::ToolUse { id, name },
..
}) => Some((*index, id.as_str(), name.as_str())),
_ => None,
})
.collect();
assert_eq!(calls, [(0, "call_0", "echo"), (1, "call_1", "echo")]);
let inputs: Vec<_> = events
.iter()
.filter_map(|event| match event {
Event::BlockDelta(delta) => match &delta.delta {
DeltaContent::InputJson(json) => Some((delta.index, json.as_str())),
_ => None,
},
_ => None,
})
.collect();
assert_eq!(
inputs,
[(0, r#"{"text":"first"}"#), (1, r#"{"text":"second"}"#)]
);
}
}
#[test]
fn partial_argument_protocol_is_not_misread_as_an_empty_complete_call() {
use crate::llm_client::scheme::Scheme;
let scheme = GeminiScheme::new();
for call in [
r#"{"name":"echo","partialArgs":[{"jsonPath":"$.text","stringValue":"part"}]}"#,
r#"{"name":"echo","willContinue":true,"args":{"text":"part"}}"#,
r#"{"name":"echo","willContinue":false}"#,
] {
let data = format!(
r#"{{"candidates":[{{"content":{{"parts":[{{"functionCall":{call}}}]}}}}]}}"#
);
assert!(
matches!(scheme.parse_sse("", &data, &mut Default::default()), Err(ClientError::Api { code: Some(code), .. }) if code == "unsupported_partial_function_call")
);
}
}
#[test]
fn test_parse_empty_data() {
let scheme = GeminiScheme::new();
assert!(scheme.parse_event("").unwrap().is_none());
assert!(scheme.parse_event("[DONE]").unwrap().is_none());
assert!(
scheme
.parse_event("", &mut GeminiStreamState::default())
.unwrap()
.is_none()
);
assert!(
scheme
.parse_event("[DONE]", &mut GeminiStreamState::default())
.unwrap()
.is_none()
);
}
}
crates/agen/src/llm_client/scheme/gemini/request.rs
+45 -3
View File
@@ -239,6 +239,8 @@ impl GeminiScheme {
let mut contents = Vec::new();
let mut pending_model_parts: Vec<GeminiPart> = Vec::new();
let mut pending_user_parts: Vec<GeminiPart> = Vec::new();
let mut call_names: std::collections::HashMap<&str, std::collections::VecDeque<&str>> =
std::collections::HashMap::new();
for item in items {
match item {
@@ -270,8 +272,12 @@ impl GeminiScheme {
}
Item::ToolCall {
name, arguments, ..
call_id,
name,
arguments,
..
} => {
call_names.entry(call_id).or_default().push_back(name);
// Flush pending user parts first
if !pending_user_parts.is_empty() {
contents.push(GeminiContent {
@@ -309,11 +315,17 @@ impl GeminiScheme {
Some(c) => format!("{summary}\n{c}"),
None => summary.clone(),
};
// Gemini responses identify the function by its declaration
// name, never by the adapter's response-local correlation ID.
let name = call_names
.get_mut(call_id.as_str())
.and_then(std::collections::VecDeque::pop_front)
.unwrap_or(call_id.as_str());
pending_user_parts.push(GeminiPart::FunctionResponse {
function_response: GeminiFunctionResponse {
name: call_id.clone(),
name: name.to_string(),
response: GeminiFunctionResponseContent {
name: call_id.clone(),
name: name.to_string(),
content: Value::String(text),
},
},
@@ -468,6 +480,36 @@ mod tests {
assert_eq!(gemini_req.contents[0].role, "user");
assert_eq!(gemini_req.contents[1].role, "model");
assert_eq!(gemini_req.contents[2].role, "user");
let body = serde_json::to_value(gemini_req).unwrap();
assert_eq!(
body["contents"][2]["parts"][0]["functionResponse"]["name"],
"get_weather"
);
}
#[test]
fn parallel_results_resolve_function_names_not_synthetic_call_ids() {
let scheme = GeminiScheme::new();
let request = Request::new()
.item(Item::tool_call("call_0", "echo", r#"{"text":"first"}"#))
.item(Item::tool_call("call_1", "echo", r#"{"text":"second"}"#))
.item(Item::tool_result("call_1", "second"))
.item(Item::tool_result("call_0", "first"))
.item(Item::tool_call("call_0", "other", "{}"))
.item(Item::tool_result("call_0", "later"));
let body = serde_json::to_value(scheme.build_request(&request, &cap())).unwrap();
let responses: Vec<_> = body["contents"]
.as_array()
.unwrap()
.iter()
.flat_map(|content| content["parts"].as_array().unwrap())
.filter_map(|part| part.get("functionResponse"))
.collect();
assert_eq!(responses.len(), 3);
for (response, name) in responses.iter().zip(["echo", "echo", "other"]) {
assert_eq!(response["name"], name);
assert_eq!(response["response"]["name"], name);
}
}
#[test]
crates/agen/src/llm_client/scheme/gemini/scheme_impl.rs
+3 -3
View File
@@ -10,7 +10,7 @@ use crate::llm_client::{
use super::GeminiScheme;
impl Scheme for GeminiScheme {
type State = ();
type State = super::events::GeminiStreamState;
fn default_base_url(&self) -> &'static str {
"https://generativelanguage.googleapis.com"
@@ -38,9 +38,9 @@ impl Scheme for GeminiScheme {
&self,
_event_type: &str,
data: &str,
_state: &mut Self::State,
state: &mut Self::State,
) -> Result<Vec<Event>, ClientError> {
Ok(self.parse_event(data)?.unwrap_or_default())
Ok(self.parse_event(data, state)?.unwrap_or_default())
}
fn default_capability(&self) -> ModelCapability {