使用 Vanna 和 Milvus 编写 SQL

Vanna 是一个开源的 Python RAG（检索增强生成）框架，用于 SQL 生成和相关功能。Milvus 是世界上最先进的开源向量数据库，专为支持嵌入相似性搜索和 AI 应用而构建。

Vanna 通过两个简单步骤工作 - 在您的数据上训练 RAG "模型"，然后提出问题，这些问题将返回可以设置为在您的数据库上运行的 SQL 查询。本指南演示了如何使用 Vanna 基于存储在数据库中的数据生成和执行 SQL 查询。

先决条件

在运行此笔记本之前，请确保您已安装以下依赖项：

$ pip install "vanna[milvus,openai]"

如果您使用的是 Google Colab，为了启用刚安装的依赖项，您可能需要重启运行时（点击屏幕顶部的"Runtime"菜单，然后从下拉菜单中选择"Restart session"）。

您需要在环境变量中设置 OPENAI_API_KEY。您可以从 OpenAI 获取 API 密钥。

import os

os.environ["OPENAI_API_KEY"] = "sk-***********"

数据准备

首先，我们需要从 Vanna 继承 Milvus_VectorStore 和 OpenAI_Chat 类，并定义一个新类 VannaMilvus，该类结合了两者的功能。

from pymilvus import MilvusClient, model
from vanna.milvus import Milvus_VectorStore
from vanna.openai import OpenAI_Chat


class VannaMilvus(Milvus_VectorStore, OpenAI_Chat):
    def __init__(self, config=None):
        Milvus_VectorStore.__init__(self, config=config)
        OpenAI_Chat.__init__(self, config=config)

我们使用必要的配置参数初始化 VannaMilvus 类。我们使用 milvus_client 实例来存储嵌入，并使用从 milvus_model 初始化的 model.DefaultEmbeddingFunction() 来生成嵌入。

关于 MilvusClient 的参数：

将 uri 设置为本地文件，例如 ./milvus.db，是最方便的方法，因为它会自动利用 Milvus Lite 将所有数据存储在此文件中。
如果您有大规模数据，可以在 docker 或 kubernetes 上设置性能更高的 Milvus 服务器。在此设置中，请使用服务器 uri，例如 http://localhost:19530，作为您的 uri。
如果您想使用 Zilliz Cloud，Milvus 的完全托管云服务，请调整 uri 和 token，它们对应于 Zilliz Cloud 中的 Public Endpoint 和 Api key。

milvus_uri = "./milvus_vanna.db"

milvus_client = MilvusClient(uri=milvus_uri)

vn_milvus = VannaMilvus(
    config={
        "api_key": os.getenv("OPENAI_API_KEY"),
        "model": "gpt-3.5-turbo",
        "milvus_client": milvus_client,
        "embedding_function": model.DefaultEmbeddingFunction(),
        "n_results": 2,  # 从 Milvus 语义搜索返回的结果数量。
    }
)

这是一个只有少量样本数据的简单示例，因此我们将 n_results 设置为 2，以确保我们搜索前 2 个最相似的结果。在实践中，当处理更大的训练数据集时，您应该将 n_results 设置为更高的值。

我们将使用一个包含几个表的示例 SQLite 数据库，其中包含一些示例数据。

import sqlite3

sqlite_path = "./my-database.sqlite"
sql_connect = sqlite3.connect(sqlite_path)
c = sql_connect.cursor()

init_sqls = """
CREATE TABLE IF NOT EXISTS Customer (
    ID INTEGER PRIMARY KEY AUTOINCREMENT,
    Name TEXT NOT NULL,
    Company TEXT NOT NULL,
    City TEXT NOT NULL,
    Phone TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS Company ( ID INTEGER PRIMARY KEY AUTOINCREMENT, Name TEXT NOT NULL, Industry TEXT NOT NULL, Location TEXT NOT NULL, EmployeeCount INTEGER NOT NULL );

CREATE TABLE IF NOT EXISTS User ( ID INTEGER PRIMARY KEY AUTOINCREMENT, Username TEXT NOT NULL UNIQUE, Email TEXT NOT NULL UNIQUE );

INSERT INTO Customer (Name, Company, City, Phone) VALUES ('John Doe', 'ABC Corp', 'New York', '123-456-7890');

INSERT INTO Customer (Name, Company, City, Phone) VALUES ('Jane Smith', 'XYZ Inc', 'Los Angeles', '098-765-4321');

INSERT INTO Company (Name, Industry, Location, EmployeeCount) VALUES ('ABC Corp', 'cutting-edge technology', 'New York', 100);

INSERT INTO User (Username, Email) VALUES ('johndoe123', 'johndoe123@example.com'); """

for sql in init_sqls.split(";"): c.execute(sql)

sql_connect.commit()

连接到 SQLite 数据库

vn_milvus.connect_to_sqlite(sqlite_path)

## 使用数据训练
我们可以在 SQLite 数据库的 DDL 数据上训练模型。我们获取 DDL 数据并将其提供给 `train` 函数。


