grammar_registry.cpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#include <entropic/inference/grammar_registry.h>
#include <entropic/types/logging.h>
 
#include <llama.h>
 
#include <fstream>
#include <sstream>
 
namespace entropic {
 
namespace {
auto logger = entropic::log::get("inference.grammar_registry");
 
std::string read_file_contents(const std::filesystem::path& path) {
    std::ifstream file(path);
    if (!file.is_open()) {
        return "";
    }
    std::ostringstream ss;
    ss << file.rdbuf();
    return ss.str();
}
 
} // anonymous namespace
 
static std::optional<GrammarEntry> build_grammar_entry(
    const std::filesystem::path& path) {
    std::string content = read_file_contents(path);
    if (content.empty()) {
        logger->warn("Empty or unreadable grammar file: {}", path.string());
        return std::nullopt;
    }
    std::string error = GrammarRegistry::validate(content);
    GrammarEntry ge;
    ge.key = path.stem().string();
    ge.gbnf_content = std::move(content);
    ge.source = "bundled";
    ge.validated = error.empty();
    ge.error = std::move(error);
    if (!ge.validated) {
        logger->warn("Bundled grammar '{}' has validation errors: {}",
                     ge.key, ge.error);
    }
    return ge;
}
 
size_t GrammarRegistry::load_bundled(
    const std::filesystem::path& grammar_dir)
{
    if (!std::filesystem::is_directory(grammar_dir)) {
        logger->warn("Grammar directory not found: {}",
                     grammar_dir.string());
        return 0;
    }
 
    size_t count = 0;
    for (const auto& entry : std::filesystem::directory_iterator(grammar_dir)) {
        if (entry.path().extension() != ".gbnf") {
            continue;
        }
        auto ge = build_grammar_entry(entry.path());
        if (!ge) {
            continue;
        }
        std::string key = ge->key;
        std::lock_guard<std::mutex> lock(registry_mutex_);
        grammars_[key] = std::move(*ge);
        ++count;
    }
 
    logger->info("Loaded {} bundled grammar(s)", count);
    return count;
}
 
bool GrammarRegistry::register_grammar(
    const std::string& key,
    const std::string& gbnf_content,
    const std::string& source)
{
    std::lock_guard<std::mutex> lock(registry_mutex_);
    if (grammars_.count(key) > 0) {
        logger->warn("Grammar key '{}' already registered", key);
        return false;
    }
 
    std::string error = validate(gbnf_content);
    GrammarEntry ge;
    ge.key = key;
    ge.gbnf_content = gbnf_content;
    ge.source = source;
    ge.validated = error.empty();
    ge.error = std::move(error);
 
    grammars_[key] = std::move(ge);
    logger->info("Registered grammar '{}' (source={}, valid={})",
                 key, source, grammars_[key].validated);
    return true;
}
 
bool GrammarRegistry::register_from_file(
    const std::string& key,
    const std::filesystem::path& path)
{
    std::string content = read_file_contents(path);
    if (content.empty()) {
        logger->error("Cannot read grammar file: {}", path.string());
        return false;
    }
 
    std::string resolved_key = key.empty()
        ? path.stem().string() : key;
    return register_grammar(resolved_key, content, "file");
}
 
bool GrammarRegistry::deregister(const std::string& key) {
    std::lock_guard<std::mutex> lock(registry_mutex_);
    auto it = grammars_.find(key);
    if (it == grammars_.end()) {
        return false;
    }
    grammars_.erase(it);
    logger->info("Deregistered grammar '{}'", key);
    return true;
}
 
std::string GrammarRegistry::get(const std::string& key) const {
    std::lock_guard<std::mutex> lock(registry_mutex_);
    auto it = grammars_.find(key);
    if (it == grammars_.end()) {
        return "";
    }
    return it->second.gbnf_content;
}
 
bool GrammarRegistry::has(const std::string& key) const {
    std::lock_guard<std::mutex> lock(registry_mutex_);
    return grammars_.count(key) > 0;
}
 
GrammarEntry GrammarRegistry::entry(const std::string& key) const {
    std::lock_guard<std::mutex> lock(registry_mutex_);
    auto it = grammars_.find(key);
    if (it == grammars_.end()) {
        return {};
    }
    return it->second;
}
 
std::vector<GrammarEntry> GrammarRegistry::list() const {
    std::lock_guard<std::mutex> lock(registry_mutex_);
    std::vector<GrammarEntry> result;
    result.reserve(grammars_.size());
    for (const auto& [k, e] : grammars_) {
        GrammarEntry meta;
        meta.key = e.key;
        meta.source = e.source;
        meta.validated = e.validated;
        meta.error = e.error;
        result.push_back(std::move(meta));
    }
    return result;
}
 
std::string GrammarRegistry::validate(const std::string& gbnf_content) {
    if (gbnf_content.empty()) {
        return "empty grammar string";
    }
 
    // Use llama_sampler_init_grammar with nullptr vocab for validation.
    // Grammar parsing is independent of vocabulary — the GBNF parser
    // only needs the grammar string. Returns nullptr on parse failure.
    llama_sampler* sampler = llama_sampler_init_grammar(
        nullptr, gbnf_content.c_str(), "root");
 
    if (sampler == nullptr) {
        return "GBNF parse failed";
    }
 
    llama_sampler_free(sampler);
    return "";
}
 
size_t GrammarRegistry::size() const {
    std::lock_guard<std::mutex> lock(registry_mutex_);
    return grammars_.size();
}
 
void GrammarRegistry::clear() {
    std::lock_guard<std::mutex> lock(registry_mutex_);
    grammars_.clear();
    logger->info("Cleared all grammars");
}
 
} // namespace entropic