chat_in.hpp

Go to the documentation of this file.

#pragma once
 
// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 Lloyal Labs
 
#include "common.hpp"
#include "tokenizer.hpp"
#include <llama/llama.h>
#include <chat.h>         // llama.cpp common library: common_chat_templates_*
#include <nlohmann/json.hpp>
#include <algorithm>
#include <exception>
#include <string>
#include <vector>
 
namespace lloyal::chat_in {
 
struct FormatInputs {
  std::string messages_json;                       
  std::string template_override = "";              
  bool add_generation_prompt = true;               
  std::string tools_json = "";                     
  std::string tool_choice = "auto";                
  bool parallel_tool_calls = false;                
  std::string reasoning_format = "none";           
  bool enable_thinking = true;                     
  std::string json_schema = "";                    
  std::string grammar = "";                        
};
 
struct FormatResult {
  // Core output
  std::string prompt;                              
  std::vector<std::string> additional_stops;       
 
  // Format awareness (all fields from common_chat_params)
  common_chat_format format = COMMON_CHAT_FORMAT_CONTENT_ONLY;  
  std::string grammar;                             
  bool grammar_lazy = false;                       
  bool thinking_forced_open = false;               
  std::vector<common_grammar_trigger> grammar_triggers; 
  std::vector<std::string> preserved_tokens;       
  std::string parser;                              
 
