diff --git a/cactus-engine/src/complete.cpp b/cactus-engine/src/complete.cpp index 51ef39b06..85273c3c9 100644 --- a/cactus-engine/src/complete.cpp +++ b/cactus-engine/src/complete.cpp @@ -25,6 +25,8 @@ namespace { std::vector> extract_schema_property_types(const std::string& schema); std::vector extract_schema_required(const std::string& schema); +std::vector extract_string_array(const std::string& json, const std::string& key); +std::string extract_json_object_field(const std::string& json, const std::string& key); std::string extract_last_user_query(const std::vector& messages) { for (auto it = messages.rbegin(); it != messages.rend(); ++it) { @@ -65,9 +67,13 @@ std::vector build_tool_constraint_specs(const std::vectorsecond); + std::string properties_object = extract_json_object_field(schema_it->second, "properties"); spec.parameter_names.reserve(properties.size()); + spec.parameter_enums.reserve(properties.size()); for (const auto& [name, _] : properties) { spec.parameter_names.push_back(name); + spec.parameter_enums.push_back( + extract_string_array(extract_json_object_field(properties_object, name), "enum")); } spec.required_parameter_names = extract_schema_required(schema_it->second); } @@ -194,25 +200,30 @@ std::vector> extract_schema_property_types(c return properties; } -std::vector extract_schema_required(const std::string& schema) { - std::vector required; - std::string key = "\"required\""; - size_t key_pos = schema.find(key); - if (key_pos == std::string::npos) return required; - size_t arr_start = schema.find('[', key_pos + key.size()); - if (arr_start == std::string::npos) return required; - size_t arr_end = schema.find(']', arr_start); - if (arr_end == std::string::npos) return required; +std::vector extract_string_array(const std::string& json, const std::string& key) { + std::vector values; + if (json.empty()) return values; + std::string pattern = "\"" + key + "\""; + size_t key_pos = json.find(pattern); + if (key_pos == std::string::npos) return values; + size_t arr_start = json.find('[', key_pos + pattern.size()); + if (arr_start == std::string::npos) return values; + size_t arr_end = json.find(']', arr_start); + if (arr_end == std::string::npos) return values; size_t pos = arr_start + 1; while (pos < arr_end) { - size_t qs = schema.find('"', pos); + size_t qs = json.find('"', pos); if (qs == std::string::npos || qs >= arr_end) break; - size_t qe = schema.find('"', qs + 1); + size_t qe = json.find('"', qs + 1); if (qe == std::string::npos || qe > arr_end) break; - required.push_back(schema.substr(qs + 1, qe - qs - 1)); + values.push_back(json.substr(qs + 1, qe - qs - 1)); pos = qe + 1; } - return required; + return values; +} + +std::vector extract_schema_required(const std::string& schema) { + return extract_string_array(schema, "required"); } static bool json_object_has_key(const std::string& obj, const std::string& key) { @@ -1141,6 +1152,10 @@ int cactus_complete( : "kept local"; } + if (prompt.options.force_tools && !prompt.tools.empty() && primary_function_calls.empty()) { + CACTUS_LOG_WARN("force_tools", "force_tools was set but no tool call was parsed from the constrained output"); + } + std::string result = construct_response_json(primary_response, primary_function_calls, time_to_first_token, total_time_ms, prefill_tps, decode_tps, prompt_tokens, completion_tokens, confidence, handoff_succeeded, diff --git a/cactus-engine/src/constraints.cpp b/cactus-engine/src/constraints.cpp index 599d3e8ed..ec6731dbd 100644 --- a/cactus-engine/src/constraints.cpp +++ b/cactus-engine/src/constraints.cpp @@ -5,18 +5,14 @@ namespace engine { constexpr float FORCE_BIAS = 500.0f; constexpr float BLOCK_BIAS = -500.0f; -constexpr float NEEDLE_BLOCK_BIAS = -1e9f; - -static bool contains_odd_quote_tags(const std::string& text) { - static constexpr const char* kQuoteTag = "<|\"|>"; - static constexpr size_t kQuoteTagLen = 5; - bool odd = false; - size_t pos = 0; - while ((pos = text.find(kQuoteTag, pos)) != std::string::npos) { - odd = !odd; - pos += kQuoteTagLen; - } - return odd; +constexpr float TRIE_BLOCK_BIAS = -1e9f; + +static constexpr const char* kEscapeTag = "<|\"|>"; +static constexpr size_t kEscapeTagLen = 5; + +static bool ends_with_escape_tag(const std::string& text) { + return text.size() >= kEscapeTagLen && + text.compare(text.size() - kEscapeTagLen, kEscapeTagLen, kEscapeTag) == 0; } void ToolCallConstrainer::add_tokens_for_string(const std::string& str, std::unordered_set& token_set) { @@ -27,36 +23,6 @@ void ToolCallConstrainer::add_tokens_for_string(const std::string& str, std::uno } } -void ToolCallConstrainer::add_tokens_for_prefix_string(const std::string& prefix, std::unordered_set& token_set) { - if (!tokenizer_) return; - const uint32_t vocab = tokenizer_->get_vocab_size(); - for (uint32_t t = 0; t < vocab; t++) { - if (tokenizer_->decode({t}).rfind(prefix, 0) == 0) { - token_set.insert(t); - } - } -} - -void ToolCallConstrainer::tokenize_function_names(bool quote_names) { - all_func_name_tokens_.clear(); - func_name_sequences_.clear(); - - for (const auto& name : function_names_) { - std::string name_to_encode = quote_names ? ("\"" + name + "\"") : name; - auto tokens = tokenizer_->encode(name_to_encode); - func_name_sequences_[name] = tokens; - for (uint32_t t : tokens) { - all_func_name_tokens_.insert(t); - } - if (quote_names) { - auto unquoted_tokens = tokenizer_->encode(name); - for (uint32_t t : unquoted_tokens) { - all_func_name_tokens_.insert(t); - } - } - } -} - void ToolCallConstrainer::init_common_tokens() { backtick_tokens_.clear(); add_tokens_for_string("`", backtick_tokens_); @@ -70,22 +36,22 @@ void ToolCallConstrainer::init_common_tokens() { add_tokens_for_string("` ", backtick_tokens_); } -void ToolCallConstrainer::needle_insert_word(NeedleTrieNode* root, const std::string& word) { +void ToolCallConstrainer::trie_insert(TrieNode* root, const std::string& word) { if (!root) return; - NeedleTrieNode* node = root; + TrieNode* node = root; for (char ch : word) { auto& child = node->children[ch]; - if (!child) child = std::make_unique(); + if (!child) child = std::make_unique(); node = child.get(); } node->is_terminal = true; } -const ToolCallConstrainer::NeedleTrieNode* ToolCallConstrainer::needle_get_trie_node( - const NeedleTrieNode* root, +const ToolCallConstrainer::TrieNode* ToolCallConstrainer::trie_seek( + const TrieNode* root, const std::string& prefix ) const { - const NeedleTrieNode* node = root; + const TrieNode* node = root; for (char ch : prefix) { if (!node) return nullptr; auto it = node->children.find(ch); @@ -95,12 +61,12 @@ const ToolCallConstrainer::NeedleTrieNode* ToolCallConstrainer::needle_get_trie_ return node; } -bool ToolCallConstrainer::needle_check_token_valid(const std::string& token_text, - const NeedleTrieNode* trie_node) const { - const NeedleTrieNode* node = trie_node; +bool ToolCallConstrainer::trie_token_valid(const std::string& token_text, + const TrieNode* node, + char terminator) const { if (!node) return false; for (char ch : token_text) { - if (ch == '"') return node->is_terminal; + if (ch == terminator) return node->is_terminal; auto it = node->children.find(ch); if (it == node->children.end()) return false; node = it->second.get(); @@ -108,28 +74,165 @@ bool ToolCallConstrainer::needle_check_token_valid(const std::string& token_text return true; } -void ToolCallConstrainer::reset_needle_constraints() { - needle_json_state_ = NeedleJsonState::FREE; - needle_buffer_.clear(); - needle_constrained_buf_.clear(); - needle_current_function_.clear(); - needle_in_arguments_ = false; - needle_arguments_depth_ = 0; - needle_nesting_depth_ = 0; - needle_in_string_value_ = false; - needle_prev_char_escape_ = false; +const ToolCallConstrainer::TrieNode* ToolCallConstrainer::current_param_trie() const { + auto it = param_tries_.find(current_function_); + return it != param_tries_.end() ? it->second.get() : nullptr; +} + +void ToolCallConstrainer::mark_trie_bias(const TrieNode* node, char terminator, + const std::vector& extra_allowed, + int32_t terminal_token, bool exact_terminator) { + if (!node) return; + + std::vector valid(token_strings_.size(), false); + bool has_valid = false; + auto mark = [&](char first_char, bool unconditional) { + auto idx_it = token_index_.find(first_char); + if (idx_it == token_index_.end()) return; + for (uint32_t token_id : idx_it->second) { + if (valid[token_id]) continue; + if (unconditional || trie_token_valid(token_strings_[token_id], node, terminator)) { + valid[token_id] = true; + has_valid = true; + } + } + }; + + for (const auto& [first_char, _] : node->children) mark(first_char, false); + if (node->is_terminal) { + if (exact_terminator) { + auto idx_it = token_index_.find(terminator); + if (idx_it != token_index_.end()) { + for (uint32_t token_id : idx_it->second) { + if (!valid[token_id] && token_strings_[token_id].size() == 1) { + valid[token_id] = true; + has_valid = true; + } + } + } + } else { + mark(terminator, false); + } + } + for (char c : extra_allowed) mark(c, true); + if (node->is_terminal && terminal_token >= 0) { + uint32_t tid = static_cast(terminal_token); + if (tid < valid.size() && !valid[tid]) { + valid[tid] = true; + has_valid = true; + } + } + if (!has_valid) return; + + for (size_t token_id = 0; token_id < valid.size(); ++token_id) { + if (!valid[token_id]) { + current_bias_[static_cast(token_id)] = TRIE_BLOCK_BIAS; + } + } +} + +void ToolCallConstrainer::reset_constraint_state() { + region_ = Region::FREE; + constraint_buffer_.clear(); + constrained_buf_.clear(); + current_function_.clear(); + in_arguments_ = false; + arguments_depth_ = 0; + nesting_depth_ = 0; + in_string_value_ = false; + prev_char_escape_ = false; + at_key_start_ = false; + await_enum_open_ = false; + active_enum_trie_ = nullptr; + emitted_keys_.clear(); + remaining_key_trie_.reset(); +} + +void ToolCallConstrainer::rebuild_remaining_keys() { + remaining_key_trie_ = std::make_unique(); + for (const auto& tool : tool_specs_) { + if (tool.name != current_function_) continue; + for (const auto& param_name : tool.parameter_names) { + if (!param_name.empty() && !emitted_keys_.count(param_name)) { + trie_insert(remaining_key_trie_.get(), param_name); + } + } + break; + } +} + +bool ToolCallConstrainer::required_satisfied() const { + for (const auto& tool : tool_specs_) { + if (tool.name != current_function_) continue; + for (const auto& req : tool.required_parameter_names) { + if (!emitted_keys_.count(req)) return false; + } + return true; + } + return true; +} + +void ToolCallConstrainer::build_constraint_tables() { + name_trie_ = std::make_unique(); + param_tries_.clear(); + enum_tries_.clear(); + for (const auto& tool : tool_specs_) { + if (!tool.name.empty()) trie_insert(name_trie_.get(), tool.name); + + if (is_needle()) { + auto param_root = std::make_unique(); + for (const auto& param_name : tool.parameter_names) { + if (!param_name.empty()) trie_insert(param_root.get(), param_name); + } + param_tries_[tool.name] = std::move(param_root); + } + + for (size_t i = 0; i < tool.parameter_names.size() && i < tool.parameter_enums.size(); ++i) { + if (tool.parameter_enums[i].empty()) continue; + auto enum_root = std::make_unique(); + for (const auto& value : tool.parameter_enums[i]) { + if (!value.empty()) trie_insert(enum_root.get(), value); + } + enum_tries_[tool.name][tool.parameter_names[i]] = std::move(enum_root); + } + } + + const uint32_t vocab_size = tokenizer_ ? tokenizer_->get_vocab_size() : 0; + if (token_strings_.size() == vocab_size) return; + + token_strings_.assign(vocab_size, ""); + token_index_.clear(); + + const uint32_t eos = tokenizer_->get_eos_token(); + const uint32_t bos = tokenizer_->get_bos_token(); + const uint32_t unk = tokenizer_->get_unk_token(); + constexpr uint32_t pad = 0; + + for (uint32_t token_id = 0; token_id < vocab_size; ++token_id) { + if (token_id == pad || token_id == eos || token_id == bos || token_id == unk) continue; + std::string decoded = tokenizer_->decode({token_id}); + if (decoded.size() == 1) { + unsigned char b = static_cast(decoded[0]); + if (b >= 0x80) { + decoded.assign({static_cast(0xC0 | (b >> 6)), + static_cast(0x80 | (b & 0x3F))}); + } + } + token_strings_[token_id] = decoded; + if (!decoded.empty()) token_index_[decoded.front()].push_back(token_id); + } } bool ToolCallConstrainer::needle_at_arg_key_start() const { - if (needle_buffer_.size() < 2) return false; - return needle_buffer_.compare(needle_buffer_.size() - 2, 2, "{\"") == 0 || - needle_buffer_.compare(needle_buffer_.size() - 2, 2, ",\"") == 0; + if (constraint_buffer_.size() < 2) return false; + return constraint_buffer_.compare(constraint_buffer_.size() - 2, 2, "{\"") == 0 || + constraint_buffer_.compare(constraint_buffer_.size() - 2, 2, ",\"") == 0; } bool ToolCallConstrainer::needle_is_value_string_start() const { - if (needle_buffer_.empty()) return false; - for (size_t i = needle_buffer_.size() - 1; i-- > 0;) { - char ch = needle_buffer_[i]; + if (constraint_buffer_.empty()) return false; + for (size_t i = constraint_buffer_.size() - 1; i-- > 0;) { + char ch = constraint_buffer_[i]; if (ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r') continue; return ch == ':'; } @@ -137,69 +240,66 @@ bool ToolCallConstrainer::needle_is_value_string_start() const { } void ToolCallConstrainer::feed_needle_char(char ch) { - if (needle_json_state_ == NeedleJsonState::IN_NAME || - needle_json_state_ == NeedleJsonState::IN_ARG_KEY) { + if (region_ == Region::IN_NAME || region_ == Region::IN_ARG_KEY) { if (ch == '"') { - if (needle_json_state_ == NeedleJsonState::IN_NAME) { - needle_current_function_ = needle_constrained_buf_; - } - needle_constrained_buf_.clear(); - needle_json_state_ = NeedleJsonState::FREE; + if (region_ == Region::IN_NAME) current_function_ = constrained_buf_; + constrained_buf_.clear(); + region_ = Region::FREE; } else { - needle_constrained_buf_.push_back(ch); + constrained_buf_.push_back(ch); } - needle_buffer_.push_back(ch); + constraint_buffer_.push_back(ch); return; } - needle_buffer_.push_back(ch); + constraint_buffer_.push_back(ch); - if (needle_in_string_value_) { - if (needle_prev_char_escape_) { - needle_prev_char_escape_ = false; + if (in_string_value_) { + if (prev_char_escape_) { + prev_char_escape_ = false; return; } if (ch == '\\') { - needle_prev_char_escape_ = true; + prev_char_escape_ = true; return; } - if (ch == '"') needle_in_string_value_ = false; + if (ch == '"') in_string_value_ = false; return; } if (ch == '{' || ch == '[') { - needle_nesting_depth_++; + nesting_depth_++; } else if (ch == '}' || ch == ']') { - needle_nesting_depth_ = std::max(0, needle_nesting_depth_ - 1); - if (ch == '}' && needle_in_arguments_ && needle_nesting_depth_ < needle_arguments_depth_) { - needle_in_arguments_ = false; + nesting_depth_ = std::max(0, nesting_depth_ - 1); + if (ch == '}' && in_arguments_ && nesting_depth_ < arguments_depth_) { + in_arguments_ = false; } } - if (needle_buffer_.size() >= 8 && !needle_in_arguments_ && - needle_buffer_.compare(needle_buffer_.size() - 8, 8, "\"name\":\"") == 0) { - needle_json_state_ = NeedleJsonState::IN_NAME; - needle_constrained_buf_.clear(); + if (constraint_buffer_.size() >= 8 && !in_arguments_ && + constraint_buffer_.compare(constraint_buffer_.size() - 8, 8, "\"name\":\"") == 0) { + region_ = Region::IN_NAME; + constrained_buf_.clear(); return; } - if (needle_buffer_.size() >= 13 && - needle_buffer_.compare(needle_buffer_.size() - 13, 13, "\"arguments\":{") == 0) { - needle_in_arguments_ = true; - needle_arguments_depth_ = needle_nesting_depth_; + if (constraint_buffer_.size() >= 13 && + constraint_buffer_.compare(constraint_buffer_.size() - 13, 13, "\"arguments\":{") == 0) { + in_arguments_ = true; + arguments_depth_ = nesting_depth_; return; } - if (needle_in_arguments_ && - needle_nesting_depth_ == needle_arguments_depth_ && + if (in_arguments_ && + nesting_depth_ == arguments_depth_ && needle_at_arg_key_start()) { - needle_json_state_ = NeedleJsonState::IN_ARG_KEY; - needle_constrained_buf_.clear(); + region_ = Region::IN_ARG_KEY; + constrained_buf_.clear(); return; } if (ch == '"' && needle_is_value_string_start()) { - needle_in_string_value_ = true; + in_string_value_ = true; } } @@ -207,61 +307,105 @@ void ToolCallConstrainer::feed_needle_text(const std::string& text) { for (char ch : text) feed_needle_char(ch); } -void ToolCallConstrainer::init_needle_constraints() { - reset_needle_constraints(); - - needle_name_trie_ = std::make_unique(); - needle_param_tries_.clear(); - for (const auto& tool : tool_specs_) { - if (!tool.name.empty()) needle_insert_word(needle_name_trie_.get(), tool.name); - - auto param_root = std::make_unique(); - for (const auto& param_name : tool.parameter_names) { - if (!param_name.empty()) needle_insert_word(param_root.get(), param_name); +void ToolCallConstrainer::feed_gemma_char(char ch) { + if (state_ != State::GEMMA_NAME && state_ != State::GEMMA_ARGS) return; + + constraint_buffer_.push_back(ch); + + if (region_ == Region::IN_NAME) { + if (ch == '{') { + current_function_ = constrained_buf_; + constrained_buf_.clear(); + region_ = Region::FREE; + state_ = State::GEMMA_ARGS; + in_arguments_ = true; + nesting_depth_ = 1; + arguments_depth_ = nesting_depth_; + at_key_start_ = true; + rebuild_remaining_keys(); + } else { + constrained_buf_.push_back(ch); } - needle_param_tries_[tool.name] = std::move(param_root); + return; } - const uint32_t vocab_size = tokenizer_ ? tokenizer_->get_vocab_size() : 0; - if (needle_token_strings_.size() != vocab_size) { - needle_token_strings_.assign(vocab_size, ""); - needle_token_index_.clear(); - - const uint32_t eos = tokenizer_->get_eos_token(); - const uint32_t bos = tokenizer_->get_bos_token(); - const uint32_t unk = tokenizer_->get_unk_token(); - constexpr uint32_t pad = 0; - - for (uint32_t token_id = 0; token_id < vocab_size; ++token_id) { - if (token_id == pad || token_id == eos || token_id == bos || token_id == unk) continue; - std::string decoded = tokenizer_->decode({token_id}); - if (decoded.size() == 1) { - unsigned char b = static_cast(decoded[0]); - if (b >= 0x80) { - decoded.assign({static_cast(0xC0 | (b >> 6)), - static_cast(0x80 | (b & 0x3F))}); + if (region_ == Region::IN_ARG_KEY) { + if (ch == ':') { + emitted_keys_.insert(constrained_buf_); + active_enum_trie_ = nullptr; + if (escape_tag_valid_) { + auto fit = enum_tries_.find(current_function_); + if (fit != enum_tries_.end()) { + auto kit = fit->second.find(constrained_buf_); + if (kit != fit->second.end()) active_enum_trie_ = kit->second.get(); } } - needle_token_strings_[token_id] = decoded; - if (!decoded.empty()) needle_token_index_[decoded.front()].push_back(token_id); + await_enum_open_ = active_enum_trie_ != nullptr; + rebuild_remaining_keys(); + constrained_buf_.clear(); + region_ = Region::FREE; + } else { + constrained_buf_.push_back(ch); + } + return; + } + + if (in_string_value_) { + if (ends_with_escape_tag(constraint_buffer_)) { + in_string_value_ = false; + active_enum_trie_ = nullptr; + constrained_buf_.clear(); + } else if (active_enum_trie_) { + constrained_buf_.push_back(ch); } + return; + } + if (ends_with_escape_tag(constraint_buffer_)) { + in_string_value_ = true; + at_key_start_ = false; + await_enum_open_ = false; + constrained_buf_.clear(); + return; + } + if (ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r') return; + + if (at_key_start_ && ch != '}' && ch != ',') { + region_ = Region::IN_ARG_KEY; + constrained_buf_.clear(); + constrained_buf_.push_back(ch); + at_key_start_ = false; + return; + } + + if (ch == '{' || ch == '[') { + nesting_depth_++; + } else if (ch == '}' || ch == ']') { + nesting_depth_--; + if (ch == '}' && nesting_depth_ < arguments_depth_) { + in_arguments_ = false; + state_ = State::GEMMA_EXPECT_END; + forced_tag_progress_ = 0; + } + } else if (ch == ',' && nesting_depth_ == arguments_depth_) { + at_key_start_ = true; } } +void ToolCallConstrainer::feed_gemma_text(const std::string& text) { + for (char ch : text) feed_gemma_char(ch); +} + void ToolCallConstrainer::tokenize_grammar_elements() { if (!tokenizer_) return; open_brace_tokens_.clear(); close_brace_tokens_.clear(); - colon_tokens_.clear(); - comma_tokens_.clear(); init_common_tokens(); gemma_call_start_tokens_.clear(); gemma_call_end_tokens_.clear(); gemma_call_prefix_tokens_.clear(); - escape_tokens_.clear(); gemma_response_start_tokens_.clear(); call_start_sequence_ = tokenizer_->encode(call_start_tag_); @@ -274,10 +418,10 @@ void ToolCallConstrainer::tokenize_grammar_elements() { add_tokens_for_string("{", open_brace_tokens_); add_tokens_for_string("}", close_brace_tokens_); - add_tokens_for_string(":", colon_tokens_); - add_tokens_for_string(",", comma_tokens_); - tokenize_function_names(false); + std::vector escape_sequence = tokenizer_->encode(kEscapeTag); + escape_tag_valid_ = escape_sequence.size() == 1; + escape_tag_token_ = escape_tag_valid_ ? escape_sequence[0] : 0; } void ToolCallConstrainer::init(Config::ModelType model_type, @@ -285,21 +429,14 @@ void ToolCallConstrainer::init(Config::ModelType model_type, Tokenizer* tokenizer) { model_type_ = model_type; tool_specs_ = tools; - function_names_.clear(); - function_names_.reserve(tool_specs_.size()); - for (const auto& tool : tool_specs_) { - function_names_.push_back(tool.name); - } tokenizer_ = tokenizer; generated_text_.clear(); current_bias_.clear(); - brace_depth_ = 0; - in_argument_string_ = false; forced_tag_progress_ = 0; const bool model_supported = is_needle() || Config::is_gemma_family(model_type_); - active_ = model_supported && !function_names_.empty() && tokenizer != nullptr; + active_ = model_supported && !tool_specs_.empty() && tokenizer != nullptr; - reset_needle_constraints(); + reset_constraint_state(); state_ = is_needle() ? State::NEEDLE_START : State::GEMMA_START; if (Config::is_gemma3_family(model_type_)) { call_start_tag_ = ""; @@ -316,9 +453,7 @@ void ToolCallConstrainer::init(Config::ModelType model_type, } tokenize_grammar_elements(); - if (is_needle()) { - init_needle_constraints(); - } + build_constraint_tables(); compute_bias(); } @@ -354,49 +489,18 @@ void ToolCallConstrainer::update(uint32_t token_id, const std::string& decoded_t case State::GEMMA_EXPECT_CALL: if (generated_text_.find("call:") != std::string::npos) { - state_ = State::GEMMA_IN_FUNC_NAME; - generated_text_.clear(); - } - break; - - case State::GEMMA_IN_FUNC_NAME: - for (const auto& name : function_names_) { - if (generated_text_.find(name) != std::string::npos) { - state_ = State::GEMMA_EXPECT_BRACE; - generated_text_.clear(); - break; - } - } - break; - - case State::GEMMA_EXPECT_BRACE: - if (generated_text_.find("{") != std::string::npos) { - state_ = State::GEMMA_IN_ARGUMENTS; - brace_depth_ = 1; - in_argument_string_ = false; + state_ = State::GEMMA_NAME; generated_text_.clear(); + region_ = Region::IN_NAME; + constrained_buf_.clear(); + constraint_buffer_.clear(); } break; - case State::GEMMA_IN_ARGUMENTS: + case State::GEMMA_NAME: + case State::GEMMA_ARGS: + feed_gemma_text(decoded_text); generated_text_.clear(); - if (contains_odd_quote_tags(decoded_text)) { - in_argument_string_ = !in_argument_string_; - } - if (!in_argument_string_) { - for (char c : decoded_text) { - if (c == '{') brace_depth_++; - else if (c == '}') { - brace_depth_--; - if (brace_depth_ == 0) { - state_ = State::GEMMA_EXPECT_END; - in_argument_string_ = false; - forced_tag_progress_ = 0; - break; - } - } - } - } break; case State::GEMMA_EXPECT_END: @@ -423,56 +527,28 @@ void ToolCallConstrainer::compute_bias() { for (uint32_t t : backtick_tokens_) { current_bias_[t] = BLOCK_BIAS; } - for (uint32_t t : gemma_response_start_tokens_) { current_bias_[t] = BLOCK_BIAS; } } - if (is_needle()) { - if (needle_json_state_ == NeedleJsonState::FREE) return; - - const NeedleTrieNode* trie_node = nullptr; - if (needle_json_state_ == NeedleJsonState::IN_NAME) { - trie_node = needle_get_trie_node(needle_name_trie_.get(), needle_constrained_buf_); - } else if (needle_json_state_ == NeedleJsonState::IN_ARG_KEY) { - auto it = needle_param_tries_.find(needle_current_function_); - if (it != needle_param_tries_.end()) { - trie_node = needle_get_trie_node(it->second.get(), needle_constrained_buf_); - } - } - if (!trie_node) return; - - std::vector valid_tokens(needle_token_strings_.size(), false); - bool has_valid = false; - auto mark_valid_tokens = [&](char first_char) { - auto idx_it = needle_token_index_.find(first_char); - if (idx_it == needle_token_index_.end()) return; - for (uint32_t token_id : idx_it->second) { - if (!valid_tokens[token_id] && - needle_check_token_valid(needle_token_strings_[token_id], trie_node)) { - valid_tokens[token_id] = true; - has_valid = true; - } - } - }; - - for (const auto& [first_char, _] : trie_node->children) { - mark_valid_tokens(first_char); - } - if (trie_node->is_terminal) { - mark_valid_tokens('"'); - } - if (!has_valid) return; - - for (size_t token_id = 0; token_id < valid_tokens.size(); ++token_id) { - if (!valid_tokens[token_id]) { - current_bias_[static_cast(token_id)] = NEEDLE_BLOCK_BIAS; - } + if (region_ != Region::FREE) { + const TrieNode* node = nullptr; + char terminator = '"'; + if (region_ == Region::IN_NAME) { + node = trie_seek(name_trie_.get(), constrained_buf_); + if (!is_needle()) terminator = '{'; + } else { + const TrieNode* root = is_needle() ? current_param_trie() : remaining_key_trie_.get(); + node = trie_seek(root, constrained_buf_); + if (!is_needle()) terminator = ':'; } + mark_trie_bias(node, terminator, {}, -1, !is_needle()); return; } + if (is_needle()) return; + switch (state_) { case State::GEMMA_START: if (forced_tag_progress_ < call_start_sequence_.size()) { @@ -498,50 +574,24 @@ void ToolCallConstrainer::compute_bias() { } break; - case State::GEMMA_IN_FUNC_NAME: - for (uint32_t t : all_func_name_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : close_brace_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : gemma_call_end_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::GEMMA_EXPECT_BRACE: - for (uint32_t t : open_brace_tokens_) { - current_bias_[t] = FORCE_BIAS; - } - for (uint32_t t : gemma_call_end_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - break; - - case State::GEMMA_IN_ARGUMENTS: - if (!in_argument_string_) { - for (uint32_t t : colon_tokens_) { - current_bias_[t] = 10.0f; - } - for (uint32_t t : comma_tokens_) { - current_bias_[t] = 8.0f; + case State::GEMMA_ARGS: + if (await_enum_open_) { + current_bias_[escape_tag_token_] = FORCE_BIAS; + for (uint32_t t : gemma_call_end_tokens_) current_bias_[t] = BLOCK_BIAS; + for (uint32_t t : gemma_call_start_tokens_) current_bias_[t] = BLOCK_BIAS; + } else if (active_enum_trie_ && in_string_value_) { + mark_trie_bias(trie_seek(active_enum_trie_, constrained_buf_), '\0', {}, + static_cast(escape_tag_token_)); + } else if (at_key_start_) { + std::vector extra; + if (required_satisfied()) extra.push_back('}'); + mark_trie_bias(remaining_key_trie_.get(), ':', extra); + } else { + for (uint32_t t : gemma_call_end_tokens_) current_bias_[t] = BLOCK_BIAS; + for (uint32_t t : gemma_call_start_tokens_) current_bias_[t] = BLOCK_BIAS; + if (nesting_depth_ == arguments_depth_ && !in_string_value_ && !required_satisfied()) { + for (uint32_t t : close_brace_tokens_) current_bias_[t] = BLOCK_BIAS; } - for (uint32_t t : escape_tokens_) { - current_bias_[t] = 5.0f; - } - for (uint32_t t : close_brace_tokens_) { - current_bias_[t] = 3.0f; - } - for (uint32_t t : open_brace_tokens_) { - current_bias_[t] = 3.0f; - } - } - for (uint32_t t : gemma_call_end_tokens_) { - current_bias_[t] = BLOCK_BIAS; - } - for (uint32_t t : gemma_call_start_tokens_) { - current_bias_[t] = BLOCK_BIAS; } break; @@ -565,10 +615,8 @@ void ToolCallConstrainer::compute_bias() { void ToolCallConstrainer::reset() { generated_text_.clear(); current_bias_.clear(); - brace_depth_ = 0; - in_argument_string_ = false; forced_tag_progress_ = 0; - reset_needle_constraints(); + reset_constraint_state(); state_ = is_needle() ? State::NEEDLE_START : State::GEMMA_START; diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 34678a10f..46bcdf415 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -281,6 +281,7 @@ struct ChatMessage { struct ToolConstraintSpec { std::string name; std::vector parameter_names; + std::vector> parameter_enums; std::vector required_parameter_names; }; @@ -510,9 +511,8 @@ class ToolCallConstrainer { DONE, GEMMA_START, GEMMA_EXPECT_CALL, - GEMMA_IN_FUNC_NAME, - GEMMA_EXPECT_BRACE, - GEMMA_IN_ARGUMENTS, + GEMMA_NAME, + GEMMA_ARGS, GEMMA_EXPECT_END, NEEDLE_START }; @@ -530,16 +530,19 @@ class ToolCallConstrainer { bool is_active() const { return active_; } private: + enum class Region { FREE, IN_NAME, IN_ARG_KEY }; + struct TrieNode { + std::unordered_map> children; + bool is_terminal = false; + }; + bool active_ = false; State state_ = State::GEMMA_START; Config::ModelType model_type_ = Config::ModelType::GEMMA4; Tokenizer* tokenizer_ = nullptr; std::vector tool_specs_; - std::vector function_names_; std::string generated_text_; - int brace_depth_ = 0; - bool in_argument_string_ = false; std::string call_start_tag_; std::string call_end_tag_; @@ -548,65 +551,65 @@ class ToolCallConstrainer { std::vector call_end_sequence_; size_t forced_tag_progress_ = 0; - std::unordered_set all_func_name_tokens_; - std::unordered_map> func_name_sequences_; - std::unordered_set gemma_call_start_tokens_; std::unordered_set gemma_call_end_tokens_; std::unordered_set gemma_response_start_tokens_; std::unordered_set gemma_call_prefix_tokens_; - std::unordered_set escape_tokens_; std::unordered_set backtick_tokens_; std::unordered_set open_brace_tokens_; std::unordered_set close_brace_tokens_; - std::unordered_set colon_tokens_; - std::unordered_set comma_tokens_; std::unordered_map current_bias_; void compute_bias(); void tokenize_grammar_elements(); void add_tokens_for_string(const std::string& str, std::unordered_set& token_set); - void add_tokens_for_prefix_string(const std::string& prefix, std::unordered_set& token_set); void advance_forced_tag(const std::vector& sequence, uint32_t token_id); - void tokenize_function_names(bool quote_names); void init_common_tokens(); bool is_needle() const { return model_type_ == Config::ModelType::NEEDLE; } - enum class NeedleJsonState { - FREE, - IN_NAME, - IN_ARG_KEY - }; - struct NeedleTrieNode { - std::unordered_map> children; - bool is_terminal = false; - }; - NeedleJsonState needle_json_state_ = NeedleJsonState::FREE; - std::string needle_buffer_; - std::string needle_constrained_buf_; - std::string needle_current_function_; - bool needle_in_arguments_ = false; - int needle_arguments_depth_ = 0; - int needle_nesting_depth_ = 0; - bool needle_in_string_value_ = false; - bool needle_prev_char_escape_ = false; - std::unique_ptr needle_name_trie_; - std::unordered_map> needle_param_tries_; - std::vector needle_token_strings_; - std::unordered_map> needle_token_index_; - - void init_needle_constraints(); - void reset_needle_constraints(); + Region region_ = Region::FREE; + std::string constraint_buffer_; + std::string constrained_buf_; + std::string current_function_; + bool in_arguments_ = false; + int arguments_depth_ = 0; + int nesting_depth_ = 0; + bool in_string_value_ = false; + bool prev_char_escape_ = false; + bool at_key_start_ = false; + bool await_enum_open_ = false; + const TrieNode* active_enum_trie_ = nullptr; + std::unordered_set emitted_keys_; + std::unique_ptr name_trie_; + std::unordered_map> param_tries_; + std::unordered_map>> enum_tries_; + std::unique_ptr remaining_key_trie_; + std::vector token_strings_; + std::unordered_map> token_index_; + uint32_t escape_tag_token_ = 0; + bool escape_tag_valid_ = false; + + void build_constraint_tables(); + void reset_constraint_state(); + void trie_insert(TrieNode* root, const std::string& word); + const TrieNode* trie_seek(const TrieNode* root, const std::string& prefix) const; + bool trie_token_valid(const std::string& token_text, const TrieNode* node, char terminator) const; + void mark_trie_bias(const TrieNode* node, char terminator, const std::vector& extra_allowed, + int32_t terminal_token = -1, bool exact_terminator = false); + const TrieNode* current_param_trie() const; + void rebuild_remaining_keys(); + bool required_satisfied() const; + void feed_needle_text(const std::string& text); void feed_needle_char(char ch); bool needle_at_arg_key_start() const; bool needle_is_value_string_start() const; - void needle_insert_word(NeedleTrieNode* root, const std::string& word); - const NeedleTrieNode* needle_get_trie_node(const NeedleTrieNode* root, const std::string& prefix) const; - bool needle_check_token_valid(const std::string& token_text, const NeedleTrieNode* trie_node) const; + + void feed_gemma_text(const std::string& text); + void feed_gemma_char(char ch); }; class Model { diff --git a/python/tests/test_tool_constraint_state.py b/python/tests/test_tool_constraint_state.py index 9fe53d807..7a1096855 100644 --- a/python/tests/test_tool_constraint_state.py +++ b/python/tests/test_tool_constraint_state.py @@ -13,20 +13,54 @@ WEIGHTS / "gemma-4-e2b-it", ) -GET_WEATHER = { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get the current weather for a location", - "parameters": { - "type": "object", - "properties": { - "location": {"type": "string", "description": "City name"}, - }, - "required": ["location"], +def _tool(name, description, properties, required): + return { + "type": "function", + "function": { + "name": name, + "description": description, + "parameters": {"type": "object", "properties": properties, "required": required}, }, - }, -} + } + + +GET_WEATHER = _tool("get_weather", "Get the current weather for a location", + {"location": {"type": "string", "description": "City name"}}, ["location"]) +SET_ALARM = _tool("set_alarm", "Set an alarm for a time", + {"time": {"type": "string", "description": "HH:MM"}}, ["time"]) +SEND = _tool("send", "Send a quick ping (no body) to a user id", + {"user": {"type": "string"}}, ["user"]) +SEND_MESSAGE = _tool("send_message", "Send a text message with a body to a contact", + {"recipient": {"type": "string"}, "body": {"type": "string"}}, ["recipient", "body"]) +SEARCH = _tool("search", "Search local notes", {"query": {"type": "string"}}, ["query"]) +SEARCH_WEB = _tool("search_web", "Search the public web", {"query": {"type": "string"}}, ["query"]) +GET = _tool("get", "Get a raw value by key", {"key": {"type": "string"}}, ["key"]) +GET_USER = _tool("get_user", "Fetch a user profile by id", {"id": {"type": "string"}}, ["id"]) +PING_SERVER = _tool("ping_server", "Check server health, takes no arguments", {}, []) +WEATHER_UNIT = _tool("get_weather", "Get weather for a city in a unit", + {"location": {"type": "string", "description": "city"}, + "unit": {"type": "string", "description": "temperature unit", + "enum": ["celsius", "fahrenheit"]}}, + ["location", "unit"]) +CREATE_EVENT = _tool("create_event", "Create a calendar event", + {"title": {"type": "string"}, "date": {"type": "string", "description": "YYYY-MM-DD"}, + "location": {"type": "string"}}, + ["title", "date"]) + + +def _forced_calls(model, tools, user): + result = cactus_complete( + model, + [ + {"role": "system", "content": "You are a helpful assistant that can use tools."}, + {"role": "user", "content": user}, + ], + {"force_tools": True, "max_tokens": 128, "temperature": 0}, + tools, + ) + assert result.get("success"), result + calls = result.get("function_calls") or [] + return [c if isinstance(c, dict) else json.loads(c) for c in calls] @pytest.fixture(scope="module") @@ -67,3 +101,124 @@ def test_turn_after_force_tools_is_not_degenerate(model): response = followup.get("response", "") assert response.strip() assert "<|tool_call>" not in response + + +def test_prefix_name_is_not_shadowed(model): + calls = _forced_calls(model, [SEND, SEND_MESSAGE], "Message my contact Alice: running late, start without me.") + assert calls + assert calls[0]["name"] == "send_message" + + +@pytest.mark.parametrize("tools, user, want", [ + ([SEARCH, SEARCH_WEB], "Search the public web for the best ramen in Osaka.", "search_web"), + ([GET, GET_USER], "Fetch the user profile for account 42.", "get_user"), +]) +def test_prefix_pairs_reach_longer_name(model, tools, user, want): + calls = _forced_calls(model, tools, user) + assert calls + assert calls[0]["name"] == want + + +def test_shorter_prefix_name_still_reachable(model): + calls = _forced_calls(model, [SEND, SEND_MESSAGE], "Ping user u_42.") + assert calls + assert calls[0]["name"] == "send" + + +def test_forced_name_is_always_a_declared_tool(model): + tools = [GET_WEATHER, SET_ALARM] + declared = {t["function"]["name"] for t in tools} + calls = _forced_calls(model, tools, "What's the weather in Tokyo right now?") + assert calls + assert calls[0]["name"] in declared + + +def test_no_argument_function_emits_empty_object(model): + calls = _forced_calls(model, [PING_SERVER], "Check whether the server is healthy.") + assert calls + assert calls[0]["name"] == "ping_server" + assert calls[0]["arguments"] == {} + + +def test_argument_keys_are_declared_parameters(model): + calls = _forced_calls(model, [GET_WEATHER], "What's the weather in Berlin?") + assert calls + assert calls[0]["name"] == "get_weather" + assert set(calls[0]["arguments"]).issubset({"location"}) + + +@pytest.mark.parametrize("user", [ + "What's the weather in Tokyo in kelvin?", + "Weather in Paris, in Kelvin please.", + "Weather in Berlin.", +]) +def test_enum_argument_restricted_to_allowed_values(model, user): + calls = _forced_calls(model, [WEATHER_UNIT], user) + assert calls + assert calls[0]["arguments"].get("unit") in ("celsius", "fahrenheit") + + +@pytest.mark.parametrize("user", [ + "Make a calendar event.", + "Add an event called Standup.", + "Schedule a dentist appointment.", +]) +def test_required_arguments_are_always_present(model, user): + calls = _forced_calls(model, [CREATE_EVENT], user) + assert calls + args = calls[0]["arguments"] + assert "title" in args and "date" in args + + +def test_enum_with_shared_prefixes_can_reach_longest_value(model): + tool = _tool("pick_plan", "pick a plan", + {"plan": {"type": "string", "enum": ["pro", "pro_plus", "pro_plus_max"]}}, ["plan"]) + calls = _forced_calls(model, [tool], "Pick the pro plus max plan.") + assert calls + assert calls[0]["arguments"].get("plan") in ("pro", "pro_plus", "pro_plus_max") + + +def test_enum_value_containing_spaces(model): + tool = _tool("move_card", "move a kanban card", + {"card": {"type": "string"}, + "column": {"type": "string", "enum": ["to do", "in progress", "done"]}}, + ["card", "column"]) + calls = _forced_calls(model, [tool], "Move the login card to in progress.") + assert calls + assert calls[0]["arguments"].get("column") in ("to do", "in progress", "done") + + +def test_single_value_enum_is_forced(model): + tool = _tool("finalize", "finalize", {"status": {"type": "string", "enum": ["confirmed"]}}, ["status"]) + calls = _forced_calls(model, [tool], "Finalize it however you want.") + assert calls + assert calls[0]["arguments"].get("status") == "confirmed" + + +def test_prefix_parameter_names_both_required_present(model): + tool = _tool("make_range", "range", {"to": {"type": "string"}, "total": {"type": "string"}}, ["to", "total"]) + calls = _forced_calls(model, [tool], "Make a range up to 10 with a total of 100.") + assert calls + assert {"to", "total"} <= set(calls[0]["arguments"]) + + +def test_many_required_arguments_all_present(model): + tool = _tool("register", "register a device", + {k: {"type": "string"} for k in ("name", "os", "owner", "ip", "room")}, + ["name", "os", "owner", "ip", "room"]) + calls = _forced_calls(model, [tool], "Register a device.") + assert calls + assert {"name", "os", "owner", "ip", "room"} <= set(calls[0]["arguments"]) + + +def test_prefix_tool_with_required_enum_on_longer_name(model): + send = _tool("send", "ping a user", {"user": {"type": "string"}}, ["user"]) + send_message = _tool("send_message", "send a message", + {"to": {"type": "string"}, "body": {"type": "string"}, + "urgency": {"type": "string", "enum": ["normal", "urgent"]}}, + ["to", "body", "urgency"]) + calls = _forced_calls(model, [send, send_message], "Send an urgent message to Alice saying the server is down.") + assert calls + assert calls[0]["name"] == "send_message" + assert calls[0]["arguments"].get("urgency") in ("normal", "urgent") + assert {"to", "body", "urgency"} <= set(calls[0]["arguments"])