diff --git a/.github/workflows/router_tests.yaml b/.github/workflows/router_tests.yaml index 6eb95f130..404b481d6 100644 --- a/.github/workflows/router_tests.yaml +++ b/.github/workflows/router_tests.yaml @@ -29,7 +29,7 @@ jobs: - name: Install Rust uses: actions-rs/toolchain@v1 with: - toolchain: 1.79.0 + toolchain: 1.83.0 override: true components: rustfmt, clippy - name: Install Protoc diff --git a/Dockerfile b/Dockerfile index eccefae58..0988daf58 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,5 +1,5 @@ # Rust builder -FROM lukemathwalker/cargo-chef:latest-rust-1.79 AS chef +FROM lukemathwalker/cargo-chef:latest-rust-1.83 AS chef WORKDIR /usr/src ARG CARGO_REGISTRIES_CRATES_IO_PROTOCOL=sparse diff --git a/router/src/health.rs b/router/src/health.rs index 5ca8e8de8..cfda0b85c 100644 --- a/router/src/health.rs +++ b/router/src/health.rs @@ -1,6 +1,6 @@ use lorax_client::{ - Batch, NextTokenChooserParameters, Request, ShardInfo, ShardedClient, - StoppingCriteriaParameters, + input_chunk, Batch, InputChunk, NextTokenChooserParameters, Request, ShardInfo, ShardedClient, + StoppingCriteriaParameters, TokenizedInputs, }; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; @@ -40,7 +40,12 @@ impl Health { let generation_liveness_request = Request { id: LIVENESS_ID, inputs: "liveness".to_string(), - tokenized_inputs: None, + tokenized_inputs: Some(TokenizedInputs { + ids: vec![75], + input_chunks: vec![InputChunk { + chunk: Some(input_chunk::Chunk::Text("liveness".to_string())), + }], + }), truncate: 10, prefill_logprobs: false, parameters: Some(NextTokenChooserParameters { @@ -66,7 +71,7 @@ impl Health { adapter_index: 0, // Block 0 is reserved for health checks blocks: vec![0], - slots: (0..16).collect(), + slots: (0..self.shard_info.block_size).collect(), cache_len: 0, chunk_len: None, }; @@ -84,15 +89,20 @@ impl Health { pub(crate) async fn check_classification(&mut self) -> bool { let classify_request = Request { id: LIVENESS_ID, - inputs: "San Francisco".to_string(), - tokenized_inputs: None, + inputs: "liveness".to_string(), + tokenized_inputs: Some(TokenizedInputs { + ids: vec![75], + input_chunks: vec![InputChunk { + chunk: Some(input_chunk::Chunk::Text("liveness".to_string())), + }], + }), truncate: 10, prefill_logprobs: false, parameters: None, stopping_parameters: None, adapter_index: 0, blocks: vec![0], - slots: (0..16).collect(), + slots: (0..self.shard_info.block_size).collect(), cache_len: 0, chunk_len: None, }; @@ -109,15 +119,20 @@ impl Health { pub(crate) async fn check_embeddings(&mut self) -> bool { let embed_request = Request { id: LIVENESS_ID, - inputs: "San Francisco".to_string(), - tokenized_inputs: None, + inputs: "liveness".to_string(), + tokenized_inputs: Some(TokenizedInputs { + ids: vec![75], + input_chunks: vec![InputChunk { + chunk: Some(input_chunk::Chunk::Text("liveness".to_string())), + }], + }), truncate: 10, prefill_logprobs: false, parameters: None, stopping_parameters: None, adapter_index: 0, blocks: vec![0], - slots: (0..16).collect(), + slots: (0..self.shard_info.block_size).collect(), cache_len: 0, chunk_len: None, }; diff --git a/router/src/infer.rs b/router/src/infer.rs index 703dacd46..83a2d4e83 100644 --- a/router/src/infer.rs +++ b/router/src/infer.rs @@ -112,12 +112,14 @@ impl ChatTemplateRenderer { // if not, we need to append the tools to the last message let text = if self.use_default_tool_template { match serde_json::to_string(&tools) { - Ok(tools_str) => format!("\n---\n{}\n{}", tools_str, tool_prompt), + // Ok(tools_str) => format!("\n---\n{}\n{}", tools_str, tool_prompt), + Ok(tools_str) => format!("\n{}\n{}", tools_str, tool_prompt), Err(e) => return Err(InferError::ToolError(e.to_string())), } } else { // if the `tools` variable is used in the template, we just append the tool_prompt - format!("\n---\n{}", tool_prompt) + // format!("\n---\n{}", tool_prompt) + format!("\n{}", tool_prompt) }; if let Some(last_message) = messages.last_mut() { if let Some(content) = &mut last_message.content { diff --git a/router/src/lib.rs b/router/src/lib.rs index c3cf2cedc..b046815a6 100644 --- a/router/src/lib.rs +++ b/router/src/lib.rs @@ -581,7 +581,7 @@ pub struct Url { } #[derive(Clone, Deserialize, Serialize, ToSchema, Default, Debug, PartialEq)] -pub(crate) struct ToolCall { +pub struct ToolCall { pub id: String, pub r#type: String, pub function: ReturnFunctionDefinition, @@ -603,6 +603,8 @@ pub struct Message { #[schema(example = "My name is David and I")] pub content: Option, #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] #[schema(example = "\"David\"")] name: Option, } @@ -642,6 +644,8 @@ pub struct TextMessage { pub role: String, #[schema(example = "My name is David and I")] pub content: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, } impl From for TextMessage { @@ -660,6 +664,7 @@ impl From for TextMessage { .join(""), None => String::new(), }, + tool_calls: value.tool_calls, } } } @@ -858,7 +863,8 @@ impl ChatCompletionRequest { } pub fn default_tool_prompt() -> String { - "\nGiven the functions available, please respond with a JSON for a function call with its proper arguments that best answers the given prompt. Respond in the format {name: function name, parameters: dictionary of argument name and its value}.Do not use variables.\n".to_string() + // "\nGiven the functions available, please respond with a JSON for a function call with its proper arguments that best answers the given prompt. Respond in the format {name: function name, parameters: dictionary of argument name and its value}.Do not use variables.\n".to_string() + "".to_string() } #[derive(Clone, Debug, Deserialize, PartialEq, Serialize, ToSchema)] @@ -951,7 +957,7 @@ pub(crate) struct FunctionDefinition { } #[derive(Clone, Debug, Deserialize, Serialize, ToSchema, Default, PartialEq)] -pub(crate) struct ReturnFunctionDefinition { +pub struct ReturnFunctionDefinition { #[serde(default)] pub description: Option, pub name: String, diff --git a/router/src/server.rs b/router/src/server.rs index 3eb24521d..941cd354a 100644 --- a/router/src/server.rs +++ b/router/src/server.rs @@ -33,6 +33,7 @@ use futures::Stream; use lorax_client::{ShardInfo, ShardedClient}; use metrics_exporter_prometheus::{Matcher, PrometheusBuilder, PrometheusHandle}; use once_cell::sync::OnceCell; +use regex::Regex; use reqwest_middleware::ClientBuilder; use reqwest_retry::{policies::ExponentialBackoff, RetryTransientMiddleware}; use serde::{Deserialize, Serialize}; @@ -210,6 +211,112 @@ async fn completions_v1( } } +fn parse_json_tool_call( + gen_text_value: Value, +) -> Result<(Option>, Option), InferError> { + let function = gen_text_value.get("function").ok_or(InferError::ToolError( + "No function found in generated text".to_string(), + ))?; + + let name = function + .get("_name") + .and_then(Value::as_str) + .ok_or(InferError::ToolError( + "No _name found in generated text".to_string(), + ))? + .to_string(); + + let mut arguments = function.clone(); + if let Value::Object(ref mut props) = arguments { + props.remove("_name"); + } + match name.as_str() { + "no_tool" => { + // parse the content message + let content_message = arguments + .get("content") + .and_then(Value::as_str) + .ok_or_else(|| { + InferError::ToolError("No `content` found in generated text".to_string()) + })? + .to_string(); + Ok((None, Some(content_message))) + } + _ => { + let arguments = serde_json::to_string(&arguments).map_err(|e| { + InferError::ToolError(format!("Failed to serialize arguments: {}", e)) + })?; + let tool_calls = vec![ToolCall { + id: "0".to_string(), + r#type: "function".to_string(), + function: ReturnFunctionDefinition { + description: None, + name, + arguments, + }, + }]; + Ok((Some(tool_calls), None)) + } + } +} + +fn parse_xml_tool_call(gen: &str) -> Result<(Option>, Option), InferError> { + let tool_call_regex = Regex::new(r"(?s)(.*?)|(.*)") + .map_err(|e| InferError::ToolError(format!("Failed to create tool call regex: {}", e)))?; + // Check for tool call matches + if let Some(captures) = tool_call_regex.captures(gen) { + // Check for complete tool call (first capture group) + let json_content = if let Some(complete_match) = captures.get(1) { + complete_match.as_str() + } + // Check for incomplete tool call (second capture group) + else if let Some(incomplete_match) = captures.get(2) { + incomplete_match.as_str() + } else { + return Ok((None, Some(gen.to_string()))); + }; + + // Parse the JSON content + let parsed_content: serde_json::Value = + serde_json::from_str(json_content.trim()).map_err(|e| { + InferError::ToolError(format!("Failed to parse tool call JSON content: {}", e)) + })?; + + // Extract name and arguments + let name = parsed_content["name"] + .as_str() + .ok_or_else(|| InferError::ToolError("Missing 'name' field in tool call".to_string()))? + .to_string(); + + // Parse the arguments field which may be a JSON string + let arguments = if let Some(args_str) = parsed_content["arguments"].as_str() { + // If arguments is a string, try to parse it as JSON + serde_json::from_str(args_str).unwrap_or(parsed_content["arguments"].clone()) + } else { + // If not a string, use the raw value + parsed_content["arguments"].clone() + }; + + // Create tool call with the extracted content + let tool_calls = vec![ToolCall { + id: "0".to_string(), + r#type: "function".to_string(), + function: ReturnFunctionDefinition { + description: None, + name, + arguments: serde_json::to_string(&arguments).map_err(|e| { + InferError::ToolError(format!("Failed to serialize arguments: {}", e)) + })?, + }, + }]; + + Ok((Some(tool_calls), None)) + } else { + // If no tool call tags are found, return the original text + Ok((None, Some(gen.to_string()))) + } +} + /// OpenAI compatible chat completions endpoint #[utoipa::path( post, @@ -319,57 +426,14 @@ async fn chat_completions_v1( let mut choice_content = vec![]; for (_, gen) in generations.iter().enumerate() { let (tool_calls, output) = if using_tools { - let gen_text_value: Value = serde_json::from_str(&gen).map_err(|e| { - InferError::ToolError(format!( - "Failed to parse generated text: {} {:?}", - e, gen - )) - })?; - let function = gen_text_value.get("function").ok_or(InferError::ToolError( - "No function found in generated text".to_string(), - ))?; - - let name = function - .get("_name") - .and_then(Value::as_str) - .ok_or(InferError::ToolError( - "No _name found in generated text".to_string(), - ))? - .to_string(); - - let mut arguments = function.clone(); - if let Value::Object(ref mut props) = arguments { - props.remove("_name"); - } - match name.as_str() { - "no_tool" => { - // parse the content message - let content_message = arguments - .get("content") - .and_then(Value::as_str) - .ok_or_else(|| { - InferError::ToolError( - "No `content` found in generated text".to_string(), - ) - })? - .to_string(); - (None, Some(content_message)) - } - _ => { - let arguments = serde_json::to_string(&arguments).map_err(|e| { - InferError::ToolError(format!("Failed to serialize arguments: {}", e)) - })?; - let tool_calls = vec![ToolCall { - id: "0".to_string(), - r#type: "function".to_string(), - function: ReturnFunctionDefinition { - description: None, - name, - arguments, - }, - }]; - (Some(tool_calls), None) - } + let tool_call_result = match serde_json::from_str::(gen) { + Ok(gen_text_value) => parse_json_tool_call(gen_text_value), + Err(_) => parse_xml_tool_call(gen), + }; + match tool_call_result { + Ok((tool_calls, output)) => (tool_calls, output), + // TODO: (magdy) How should we tell the user that the tool call failed? + Err(_) => (None, Some(gen.clone())), } } else { (None, Some(gen.clone())) @@ -435,7 +499,8 @@ pub(crate) fn prepare_chat_input( messages, Some((updated_tools, tool_prompt.into())), )?; - return Ok((inputs, grammar, tool_schema.is_some())); + // return Ok((inputs, grammar, tool_schema.is_some())); + return Ok((inputs, grammar, true)); } // if no response_format or tools are set simply apply the chat template to generate inputs diff --git a/router/src/tool_grammar.rs b/router/src/tool_grammar.rs index 6a1f604ba..2ecadda0e 100644 --- a/router/src/tool_grammar.rs +++ b/router/src/tool_grammar.rs @@ -1,8 +1,5 @@ use crate::infer::InferError; -use crate::{ - FunctionDefinition, FunctionRef, FunctionsMap, JsonSchemaTool, Properties, Tool, ToolChoice, - ToolType, -}; +use crate::{FunctionRef, FunctionsMap, JsonSchemaTool, Properties, Tool, ToolChoice, ToolType}; use serde_json::{json, Map, Value}; use std::collections::HashMap; @@ -29,27 +26,27 @@ impl ToolGrammar { let tool_choice = tool_choice.0.unwrap_or(ToolType::OneOf); - let mut tools = tools.clone(); - - // add the no_tool function to the tools - let no_tool = Tool { - r#type: "function".to_string(), - function: FunctionDefinition { - name: "no_tool".to_string(), - description: Some("Open ened response with no specific tool selected".to_string()), - parameters: json!({ - "type": "object", - "properties": { - "content": { - "type": "string", - "description": "The response content", - } - }, - "required": ["content"] - }), - }, - }; - tools.push(no_tool); + // let mut tools = tools.clone(); + + // // add the no_tool function to the tools + // let no_tool = Tool { + // r#type: "function".to_string(), + // function: FunctionDefinition { + // name: "no_tool".to_string(), + // description: Some("Open ened response with no specific tool selected".to_string()), + // parameters: json!({ + // "type": "object", + // "properties": { + // "content": { + // "type": "string", + // "description": "The response content", + // } + // }, + // "required": ["content"] + // }), + // }, + // }; + // tools.push(no_tool); // if tools are provided and no tool_choice we default to the OneOf let tools_to_use = match tool_choice { @@ -106,7 +103,7 @@ impl ToolGrammar { }) .collect(); - let tool_schema = JsonSchemaTool { + let _tool_schema = JsonSchemaTool { functions_map: FunctionsMap { functions }, properties: Properties { function: tools_to_use @@ -118,6 +115,7 @@ impl ToolGrammar { }, }; - Ok((tools, Some(tool_schema))) + // Ok((tools, Some(tool_schema))) + Ok((tools, None)) } } diff --git a/rust-toolchain.toml b/rust-toolchain.toml index b6ffc9d2c..80afd2d32 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,3 +1,3 @@ [toolchain] -channel = "1.79.0" -components = ["rustfmt", "clippy"] \ No newline at end of file +channel = "1.83.0" +components = ["rustfmt", "clippy"] diff --git a/server/lorax_server/layers/fp8.py b/server/lorax_server/layers/fp8.py index f03d2974a..c11f23c15 100644 --- a/server/lorax_server/layers/fp8.py +++ b/server/lorax_server/layers/fp8.py @@ -14,7 +14,7 @@ def apply_fp8_linear( input_scale_ub: Optional[torch.Tensor] = None, qbias: Optional[torch.Tensor] = None, ) -> torch.Tensor: - qinput, x_scale = ops.scaled_fp8_quant(input, input_scale, scale_ub=input_scale_ub, use_per_token_if_dynamic=False) + qinput, x_scale = ops.scaled_fp8_quant(input, input_scale, scale_ub=input_scale_ub, use_per_token_if_dynamic=True) output = ops.cutlass_scaled_mm( qinput, qweight, out_dtype=input.dtype, scale_a=x_scale, scale_b=weight_scale, bias=qbias