Merge branch 'fix_tool_call'

src/chat_template/mod.rs

@@ -102,6 +102,7 @@ impl<'a> ChatTemplate<'a> {
     pub fn apply_chat_template(&self, messages: &ChatCompletionParameters) -> Result<String> {
         let context = context! {
             messages => &messages.messages,
+            tools => &messages.tools.as_ref(),
             add_generation_prompt => true,
         };
         let template = self

src/models/deepseek_ocr/generate.rs

@@ -127,9 +127,10 @@ impl GenerateModel for DeepseekOCRGenerateModel {
             images_seq_mask = None;
             images_spatial_crop_t = None;
         }
+        let num_token = generate.len() as u32;
         let res = self.tokenizer.token_decode(generate)?;
         self.deepseekocr_model.clear_kv_cache();
-        let response = build_completion_response(res, &self.model_name);
+        let response = build_completion_response(res, &self.model_name, Some(num_token));
         Ok(response)
     }
 

src/models/hunyuan_ocr/generate.rs

@@ -119,9 +119,10 @@ impl<'a> GenerateModel for HunyuanOCRGenerateModel<'a> {
             pixel_values = None;
             image_grid_thw = None;
         }
+        let num_token = generate.len() as u32;
         let res = self.tokenizer.token_decode(generate)?;
         self.hunyuan_vl.clear_kv_cache();
-        let response = build_completion_response(res, &self.model_name);
+        let response = build_completion_response(res, &self.model_name, Some(num_token));
         Ok(response)
     }
 

src/models/minicpm4/generate.rs

@@ -78,9 +78,10 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> {
             seq_len = 1;
             input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
         }
+        let num_token = generate.len() as u32;
         let res = self.tokenizer.token_decode(generate)?;
         self.minicpm.clear_kv_cache();
-        let response = build_completion_response(res, &self.model_name);
+        let response = build_completion_response(res, &self.model_name, Some(num_token));
         Ok(response)
     }
     fn generate_stream(

src/models/paddleocr_vl/generate.rs

@@ -104,9 +104,10 @@ impl<'a> GenerateModel for PaddleOCRVLGenerateModel<'a> {
             pixel_values = None;
             image_grid_thw = None;
         }
+        let num_token = generate.len() as u32;
         let res = self.tokenizer.token_decode(generate)?;
         self.paddleocr_vl.clear_kv_cache();
-        let response = build_completion_response(res, &self.model_name);
+        let response = build_completion_response(res, &self.model_name, Some(num_token));
         Ok(response)
     }
 

src/models/qwen2_5vl/generate.rs

@@ -118,9 +118,10 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> {
             pixel_values = None;
             pixel_values_video = None;
         }
+        let num_token = generate.len() as u32;
         let res = self.tokenizer.token_decode(generate)?;
         self.qwen2_5_vl.clear_kv_cache();
-        let response = build_completion_response(res, &self.model_name);
+        let response = build_completion_response(res, &self.model_name, Some(num_token));
         Ok(response)
     }
 
