大模型 agent 执行sql查询

转载，来源：How to call functions with chat models

1.定义client

使用了阿里云的qwen-max 模型。

from openai import OpenAI
import json
import os
import sqlite3 as sl

client = OpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"), # 如果您没有配置环境变量，请在此处用您的API Key进行替换
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",  # 填写DashScope SDK的base_url
)

2. 准备数据库信息

使用了SQLite sample Database，下载地址：
Chinook sample database
这部分主要是准备表结构信息，加入到tools中。

import sqlite3 as sl
conn = sl.connect("./chinook.db")
def get_table_names(conn):
    """返回表名列表"""
    table_names = []
    tables = conn.execute("SELECT name FROM sqlite_master WHERE type='table';")
    for table in tables.fetchall():
        table_names.append(table[0])
    return table_names


def get_column_names(conn, table_name):
    """返回字段名"""
    column_names = []
    columns = conn.execute(f"PRAGMA table_info('{
     table_name}');").fetchall()