huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-12 04:44:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: improve tool choice syntax and response parsing/errors

Browse Source
This commit is contained in:
drbh 2024-07-18 17:41:17 +00:00
parent 35f8a88db5
commit 21dc6776b1
3 changed files with 189 additions and 157 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

225
router/src/infer/mod.rs
View File

@ -7,7 +7,7 @@ pub(crate) use health::HealthCheck;
use crate::validation::{ValidGenerateRequest, Validation, ValidationError}; use crate::validation::{ValidGenerateRequest, Validation, ValidationError};
use crate::{ use crate::{
ChatTemplateInputs, ChatTemplateVersions, FinishReason, GenerateRequest, HubProcessorConfig, ChatTemplateInputs, ChatTemplateVersions, FinishReason, GenerateRequest, HubProcessorConfig,
HubTokenizerConfig, Message, MessageChunk, PrefillToken, TextMessage, Token, HubTokenizerConfig, Message, MessageChunk, PrefillToken, TextMessage, Token, ToolChoice,
}; };
use crate::{ use crate::{
FunctionRef, FunctionsMap, GrammarType, Properties, TokenizerConfigToken, Tool, ToolType, Tools, FunctionRef, FunctionsMap, GrammarType, Properties, TokenizerConfigToken, Tool, ToolType, Tools,
@ -332,132 +332,131 @@ impl ChatTemplate {
pub struct ToolGrammar {} pub struct ToolGrammar {}
impl ToolGrammar { impl ToolGrammar {
// find a tool by name
fn find_tool_by_name(tools: &[Tool], name: &str) -> Result<Tool, InferError> {
tools
.iter()
.find(|tool| tool.function.name == name)
.cloned()
.ok_or_else(|| InferError::ToolError(format!("Tool with name {} not found", name)))
}
pub fn apply( pub fn apply(
tools: Option<Vec<Tool>>, tools: Option<Vec<Tool>>,
tool_choice: Option<ToolType>, tool_choice: ToolChoice,
) -> Result<Option<Tools>, InferError> { ) -> Result<Option<Tools>, InferError> {
if let Some(req_tools) = tools { // if no tools are provided, we return None
let tool_choice = tool_choice let tools = match tools {
.map(|t| match t { Some(tools) if !tools.is_empty() => tools,
ToolType::FunctionName(name) if name == "auto" => ToolType::OneOf, _ => return Ok(None),
_ => t, };
})
.unwrap_or_default();
let tools_to_use = match tool_choice { let tool_choice = tool_choice.0.unwrap_or(ToolType::OneOf);
ToolType::FunctionName(name) => {
vec![req_tools
.iter()
.find(|tool| tool.function.name == *name)
.unwrap_or_else(|| panic!("Tool with name {} not found", name))
.clone()]
}
ToolType::Function { function } => {
let tool = req_tools
.iter()
.find(|tool| tool.function.name == function.name)
.unwrap_or_else(|| panic!("Tool with name {} not found", function.name))
.clone();
vec![tool]
}
ToolType::OneOf => req_tools.to_owned(),
};
// adds the error notification function for LLM feedback if required // if tools are provided and no tool_choice we default to the OneOf
let mut text_response_properties = Map::new(); let tools_to_use = match tool_choice {
text_response_properties.insert( ToolType::FunctionName(name) => {
"error".to_string(), vec![Self::find_tool_by_name(&tools, &name)?]
serde_json::json!({ }
"type": "string", ToolType::Function { function } => {
"description": "The error or issue to notify" vec![Self::find_tool_by_name(&tools, &function.name)?]
}), }
); ToolType::OneOf => tools,
text_response_properties.insert( ToolType::NoTool => return Ok(None),
"_name".to_string(), };
serde_json::json!({
"type": "string",
"const": "notify_error"
}),
);
let functions: HashMap<String, serde_json::Value> = tools_to_use // adds the error notification function for LLM feedback if required
.iter() let mut text_response_properties = Map::new();
.map(|tool| { text_response_properties.insert(
let func = tool.function.clone(); "error".to_string(),
serde_json::json!({
"type": "string",
"description": "The error or issue to notify"
}),
);
text_response_properties.insert(
"_name".to_string(),
serde_json::json!({
"type": "string",
"const": "notify_error"
}),
);
// Clone the existing parameters, which are expected to be a JSON object let functions: HashMap<String, serde_json::Value> = tools_to_use
let mut params = if let Value::Object(params) = &func.arguments { .iter()
params.clone() .map(|tool| {
} else { let func = tool.function.clone();
Map::new()
};
// Insert the function's description at the top level, outside of properties // Clone the existing parameters, which are expected to be a JSON object
params.insert( let mut params = if let Value::Object(params) = &func.arguments {
"description".to_string(), params.clone()
Value::String(func.description.clone().unwrap_or_default()), } else {
); Map::new()
};
// Ensure 'properties' exists and is an object // Insert the function's description at the top level, outside of properties
let properties = params params.insert(
.entry("properties".to_string()) "description".to_string(),
.or_insert_with(|| json!({})) Value::String(func.description.clone().unwrap_or_default()),
.as_object_mut() );
.unwrap();
// Insert the constant for the function name inside 'properties' // Ensure 'properties' exists and is an object
properties.insert( let properties = params
"_name".to_string(), .entry("properties".to_string())
json!({ .or_insert_with(|| json!({}))
"type": "string", .as_object_mut()
"const": func.name.clone(), .unwrap();
// "description": "The name of the function"
}),
);
// Check if 'required' exists, and it is an array. If not, create an empty array. // Insert the constant for the function name inside 'properties'
let required = params properties.insert(
.entry("required".to_string()) "_name".to_string(),
.or_insert_with(|| json!([])) json!({
.as_array_mut() "type": "string",
.unwrap(); "const": func.name.clone(),
// "description": "The name of the function"
// Add 'name' to the 'required' array if it is not already present
if !required.iter().any(|r| r == "_name") {
required.push(json!("_name"));
}
(func.name, Value::Object(params))
})
.chain([(
"notify_error".to_string(),
serde_json::json!({
"properties": text_response_properties,
"required": ["error", "_name"],
"type": "object"
}), }),
)]) );
.collect();
let tools = Tools { // Check if 'required' exists, and it is an array. If not, create an empty array.
functions_map: FunctionsMap { functions }, let required = params
properties: Properties { .entry("required".to_string())
function: tools_to_use .or_insert_with(|| json!([]))
.iter() .as_array_mut()
.map(|tool| FunctionRef { .unwrap();
ref_path: format!("#/$functions/{}", tool.function.name.clone()),
})
.chain(std::iter::once(FunctionRef {
ref_path: "#/$functions/notify_error".to_string(),
}))
.collect(),
},
};
return Ok(Some(tools)); // Add 'name' to the 'required' array if it is not already present
} if !required.iter().any(|r| r == "_name") {
// Err(InferError::ToolError("No tools provided".to_string())) required.push(json!("_name"));
Ok(None) }
(func.name, Value::Object(params))
})
.chain([(
"notify_error".to_string(),
serde_json::json!({
"properties": text_response_properties,
"required": ["error", "_name"],
"type": "object"
}),
)])
.collect();
let tools = Tools {
functions_map: FunctionsMap { functions },
properties: Properties {
function: tools_to_use
.iter()
.map(|tool| FunctionRef {
ref_path: format!("#/$functions/{}", tool.function.name.clone()),
})
.chain(std::iter::once(FunctionRef {
ref_path: "#/$functions/notify_error".to_string(),
}))
.collect(),
},
};
Ok(Some(tools))
} }
} }

71
router/src/lib.rs
View File

@ -824,7 +824,7 @@ pub(crate) struct ChatRequest {
/// A specific tool to use. If not provided, the model will default to use any of the tools provided in the tools parameter. /// A specific tool to use. If not provided, the model will default to use any of the tools provided in the tools parameter.
#[serde(default)] #[serde(default)]
#[schema(nullable = true, example = "null")] #[schema(nullable = true, example = "null")]
pub tool_choice: Option<ToolType>, pub tool_choice: ToolChoice,
/// Response format constraints for the generation. /// Response format constraints for the generation.
/// ///
@ -840,16 +840,13 @@ fn default_tool_prompt() -> Option<String> {
) )
} }
#[derive(Clone, Default, Debug, Deserialize, PartialEq, Serialize, ToSchema)] #[derive(Clone, Debug, Deserialize, PartialEq, Serialize, ToSchema)]
#[serde(untagged)] #[serde(untagged)]
pub enum ToolType { pub enum ToolType {
#[default]
#[serde(alias = "auto")]
OneOf, OneOf,
FunctionName(String), FunctionName(String),
Function { Function { function: FunctionName },
function: FunctionName, NoTool,
},
} }
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)] #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
@ -857,27 +854,26 @@ pub struct FunctionName {
pub name: String, pub name: String,
} }
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
#[serde(from = "ToolTypeDeserializer")] #[serde(from = "ToolTypeDeserializer")]
pub struct ToolChoice(pub Option<ToolType>); pub struct ToolChoice(pub Option<ToolType>);
#[derive(Deserialize)] #[derive(Deserialize)]
#[serde(untagged)] #[serde(untagged)]
enum ToolTypeDeserializer { enum ToolTypeDeserializer {
None(Option<String>), String(String),
Some(ToolType), ToolType(ToolType),
} }
impl From<ToolTypeDeserializer> for ToolChoice { impl From<ToolTypeDeserializer> for ToolChoice {
fn from(value: ToolTypeDeserializer) -> Self { fn from(value: ToolTypeDeserializer) -> Self {
match value { match value {
ToolTypeDeserializer::None(opt) => match opt.as_deref() { ToolTypeDeserializer::String(s) => match s.as_str() {
Some("none") => ToolChoice(None), "none" => ToolChoice(Some(ToolType::NoTool)),
Some("auto") => ToolChoice(Some(ToolType::OneOf)), "auto" => ToolChoice(Some(ToolType::OneOf)),
Some(s) => ToolChoice(Some(ToolType::FunctionName(s.to_string()))), _ => ToolChoice(Some(ToolType::FunctionName(s))),
None => ToolChoice(Some(ToolType::OneOf)),
}, },
ToolTypeDeserializer::Some(tool_type) => ToolChoice(Some(tool_type)), ToolTypeDeserializer::ToolType(tool_type) => ToolChoice(Some(tool_type)),
} }
} }
} }
@ -1375,4 +1371,47 @@ mod tests {
r#"{"role":"assistant","tool_calls":[{"id":"0","type":"function","function":{"description":null,"name":"myfn","arguments":{"format":"csv"}}}]}"# r#"{"role":"assistant","tool_calls":[{"id":"0","type":"function","function":{"description":null,"name":"myfn","arguments":{"format":"csv"}}}]}"#
); );
} }
#[test]
fn tool_deserialize() {
// Test ToolCall deserialization
let json = r#"{"id":"0","type":"function","function":{"description":null,"name":"myfn","arguments":{"format":"csv"}}}"#;
let tool: ToolCall = serde_json::from_str(json).unwrap();
assert_eq!(
tool,
ToolCall {
id: "0".to_string(),
r#type: "function".to_string(),
function: FunctionDefinition {
description: None,
name: "myfn".to_string(),
arguments: json!({
"format": "csv"
}),
},
}
);
// Test ToolChoice deserialization with "auto"
let auto_json = r#""auto""#;
let auto_choice: ToolChoice = serde_json::from_str(auto_json).unwrap();
assert_eq!(auto_choice, ToolChoice(Some(ToolType::OneOf)));
// Test ToolChoice deserialization with "none"
let none_json = r#""none""#;
let none_choice: ToolChoice = serde_json::from_str(none_json).unwrap();
assert_eq!(none_choice, ToolChoice(None));
// Test ToolChoice deserialization with a specific function name
let function_json = r#""my_function""#;
let function_choice: ToolChoice = serde_json::from_str(function_json).unwrap();
assert_eq!(
function_choice,
ToolChoice(Some(ToolType::FunctionName("my_function".to_string())))
);
// Test ToolChoice deserialization with no value (should default to OneOf)
let default_json = r#"null"#;
let default_choice: ToolChoice = serde_json::from_str(default_json).unwrap();
assert_eq!(default_choice, ToolChoice(Some(ToolType::OneOf)));
}
} }

