Skip to content

Commit 041eec9

Browse files
committed
fix: some provider may not return funciton id / function name in the first chunk during tool_call stream
1 parent 7796784 commit 041eec9

File tree

1 file changed

+33
-37
lines changed

1 file changed

+33
-37
lines changed

async-openai/src/tools.rs

Lines changed: 33 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
//! This module provides functionality for managing and executing tools in an async OpenAI context.
22
//! It defines traits and structures for tool management, execution, and streaming.
33
use std::{
4-
collections::{hash_map::Entry, BTreeMap, HashMap},
4+
collections::{BTreeMap, HashMap},
55
future::Future,
66
pin::Pin,
77
sync::Arc,
@@ -14,7 +14,7 @@ use serde_json::json;
1414
use crate::types::{
1515
ChatCompletionMessageToolCall, ChatCompletionMessageToolCallChunk,
1616
ChatCompletionRequestToolMessage, ChatCompletionTool, ChatCompletionToolType, FunctionCall,
17-
FunctionCallStream, FunctionObject,
17+
FunctionObject,
1818
};
1919

2020
/// A trait defining the interface for tools that can be used with the OpenAI API.
@@ -92,7 +92,7 @@ impl<T: Tool> ToolDyn for T {
9292
// If the tool doesn't require arguments (T::Args is unit type),
9393
// we can safely ignore the provided arguments string
9494
match serde_json::from_str::<T::Args>(&args)
95-
.or_else(|e| serde_json::from_str::<T::Args>(&"null").map_err(|_| e))
95+
.or_else(|e| serde_json::from_str::<T::Args>("null").map_err(|_| e))
9696
{
9797
Ok(args) => T::call(self, args)
9898
.await
@@ -166,9 +166,7 @@ impl ToolManager {
166166
Ok(Err(e)) => {
167167
format!("Tool call failed: {}", e)
168168
}
169-
Err(_) => {
170-
format!("Tool call failed: runtime error")
171-
}
169+
Err(_) => "Tool call failed: runtime error".to_string(),
172170
};
173171
outputs.push(ChatCompletionRequestToolMessage {
174172
content: output.into(),
@@ -192,37 +190,26 @@ impl ToolCallStreamManager {
192190

193191
/// Processes a single streaming tool call chunk and merges it with existing data.
194192
pub fn process_chunk(&mut self, chunk: ChatCompletionMessageToolCallChunk) {
195-
match self.0.entry(chunk.index) {
196-
Entry::Occupied(mut o) => {
197-
if let Some(FunctionCallStream {
198-
name: _,
199-
arguments: Some(arguments),
200-
}) = chunk.function
201-
{
202-
o.get_mut().function.arguments.push_str(&arguments);
203-
}
204-
}
205-
Entry::Vacant(o) => {
206-
let ChatCompletionMessageToolCallChunk {
207-
index: _,
208-
id: Some(id),
209-
r#type: _,
210-
function:
211-
Some(FunctionCallStream {
212-
name: Some(name),
213-
arguments: Some(arguments),
214-
}),
215-
} = chunk
216-
else {
217-
tracing::error!("Tool call chunk is not complete: {:?}", chunk);
218-
return;
219-
};
220-
let tool_call = ChatCompletionMessageToolCall {
221-
id,
193+
let tool_call =
194+
self.0
195+
.entry(chunk.index)
196+
.or_insert_with(|| ChatCompletionMessageToolCall {
197+
id: "".to_string(),
222198
r#type: ChatCompletionToolType::Function,
223-
function: FunctionCall { name, arguments },
224-
};
225-
o.insert(tool_call);
199+
function: FunctionCall {
200+
name: "".to_string(),
201+
arguments: "".to_string(),
202+
},
203+
});
204+
if let Some(id) = chunk.id {
205+
tool_call.id = id;
206+
}
207+
if let Some(function) = chunk.function {
208+
if let Some(name) = function.name {
209+
tool_call.function.name = name;
210+
}
211+
if let Some(arguments) = function.arguments {
212+
tool_call.function.arguments.push_str(&arguments);
226213
}
227214
}
228215
}
@@ -239,6 +226,15 @@ impl ToolCallStreamManager {
239226

240227
/// Returns all completed tool calls as a vector.
241228
pub fn finish_stream(self) -> Vec<ChatCompletionMessageToolCall> {
242-
self.0.into_values().collect()
229+
self.0
230+
.into_values()
231+
.filter(|tool_call| {
232+
let is_complete = !tool_call.id.is_empty() && !tool_call.function.name.is_empty();
233+
if !is_complete {
234+
tracing::error!("Tool call is not complete: {:?}", tool_call);
235+
}
236+
is_complete
237+
})
238+
.collect()
243239
}
244240
}

0 commit comments

Comments
 (0)