@@ -167,6 +168,8 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> {
             let image_grid_thw = image_grid_thw.as_ref();
             let mut pixel_values_video = pixel_values_video.as_ref();
             let video_grid_thw = video_grid_thw.as_ref();
+            let mut tool_call_id = None;
+            let mut tool_call_content = String::new();
             for _ in 0..sample_len {
                 let logits = self.qwen2_5_vl.forward(
                     &input_ids,
@@ -203,8 +206,56 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> {
                     continue;
                 }
                 error_tokens.clear();
-                let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None);
-                yield Ok(chunk);
+                // 处理特殊标记和工具调用
+                match decoded_token.as_str() {
+                    "<tool_call>" => {
+                        // 开始工具调用
+                        tool_call_id = Some(uuid::Uuid::new_v4().to_string());
+                        seqlen_offset += seq_len;
+                        seq_len = 1;
+                        input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+                        cache_position = Tensor::from_vec(vec![seqlen_offset as u32], 1, &self.device)?;
+                        pixel_values = None;
+                        pixel_values_video = None;
+                        continue;
+                    }
+                    "</tool_call>" => {
+                        // 结束工具调用
+                        let chunk = build_completion_chunk_response(
+                            decoded_token, 
+                            &self.model_name, 
+                            tool_call_id.clone(), 
+                            Some(tool_call_content.clone())
+                        );
+                        tool_call_id = None;
+                        tool_call_content = String::new();
+                        yield Ok(chunk);
+                    }
+                    _ => {
+                        if tool_call_id.is_some() {
+                            // 在工具调用过程中，收集工具调用内容
+                            tool_call_content.push_str(&decoded_token);
+                            seqlen_offset += seq_len;
+                            seq_len = 1;
+                            input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+                            cache_position = Tensor::from_vec(vec![seqlen_offset as u32], 1, &self.device)?;
+                            pixel_values = None;
+                            pixel_values_video = None;
+                            continue;
+                        } else {
+                            // 正常文本输出
+                            let chunk = build_completion_chunk_response(
+                                decoded_token, 
+                                &self.model_name, 
+                                None, 
+                                None
+                            );
+                            yield Ok(chunk);
+                        }
+                    }
+                }          
+                // let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None);
+                // yield Ok(chunk);
                 if next_token == self.endoftext_id || next_token == self.im_end_id {
                     break;
                 }

src/models/qwen2_5vl/processor.rs

@@ -72,6 +72,9 @@ impl Qwen2_5VLProcessor {
                     if let ChatMessageContentPart::Image(img_part) = part {
                         let img_url = img_part.image_url;
                         vision_map.get_mut("image").unwrap().push(img_url.url);
+                    // } else if let ChatMessageContentPart::Video(video_part) = part {
+                    //     let video_url = video_part.video_url;
+                    //     vision_map.get_mut("video").unwrap().push(video_url.url);
                     }
                 }
             }

src/models/qwen3vl/generate.rs

@@ -120,9 +120,10 @@ impl<'a> GenerateModel for Qwen3VLGenerateModel<'a> {
             pixel_values = None;
             pixel_values_video = None;
         }
+        let num_token = generate.len() as u32;
         let res = self.tokenizer.token_decode(generate)?;
         self.qwen3_vl.clear_kv_cache();
-        let response = build_completion_response(res, &self.model_name);
+        let response = build_completion_response(res, &self.model_name, Some(num_token));
         Ok(response)
     }
 
