Tool Use中的JSON模式强制：利用Context-Free Grammar（CFG）约束采样保证API调用正确性

大家好，今天我们来探讨一个非常关键且实用的主题：在Tool Use中，如何利用JSON模式强制和上下文无关文法（CFG）约束采样来保证API调用的正确性。在大型语言模型（LLM）驱动的智能体（Agent）应用中，让智能体学会使用工具（Tool Use）是增强其能力的关键。而工具通常以API的形式暴露，因此，如何确保智能体生成的API调用是正确的、符合规范的，就变得至关重要。

1. Tool Use的挑战与JSON模式

Tool Use涉及的核心问题是：如何让LLM理解工具的功能，并根据给定的上下文生成符合API规范的请求。这其中面临着诸多挑战：

API规范复杂性： 现实世界中的API往往非常复杂，包含多种参数、不同的数据类型、以及复杂的依赖关系。
LLM理解偏差： LLM虽然强大，但对API规范的理解可能存在偏差，导致生成的请求不符合规范。
推理能力限制： LLM在复杂推理场景下，可能难以准确选择合适的工具和参数。

为了应对这些挑战，JSON模式提供了一种有效的方法来描述API的请求和响应格式。JSON模式是一个标准化的格式，用于描述JSON数据的结构和数据类型。通过JSON模式，我们可以明确地定义API请求中每个参数的名称、类型、取值范围等信息。

例如，假设我们有一个天气查询API，其请求需要包含城市名称（city）和单位（unit，可以是摄氏度"celsius"或华氏度"fahrenheit"）。我们可以用以下JSON模式来描述这个API的请求：

{
  "type": "object",
  "properties": {
    "city": {
      "type": "string",
      "description": "要查询天气的城市名称"
    },
    "unit": {
      "type": "string",
      "enum": ["celsius", "fahrenheit"],
      "description": "温度单位，摄氏度或华氏度"
    }
  },
  "required": ["city", "unit"]
}

这个JSON模式明确地规定了：

请求必须是一个JSON对象（"type": "object"）。
请求包含两个属性：city和unit（"properties": { ... }）。
city属性必须是一个字符串（"type": "string"），并且描述了它的含义（"description": "要查询天气的城市名称"）。
unit属性必须是一个字符串，并且只能取"celsius"或"fahrenheit"这两个值（"enum": ["celsius", "fahrenheit"]），并且描述了它的含义（"description": "温度单位，摄氏度或华氏度"）。
city和unit属性都是必需的（"required": ["city", "unit"]）。

有了JSON模式，我们就可以利用它来强制LLM生成的API请求符合规范。

2. Context-Free Grammar (CFG) 的引入

仅仅使用JSON模式还不够，因为LLM仍然可能生成不符合JSON模式的文本。为了更严格地约束LLM的生成过程，我们引入上下文无关文法（CFG）。

CFG是一种形式文法，用于描述语言的语法结构。它可以用来定义JSON模式所描述的JSON对象的合法生成规则。通过CFG，我们可以精确地控制LLM生成的文本，确保它总是符合JSON模式的规范。

将JSON模式转换为CFG是一个关键步骤。下面是一个将上述天气查询API的JSON模式转换为CFG的示例：

root ::= object
object ::= "{" ws properties ws "}"
properties ::= pair ("," ws pair)*
pair ::= string ":" ws value
value ::= string | enum
string ::= """ chars """
chars ::= char*
char ::= [a-zA-Z0-9 ]  // 允许字母、数字和空格
enum ::= "celsius" | "fahrenheit"
ws ::= [ tnr]*  // 允许空白字符

这个CFG定义了以下规则：

root是文法的起始符号，表示一个JSON对象。
object定义了JSON对象的结构，包含一个左花括号{，一些属性properties，和一个右花括号}。
properties定义了JSON对象中的属性，可以是一个pair，也可以是多个pair用逗号,分隔。
pair定义了一个键值对，包含一个字符串string作为键，一个冒号:，和一个value作为值。
value定义了值的类型，可以是字符串string，也可以是枚举类型enum。
string定义了一个字符串，由双引号"包裹一些字符chars。
chars定义了字符串中的字符，可以是零个或多个char。
char定义了允许的字符，这里限制为字母、数字和空格。
enum定义了枚举类型，只能取"celsius"或"fahrenheit"这两个值。
ws定义了空白字符，可以包含空格、制表符、换行符和回车符。