  // Carried through for chat_out pairing
  common_reasoning_format reasoning_format = COMMON_REASONING_FORMAT_NONE; 
};
 
inline FormatResult format(const llama_model *model, const FormatInputs& inputs) {
  FormatResult result;
 
  try {
    using json = nlohmann::ordered_json;
    json messages_array = json::parse(inputs.messages_json);
 
    // Initialize templates from model (or override)
    common_chat_templates_ptr tmpls = common_chat_templates_init(model, inputs.template_override);
    if (!tmpls) {
      LLOYAL_LOG_DEBUG("[chat_in::format] Template init failed, using fallback");
      goto fallback;
    }
 
    {
      // Parse messages
      std::vector<common_chat_msg> messages = common_chat_msgs_parse_oaicompat(messages_array);
 
      // Build full template inputs
      common_chat_templates_inputs tmpl_inputs;
      tmpl_inputs.messages = messages;
      tmpl_inputs.add_generation_prompt = inputs.add_generation_prompt;
      tmpl_inputs.use_jinja = true;
 
      // Tools
      if (!inputs.tools_json.empty()) {
        json tools_array = json::parse(inputs.tools_json);
        tmpl_inputs.tools = common_chat_tools_parse_oaicompat(tools_array);
        tmpl_inputs.tool_choice = common_chat_tool_choice_parse_oaicompat(inputs.tool_choice);
        tmpl_inputs.parallel_tool_calls = inputs.parallel_tool_calls;
      }
 
      // Reasoning
      tmpl_inputs.reasoning_format = common_reasoning_format_from_name(inputs.reasoning_format);
      tmpl_inputs.enable_thinking = inputs.enable_thinking;
 
      // Structured output
      tmpl_inputs.json_schema = inputs.json_schema;
      tmpl_inputs.grammar = inputs.grammar;
 
      // Apply template
      common_chat_params params = common_chat_templates_apply(tmpls.get(), tmpl_inputs);
 
      // Implicit empty system prompt stripping: if messages[0] is {system, ""},
      // strip the resulting empty system block from the output. This lets callers
      // suppress template auto-injection (e.g. SmolLM2/ChatML) by prepending an
      // empty system message — the library completes the intent by removing the
      // rendered empty block, leaving only the user+assistant portion.
      if (!messages.empty() && messages[0].role == "system" && messages[0].content.empty()) {
        common_chat_msg sys_msg;
        sys_msg.role = "system";
        sys_msg.content = "";
 
        common_chat_templates_inputs sys_inputs;
        sys_inputs.messages = {sys_msg};
        sys_inputs.add_generation_prompt = false;
        sys_inputs.use_jinja = true;
        auto sys_params = common_chat_templates_apply(tmpls.get(), sys_inputs);
 
        const auto& sys_prefix = sys_params.prompt;
        if (!sys_prefix.empty() &&
            params.prompt.size() >= sys_prefix.size() &&
            params.prompt.substr(0, sys_prefix.size()) == sys_prefix) {
          params.prompt = params.prompt.substr(sys_prefix.size());
          LLOYAL_LOG_DEBUG("[chat_in::format] Stripped empty system prefix (%zu bytes)", sys_prefix.size());
        }
      }
 
      // Populate ALL result fields from common_chat_params
      result.prompt = params.prompt;
      result.additional_stops = params.additional_stops;
      result.format = params.format;
      result.grammar = params.grammar;
      result.grammar_lazy = params.grammar_lazy;
      result.thinking_forced_open = params.thinking_forced_open;
      result.grammar_triggers = params.grammar_triggers;
      result.preserved_tokens = params.preserved_tokens;
      result.parser = params.parser;
 
      // Carry reasoning_format through for chat_out pairing
      result.reasoning_format = tmpl_inputs.reasoning_format;
 
      LLOYAL_LOG_DEBUG(
          "[chat_in::format] Successfully formatted with format=%d, %zu stop tokens, grammar=%zu bytes",
          static_cast<int>(result.format),
          result.additional_stops.size(),
          result.grammar.size());
      return result;
    }
 
  } catch (const std::exception &e) {
    LLOYAL_LOG_DEBUG("[chat_in::format] Template processing failed: %s", e.what());
  }
 
fallback:
  // Fallback to simple "role: content" format
  try {
    using json = nlohmann::ordered_json;
    json messages = json::parse(inputs.messages_json);
    std::string fallback_prompt;
    for (const auto &msg : messages) {
      if (msg.contains("role") && msg.contains("content")) {
        std::string role = msg["role"].get<std::string>();
        std::string content;
        const auto& c = msg["content"];
        if (c.is_null()) {
          content = "";
        } else if (c.is_string()) {
          content = c.get<std::string>();
        } else {
          content = c.dump();
        }
        fallback_prompt += role + ": " + content + "\n";
      }
    }
 
    result.prompt = fallback_prompt;
    result.additional_stops = {};
 
    LLOYAL_LOG_DEBUG(
        "[chat_in::format] Using fallback format (%zu bytes)",
        fallback_prompt.size());
    return result;
 
  } catch (const std::exception &e) {
    LLOYAL_LOG_DEBUG(
        "[chat_in::format] ERROR: Failed to parse messages JSON: %s",
        e.what());
    result.prompt = "";
    result.additional_stops = {};
    return result;
  }
}
 
inline bool validate(const std::string &template_str) {
  try {
    bool isValid = common_chat_verify_template(template_str, /* use_jinja */ true);
    LLOYAL_LOG_DEBUG("[chat_in::validate] Template validation: %s",
                     isValid ? "valid" : "invalid");
    return isValid;
  } catch (const std::exception &e) {
    LLOYAL_LOG_DEBUG("[chat_in::validate] ERROR: %s", e.what());
    return false;
  }
}
 
inline std::vector<llama_token> fallback_to_eog(const llama_model* model) {
  if (model == nullptr) {
    return {};
  }
  const llama_vocab* vocab = llama_model_get_vocab(model);
  llama_token eot = llama_vocab_eot(vocab);
  if (eot == LLAMA_TOKEN_NULL) {
    eot = llama_vocab_eos(vocab);
  }
  if (eot != LLAMA_TOKEN_NULL) {
    return {eot};
  }
  return {};
}
 
inline std::string get_token_safe(const llama_model *model, llama_token token) {
  if (!model || token == LLAMA_TOKEN_NULL) {
    return "";
  }
  return lloyal::tokenizer::detokenize(model, token);
}
 
inline std::vector<llama_token> get_turn_separator(const llama_model* model) {
  using json = nlohmann::ordered_json;
 
  if (!model) return {};
 
  // Collision-resistant sentinels
  const std::string SENTINEL = "\x1F__LLOYAL_SEP__\x1F";
  const std::string SENTINEL2 = "\x1F__LLOYAL_SEP2__\x1F";
 
  try {
    // Initialize templates from model
    common_chat_templates_ptr tmpls = common_chat_templates_init(model, "");
    if (!tmpls) {
      return fallback_to_eog(model);
    }
 
    // 3-message probe: captures REAL assistant→user boundary
    std::vector<common_chat_msg> messages = {
      {.role = "user", .content = "X"},
      {.role = "assistant", .content = SENTINEL},
      {.role = "user", .content = SENTINEL2}
    };
 
    common_chat_templates_inputs inputs;
    inputs.messages = messages;
    inputs.add_generation_prompt = false;  // Don't add assistant prompt at end
    inputs.use_jinja = true;
 
    auto params = common_chat_templates_apply(tmpls.get(), inputs);
    const std::string& formatted = params.prompt;
 
    // Extract substring between sentinels
    size_t sep_start = formatted.rfind(SENTINEL);
    if (sep_start == std::string::npos) {
      return fallback_to_eog(model);
    }
    sep_start += SENTINEL.length();
 
    size_t sep_end = formatted.find(SENTINEL2, sep_start);
    if (sep_end == std::string::npos) {
      return fallback_to_eog(model);
    }
 
    std::string between = formatted.substr(sep_start, sep_end - sep_start);
    if (between.empty()) {
      return fallback_to_eog(model);
    }
 
    // Tokenize with parse_special=true
    const auto* vocab = llama_model_get_vocab(model);
    std::vector<llama_token> tokens = lloyal::tokenizer::tokenize(vocab, between, false, true);
    if (tokens.empty()) {
      return fallback_to_eog(model);
    }
 
    // Extract: everything up to and including EOG + trailing whitespace
    std::vector<llama_token> separator;
    bool found_eog = false;
 
    for (auto tok : tokens) {
      if (!found_eog) {
        separator.push_back(tok);
        if (lloyal::tokenizer::is_eog(model, tok)) {
          found_eog = true;
        }
      } else {
        // After EOG, only keep whitespace tokens
        std::string text = lloyal::tokenizer::detokenize(model, tok);
        bool is_whitespace = !text.empty() && std::all_of(text.begin(), text.end(),
            [](unsigned char c) { return c == ' ' || c == '\n' || c == '\r' || c == '\t'; });
        if (is_whitespace) {
          separator.push_back(tok);
        } else {
          break;  // Non-whitespace = next message opener, stop
        }
      }
    }
 
    if (separator.empty() || !found_eog) {
      return fallback_to_eog(model);
    }
 
    return separator;
 
  } catch (const std::exception& e) {
    LLOYAL_LOG_DEBUG("[chat_in::get_turn_separator] Error: %s", e.what());
    return fallback_to_eog(model);
  }
}
 
} // namespace lloyal::chat_in