@@ -171,19 +172,21 @@ impl<'a> GenerateModel for Qwen3VLGenerateModel<'a> {
             let image_grid_thw = image_grid_thw.as_ref();
             let mut pixel_values_video = pixel_values_video.as_ref();
             let video_grid_thw = video_grid_thw.as_ref();
+            let mut tool_call_id = None;
+            let mut tool_call_content = String::new();
             for _ in 0..sample_len {
-            let logits = self.qwen3_vl.forward(
-                &input_ids,
-                pixel_values,
-                image_grid_thw,
-                pixel_values_video,
-                video_grid_thw,
-                Some(&cache_position),
-                seqlen_offset,
-            )?;
-            let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
-            let next_token = logit_processor.sample(&logits)?;
-            let mut decode_ids = Vec::new();
+                let logits = self.qwen3_vl.forward(
+                    &input_ids,
+                    pixel_values,
+                    image_grid_thw,
+                    pixel_values_video,
+                    video_grid_thw,
+                    Some(&cache_position),
+                    seqlen_offset,
+                )?;
+                let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
+                let next_token = logit_processor.sample(&logits)?;
+                let mut decode_ids = Vec::new();
                 if !error_tokens.is_empty() {
                     decode_ids.extend_from_slice(&error_tokens);
                 }
@@ -203,8 +206,55 @@ impl<'a> GenerateModel for Qwen3VLGenerateModel<'a> {
                     continue;
                 }
                 error_tokens.clear();
-                let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None);
-                yield Ok(chunk);
+
+                // 处理特殊标记和工具调用
+                match decoded_token.as_str() {
+                    "<tool_call>" => {
+                        // 开始工具调用
+                        tool_call_id = Some(uuid::Uuid::new_v4().to_string());
+                        seqlen_offset += seq_len;
+                        seq_len = 1;
+                        input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+                        cache_position = Tensor::from_vec(vec![seqlen_offset as u32], 1, &self.device)?;
+                        pixel_values = None;
+                        pixel_values_video = None;
+                        continue;
+                    }
+                    "</tool_call>" => {
+                        // 结束工具调用
+                        let chunk = build_completion_chunk_response(
+                            decoded_token,
+                            &self.model_name,
+                            tool_call_id.clone(),
+                            Some(tool_call_content.clone())
+                        );
+                        tool_call_id = None;
+                        tool_call_content = String::new();
+                        yield Ok(chunk);
+                    }
+                    _ => {
+                        if tool_call_id.is_some() {
+                            // 在工具调用过程中，收集工具调用内容
+                            tool_call_content.push_str(&decoded_token);
+                            seqlen_offset += seq_len;
+                            seq_len = 1;
+                            input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+                            cache_position = Tensor::from_vec(vec![seqlen_offset as u32], 1, &self.device)?;
+                            pixel_values = None;
+                            pixel_values_video = None;
+                            continue;
+                        } else {
+                            // 正常文本输出
+                            let chunk = build_completion_chunk_response(
+                                decoded_token,
+                                &self.model_name,
+                                None,
+                                None
+                            );
+                            yield Ok(chunk);
+                        }
+                    }
+                }
                 if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 {
                     break;
                 }

src/models/qwen3vl/model.rs

@@ -730,7 +730,6 @@ impl Qwen3VLTextModel {
         deepstack_visual_embeds: Option<Vec<Tensor>>,
     ) -> Result<Tensor> {
         let (b_size, seq_len, _) = inputs_embeds.dims3()?;
-
         let position_ids = match position_ids {
             Some(ids) => ids.clone(),
             None => Tensor::arange(

src/utils/mod.rs

@@ -11,7 +11,7 @@ use aha_openai_dive::v1::resources::{
         ChatCompletionParameters, ChatCompletionResponse, ChatMessage, ChatMessageContent,
         ChatMessageContentPart, DeltaChatMessage, DeltaFunction, DeltaToolCall, Function, ToolCall,
     },
-    shared::FinishReason,
+    shared::{FinishReason, Usage},
 };
 use anyhow::Result;
 use candle_core::{DType, Device};
@@ -137,8 +137,24 @@ pub fn ceil_by_factor(num: f32, factor: u32) -> u32 {
     ceil * factor
 }
 
-pub fn build_completion_response(res: String, model_name: &str) -> ChatCompletionResponse {
+pub fn build_completion_response(res: String, model_name: &str, num_tokens: Option<u32>) -> ChatCompletionResponse {
     let id = uuid::Uuid::new_v4().to_string();
+    let usage = match num_tokens {
+        Some(num) => {
+            Some(Usage {
+                input_tokens: None,
+                input_tokens_details: None,
+                output_tokens: None,
+                output_tokens_details: None,
+                prompt_tokens: None,
+                completion_tokens: None,
+                total_tokens: num,
+                prompt_tokens_details: None,
+                completion_tokens_details: None
+            })
+        },
+        None => None,
+    };
     let mut response = ChatCompletionResponse {
         id: Some(id),
         choices: vec![],
@@ -147,7 +163,7 @@ pub fn build_completion_response(res: String, model_name: &str) -> ChatCompletio
         service_tier: None,
         system_fingerprint: None,
         object: "chat.completion".to_string(),
-        usage: None,
+        usage
     };
     let choice = if res.contains("<tool_call>") {
         let mes: Vec<&str> = res.split("<tool_call>").collect();

tests/test_deepseek_ocr.rs

@@ -41,8 +41,14 @@ fn deepseek_ocr_generate() -> Result<()> {
     let i_start = Instant::now();
     let res = model.generate(mes)?;
     let i_duration = i_start.elapsed();
-    println!("Time elapsed in generate is: {:?}", i_duration);
     println!("generate: \n {:?}", res);
+    if res.usage.is_some() {
+        let num_token = res.usage.as_ref().unwrap().total_tokens;
+        let duration_secs = i_duration.as_secs_f64();
+        let tps = num_token as f64 / duration_secs;
+        println!("Tokens per second (TPS): {:.2}", tps);
+    }
+    println!("Time elapsed in generate is: {:?}", i_duration);    
     Ok(())
 }
 

tests/test_gelab_zero.rs

@@ -19,11 +19,39 @@ fn gelab_zero_generate() -> Result<()> {
                 "content": [    
                     {
                         "type": "text", 
-                        "text": "Hello, GELab-Zero!"
+                        "text": "Hello, GELab-Zero!, 现在几点了"
                     }
                 ]
             }
-        ]
+        ],
+        "tools": [
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_current_time",
+                    "description": "当你想知道现在的时间时非常有用。",
+                    "parameters": {}
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_current_weather",
+                    "description": "当你想查询指定城市的天气时非常有用。",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "location": {
+                                "type": "string",
+                                "description": "城市或县区，比如北京市、杭州市、余杭区等。"
+                            }
+                        },
+                        "required": ["location"]
+                    }
+                }
+            }
+        ],
+        "tool_choice": null
     }
     "#;
     let mes: ChatCompletionParameters = serde_json::from_str(message)?;
@@ -36,6 +64,12 @@ fn gelab_zero_generate() -> Result<()> {
     let res = qwen3vl.generate(mes)?;
     let i_duration = i_start.elapsed();
     println!("generate: \n {:?}", res);
+    if res.usage.is_some() {
+        let num_token = res.usage.as_ref().unwrap().total_tokens;
+        let duration_secs = i_duration.as_secs_f64();
+        let tps = num_token as f64 / duration_secs;
+        println!("Tokens per second (TPS): {:.2}", tps);
+    }
     println!("Time elapsed in generate is: {:?}", i_duration);
     Ok(())
 }

