1
1
//! This module provides functionality for managing and executing tools in an async OpenAI context.
2
2
//! It defines traits and structures for tool management, execution, and streaming.
3
3
use std:: {
4
- collections:: { hash_map :: Entry , BTreeMap , HashMap } ,
4
+ collections:: { BTreeMap , HashMap } ,
5
5
future:: Future ,
6
6
pin:: Pin ,
7
7
sync:: Arc ,
@@ -14,7 +14,7 @@ use serde_json::json;
14
14
use crate :: types:: {
15
15
ChatCompletionMessageToolCall , ChatCompletionMessageToolCallChunk ,
16
16
ChatCompletionRequestToolMessage , ChatCompletionTool , ChatCompletionToolType , FunctionCall ,
17
- FunctionCallStream , FunctionObject ,
17
+ FunctionObject ,
18
18
} ;
19
19
20
20
/// 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 {
92
92
// If the tool doesn't require arguments (T::Args is unit type),
93
93
// we can safely ignore the provided arguments string
94
94
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) )
96
96
{
97
97
Ok ( args) => T :: call ( self , args)
98
98
. await
@@ -166,9 +166,7 @@ impl ToolManager {
166
166
Ok ( Err ( e) ) => {
167
167
format ! ( "Tool call failed: {}" , e)
168
168
}
169
- Err ( _) => {
170
- format ! ( "Tool call failed: runtime error" )
171
- }
169
+ Err ( _) => "Tool call failed: runtime error" . to_string ( ) ,
172
170
} ;
173
171
outputs. push ( ChatCompletionRequestToolMessage {
174
172
content : output. into ( ) ,
@@ -192,37 +190,26 @@ impl ToolCallStreamManager {
192
190
193
191
/// Processes a single streaming tool call chunk and merges it with existing data.
194
192
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 ( ) ,
222
198
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) ;
226
213
}
227
214
}
228
215
}
@@ -239,6 +226,15 @@ impl ToolCallStreamManager {
239
226
240
227
/// Returns all completed tool calls as a vector.
241
228
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 ( )
243
239
}
244
240
}
0 commit comments