yugasun
Published on

TRAG - 基于类型定义的检索增强生成

Authors
  • avatar
    Name
    Yuga Sun
    Twitter

1. 前言

TRAG:Type-based Retrival-Augmented Generation 基于类型的检索增强生成

大语言模型(以下简称 LLM)本身是一个很好的聊天助手,它可以为我们提供现实中各种场景的对话回复,交互流程是:用户输入问题,LLM 回答输出自然语言。如果是单纯的自然语言,很难作为应用程序的输入。 比如,我们要设计一款智能用户服务,可以通过对话方式来管理用户,我们需要添加一个用户,为了生成接口参数,我们可以通过 Prompt 来引导 LLM,Prompt 如下:

帮忙添加一个名称为 `Jack` 的测试账号,以 JSON 的形式响应

LLM 可能得输出结果如下:

// 结果1
{"name": "Jack", "type": "test"}

// 结果2
{
    "username": "Jack",
    "type": "test"
}

// 结果3
{
  "account": {
    "username": "Jack",
    "type": "test_user",
  }
}

// 结果x

我们可以发现生成的结构可能不太稳定,这就导致我们的程序没法直接使用,为了解决这个问题, OpenAI 推出了 Function Call 的功能,即函数调用。调用 GPT API 时,需要先定义函数,如下:

{
    "name": "addUser",
    "description": "添加用户",
    "parameters": {
        "type": "object",
        "properties": {
            "username": {
                "type": "string",
                "description": "用户名称"
            },
            "type": {
                "type": "string",
                "description": "用户类型",
                "enum": ["admin", "test", "normal"]
            }
        },
        "required": ["username"]
    }
}

添加函数定义后,GPT API 的返回中,就会携带函数调用相关参数,比如:

id='call_62136354',
function=Function(
    arguments='{"username":"Jack","type":"test"}',
    name='addUser'),
type='function'

此时我们只需要执行定义好的 addUser 函数就可以了。

Function Call 功能,极大的扩展了 LLM 的应用场景,ChatGPT 官方的应用市场 GPTs 就是基于它实现的。 但是,Function Call 功能同样存在问题,比如:

  1. 目前只有少数开源大模型支持
  2. Function Call 目前只支持单个函数调用
  3. 函数声明不够友好,内部黑盒,强依赖大模型;当返回不符合预期,难以排查

那么如何能够让 LLM 返回,严格遵守我们的应用程序输入呢?

2. TRAG 原理

我们在设计大型软件应用时,往往会选择强类型语言(强类型语言是一种编程语言,其特点是数据类型在使用时必须严格遵守,并且在未经显式转换的情况下,不允许自动将一种数据类型转换为另一种数据类型。),通过一致和严格的类型声明,来约束系统代码,从而提升系统的健壮性和维护性。 同样,我们也可以通过给 Prompt 添加类型声明,来约束 LLM 的返回结果。以下以 TypeScript 语言作为示例。 首先我们声明一个语义类型声明(Schema):

interface SentimentResponse {
    sentiment: "negative" | "neutral" | "positive"; // The sentiment of the text
}

然后结合类型声明,定制 Prompt 如下:

You are a service that translates user requests into JSON objects of type "SentimentResponse" according to the following TypeScript definitions:
"""
interface SentimentResponse {
    sentiment: "negative" | "neutral" | "positive"; // The sentiment of the text
}
"""

The following is a user request:
"""
It's very rainy outside
"""
The following is the user request translated into a JSON object with 2 spaces of indentation and no properties with the value undefined:

以上提示词,主要含有三个部分:

  1. 用来指定参考 TypeScript 类型定义
  2. 用户的请求
  3. 指定输出 JSON 结构,并且属性值不能是 undefined(未定义)

将拼接后的提示词输入给 LLM 后,它将严格按照类型定义进行返回,如下:

{
    "sentiment": "neutral"
}

有时为了验证返回结果有效性,我们还可以通过 zod 进行,可以稍微修改下上面个的类型声明即可:

import { z } from 'zod';

export const SentimentResponse = z.object({
  sentiment: z.enum(['negative', 'neutral', 'positive']),
});

export const SentimentSchema = {
  SentimentResponse,
};

然后可以针对返回的 JSON 进行验证:

result = {
    "sentiment": "neutral"
}
SentimentResponse.safeParse(result)

如果验证不通过,可以再一次将返回 JSON 结果跟之前的 Prompt 进行结合,来让 LLM 进行类型修正,这样就可以做到严格按照应用程序的类型定义进行返回了。

3. TRAG 执行流程

TRAG 的整个执行流程如下:

TRAG

以上流程中,红色部分为输出产物纠错过程。

TRAG 架构包括两个核心模块:

  1. Parser:主要用来处理提示词,通过将系统提示词、类型定义和用户问题进行组合,输出固定格式的 Prompt
  2. Validator: 用来进行生成结果的类型检查

如果生成结构未通过类型检查,则会将不符合的结果输入到 Parser 的 Repairer,Repairer 会将之前的系统提示词、类型定义、用户问题和异常 JSON 一起拼接成新的提示词,重新输入给 LLM,进行重试,直到返回正确的结果。

3. TypeChat 简介

基于以上原理,TypeScript 之父 Anders Hejlsberg (同时也是 C#之父!!!) 开源了新项目 TypeChat。 TypeChat 是一个工具库,可以轻松地使用类型构建自然语言界面。使用时,只需定义代表自然语言应用程序中支持的意图的类型即可。这可以像用于对情绪进行分类的界面一样简单,也可以像购物车或音乐应用程序的类型一样复杂。 例如,要向架构添加其他意图,开发人员可以将其他类型添加到可区分联合中。要使架构具有层次结构,开发人员可以使用“元架构”根据用户输入选择一个或多个子架构。 定义类型后,TypeChat 将通过以下方式处理其余部分:

  1. 使用类型构建 LLM 提示。
  2. 验证 LLM 响应是否符合架构。如果验证失败,则通过进一步的语言模型交互修复不符合要求的输出。
  3. 简洁地总结(不使用 LLM)实例并确认它符合用户意图。

比如,对于上面提到的示例,如果使用 TypeChat,代码可以简化为:

import { createJsonTranslator, createLanguageModel, createZodJsonValidator } from '../../src/parser';
import { processRequests } from '../interactive';
import { logger as log } from '../logger';
import { SentimentSchema } from './schema';

async function main() {
    // 创建 LLM 实例
    const model = createLanguageModel(process.env);
    // 定义类型检查器
    const validator = createZodJsonValidator(SentimentSchema, 'SentimentResponse');
    // 定义 JSON 编译器
    const translator = createJsonTranslator(model, validator);

    const question = "It's very rainy outside"

    // 发送请求
    const response = await translator.translate(question);
    // 打印结果
    console.log(JSON.stringify(response.data))
}

TypeChat 就介绍介绍到这里,感兴趣的朋友可以通过阅读官方源码,进行深入研究。

4. 后记

TRAG 架构是笔者在设计 ChatBI 系统时探索的一种创新实现方案。传统的 RAG 方法在稳定输出用户所需的图表类型方面效果不佳,最终通过 TRAG 架构成功解决了这一问题。未来,笔者还会将 TRAG 应用于通过自然语言操作 RESTful API 服务的场景,从目前的测试结果来看,预期能够顺利实现。 TRAG 架构有效弥合了自然语言与结构化数据之间的鸿沟,使得 LLM 可以更轻松地集成到各类应用程序中,助力现有系统的智能化改造,加速推进业务升级。