Add Tool Calling and Reasoning Token Metadata to genai_config.json with Fallback Map

src/config.cpp

@@ -1181,6 +1181,14 @@ struct Model_Element : JSON::Element {
       v_.left_context_samples = SafeDoubleToInt(JSON::Get<double>(value), name);
     } else if (name == "right_context_samples") {
       v_.right_context_samples = SafeDoubleToInt(JSON::Get<double>(value), name);
+    } else if (name == "tool_call_start_token_id") {
+      v_.tool_call_start_token_id = SafeDoubleToInt(JSON::Get<double>(value), name);
+    } else if (name == "tool_call_end_token_id") {
+      v_.tool_call_end_token_id = SafeDoubleToInt(JSON::Get<double>(value), name);
+    } else if (name == "reasoning_start_token_id") {
+      v_.reasoning_start_token_id = SafeDoubleToInt(JSON::Get<double>(value), name);
+    } else if (name == "reasoning_end_token_id") {
+      v_.reasoning_end_token_id = SafeDoubleToInt(JSON::Get<double>(value), name);
     } else {
       throw JSON::unknown_value_error{};
     }

src/config.h

@@ -137,6 +137,12 @@ struct Config {
     int video_token_id{};
     int vision_start_token_id{};
 
+    // Tool-calling and reasoning token IDs (used for efficient token-level detection)
+    int tool_call_start_token_id{-1};
+    int tool_call_end_token_id{-1};
+    int reasoning_start_token_id{-1};
+    int reasoning_end_token_id{-1};
+
     int vocab_size{};
     int context_length{};
 

src/models/model.cpp

@@ -6,11 +6,13 @@
 #include <algorithm>
 #include <array>
 #include <climits>
+#include <functional>
 #include <random>
 #include <set>
 #include <string>
 #include <string_view>
 #include <thread>
+#include <unordered_map>
 
 #include "../generators.h"
 #include "../search.h"
@@ -306,6 +308,11 @@ Tokenizer::Tokenizer(Config& config) : bos_token_id_{config.model.bos_token_id},
   // Resolve tokenizer_dir (may be empty, relative, absolute, or "package:"-scheme).
   const fs::path tokenizer_dir = config.ResolvePath(config.model.tokenizer_dir);
   CheckResult(OrtxCreateTokenizerWithOptions(tokenizer_.Address(), tokenizer_dir.string().c_str(), keys, values, 2));
+
+  // TODO: Once ORT Extensions supports an "additional_special_tokens" option, pass the generation
+  // tags (tool_calling/reasoning tokens) here so that models which don't already mark them as
+  // special in their tokenizer_config.json will still get correct skip_special_tokens behavior.
+  // This is needed for the FL SDK's dual-stream special token detection in OnnxChatGenerator::Decode().
 }
 
 std::unique_ptr<TokenizerStream> Tokenizer::CreateStream() const {
@@ -820,6 +827,64 @@ bool Model::IsPruned() const {
   return logits_shape[1] == 1;
 }
 
+namespace {
+
+// Fallback map for models whose genai_config.json doesn't yet have token IDs in the model section.
+// Keyed by model.type string from genai_config.json.
+// Inner map: tag_name -> token string (used for vocab lookup to get the ID).
+const std::string* GetFallbackTag(const std::string& model_type, const std::string& tag_name) {
+  static const std::unordered_map<std::string, std::unordered_map<std::string, std::string>> fallback_map = {
+      {"qwen2", {{"tool_call_start", "<tool_call>"}, {"tool_call_end", "</tool_call>"}}},
+      {"qwen3", {{"tool_call_start", "<tool_call>"}, {"tool_call_end", "</tool_call>"}, {"reasoning_start", "<think>"}, {"reasoning_end", "</think>"}}},
+      {"phi3", {{"tool_call_start", "<tool_call>"}, {"tool_call_end", "</tool_call>"}}},
+      {"gptoss", {{"tool_call_start", "<|start|>"}, {"tool_call_end", "<|call|>"}}},
+  };
+  auto type_it = fallback_map.find(model_type);
+  if (type_it == fallback_map.end()) return nullptr;
+  auto tag_it = type_it->second.find(tag_name);
+  if (tag_it == type_it->second.end()) return nullptr;
+  return &tag_it->second;
+}
+
+}  // namespace
+
+void Model::InitTagIdCache() const {
+  auto tokenizer = CreateTokenizer();
+
+  static const char* tag_names[] = {"tool_call_start", "tool_call_end", "reasoning_start", "reasoning_end"};
+
+  auto get_config_id = [&](const std::string& name) -> int32_t {
+    if (name == "tool_call_start") return config_->model.tool_call_start_token_id;
+    if (name == "tool_call_end") return config_->model.tool_call_end_token_id;
+    if (name == "reasoning_start") return config_->model.reasoning_start_token_id;
+    if (name == "reasoning_end") return config_->model.reasoning_end_token_id;
+    return -1;
+  };
+
+  for (const auto* tag_name : tag_names) {
+    int32_t id = get_config_id(tag_name);
+    if (id >= 0) {
+      tag_id_cache_[tag_name] = id;
+      continue;
+    }
+
+    // Fallback: look up the token string in the vocabulary to get its ID.
+    const auto* fallback_str = GetFallbackTag(config_->model.type, tag_name);
+    if (fallback_str && !fallback_str->empty()) {
+      int32_t fallback_id = tokenizer->TokenToTokenId(fallback_str->c_str());
+      if (fallback_id >= 0) {
+        tag_id_cache_[tag_name] = fallback_id;
+      }
+    }
+  }
+}
+
+int32_t Model::GetTagId(const std::string& tag_name) const {
+  std::call_once(tag_id_cache_flag_, [this]() { InitTagIdCache(); });
+  auto it = tag_id_cache_.find(tag_name);
+  return (it != tag_id_cache_.end()) ? it->second : -1;
+}
+
 std::shared_ptr<Model> CreateModel(OrtEnv& ort_env, const char* config_path, const RuntimeSettings* settings /*= nullptr*/) {
   std::string config_overlay;
   if (settings) {

src/models/model.h

@@ -168,6 +168,11 @@ struct Model : std::enable_shared_from_this<Model>, LeakChecked<Model>, External
 
   bool IsPruned() const;
 
+  // Returns the token ID for the given tag, or -1 if the model doesn't define the tag.
+  // Checks genai_config.json model section first, then encodes the model-type fallback string.
+  // Known tag names: "tool_call_start", "tool_call_end", "reasoning_start", "reasoning_end".
+  int32_t GetTagId(const std::string& tag_name) const;
+
   std::unique_ptr<Config> config_;
   std::unique_ptr<OrtSessionOptions> session_options_;
 
@@ -189,8 +194,13 @@ struct Model : std::enable_shared_from_this<Model>, LeakChecked<Model>, External
 
  protected:
   void CreateSessionOptions();
+  void InitTagIdCache() const;
 
   std::map<std::string, std::unique_ptr<OrtSessionOptions>> pipeline_session_options_;
+
+  // Cached tag token IDs (lazily populated on first GetTagId call).
+  mutable std::unordered_map<std::string, int32_t> tag_id_cache_;
+  mutable std::once_flag tag_id_cache_flag_;
 };
 
 }  // namespace Generators

src/ort_genai.h

@@ -258,6 +258,12 @@ struct OgaModel : OgaAbstract {
     return p;
   }
 
+  int32_t GetTagId(const char* tag_name) const {
+    int32_t id;
+    OgaCheckResult(OgaModelGetTagId(this, tag_name, &id));
+    return id;
+  }
+
   static void operator delete(void* p) { OgaDestroyModel(reinterpret_cast<OgaModel*>(p)); }
 };
 