```python
# 如果存在训练数据，我们应该在训练前删除它。
existing_training_data = vn_milvus.get_training_data()
if len(existing_training_data) > 0:
    for _, training_data in existing_training_data.iterrows():
        vn_milvus.remove_training_data(training_data["id"])

# 获取 SQLite 数据库的 DDL
df_ddl = vn_milvus.run_sql("SELECT type, sql FROM sqlite_master WHERE sql is not null")

# 在 DDL 数据上训练模型
for ddl in df_ddl["sql"].to_list():
    vn_milvus.train(ddl=ddl)

Adding ddl: CREATE TABLE Customer ( ID INTEGER PRIMARY KEY AUTOINCREMENT, Name TEXT NOT NULL, Company TEXT NOT NULL, City TEXT NOT NULL, Phone TEXT NOT NULL ) Adding ddl: CREATE TABLE sqlite_sequence(name,seq) Adding ddl: CREATE TABLE Company ( ID INTEGER PRIMARY KEY AUTOINCREMENT, Name TEXT NOT NULL, Industry TEXT NOT NULL, Location TEXT NOT NULL, EmployeeCount INTEGER NOT NULL ) Adding ddl: CREATE TABLE User ( ID INTEGER PRIMARY KEY AUTOINCREMENT, Username TEXT NOT NULL UNIQUE, Email TEXT NOT NULL UNIQUE )

除了在 DDL 数据上训练外，我们还可以在数据库的文档和 SQL 查询上进行训练。

# 添加有关您的业务术语或定义的文档。
vn_milvus.train(
    documentation="ABC Corp 专门从事尖端技术解决方案和创新。"
)
vn_milvus.train(
    documentation="XYZ Inc 是制造和供应链管理的全球领导者。"
)

# 您还可以添加 SQL 查询来训练您的模型
vn_milvus.train(sql="SELECT * FROM Customer WHERE Name = 'John Doe'")

Adding documentation.... Adding documentation.... Using model gpt-3.5-turbo for 65.0 tokens (approx) Question generated with sql: What are the details of the customer named John Doe? Adding SQL...

'595b185c-e6ad-47b0-98fd-0e93ef9b6a0a-sql'

现在让我们看看训练数据。

training_data = vn_milvus.get_training_data()
training_data

&lt;table&gt;
  &lt;thead&gt;
    &lt;tr&gt;
      &lt;th&gt;#</th>
      &lt;th&gt;id</th>
      &lt;th&gt;question</th>
      &lt;th&gt;content</th>
    </tr>
  </thead>
  &lt;tbody&gt;
    &lt;tr&gt;
      &lt;th&gt;0</th>
      &lt;td&gt;595b185c-e6ad-47b0-98fd-0e93ef9b6a0a-sql</td>
      &lt;td&gt;What are the details of the customer named Joh...</td>
      &lt;td&gt;```sql
SELECT * FROM Customer WHERE Name = 'John Doe'</td>
    </tr>
    &lt;
```tr&gt;
      &lt;th&gt;0</th>
      &lt;td&gt;25f4956c-e370-4097-994f-996f22d145fa-ddl</td>
      &lt;td&gt;None</td>
      &lt;td&gt;CREATE TABLE Company (\n    ID INTEGER PRIMARY...</td>
    </tr>
    &lt;tr&gt;
      &lt;th&gt;1</th>
      &lt;td&gt;b95ecc66-f65b-49dc-a9f1-c1842ad230ff-ddl</td>
      &lt;td&gt;None</td>
      &lt;td&gt;CREATE TABLE Customer (\n    ID INTEGER PRIMAR...</td>
    </tr>
    &lt;tr&gt;
      &lt;th&gt;2</th>
      &lt;td&gt;fcc73d15-30a5-4421-9d73-b8c3b0ed5305-ddl</td>
      &lt;td&gt;None</td>
      &lt;td&gt;CREATE TABLE sqlite_sequence(name,seq)</td>
    </tr>
    &lt;tr&gt;
      &lt;th&gt;3</th>
      &lt;td&gt;feae618c-5910-4f6f-8b4b-6cc3e03aec06-ddl</td>
      &lt;td&gt;None</td>
      &lt;td&gt;CREATE TABLE User (\n    ID INTEGER PRIMARY KE...</td>
    </tr>
    &lt;tr&gt;
      &lt;th&gt;0</th>
      &lt;td&gt;79a48db1-ba1f-4fd5-be99-74f2ca2eaeeb-doc</td>
      &lt;td&gt;None</td>
      &lt;td&gt;XYZ Inc is a global leader in manufacturing an...</td>
    </tr>
    &lt;tr&gt;
      &lt;th&gt;1</th>
      &lt;td&gt;9f9df1b8-ae62-4823-ad28-d7e0f2d1f4c0-doc</td>
      &lt;td&gt;None</td>
      &lt;td&gt;ABC Corp specializes in cutting-edge technolog...</td>
    </tr>
  </tbody>
</table>
</div>



## 生成 SQL 查询

现在我们可以提出问题，Vanna 将生成 SQL 查询。


```python
sql_query = vn_milvus.generate_sql("ABC Corp 有多少员工？")
print(sql_query)

SELECT EmployeeCount FROM Company WHERE Name = 'ABC Corp'

sql_query = vn_milvus.generate_sql("ABC Corp 的客户是谁？")
print(sql_query)

SELECT * FROM Customer WHERE Company = 'ABC Corp'

执行 SQL 查询

现在我们可以执行生成的 SQL 查询。

result = vn_milvus.run_sql(sql_query)
print(result)

ID Name Company City Phone 0 1 John Doe ABC Corp New York 123-456-7890

我们还可以使用 ask 函数，它将生成 SQL 查询并执行它。

result = vn_milvus.ask("XYZ Inc 的客户是谁？")
print(result)

SELECT * FROM Customer WHERE Company = 'XYZ Inc' ID Name Company City Phone 0 2 Jane Smith XYZ Inc Los Angeles 098-765-4321

result = vn_milvus.ask("有多少客户？")
print(result)

SELECT COUNT() FROM Customer COUNT() 0 2

太好了！我们已经成功使用 Vanna 和 Milvus 构建了一个文本到 SQL 的生成系统。

快速部署

要了解如何使用此教程启动在线演示，请参阅示例应用程序。

先决条件​

数据准备​

连接到 SQLite 数据库

执行 SQL 查询​

快速部署​

先决条件

数据准备

执行 SQL 查询

快速部署