text2sql框架-vanna使用总结

1. 说明

Vanna 是一个开源的基于大模型和 RAG 的 Text2SQL 工具

官方地址是：https://vanna.ai/

2. 项目下载

vanna 提供了 vanna-flask 项目，通过 Python 的 Flash 框架提供 Web 界面，项目地址为：https://github.com/vanna-ai/vanna-flask

3. 项目依赖

https://vanna.ai/docs/mysql-other-llm-chromadb/

在 Vann 的以上链接可以选择准备使用的模型、向量数据库、需要查询的数据库

以上链接对应的情况如下：

内容	选择
模型	Other LLM
向量数据库	ChromaDB
需要查询的数据库	MySQL

3.1. 模型

Vanna 推荐通过 Vanna.Ai 免费使用 OpenAI（需要注册并通过公网调用 Vanna 接口）

也支持通过 API key 使用 OpenAI 等在线模型

准备使用局域网搭建的 qwen 模型，因此选择 “Other LLM”

3.2. 向量数据库

向量数据库用于保存训练数据库

Vanna 推荐使用 Vanna 提供的向量数据库 pgvector（需要注册并通过公网调用 Vanna 接口）

准备使用本地的 ChromaDB，不需要安装软件，只需要安装对应的 Python 依赖包，因此选择 “ChromaDB”

3.3. 文档中关于模型与向量数据库的使用方式

对应以上网页的 “Setup” 部分，提供了使用 ChromaDB 与其他模型的代码示例

定义 MyCustomLLM 类，用于实现自定义模型访问，父类使用 VannaBase

定义 MyVanna 类，用于对向量数据库与模型使用类进行初始化，父类使用 ChromaDB_VectorStore、MyCustomLLM

VannaBase 类是 Vanna 提供的访问其他模型的基类，ChromaDB_VectorStore 类是 Vanna 提供的访问 ChromaDB 的类

访问模型的代码参考：https://github.com/vanna-ai/vanna/blob/main/src/vanna/mistral/mistral.py，需要继承 VannaBase 类，访问模型的代码在 submit_prompt 方法中实现，如下所示：

    def submit_prompt(self, prompt, **kwargs) -> str:
        chat_response = self.client.chat(
            model=self.model,
            messages=prompt,
        )

        return chat_response.choices[0].message.content

4. Python 依赖包安装

下载 vanna-flask 项目后，在根目录执行以下命令安装 Python 依赖包

pip install -r requirements.txt
pip install vanna[chromadb,mysql]

假如代码中有其他依赖的包，也需要安装

5. 代码修改

需要修改 app.py 代码，对模型与向量数据库的访问进行调整

5.1. 注释代码

注释以下访问不准备使用的模型的代码：

from vanna.local import LocalContext_OpenAI
vn = LocalContext_OpenAI()
from vanna.remote import VannaDefault
vn = VannaDefault(model=os.environ['VANNA_MODEL'], api_key=os.environ['VANNA_API_KEY'])

vn.connect_to_snowflake(
    account=os.environ['SNOWFLAKE_ACCOUNT'],
    username=os.environ['SNOWFLAKE_USERNAME'],
    password=os.environ['SNOWFLAKE_PASSWORD'],
    database=os.environ['SNOWFLAKE_DATABASE'],
    warehouse=os.environ['SNOWFLAKE_WAREHOUSE'],
)

5.2. 自定义模型与向量数据库配置

ChromaDB_VectorStore、MyCustomLLM、MyVanna 类都通过 config 接收使用的配置参数

在创建 MyVanna 类时，可通过 config 传递配置参数

如下所示：

vn = MyVanna(config={'model_url': MODEL_URL, 'model': MODEL, 'api_key': API_KEY, 'path': CHROMADB_PATH})

可在 environment.py 中定义使用的模型与向量数据库相关参数，格式如下：

MODEL_URL = "模型接口的 URL，http://xxx"
MODEL = "模型名称"
API_KEY = "模型提供的 API KEY"

CHROMADB_PATH = "保存 ChromaDB 数据库文件的目录路径"

在 app.py 中引入相关参数：

from environment import MODEL_URL, API_KEY, MODEL, CHROMADB_PATH

5.3. 模型访问相关

5.3.1. 自定义 ChatMessage 类

参考 vanna 示例代码 https://github.com/vanna-ai/vanna/blob/main/src/vanna/mistral/mistral.py

需要实现 VannaBase 的子类，system_message、user_message、assistant_message 方法返回对象为 ChatMessage，引用方式为 “from mistralai.models.chat_completion import ChatMessage”

由于以上 ChatMessage 类的功能比较简单，为了避免引入以上 mistralai.models.chat_completion，因此在项目中增加 chat_message.py 文件，实现 ChatMessage 类：

class ChatMessage:
    role: str
    content: str

    def __init__(self, role, content):
        self.role = role
        self.content = content

    def __str__(self):
        return f"ChatMessage(role={self.role}, content={self.content})"

    def __repr__(self):
        return f"ChatMessage(role={self.role}, content={self.content})"

5.3.2. 指定模型访问方式

参考以上 vanna 示例代码，在 VannaBase 的子类 MyCustomLLM 的 submit_prompt 方法中实现对模型的访问：

以下使用的接口格式与 OpenAI 的相同，在 submit_prompt 方法中请求模型，并返回模型的返回值

