VannaAI 介绍及使用 - 第三篇_python vanna

Vanna的多样化使用

前言

前面两篇，对Vanna的一些原理和基本使用进行的总结，本篇将根据Vanna提供的多样化配置进行体验尝试，并共同来探索一下Vanna源码，有不正之处请友友们指正。

先做一些说明：

1.我本地目前使用的是 python3.12版本，Vanna是0.5.5版本，不同版本的包可能源码会有些许区别，但主体是一致的。

2.因为Vanna默认是使用chatGPT3.5，我自己本地测试连接其他LLM是通过代理服务实现的，有条件的可以自己部署大模型尝试，本文仅作为自主体验的记录，仅供参考。

3.本篇会先对Vanna一些常用的源码进行说明， 并在此基础上进行一些改造尝试，有误之处请大家多多指正，抱拳抱拳~

一、自定义测试

1.大模型准备 

这里需要先推荐一个代理服务的使用教程链接KIMI AI Free 服务

 我本地试验是用的自己的云服务器，通过链接中的教程，部署了kimi服务。

 2.自定义初始化配置

from vanna.openai import OpenAI_Chat
from vanna.vannadb import VannaDB_VectorStore
from openai import OpenAI
 
# 使用自己的vanna模型及api key
MY_VANNA_MODEL = "my_vanna_model"
MY_VANNA_API_KEY = "my_vanna_api_key"
 
# 使用自己的llm模型及所需配置
refresh_token='refresh_token'
my_llm_base_url='http://127.0.0.1:8000/v1' #IP和端口号改为你自己的，我这里实际用的我的云服务器的地址
my_llm_name='kimi'
 
# 初始化客户端
client = OpenAI(
    api_key=refresh_token,
    base_url=my_llm_base_url,
    default_headers={"Authorization": f"Bearer {refresh_token}", "Content-Type": "application/json"}
)
 
class MyVanna(VannaDB_VectorStore, OpenAI_Chat):
    def __init__(self, client=None, config=None):
        VannaDB_VectorStore.__init__(self, vanna_model=MY_VANNA_MODEL, vanna_api_key=MY_VANNA_API_KEY, config=config)
        OpenAI_Chat.__init__(self, client=client, config=config)
 
# 使用自定义的大模型及vanna提供的向量库
vn = MyVanna(client=client, config={"model": my_llm_name, })

有朋友肯定会问，为什么我这里的初始化这么复杂，官网的初始化代码很简单嘛，这里做一些说明：

1. 我这里使用的是openai的自定义配置，但是呢实际我是去请求我云服务器的kimi服务，他们的请求体是差不多的，所以可以直接套用，但是呢要将所需的参数都提前定义，一个是Vanna的apikey和model，在Vanna官网根据自己的账号获取，还有一个就是kimi大模型服务需要的token，接口url和model。
2. 这里我比官网多了一个client客户端的初始化，实际Vanna源码中也是有client初始化的代码，但是我这里自定义了一个，满足需要就行了，client的作用主要就是传递相关参数，这里会用到的就是api_key，base_url，default_headers，为了方便测试，这里就提前定义好。

3. 数据库连接

以下是数据库连接方法，这里以mysql为例。Vanna还提供其他的数据库连接方式，都是以connect_to_数据库的方式命名。

# 初始化数据库连接
vn.connect_to_mysql(host="localhost", dbname="test", user="root", password="123456", port=3306)

 4.提问验证

其实这个时候就可以问问题了，我先问个‘你是谁’看看，然后执行。

vn.ask(question="你是谁")

 回答结果：

 连接成功啦，回答也是OK的，最后一行报错是因为当前kimi的回答无法作为sql去执行，这个没关系，我们换一个能生成sql的问题。

vn.ask(question="请帮我查一下student_info表score大于92的数据")

回答结果：

可以看到kimi是正常生成了sql的，并且Vanna通过sql也从数据库查到了数据，并让kimi生成相应的图形代码，展示出了图形。（这里我是提前准备了数据库的测试数据的，第二篇文章有建表语句和插入数据语句） 

5.训练数据 

刚才我们的提问是很详细具体的，我们加入一些训练数据，以实现能满足多样化的提问。

vn.train(ddl="CREATE TABLE student_info (id INT NOT NULL AUTO_INCREMENT COMMENT '唯一标识符',name VARCHAR(100) NOT NULL COMMENT '学生姓名',age INT NOT NULL COMMENT '年龄',province VARCHAR(100) NOT NULL COMMENT '省份',score DECIMAL(5,2) NOT NULL COMMENT '分数',PRIMARY KEY (id)) ENGINE=InnoDB AUTO_INCREMENT=21 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci COMMENT='存储学生基本信息及成绩';")
 
vn.train(question="请查询student_info表中score大于90的数据", sql="SELECT * FROM student_info WHERE score > 90;")
vn.train(question='请查询分数小于85的省份', sql="SELECT DISTINCT province FROM student_info WHERE score < 85;")
#自己再加一些样例

然后我们再继续进行一些数据相关的问题。这里我就不做过多演示了，小伙伴们可以自己多多尝试。

vn.ask(question="请统计各省份的分数情况")

 从输出的内容我们看到除了生成的sql，查询到的数据，以及图表和代码，最前面还有一大堆东西，这些就是prompt提示词，这个我把它提取出来格式化一下。

prompt = [{'role': 'system',
                        #你是一位SQL专家，请帮助生成一个SQL查询来回答这个问题。
           'content': "You are a SQL expert. Please help to generate a SQL query to answer the question. "
                      # 你的回答应该仅仅基于给定的上下文，并遵循回答指南和格式说明。
                      "Your response should ONLY be based on the given context and follow the response guidelines and format instructions."
                      # 我们的ddl建表语句
                      " \n===Tables \nCREATE TABLE student_info "
                      "(id INT NOT NULL AUTO_INCREMENT COMMENT '唯一标识符',"
                      "name VARCHAR(100) NOT NULL COMMENT '学生姓名',"
                      "age INT NOT NULL COMMENT '年龄',"
                      "province VARCHAR(100) NOT NULL COMMENT '省份',"
                      "score DECIMAL(5,2) NOT NULL COMMENT '分数',"
                      "PRIMARY KEY (id)) "
                      "ENGINE=InnoDB AUTO_INCREMENT=21 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci COMMENT='存储学生基本信息及成绩';"
                        # 回答指南
                        # 如果提供的上下文足够，针对问题生成一个有效的SQL查询，无需任何解释。
                        # 如果提供的上下文几乎足够，但需要了解特定列中的特定字符串知识，请生成一个中间SQL查询以查找该列中的不同字符串。在查询前加上注释，标记为intermediate_sql。
                        # 如果提供的上下文不足，请解释为什么无法生成查询。
                        # 请使用最相关的表。
                        # 如果问题之前已经被提出并回答过，请完全按照之前的答案重复回答。
                      "\n\n===Response Guidelines "
                      "\n1. If the provided context is sufficient, please generate a valid SQL query without any explanations for the question. "
                      "\n2. If the provided context is almost sufficient but requires knowledge of a specific string in a particular column,"