Accumulate streamed tool calls across chunks in Ollama streaming
Ollama >=0.8 can stream tool_calls incrementally across NDJSON chunks; chat_with_tools_stream did `tool_calls = Some(tcs)` per chunk, so only the last chunk's calls survived assembly and earlier calls were silently dropped. Append into the accumulator instead. - ollama: append_streamed_tool_calls helper + tests covering two calls arriving in separate chunks and the single-chunk batch case. - llamacpp: the SSE delta assembly was already correct (per-index BTreeMap, same-index argument fragments concatenate, distinct indexes accumulate); extracted it into apply_tool_call_deltas / finalize_tool_calls and added tests pinning that behavior. - llm_client: new shared strip_think_blocks (moved from ollama's private extract_final_answer, which now delegates) so the tool-calling final content paths can reuse it; unit tests for tagged/plain/unclosed/empty cases. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
1 parent
8e4f91561b
commit
592dfcb42c
3 files changed
+223
-66
No files matched your search
+52
-14
@@ -360,18 +360,7 @@ impl OllamaClient {
|
||||
/// Extract final answer from thinking model output
|
||||
/// Handles <think>...</think> tags and takes everything after
|
||||
fn extract_final_answer(&self, response: &str) -> String {
|
||||
let response = response.trim();
|
||||
|
||||
// Look for </think> tag and take everything after it
|
||||
if let Some(pos) = response.find("</think>") {
|
||||
let answer = response[pos + 8..].trim();
|
||||
if !answer.is_empty() {
|
||||
return answer.to_string();
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: return the whole response trimmed
|
||||
response.to_string()
|
||||
crate::ai::llm_client::strip_think_blocks(response)
|
||||
}
|
||||
|
||||
async fn try_generate(
|
||||
@@ -846,11 +835,14 @@ Analyze the image and use specific details from both the visual content and the
|
||||
if !chunk.message.role.is_empty() {
|
||||
role = chunk.message.role;
|
||||
}
|
||||
// Ollama only attaches tool_calls on the final chunk.
|
||||
// Ollama ≥0.8 can stream tool_calls incrementally
|
||||
// across chunks (older servers attach them all to
|
||||
// one chunk) — append rather than overwrite so
|
||||
// calls from earlier chunks survive.
|
||||
if let Some(tcs) = chunk.message.tool_calls
|
||||
&& !tcs.is_empty()
|
||||
{
|
||||
tool_calls = Some(tcs);
|
||||
append_streamed_tool_calls(&mut tool_calls, tcs);
|
||||
}
|
||||
if chunk.done {
|
||||
prompt_eval_count = chunk.prompt_eval_count;
|
||||
@@ -1329,8 +1321,20 @@ struct OllamaEmbedResponse {
|
||||
embeddings: Vec<Vec<f32>>,
|
||||
}
|
||||
|
||||
/// Accumulate tool calls streamed across NDJSON chunks. Ollama ≥0.8 may
|
||||
/// emit each tool call on its own chunk; replacing the accumulator on every
|
||||
/// chunk would keep only the last call, so extend instead.
|
||||
fn append_streamed_tool_calls(
|
||||
acc: &mut Option<Vec<crate::ai::llm_client::ToolCall>>,
|
||||
new: Vec<crate::ai::llm_client::ToolCall>,
|
||||
) {
|
||||
acc.get_or_insert_with(Vec::new).extend(new);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::append_streamed_tool_calls;
|
||||
use crate::ai::llm_client::{ToolCall, ToolCallFunction};
|
||||
|
||||
#[test]
|
||||
fn generate_photo_description_prompt_is_concise() {
|
||||
@@ -1341,4 +1345,38 @@ mod tests {
|
||||
Focus on the people, location, and activity.";
|
||||
assert!(prompt.len() < 200, "Prompt should be concise");
|
||||
}
|
||||
|
||||
fn call(name: &str) -> ToolCall {
|
||||
ToolCall {
|
||||
id: None,
|
||||
function: ToolCallFunction {
|
||||
name: name.to_string(),
|
||||
arguments: serde_json::json!({}),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn streamed_tool_calls_across_chunks_accumulate() {
|
||||
// Two tool calls arriving in two separate stream chunks must BOTH
|
||||
// survive assembly — the old `tool_calls = Some(tcs)` kept only the
|
||||
// last chunk's calls.
|
||||
let mut acc: Option<Vec<ToolCall>> = None;
|
||||
append_streamed_tool_calls(&mut acc, vec![call("get_sms_messages")]);
|
||||
append_streamed_tool_calls(&mut acc, vec![call("reverse_geocode")]);
|
||||
|
||||
let calls = acc.expect("tool calls accumulated");
|
||||
assert_eq!(calls.len(), 2);
|
||||
assert_eq!(calls[0].function.name, "get_sms_messages");
|
||||
assert_eq!(calls[1].function.name, "reverse_geocode");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn streamed_tool_calls_single_chunk_batch_kept_intact() {
|
||||
// Older Ollama servers attach all calls to one chunk — unchanged.
|
||||
let mut acc: Option<Vec<ToolCall>> = None;
|
||||
append_streamed_tool_calls(&mut acc, vec![call("a"), call("b")]);
|
||||
let calls = acc.expect("tool calls accumulated");
|
||||
assert_eq!(calls.len(), 2);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user