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

fix: simplify tool grammar logic and improve schema

Browse Source
This commit is contained in:
drbh 2024-08-26 17:59:21 +00:00
parent 8b45d82897
commit 1f72dcf062
2 changed files with 23 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
router/src/infer/tool_grammar.rs
View File

@ -43,13 +43,9 @@ impl ToolGrammar {
"error": { "error": {
"type": "string", "type": "string",
"description": "The error or issue to notify" "description": "The error or issue to notify"
},
"_name": {
"type": "string",
"const": "notify_error"
} }
}, },
"required": ["error", "_name"] "required": ["error"]
}), }),
}, },
}; };
@ -72,48 +68,43 @@ impl ToolGrammar {
.map(|tool| { .map(|tool| {
let func = tool.function.clone(); let func = tool.function.clone();
// Clone the existing parameters, which are expected to be a JSON object let mut params = Map::new();
let mut params = if let Value::Object(params) = &func.arguments {
params.clone()
} else {
Map::new()
};
// Insert the function's description at the top level, outside of properties
params.insert( params.insert(
"description".to_string(), "description".to_string(),
Value::String(func.description.clone().unwrap_or_default()), Value::String(func.description.unwrap_or_default()),
); );
// Ensure 'properties' exists and is an object let mut properties = Map::new();
let properties = params let mut required = vec![Value::String("_name".to_string())];
.entry("properties".to_string())
.or_insert_with(|| json!({}))
.as_object_mut()
.unwrap();
// Insert the constant for the function name inside 'properties'
properties.insert( properties.insert(
"_name".to_string(), "_name".to_string(),
json!({ json!({
"type": "string", "type": "string",
"const": func.name.clone(), "const": func.name.clone(),
// "description": "The name of the function"
}), }),
); );
// Check if 'required' exists, and it is an array. If not, create an empty array. if let Value::Object(args) = func.arguments {
let required = params if let Some(Value::Object(props)) = args.get("properties") {
.entry("required".to_string()) properties.extend(props.clone());
.or_insert_with(|| json!([])) }
.as_array_mut() if let Some(Value::Array(reqs)) = args.get("required") {
.unwrap(); required.extend(reqs.clone());
}
// Add 'name' to the 'required' array if it is not already present params.insert(
if !required.iter().any(|r| r == "_name") { "additionalProperties".to_string(),
required.push(json!("_name")); Value::Bool(
args.get("additionalProperties").and_then(|v| v.as_str())
== Some("true"),
),
);
} }
params.insert("properties".to_string(), Value::Object(properties));
params.insert("required".to_string(), Value::Array(required));
(func.name, Value::Object(params)) (func.name, Value::Object(params))
}) })
.collect(); .collect();
@ -126,9 +117,6 @@ impl ToolGrammar {
.map(|tool| FunctionRef { .map(|tool| FunctionRef {
ref_path: format!("#/$functions/{}", tool.function.name.clone()), ref_path: format!("#/$functions/{}", tool.function.name.clone()),
}) })
.chain(std::iter::once(FunctionRef {
ref_path: "#/$functions/notify_error".to_string(),
}))
.collect(), .collect(),
}, },
}; };

2
router/src/server.rs
View File

@ -2689,6 +2689,6 @@ mod tests {
assert!(result.is_ok()); assert!(result.is_ok());
let (inputs, _grammar, using_tools) = result.unwrap(); let (inputs, _grammar, using_tools) = result.unwrap();
assert_eq!(using_tools, true); assert_eq!(using_tools, true);
assert_eq!(inputs, "<s>[AVAILABLE_TOOLS] [{\"type\": \"function\", \"function\": {\"arguments\": {\"properties\":{\"format\":{\"description\":\"The temperature unit to use. Infer this from the users location.\",\"enum\":[\"celsius\",\"fahrenheit\"],\"type\":\"string\"},\"location\":{\"description\":\"The city and state, e.g. San Francisco, CA\",\"type\":\"string\"}},\"required\":[\"location\",\"format\"],\"type\":\"object\"}, \"description\": \"Get the current weather\", \"name\": \"get_current_weather\"}}, {\"type\": \"function\", \"function\": {\"arguments\": {\"properties\":{\"_name\":{\"const\":\"notify_error\",\"type\":\"string\"},\"error\":{\"description\":\"The error or issue to notify\",\"type\":\"string\"}},\"required\":[\"error\",\"_name\"],\"type\":\"object\"}, \"description\": \"Notify an error or issue\", \"name\": \"notify_error\"}}][/AVAILABLE_TOOLS][INST] What is the weather like in New York?\n---\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.[/INST]".to_string()); assert_eq!(inputs, "<s>[AVAILABLE_TOOLS] [{\"type\": \"function\", \"function\": {\"arguments\": {\"properties\":{\"format\":{\"description\":\"The temperature unit to use. Infer this from the users location.\",\"enum\":[\"celsius\",\"fahrenheit\"],\"type\":\"string\"},\"location\":{\"description\":\"The city and state, e.g. San Francisco, CA\",\"type\":\"string\"}},\"required\":[\"location\",\"format\"],\"type\":\"object\"}, \"description\": \"Get the current weather\", \"name\": \"get_current_weather\"}}, {\"type\": \"function\", \"function\": {\"arguments\": {\"properties\":{\"error\":{\"description\":\"The error or issue to notify\",\"type\":\"string\"}},\"required\":[\"error\"],\"type\":\"object\"}, \"description\": \"Notify an error or issue\", \"name\": \"notify_error\"}}][/AVAILABLE_TOOLS][INST] What is the weather like in New York?\n---\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.[/INST]".to_string());
} }
} }