50
router/src/server.rs
View File

@ -1192,39 +1192,33 @@ async fn chat_completions(
.as_secs(); .as_secs();
let (tool_calls, output) = if tool_grammar.is_some() { let (tool_calls, output) = if tool_grammar.is_some() {
// gen_text should be valid json let gen_text_value: Value = serde_json::from_str(&generation.generated_text)
let gen_text_value: Value = .map_err(|e| InferError::ToolError(e.to_string()))?;
serde_json::from_str(&generation.generated_text).map_err(|e| {
( let function = gen_text_value.get("function").ok_or(InferError::ToolError(
StatusCode::UNPROCESSABLE_ENTITY, "No function found in generated text".to_string(),
Json(ErrorResponse { ))?;
error: e.to_string(),
error_type: "Input validation error".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");
}
let tool_calls = vec![ToolCall { let tool_calls = vec![ToolCall {
id: "0".to_string(), id: "0".to_string(),
r#type: "function".to_string(), r#type: "function".to_string(),
function: FunctionDefinition { function: FunctionDefinition {
description: None, description: None,
name: gen_text_value name,
.get("function") arguments,
.and_then(|f| f.get("_name"))
.and_then(|name| name.as_str())
.unwrap_or("default_function_name")
.to_string(),
// Serialize the JSON object obtained from "function" to an escaped JSON string
arguments: gen_text_value
.get("function")
.map(|f| {
let mut f_cloned = f.clone();
if let Value::Object(ref mut props) = f_cloned {
props.remove("_name");
}
f_cloned
})
.unwrap_or_default(),
}, },
}]; }];
(Some(tool_calls), None) (Some(tool_calls), None)