tests/test_hunyuan_ocr.rs

@@ -40,8 +40,15 @@ fn hunyuan_ocr_generate() -> Result<()> {
     let i_start = Instant::now();
     let res = model.generate(mes)?;
     let i_duration = i_start.elapsed();
-    println!("Time elapsed in generate is: {:?}", i_duration);
     println!("generate: \n {:?}", res);
+    if res.usage.is_some() {
+        let num_token = res.usage.as_ref().unwrap().total_tokens;
+        let duration_secs = i_duration.as_secs_f64();
+        let tps = num_token as f64 / duration_secs;
+        println!("Tokens per second (TPS): {:.2}", tps);
+    }
+    println!("Time elapsed in generate is: {:?}", i_duration);
+    
     Ok(())
 }
 

tests/test_minicpm4.rs

@@ -33,8 +33,14 @@ fn minicpm_generate() -> Result<()> {
 
     let i_start = Instant::now();
     let result = model.generate(mes)?;
-    println!("generate: \n {:?}", result);
     let i_duration = i_start.elapsed();
+    println!("generate: \n {:?}", result);
+    if result.usage.is_some() {
+        let num_token = result.usage.as_ref().unwrap().total_tokens;
+        let duration_secs = i_duration.as_secs_f64();
+        let tps = num_token as f64 / duration_secs;
+        println!("Tokens per second (TPS): {:.2}", tps);
+    }
     println!("Time elapsed in generate is: {:?}", i_duration);
 
     Ok(())

tests/test_paddleocr_vl.rs

@@ -41,8 +41,14 @@ fn paddleocr_vl_generate() -> Result<()> {
     let i_start = Instant::now();
     let res = model.generate(mes)?;
     let i_duration = i_start.elapsed();
-    println!("Time elapsed in generate is: {:?}", i_duration);
     println!("generate: \n {:?}", res);
+    if res.usage.is_some() {
+        let num_token = res.usage.as_ref().unwrap().total_tokens;
+        let duration_secs = i_duration.as_secs_f64();
+        let tps = num_token as f64 / duration_secs;
+        println!("Tokens per second (TPS): {:.2}", tps);
+    }
+    println!("Time elapsed in generate is: {:?}", i_duration);    
     Ok(())
 }
 

tests/test_qwen2_5vl.rs

@@ -46,8 +46,14 @@ fn qwen2_5vl_generate() -> Result<()> {
 
     let i_start = Instant::now();
     let result = model.generate(mes)?;    
-    println!("generate: \n {:?}", result);
     let i_duration = i_start.elapsed();
+    println!("generate: \n {:?}", result);
+    if result.usage.is_some() {
+        let num_token = result.usage.as_ref().unwrap().total_tokens;
+        let duration_secs = i_duration.as_secs_f64();
+        let tps = num_token as f64 / duration_secs;
+        println!("Tokens per second (TPS): {:.2}", tps);
+    }
     println!("Time elapsed in generate is: {:?}", i_duration);
 
     Ok(())

tests/test_qwen3vl.rs

@@ -27,7 +27,7 @@ fn qwen3vl_generate() -> Result<()> {
                     },              
                     {
                         "type": "text", 
-                        "text": "视频里发生了什么"
+                        "text": "视频中发生了什么？, 现在几点了"
                     }
                 ]
             }
@@ -44,6 +44,12 @@ fn qwen3vl_generate() -> Result<()> {
     let res = qwen3vl.generate(mes)?;
     let i_duration = i_start.elapsed();
     println!("generate: \n {:?}", res);
+    if res.usage.is_some() {
+        let num_token = res.usage.as_ref().unwrap().total_tokens;
+        let duration_secs = i_duration.as_secs_f64();
+        let tps = num_token as f64 / duration_secs;
+        println!("Tokens per second (TPS): {:.2}", tps);
+    }
     println!("Time elapsed in generate is: {:?}", i_duration);
     Ok(())
 }

tests/test_robo_brain.rs

@@ -34,8 +34,14 @@ fn robo_brain_generate() -> Result<()> {
 
     let i_start = Instant::now();
     let result = model.generate(mes)?;
-    println!("generate: \n {:?}", result);
     let i_duration = i_start.elapsed();
+    println!("generate: \n {:?}", result);    
+    if result.usage.is_some() {
+        let num_token = result.usage.as_ref().unwrap().total_tokens;
+        let duration_secs = i_duration.as_secs_f64();
+        let tps = num_token as f64 / duration_secs;
+        println!("Tokens per second (TPS): {:.2}", tps);
+    }
     println!("Time elapsed in generate is: {:?}", i_duration);
 
     Ok(())