使用ChatGLM3自定义工具实现大模型查询MySQL数据库

news2024/11/26 11:19:49

ChatGLM3-6B 采用了全新设计的 Prompt 格式，除正常的多轮对话外。同时原生支持工具调用（Function Call）、代码执行（Code Interpreter）和 Agent 任务等复杂场景。

什么是工具调用

大模型虽然强大，但是由于训练的时间和语料限制。大模型通常会存在以下问题：

只能获取训练数据集中有的事件和内容，这意味着大模型不具备访问最新资料的能力。
ChatGLM3-6B模型以通用语料训练为主，因此，缺少专业领域的知识。
ChatGLM3-6B体量较小，虽然拥有较强的数学，英语等能力，但仍然无法与GPT4等大模型进行抗衡，因此，会出现数学计算不准确等问题。

简而言之，就是给大模型加上“四肢“，让大模型学会使用工具，很多问题将迎刃而解。

如何进行工具调用

ChatGLM3自带了查询某地实时天气的工具，如下图所示：
在这里插入图片描述

如何进行自定义工具：

（1）了解工具注册流程

在ChatGLM3项目的tool_register.py中定义了register_tool工具注册的函数，这个 Python 函数 register_tool（注册工具）的目的是将另一个函数作为参数（用 func: 可调用参数表示），注册它并收集它的元数据。感兴趣的可自行去看源码。

（2）编写自定义工具代码

我们这里想要大模型能够返回给我一些想要看的数据库数据，代码如下：

import pymysql
import pandas as pd
  
def getDataFromMysql(table_name):
    # 创建连接
    conn = pymysql.connect(
        host="10.8.xxx.xxx", 
        port=3306, 
        user='root', 
        passwd='123456', 
        db='chat', 
        charset='utf8mb4'
    )
    # 创建游标
    cursor = conn.cursor()
    cursor.execute("select * from "+table_name)
    # 获取前n行数据
    datas = cursor.fetchmany(5)
    cols_info = cursor.description       # 获取行相关信息
    cols = [col[0] for col in cols_info] # 处理保留列名
    # 关闭游标
    cursor.close()
    # 关闭连接
    conn.close()
    df = pd.DataFrame(datas,columns=cols)
    return df
   
@register_tool
def get_mysqlData(
        table_name: Annotated[str, 'The name of the table to be queried', True],
) -> str:
    """
    Get the mysql data for `table_name`
    """
	data = getDataFromMysql(table_name)
	return data