class MyCustomLLM(VannaBase):
    def __init__(self, config=None):
        super().__init__(config)

    def system_message(self, message: str) -> any:
        return ChatMessage(role="system", content=message)

    def user_message(self, message: str) -> any:
        return ChatMessage(role="user", content=message)

    def assistant_message(self, message: str) -> any:
        return ChatMessage(role="assistant", content=message)

    def submit_prompt(self, prompt, **kwargs) -> str:
        messages = []

        for prompt_data in prompt:
            messages.append({"role": prompt_data.role, "content": prompt_data.content})

        payload = {
            "model": self.config['model'],
            "messages": messages,
            "stream": False
        }

        print(f"payload: {json.dumps(payload, ensure_ascii=False)}")

        headers = {
            "Content-Type": "application/json",
            "Authorization": "Bearer" + self.config['api_key']
        }

        response = requests.post(self.config['model_url'], json=payload, headers=headers, verify=False)
        rsp_data = response.json()
        print(f"rsp_data: {rsp_data}")
        return rsp_data["choices"][0]["message"]["content"]

6. 启动 Web 服务

执行 python app.py 启动 vanna 的 Web 服务（或者使用 IDE 的启动功能）

6.1. 访问网页

启动 vanna 的 Web 服务成功后，在日志中会出现以下内容，即 Web 服务的访问地址

 * Running on http://127.0.0.1:5000

6.2. 问题 - 网页白屏

6.2.1. 现象

打开以上网页后白屏，查看浏览器 Console 中有以下报错：

index-d29524f4.js:1 Failed to load module script: Expected a JavaScript module script but the server responded with a MIME type of "text/plain". Strict MIME type checking is enforced for module scripts per HTML spec.

6.2.2. 解决

参考以下链接的说明

https://stackoverflow.com/questions/59355194/python-flask-error-failed-to-load-module-script-strict-mime-type-checking-i

在 app.py 中 import flask 之前添加以下代码：

# fix windows registry stuff
import mimetypes
mimetypes.add_type('application/javascript', '.js')
mimetypes.add_type('text/css', '.css')

重新启动 Web 服务，并清空浏览器缓存，解决问题

6.3. 问题 - 网页打开较慢

6.3.1. 现象

static\index.html 文件中引用了 plotly-latest.min.js：

<script src="https://cdn.plot.ly/plotly-latest.min.js" type="text/javascript"></script>

从对应的网站下载 plotly-latest.min.js 文件可能很慢，导致本地的 vanna 网页打开很慢

6.3.2. 解决

可先下载以上 js 文件，保存到本地项目的 static 目录中

修改 static\index.html 文件中通过引用本地的 plotly-latest.min.js 文件：

<script src="/assets/plotly-latest.min.js" type="text/javascript"></script>

6.4. 问题 - 从 awx 下载 onnx 失败

6.4.1. 现象

启动 Web 服务器打开网页，python 代码出现报错，报错内容较长，关键提示如下，提示从 awx 下载 onnx 失败，说明 SSL 证书验证失败

requests.exceptions.SSLError: HTTPSConnectionPool(host='chroma-onnx-models.s3.amazonaws.com', port=443): Max retries exceeded with url: /all-MiniLM-L6-v2/onnx.tar.gz (Caused by SSLError(SSLCertVerificationError(1, '[SSL: CERTIFICATE_VERIFY_FAILED] certificate verify failed: unable to get local issuer certificate (_ssl.c:992)')))

6.4.2. 解决

修改 app.py 代码

增加以下引用

import urllib3
from urllib3 import exceptions

在 “app = Flask” 代码之前增加以下代码，取消 SSL 证书验证：

urllib3.disable_warnings()

问题解决，从控制台可以看到开始下载对应文件

.cache\chroma\onnx_models\all-MiniLM-L6-v2\onnx.tar.gz:  37%|███▋      | 29.4M/79.3M [00:21<00:35, 1.45MiB/s]

7. 连接 MySQL 数据库

7.1. 安装 Python 依赖包

pip install PyMySQL

7.2. 增加连接 MySQL 数据库代码

在 “vn = MyVanna()” 代码后增加连接 MySQL 数据库的代码，如下所示

vn.connect_to_mysql(host='127.0.0.1', dbname='dbname', user='user', password='password', port=3306)

s

8. 训练

需要对 vanna 执行训练操作，使 vanna 能够读取到数据库表信息

8.1. 文档说明

参考 Vanna 文档 https://vanna.ai/docs/train/，“Training Plan” 部分：

# The information schema query may need some tweaking depending on your database. This is a good starting point.
df_information_schema = vn.run_sql("SELECT * FROM INFORMATION_SCHEMA.COLUMNS")

# This will break up the information schema into bite-sized chunks that can be referenced by the LLM
plan = vn.get_training_plan_generic(df_information_schema)
plan

# If you like the plan, then uncomment this and run it to train
vn.train(plan=plan)

可增加以上代码使 Vanna 读取到当前数据库的数据库的字段信息，可在 app 对象定义之前或之后执行

8.2. 训练内容优化

在使用 MySQL 数据库时，将以上通过 vn.run_sql 执行的 sql 语句修改为如下，可以使得仅查询当前 schema 的字段信息，且能够查询到对应表的注释：

SELECT c.*, t.table_comment FROM INFORMATION_SCHEMA.COLUMNS c, INFORMATION_SCHEMA.TABLES t where c.table_schema=database() and c.table_schema=t.table_schema and c.table_name=t.table_name

8.3. 实际效果

vanna 在使用训练数据中的数据库表信息时，只会将部分表信息发送给模型，因此模型无法针对所有的表做出正确的回答