From ca1a640da206359f36d0aa7b7bc1c97ac068a213 Mon Sep 17 00:00:00 2001 From: ochafik Date: Thu, 2 May 2024 03:20:00 +0100 Subject: [PATCH] server: tool call grammar-constraints fix --- common/json-schema-to-grammar.cpp | 138 ++++++++++++++++++++++-------- common/json-schema-to-grammar.h | 1 + examples/server/server.cpp | 2 +- examples/server/utils.hpp | 28 +++++- 4 files changed, 127 insertions(+), 42 deletions(-) diff --git a/common/json-schema-to-grammar.cpp b/common/json-schema-to-grammar.cpp index 0f8f1b1d4..393153b40 100644 --- a/common/json-schema-to-grammar.cpp +++ b/common/json-schema-to-grammar.cpp @@ -11,6 +11,9 @@ using json = nlohmann::ordered_json; +const char * DOTALL = "[\\U00000000-\\U0010FFFF]"; +const char * DOT = "[^\\x0A\\x0D]"; + template static std::string join(Iterator begin, Iterator end, const std::string & separator); @@ -198,6 +201,29 @@ static std::string format_literal(const std::string & literal) { } +/* + not_literal('a') -> '[^a]' + not_literal('abc') -> '([^a] | "a" ([^b] | "b" ([^c])?)?)?' +*/ +static std::string not_literal(const std::string & literal, bool dotall = true) { + assert(literal.size() > 0); + std::stringstream out; + std::function recurse = [&](size_t i) { + const auto & c = literal[i]; + out << "[^" << c << "]"; + if (i < literal.size() - 1) { + out << " | " << format_literal(std::to_string(c)) << " ("; + recurse(i + 1); + out << ")?"; + } + }; + out << "("; + recurse(0); + out << ")" << (dotall ? DOTALL : DOT) << "*"; + return out.str(); +} + + class SchemaConverter { private: std::function _fetch_json; @@ -208,22 +234,6 @@ private: std::vector _errors; std::vector _warnings; - std::string _add_rule(const std::string & name, const std::string & rule) { - std::string esc_name = regex_replace(name, INVALID_RULE_CHARS_RE, "-"); - if (_rules.find(esc_name) == _rules.end() || _rules[esc_name] == rule) { - _rules[esc_name] = rule; - return esc_name; - } else { - int i = 0; - while (_rules.find(esc_name + std::to_string(i)) != _rules.end() && _rules[esc_name + std::to_string(i)] != rule) { - i++; - } - std::string key = esc_name + std::to_string(i); - _rules[key] = rule; - return key; - } - } - std::string _generate_union_rule(const std::string & name, const std::vector & alt_schemas) { std::vector rules; for (size_t i = 0; i < alt_schemas.size(); i++) { @@ -256,11 +266,11 @@ private: auto get_dot = [&]() { std::string rule; if (_dotall) { - rule = "[\\U00000000-\\U0010FFFF]"; + rule = DOTALL; } else { - rule = "[^\\x0A\\x0D]"; + rule = DOT; } - return _add_rule("dot", rule); + return add_rule("dot", rule); }; // Joins the sequence, merging consecutive literals together. @@ -377,7 +387,7 @@ private: if (!sub_is_literal) { std::string & sub_id = sub_rule_ids[sub]; if (sub_id.empty()) { - sub_id = _add_rule(name + "-" + std::to_string(sub_rule_ids.size()), sub); + sub_id = add_rule(name + "-" + std::to_string(sub_rule_ids.size()), sub); } sub = sub_id; } @@ -423,7 +433,7 @@ private: } return join_seq(); }; - return _add_rule(name, "\"\\\"\" " + to_rule(transform()) + " \"\\\"\" space"); + return add_rule(name, "\"\\\"\" " + to_rule(transform()) + " \"\\\"\" space"); } std::string _resolve_ref(const std::string & ref) { @@ -451,7 +461,7 @@ private: const auto &prop_schema = kv.second; std::string prop_rule_name = visit(prop_schema, name + (name.empty() ? "" : "-") + prop_name); - prop_kv_rule_names[prop_name] = _add_rule( + prop_kv_rule_names[prop_name] = add_rule( name + (name.empty() ? "" : "-") + prop_name + "-kv", format_literal(json(prop_name).dump()) + " space \":\" space " + prop_rule_name ); @@ -464,7 +474,7 @@ private: if (additional_properties.is_object() || (additional_properties.is_boolean() && additional_properties.get())) { std::string sub_name = name + (name.empty() ? "" : "-") + "additional"; std::string value_rule = visit(additional_properties.is_object() ? additional_properties : json::object(), sub_name + "-value"); - std::string kv_rule = _add_rule(sub_name + "-kv", _add_primitive("string", PRIMITIVE_RULES.at("string")) + " \":\" space " + value_rule); + std::string kv_rule = add_rule(sub_name + "-kv", _add_primitive("string", PRIMITIVE_RULES.at("string")) + " \":\" space " + value_rule); prop_kv_rule_names["*"] = kv_rule; optional_props.push_back("*"); } @@ -491,7 +501,7 @@ private: std::string k = ks[0]; std::string kv_rule_name = prop_kv_rule_names[k]; if (k == "*") { - res = _add_rule( + res = add_rule( name + (name.empty() ? "" : "-") + "additional-kvs", kv_rule_name + " ( \",\" space " + kv_rule_name + " )*" ); @@ -501,7 +511,7 @@ private: res = kv_rule_name; } if (ks.size() > 1) { - res += " " + _add_rule( + res += " " + add_rule( name + (name.empty() ? "" : "-") + k + "-rest", get_recursive_refs(std::vector(ks.begin() + 1, ks.end()), true) ); @@ -527,7 +537,7 @@ private: } std::string _add_primitive(const std::string & name, const BuiltinRule & rule) { - auto n = _add_rule(name, rule.content); + auto n = add_rule(name, rule.content); for (const auto & dep : rule.deps) { BuiltinRule dep_rule; auto it = PRIMITIVE_RULES.find(dep); @@ -615,6 +625,22 @@ public: visit_refs(schema); } + std::string add_rule(const std::string & name, const std::string & rule) { + std::string esc_name = regex_replace(name, INVALID_RULE_CHARS_RE, "-"); + if (_rules.find(esc_name) == _rules.end() || _rules[esc_name] == rule) { + _rules[esc_name] = rule; + return esc_name; + } else { + int i = 0; + while (_rules.find(esc_name + std::to_string(i)) != _rules.end() && _rules[esc_name + std::to_string(i)] != rule) { + i++; + } + std::string key = esc_name + std::to_string(i); + _rules[key] = rule; + return key; + } + } + std::string _generate_constant_rule(const json & value) { return format_literal(value.dump()); } @@ -625,24 +651,24 @@ public: std::string rule_name = is_reserved_name(name) ? name + "-" : name.empty() ? "root" : name; if (schema.contains("$ref")) { - return _add_rule(rule_name, _resolve_ref(schema["$ref"])); + return add_rule(rule_name, _resolve_ref(schema["$ref"])); } else if (schema.contains("oneOf") || schema.contains("anyOf")) { std::vector alt_schemas = schema.contains("oneOf") ? schema["oneOf"].get>() : schema["anyOf"].get>(); - return _add_rule(rule_name, _generate_union_rule(name, alt_schemas)); + return add_rule(rule_name, _generate_union_rule(name, alt_schemas)); } else if (schema_type.is_array()) { std::vector schema_types; for (const auto & t : schema_type) { schema_types.push_back({{"type", t}}); } - return _add_rule(rule_name, _generate_union_rule(name, schema_types)); + return add_rule(rule_name, _generate_union_rule(name, schema_types)); } else if (schema.contains("const")) { - return _add_rule(rule_name, _generate_constant_rule(schema["const"])); + return add_rule(rule_name, _generate_constant_rule(schema["const"])); } else if (schema.contains("enum")) { std::vector enum_values; for (const auto & v : schema["enum"]) { enum_values.push_back(_generate_constant_rule(v)); } - return _add_rule(rule_name, join(enum_values.begin(), enum_values.end(), " | ")); + return add_rule(rule_name, join(enum_values.begin(), enum_values.end(), " | ")); } else if ((schema_type.is_null() || schema_type == "object") && (schema.contains("properties") || (schema.contains("additionalProperties") && schema["additionalProperties"] != true))) { @@ -660,7 +686,7 @@ public: properties.emplace_back(prop.key(), prop.value()); } } - return _add_rule(rule_name, + return add_rule(rule_name, _build_object_rule( properties, required, name, schema.contains("additionalProperties") ? schema["additionalProperties"] : json())); @@ -691,7 +717,7 @@ public: add_component(t, true); } } - return _add_rule(rule_name, _build_object_rule(properties, required, hybrid_name, json())); + return add_rule(rule_name, _build_object_rule(properties, required, hybrid_name, json())); } else if ((schema_type.is_null() || schema_type == "array") && (schema.contains("items") || schema.contains("prefixItems"))) { json items = schema.contains("items") ? schema["items"] : schema["prefixItems"]; if (items.is_array()) { @@ -703,14 +729,14 @@ public: rule += visit(items[i], name + (name.empty() ? "" : "-") + "tuple-" + std::to_string(i)); } rule += " \"]\" space"; - return _add_rule(rule_name, rule); + return add_rule(rule_name, rule); } else { std::string item_rule_name = visit(items, name + (name.empty() ? "" : "-") + "item"); int min_items = schema.contains("minItems") ? schema["minItems"].get() : 0; json max_items_json = schema.contains("maxItems") ? schema["maxItems"] : json(); int max_items = max_items_json.is_number_integer() ? max_items_json.get() : std::numeric_limits::max(); - return _add_rule(rule_name, "\"[\" space " + build_repetition(item_rule_name, min_items, max_items, "\",\" space") + " \"]\" space"); + return add_rule(rule_name, "\"[\" space " + build_repetition(item_rule_name, min_items, max_items, "\",\" space") + " \"]\" space"); } } else if ((schema_type.is_null() || schema_type == "string") && schema.contains("pattern")) { return _visit_pattern(schema["pattern"], rule_name); @@ -718,14 +744,14 @@ public: return _add_primitive(rule_name == "root" ? "root" : schema_format, PRIMITIVE_RULES.at("uuid")); } else if ((schema_type.is_null() || schema_type == "string") && STRING_FORMAT_RULES.find(schema_format + "-string") != STRING_FORMAT_RULES.end()) { auto prim_name = schema_format + "-string"; - return _add_rule(rule_name, _add_primitive(prim_name, STRING_FORMAT_RULES.at(prim_name))); + return add_rule(rule_name, _add_primitive(prim_name, STRING_FORMAT_RULES.at(prim_name))); } else if (schema_type == "string" && (schema.contains("minLength") || schema.contains("maxLength"))) { std::string char_rule = _add_primitive("char", PRIMITIVE_RULES.at("char")); int min_len = schema.contains("minLength") ? schema["minLength"].get() : 0; int max_len = schema.contains("maxLength") ? schema["maxLength"].get() : std::numeric_limits::max(); - return _add_rule(rule_name, "\"\\\"\" " + build_repetition(char_rule, min_len, max_len) + " \"\\\"\" space"); + return add_rule(rule_name, "\"\\\"\" " + build_repetition(char_rule, min_len, max_len) + " \"\\\"\" space"); } else if (schema.empty() || schema_type == "object") { - return _add_rule(rule_name, _add_primitive("object", PRIMITIVE_RULES.at("object"))); + return add_rule(rule_name, _add_primitive("object", PRIMITIVE_RULES.at("object"))); } else { if (!schema_type.is_string() || PRIMITIVE_RULES.find(schema_type.get()) == PRIMITIVE_RULES.end()) { _errors.push_back("Unrecognized schema: " + schema.dump()); @@ -762,3 +788,39 @@ std::string json_schema_to_grammar(const json & schema) { converter.check_errors(); return converter.format_grammar(); } + +std::string tool_call_grammar(const json & tools) { + SchemaConverter converter([](const std::string &) { return json::object(); }, /* dotall= */ false); + + std::vector tool_rules; + + for (const auto & tool : tools) { + const auto & function = tool["function"]; + std::string name = function["name"]; + std::string description = function["description"]; + auto parameters_copy = function["parameters"]; + converter.resolve_refs(parameters_copy, name); + + tool_rules.push_back(converter.visit(json { + {"type", "object"}, + {"description", description}, + {"properties", json { + {"name", json {{"const", name}}}, + {"arguments", parameters_copy}, + }}, + {"required", json::array({"name", "arguments"})}, + }, name + "-tool-call")); + } + + converter.add_rule( + "root", + not_literal("") + " | " + + converter.add_rule( + "tool_call", + "\"\" " + + join(tool_rules.begin(), tool_rules.end(), " | ") + + " \"\"")); + + converter.check_errors(); + return converter.format_grammar(); +} diff --git a/common/json-schema-to-grammar.h b/common/json-schema-to-grammar.h index e1abed303..5594051b1 100644 --- a/common/json-schema-to-grammar.h +++ b/common/json-schema-to-grammar.h @@ -1,4 +1,5 @@ #pragma once #include "json.hpp" +std::string tool_call_grammar(const nlohmann::ordered_json & tools); std::string json_schema_to_grammar(const nlohmann::ordered_json& schema); diff --git a/examples/server/server.cpp b/examples/server/server.cpp index f60530cf3..da45ab11b 100644 --- a/examples/server/server.cpp +++ b/examples/server/server.cpp @@ -3031,7 +3031,7 @@ int main(int argc, char ** argv) { chat.push_back({{"role", "assistant"}, {"content", "Hi there"}}); chat.push_back({{"role", "user"}, {"content", "How are you?"}}); - const std::string chat_example = format_chat(ctx_server.model, sparams.chat_template, chat); + const std::string chat_example = format_chat(ctx_server.model, sparams.chat_template, chat, ""); LOG_INFO("chat template", { {"chat_example", chat_example}, diff --git a/examples/server/utils.hpp b/examples/server/utils.hpp index 1a2212502..336307a56 100644 --- a/examples/server/utils.hpp +++ b/examples/server/utils.hpp @@ -4,6 +4,7 @@ #include "common.h" #include "json.hpp" +#include "json-schema-to-grammar.h" #include #include @@ -122,7 +123,7 @@ inline bool verify_custom_template(const std::string & tmpl) { } // Format given chat. If tmpl is empty, we take the template from model metadata -inline std::string format_chat(const struct llama_model * model, const std::string & tmpl, const std::vector & messages) { +inline std::string format_chat(const struct llama_model * model, const std::string & tmpl, const std::vector & messages, const std::string & tools_tag) { size_t alloc_size = 0; // vector holding all allocated string to be passed to llama_chat_apply_template std::vector str(messages.size() * 2); @@ -137,6 +138,20 @@ inline std::string format_chat(const struct llama_model * model, const std::stri chat[i].content = str[i*2 + 1].c_str(); } + if (!tools_tag.empty()) { + alloc_size += tools_tag.size(); + if (chat.empty()) { + str.resize(2); + str[0] = "user"; + str[1] = tools_tag; + chat.push_back({str[0].c_str(), str[1].c_str()}); + } else { + auto & content = str[str.size() - 1]; + content += tools_tag; + chat[chat.size() - 1].content = content.c_str(); + } + } + const char * ptr_tmpl = tmpl.empty() ? nullptr : tmpl.c_str(); std::vector buf(alloc_size * 2); @@ -372,8 +387,15 @@ static json oaicompat_completion_params_parse( llama_params["temperature"] = json_value(body, "temperature", 0.0); llama_params["top_p"] = json_value(body, "top_p", 1.0); + std::string tools_tag; + if (body.contains("tools") && body["tools"].is_array()) { + const auto & tools = body["tools"]; + llama_params["grammar"] = tool_call_grammar(tools); + tools_tag = (std::stringstream() << "\n\n" << tools.dump(2) << "").str(); + } + // Apply chat template to the list of messages - llama_params["prompt"] = format_chat(model, chat_template, body["messages"]); + llama_params["prompt"] = format_chat(model, chat_template, body["messages"], tools_tag); // Handle "stop" field if (body.contains("stop") && body["stop"].is_string()) { @@ -408,7 +430,7 @@ static json oaicompat_completion_params_parse( } // Params supported by OAI but unsupported by llama.cpp - static const std::vector unsupported_params { "tools", "tool_choice" }; + static const std::vector unsupported_params { "tool_choice" }; for (auto & param : unsupported_params) { if (body.contains(param)) { throw std::runtime_error("Unsupported param: " + param);