这个CFG非常简化，没有包含city这个字段。一个更完整的CFG，能够更好地约束生成，如下：

root ::= object
object ::= "{" ws properties ws "}"
properties ::= city_pair "," ws unit_pair
city_pair ::= ""city"" ":" ws string
unit_pair ::= ""unit"" ":" ws enum
string ::= """ chars """
chars ::= char+  // 至少一个字符
char ::= [a-zA-Z0-9 ]  // 允许字母、数字和空格
enum ::= ""celsius"" | ""fahrenheit""
ws ::= [ tnr]*  // 允许空白字符

这个CFG的关键改进在于：

properties被精确地定义为city_pair和unit_pair的组合，保证了city和unit字段的顺序。
city_pair和unit_pair分别定义了city和unit键值对的结构。
chars被修改为至少包含一个字符（char+），避免生成空字符串。
enum直接定义了"celsius"和"fahrenheit"这两个枚举值，避免了LLM生成其他值。

更进一步，为了更精确地控制city的内容，我们可以引入一个city_value的规则，并使用一个城市名称的列表来约束city的取值：

root ::= object
object ::= "{" ws properties ws "}"
properties ::= city_pair "," ws unit_pair
city_pair ::= ""city"" ":" ws city_value
unit_pair ::= ""unit"" ":" ws enum
city_value ::= """ city_name """
city_name ::= "Beijing" | "Shanghai" | "Guangzhou" | "Shenzhen" // 城市列表
enum ::= ""celsius"" | ""fahrenheit""
ws ::= [ tnr]*  // 允许空白字符

现在，city的值只能从"Beijing"、"Shanghai"、"Guangzhou"和"Shenzhen"这四个城市中选择。

3. 利用CFG约束采样

有了CFG，我们就可以利用它来约束LLM的采样过程，保证生成的文本总是符合CFG的语法规则。这可以通过多种方法实现，其中一种常用的方法是使用“指导解码”（Guided Decoding）。

指导解码的核心思想是：在LLM生成每个token时，只允许它选择那些能够扩展当前语法树的token。具体来说，我们可以维护一个当前语法树的状态，并根据CFG的规则，计算出下一步可以生成的token的集合。然后，我们只允许LLM从这个集合中进行采样。

下面是一个简化的Python代码示例，演示如何使用CFG约束采样：

import re

class CFG:
    def __init__(self, grammar_string):
        self.grammar = {}
        for line in grammar_string.strip().split('n'):
            nonterminal, production = line.split('::=')
            nonterminal = nonterminal.strip()
            productions = [p.strip() for p in production.split('|')]
            self.grammar[nonterminal] = productions
        self.start_symbol = list(self.grammar.keys())[0]

    def get_possible_tokens(self, current_state):
        """
        根据当前状态，返回下一步可以生成的token集合。
        current_state: 当前已经生成的文本，例如 "" 或 "{" 或 "{"city":"
        """
        # 简化实现，假设current_state已经是完整的token
        # 实际应用中需要更复杂的逻辑来解析current_state
        possible_tokens = set()
        for nonterminal, productions in self.grammar.items():
            for production in productions:
                if production.startswith(current_state):
                    # 找到以current_state开头的production
                    remaining = production[len(current_state):].strip()
                    if remaining:
                        # 取第一个token
                        first_token = remaining.split()[0]
                        if first_token not in self.grammar:
                            # 如果第一个token是终结符
                            possible_tokens.add(first_token)
                        else:
                            # 如果第一个token是非终结符，需要进一步展开
                            # (这里简化处理，实际需要递归展开)
                            possible_tokens.add(first_token)
        return possible_tokens

# 示例CFG
grammar_string = """
root ::= object
object ::= "{" ws properties ws "}"
properties ::= city_pair "," ws unit_pair
city_pair ::= "\"city\"" ":" ws city_value
unit_pair ::= "\"unit\"" ":" ws enum
city_value ::= "\"Beijing\"" | "\"Shanghai\"" | "\"Guangzhou\"" | "\"Shenzhen\""
enum ::= "\"celsius\"" | "\"fahrenheit\""
ws ::= [ \t\n\r]*
"""

cfg = CFG(grammar_string)

