[feature] Support tool function calling with Structural Tag with xgrammar

- ensure the tool function will be called in expected format using xgrammar
- modify RequestResponseFormat: add structural tag according to the tools when building response format
- the tool function calling is now constrained by format: <function=function_name>parameters</function>
- tools call list will be parsed according to the calling format when processing the response
- also expose the Structural Tag api of xgrammar to RequestResponseFormat

Author: Irfnfnkemed

cpp/serve/config.cc

@@ -10,11 +10,14 @@
 
 #include <limits>
 #include <random>
+#include <string>
+#include <vector>
 
 #include "../json_ffi/openai_api_protocol.h"
 #include "../support/json_parser.h"
 #include "../support/utils.h"
 #include "data.h"
+#include "tvm/runtime/container/array.h"
 
 namespace mlc {
 namespace llm {
@@ -42,13 +45,43 @@ Result<ResponseFormat> ResponseFormat::FromJSON(const picojson::object& config)
   ResponseFormat res;
   res.type = json::LookupOrDefault<std::string>(config, "type", "text");
 
+  if (res.type != "text" && res.type != "function" && res.type != "json_object" &&
+      res.type != "structural_tag") {
+    return TResult::Error("Uknonwn response_format type " + res.type);
+  }
+
   std::optional<std::string> schema = json::LookupOptional<std::string>(config, "schema");
   if (schema.has_value()) {
     res.schema = schema.value();
   }
 
-  if (res.type != "text" && res.type != "function" && res.type != "json_object") {
-    return TResult::Error("Uknonwn response_format type " + res.type);
+  if (auto tags_obj = json::LookupOptional<picojson::array>(config, "tags")) {
+    auto tags = Array<Array<String>>();
+    for (auto tag_obj : tags_obj.value()) {
+      Array<String> tag = Array<String>();
+      std::optional<std::string> begin =
+          json::LookupOptional<std::string>(tag_obj.get<picojson::object>(), "begin");
+      std::optional<std::string> schema =
+          json::LookupOptional<std::string>(tag_obj.get<picojson::object>(), "schema");
+      std::optional<std::string> end =
+          json::LookupOptional<std::string>(tag_obj.get<picojson::object>(), "end");
+      if (!(begin.has_value() && schema.has_value() && end.has_value())) {
+        return TResult::Error("Miss tag attribute.");
+      }
+      tag.push_back(begin.value());
+      tag.push_back(schema.value());
+      tag.push_back(end.value());
+      tags.push_back(std::move(tag));
+    }
+    res.tags = tags;
+  }
+
+  if (auto triggers_obj = json::LookupOptional<picojson::array>(config, "triggers")) {
+    auto triggers = Array<String>();
+    for (auto trigger : triggers_obj.value()) {
+      triggers.push_back(trigger.get<std::string>());
+    }
+    res.triggers = triggers;
   }
 
   return TResult::Ok(res);
@@ -60,6 +93,24 @@ picojson::object ResponseFormat::AsJSON() const {
   if (schema.defined()) {
     config["schema"] = picojson::value(schema.value().operator std::string());
   }
+  if (tags.defined()) {
+    picojson::array tags_obj = picojson::array();
+    for (auto tag : tags.value()) {
+      picojson::array tag_obj = picojson::array();
+      tag_obj.emplace_back(tag[0]);
+      tag_obj.emplace_back(tag[1]);
+      tag_obj.emplace_back(tag[2]);
+      tags_obj.emplace_back(tag_obj);
+    }
+    config["tags"] = picojson::value(tags_obj);
+  }
+  if (triggers.defined()) {
+    picojson::array trigger_obj = picojson::array();
+    for (std::string trigger : triggers.value()) {
+      trigger_obj.emplace_back(trigger);
+    }
+    config["triggers"] = picojson::value(trigger_obj);
+  }
   return config;
 }
 

cpp/serve/config.h

@@ -14,6 +14,7 @@
 
 #include "../metadata/model.h"
 #include "../support/result.h"
+#include "tvm/runtime/container/optional.h"
 
 namespace mlc {
 namespace llm {
@@ -28,6 +29,8 @@ using namespace tvm::runtime;
 struct ResponseFormat {
   String type = "text";
   Optional<String> schema = NullOpt;
+  Optional<Array<Array<String>>> tags = NullOpt;
+  Optional<Array<String>> triggers = NullOpt;
   /*!
    * \brief Create debug config from JSON.
    * \param config_json The json string for generation config

cpp/serve/engine.cc

@@ -19,8 +19,10 @@
 #include <functional>
 #include <numeric>
 #include <optional>
+#include <string>
 #include <tuple>
 #include <unordered_set>
+#include <utility>
 
 #include "../support/json_parser.h"
 #include "../support/result.h"
@@ -35,6 +37,7 @@
 #include "request.h"
 #include "request_state.h"
 #include "sampler/sampler.h"
+#include "xgrammar/grammar.h"
 
 namespace mlc {
 namespace llm {
@@ -978,12 +981,24 @@ class EngineImpl : public Engine {
   std::optional<xgrammar::CompiledGrammar> GetGrammarFromResponseFormat(
       const ResponseFormat& response_format) {
     // TODO: add other grammar type
-    if (response_format.type != "json_object") {
+    if (response_format.type == "text") {
       return std::nullopt;
-    } else if (!response_format.schema) {
-      return grammar_compiler_.CompileBuiltinJSONGrammar();
+    } else if (response_format.type == "json_object") {
+      if (!response_format.schema) {
+        return grammar_compiler_.CompileBuiltinJSONGrammar();
+      } else {
+        return grammar_compiler_.CompileJSONSchema(response_format.schema.value());
+      }
     } else {
-      return grammar_compiler_.CompileJSONSchema(response_format.schema.value());
+      std::vector<xgrammar::StructuralTagItem> tags;
+      std::vector<std::string> triggers;
+      for (auto tag : response_format.tags.value()) {
+        tags.emplace_back(xgrammar::StructuralTagItem{tag[0], tag[1], tag[2]});
+      }
+      for (auto trigger : response_format.triggers.value()) {
+        triggers.emplace_back(trigger);
+      }
+      return grammar_compiler_.CompileStructuralTag(std::move(tags), std::move(triggers));
     }
   }
 

cpp/serve/logit_processor.cc

@@ -12,6 +12,8 @@
 #include <tvm/runtime/registry.h>
 #include <tvm/runtime/threading_backend.h>
 
+#include <cstdio>
+
 namespace mlc {
 namespace llm {
 namespace serve {

python/mlc_llm/protocol/openai_api_protocol.py

@@ -86,12 +86,33 @@ class ModelResponse(BaseModel):
 
 
 class RequestResponseFormat(BaseModel):
-    type: Literal["text", "json_object"] = "text"
-    json_schema: Optional[str] = Field(default=None, alias="schema")
+    type: Literal["text", "json_object", "structural_tag"] = "text"
     """This field is named json_schema instead of schema because BaseModel defines a method called
     schema. During construction of RequestResponseFormat, key "schema" still should be used:
     `RequestResponseFormat(type="json_object", schema="{}")`
     """
+    json_schema: Optional[str] = Field(default=None, alias="schema")
+
+    """These field are only used for type="structural_tag"."""
+    tags: Optional[List[Dict[str, str]]] = Field(default=None, alias="tags")
+    triggers: Optional[List[str]] = Field(default=None, alias="triggers")
+
+    @model_validator(mode="after")
+    def check_request_response_format(self) -> "RequestResponseFormat":
+        """Check if the RequestResponseFormat is valid."""
+        if self.type == "structural_tag":
+            if self.tags is None or self.triggers is None:
+                raise ValueError("structural_tag type must contain keys 'tags' and 'triggers'.")
+            for tag in self.tags:
+                if set(tag.keys()) != {"begin", "schema", "end"}:
+                    raise ValueError(
+                        f"Each tag must contain exactly 'begin', 'schema' and 'end' keys. Got keys: {list(tag.keys())}."
+                    )
+        elif self.tags is not None or self.triggers is not None:
+            raise Warning(
+                "'tags' and 'triggers' attributes should be used when type='structural_tag'"
+            )
+        return self
 
 
 class CompletionRequest(BaseModel):

python/mlc_llm/serve/engine.py

@@ -976,6 +976,12 @@ async def _chat_completion(  # pylint: disable=too-many-arguments,too-many-local
         if request_id is None:
             request_id = f"chatcmpl-{engine_utils.random_uuid()}"
 
+        tools = (
+            [openai_api_protocol.ChatTool.model_validate(tool) for tool in tools]
+            if tools is not None
+            else None
+        )
+
         chatcmpl_generator = self._handle_chat_completion(
             openai_api_protocol.ChatCompletionRequest(
                 messages=[
@@ -1207,6 +1213,10 @@ async def _handle_chat_completion(
         e : BadRequestError
             BadRequestError is raised when the request is invalid.
         """
+        request.response_format = engine_base.set_structural_tag_from_tools(
+            request.tools, request.response_format
+        )
+
         (
             prompts,
             generation_cfg,
@@ -1764,6 +1774,10 @@ def _handle_chat_completion(
         e : BadRequestError
             BadRequestError is raised when the request is invalid.
         """
+        request.response_format = engine_base.set_structural_tag_from_tools(
+            request.tools, request.response_format
+        )
+
         (
             prompts,
             generation_cfg,

python/mlc_llm/serve/engine_base.py

@@ -7,6 +7,7 @@
 import json
 import numbers
 import queue
+import re
 import sys
 import threading
 from dataclasses import dataclass
@@ -1146,29 +1147,52 @@ def create_completion_suffix_response(
     return response
 
 
-def convert_function_str_to_json(stringified_calls: str) -> List[Union[Dict, None]]:
+def convert_function_str_to_json(stringified_calls: str):
     """Convert a (possibly list) of function call string to a list of json objects.
     Return None for invalid function call string."""
+    function_calls_json = []
+    for call in re.finditer(r"<function=(.*?)>(.*?)</function>", stringified_calls, re.DOTALL):
+        function_name = call.group(1)
+        params_str = call.group(2).strip()
+        params = ast.literal_eval(params_str)
+        function_calls_json.append({"name": function_name, "arguments": params})
 
-    def parse_function_call(call_str: str):
-        node = ast.parse(call_str, mode="eval")
-        call_node = node.body
-        if isinstance(call_node, ast.Call) and isinstance(call_node.func, ast.Name):
-            name = call_node.func.id
-            arguments = {}
-            for keyword in call_node.keywords:
-                arguments[keyword.arg] = ast.literal_eval(keyword.value)
-            return {"name": name, "arguments": arguments}
-        return None
+    return function_calls_json
 
-    if (
-        stringified_calls[0] == "[" and stringified_calls[-1] == "]"
-    ):  # hacky way to check if string list
-        calls = ast.literal_eval(stringified_calls)
+
+def set_structural_tag_from_tools(
+    tools: Optional[List[openai_api_protocol.ChatTool]],
+    response_format: Optional[openai_api_protocol.RequestResponseFormat],
+):
+    """Add the corresponding structural tag to the response format according to the tools to ensure valid function calling.
+    Return the updated response format.
+    """
+    if tools is None:
+        return response_format
     else:
-        calls = [stringified_calls]
-    function_calls_json = [parse_function_call(call_str) for call_str in calls]
-    return function_calls_json
+        if response_format is None or response_format.type == "text":
+            response_format = openai_api_protocol.RequestResponseFormat.model_validate(
+                {"type": "structural_tag", "tags": [], "triggers": []}
+            )
+        elif response_format.type == "json_object":
+            response_format.tags = []
+            response_format.triggers = []
+
+        response_format.triggers.append("<function=")
+        for tool in tools:
+            schema = {
+                "properties": tool.function.parameters["properties"],
+                "required": tool.function.parameters["required"],
+                "type": tool.function.parameters["type"],
+            }
+            response_format.tags.append(
+                {
+                    "begin": f"<function={tool.function.name}>",
+                    "schema": json.dumps(schema),
+                    "end": "</function>",
+                }
+            )
+        return response_format
 
 
 def process_function_call_output(