src/ort_genai_c.cpp

@@ -378,6 +378,13 @@ OgaResult* OGA_API_CALL OgaModelGetDeviceType(const OgaModel* model, const char*
   OGA_CATCH
 }
 
+OgaResult* OGA_API_CALL OgaModelGetTagId(const OgaModel* model, const char* tag_name, int32_t* out) {
+  OGA_TRY
+  *out = model->GetTagId(tag_name);
+  return nullptr;
+  OGA_CATCH
+}
+
 OgaResult* OGA_API_CALL OgaCreateGeneratorParams(const OgaModel* model, OgaGeneratorParams** out) {
   OGA_TRY
   auto params = std::make_shared<Generators::GeneratorParams>(*model);

src/ort_genai_c.h

@@ -415,6 +415,21 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaModelGetType(const OgaModel* model, const
  */
 OGA_EXPORT OgaResult* OGA_API_CALL OgaModelGetDeviceType(const OgaModel* model, const char** out);
 
+/**
+ * \brief Returns a tag token ID for this model by name.
+ *
+ * Checks the genai_config.json model section first, then falls back to encoding the
+ * model-type-specific fallback token string via the tokenizer vocabulary.
+ * Known tag names: "tool_call_start", "tool_call_end", "reasoning_start", "reasoning_end".
+ * Returns -1 if the model doesn't define the requested tag.
+ *
+ * \param[in] model The model to query.
+ * \param[in] tag_name The name of the tag to retrieve.
+ * \param[out] out The tag token ID.
+ * \return OgaResult containing the error message if the call failed.
+ */
+OGA_EXPORT OgaResult* OGA_API_CALL OgaModelGetTagId(const OgaModel* model, const char* tag_name, int32_t* out);
+
 /**
  * \brief Destroys the given config
  * \param[in] config The config to be destroyed.

test/c_api_tests.cpp

@@ -1817,6 +1817,72 @@ TEST(CAPITests, ParakeetTdtTranscribeLong) {
   EXPECT_FALSE(transcription.empty());
 }
 
+// Test that GetTagId returns -1 for unknown model types (not in config, not in fallback map)
+TEST(CAPITests, TagId_Unknown) {
+  // tiny-random-gpt2 model has type "gpt2" which is NOT in the fallback map → -1
+  auto model = OgaModel::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
+
+  EXPECT_EQ(model->GetTagId("tool_call_start"), -1);
+  EXPECT_EQ(model->GetTagId("tool_call_end"), -1);
+  EXPECT_EQ(model->GetTagId("reasoning_start"), -1);
+  EXPECT_EQ(model->GetTagId("reasoning_end"), -1);
+}
+
+TEST(CAPITests, TagId_FromConfig) {
+  // Create a temporary model directory with tool_call/reasoning token IDs in model section
+  auto temp_dir = std::filesystem::temp_directory_path() / "oga_test_tool_tags";
+  std::filesystem::remove_all(temp_dir);  // Clean up any leftover from a previous failed run
+  std::filesystem::create_directories(temp_dir);
+
+  // Copy minimal model files from tiny-random-gpt2
+  std::string src_dir = MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32";
+  for (const auto& entry : std::filesystem::directory_iterator(src_dir)) {
+    if (entry.path().filename() != "genai_config.json") {
+      std::filesystem::copy_file(entry.path(), temp_dir / entry.path().filename(),
+                                 std::filesystem::copy_options::overwrite_existing);
+    }
+  }
+
+  // Write genai_config.json with token IDs in model section
+  {
+    std::ofstream f((temp_dir / "genai_config.json").string());
+    f << R"({
+  "model": {
+    "type": "gpt2",
+    "pad_token_id": 98,
+    "bos_token_id": 98,
+    "eos_token_id": 98,
+    "vocab_size": 1000,
+    "context_length": 512,
+    "tool_call_start_token_id": 151657,
+    "tool_call_end_token_id": 151658,
+    "reasoning_start_token_id": 151659,
+    "reasoning_end_token_id": 151660,
+    "decoder": {
+      "session_options": { "provider_options": [] },
+      "filename": "past.onnx",
+      "num_key_value_heads": 4,
+      "head_size": 8,
+      "num_hidden_layers": 5,
+      "inputs": { "past_names": "past_%d" },
+      "outputs": { "present_names": "present_%d" }
+    }
+  }
+})";
+  }
+
+  auto model = OgaModel::Create(temp_dir.string().c_str());
+
+  // GetTagId returns configured IDs from model section
+  EXPECT_EQ(model->GetTagId("tool_call_start"), 151657);
+  EXPECT_EQ(model->GetTagId("tool_call_end"), 151658);
+  EXPECT_EQ(model->GetTagId("reasoning_start"), 151659);
+  EXPECT_EQ(model->GetTagId("reasoning_end"), 151660);
+
+  // Cleanup
+  std::filesystem::remove_all(temp_dir);
+}
+
 // Regression test for MSRC: malformed audio buffers smaller than the minimum valid
 // audio header size must be rejected with an error, not cause a crash.
 TEST(CAPITests, LoadAudiosFromBuffersRejectsEmptyBuffer) {