# 模拟LLM生成过程
current_state = ""
generated_text = ""
for _ in range(100): # 限制生成长度
    possible_tokens = cfg.get_possible_tokens(current_state)
    if not possible_tokens:
        break # 没有可生成的token，停止生成

    # 模拟LLM从possible_tokens中选择一个token
    # 实际应用中，LLM会根据概率分布选择token
    next_token = list(possible_tokens)[0] # 这里简单地选择第一个token
    generated_text += next_token
    current_state = next_token # 更新当前状态

print(generated_text)

这个代码示例只是一个简化的演示，实际应用中需要更复杂的逻辑来处理以下问题：

状态维护： 需要更精确地维护当前语法树的状态，例如使用栈来跟踪非终结符的展开过程。
tokenization： 需要对文本进行tokenization，将文本分解为token序列。
LLM集成： 需要将CFG约束采样集成到LLM的解码过程中，例如修改LLM的logits，使其只在possible_tokens对应的位置上有非零概率。
错误处理： 需要处理生成过程中可能出现的错误，例如当possible_tokens为空时，需要采取回溯或其他策略。

4. 实际应用中的考虑

在实际应用中，还需要考虑以下几个方面：

CFG的复杂性： CFG的复杂性直接影响到约束采样的效率。需要权衡CFG的精确性和效率，选择合适的CFG设计。
LLM的选择： 不同的LLM对CFG约束采样的支持程度不同。需要选择适合的LLM，并了解其API和使用方法。
动态CFG： 在某些场景下，API的规范可能会动态变化。需要设计一种机制，能够动态地更新CFG，以适应API的变化。
混合方法： 可以将CFG约束采样与其他方法结合使用，例如使用JSON模式验证器来验证生成的API请求，以提高API调用的正确性。

方法	优点	缺点	适用场景
JSON模式验证	简单易用，能够验证生成的JSON是否符合规范	只能事后验证，无法指导生成过程	对API调用正确性要求不高，或者LLM生成能力较强的情况
CFG约束采样	能够指导生成过程，保证生成的文本总是符合规范	CFG设计复杂，约束过于严格可能影响生成的多样性	对API调用正确性要求高，或者LLM生成能力较弱的情况
混合方法	结合了JSON模式验证和CFG约束采样的优点，既能指导生成过程，又能事后验证，提高API调用的正确性	实现复杂，需要权衡两种方法的优缺点	对API调用正确性要求很高，并且需要兼顾生成的多样性的情况

5. 代码示例：使用transformers库进行CFG约束采样（简化版）

下面是一个使用transformers库进行CFG约束采样的简化版代码示例。这个示例使用了transformers库的Constraint API，可以用来约束LLM的生成过程。

from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessorList, Constraint

# 假设我们已经定义了一个CFG，并将其转换为一个约束列表
# 这里简化起见，直接使用一个简单的约束
constraints = [Constraint(permitted_token_ids=[1, 2, 3])] # 假设token ID 1, 2, 3 是允许的

# 加载tokenizer和模型
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

# 创建LogitsProcessorList
logits_processor = LogitsProcessorList()
logits_processor.append(Constraint(permitted_token_ids=[1, 2, 3])) # 添加约束

# 生成文本
input_text = "The weather is"
input_ids = tokenizer.encode(input_text, return_tensors="pt")

output = model.generate(
    input_ids,
    max_length=20,
    logits_processor=logits_processor
)

generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
print(generated_text)

这个示例只是一个非常简化的演示，实际应用中需要更复杂的逻辑来将CFG转换为约束列表，并将其集成到transformers库的生成过程中。需要注意的是，transformers库的Constraint API可能不够灵活，无法满足所有CFG约束采样的需求。

6. 总结：保证Tool Use的关键

本文探讨了在Tool Use中，如何利用JSON模式强制和上下文无关文法（CFG）约束采样来保证API调用的正确性。我们介绍了JSON模式和CFG的概念，以及如何将JSON模式转换为CFG。我们还讨论了如何利用CFG约束LLM的采样过程，并提供了一些代码示例。

要确保智能体能够正确使用工具，需要仔细设计JSON模式和CFG，并选择合适的LLM和约束采样方法。通过结合JSON模式验证、CFG约束采样和其他技术，我们可以构建出更加健壮和可靠的Tool Use系统。