Vanna-ai 大模型开源项目 基于RAG的TextToSql框架 安装和使用教程

github项目地址：vanna-ai/vanna: 🤖 与 SQL 数据库聊天📊。通过 LLMs使用RAG进行准确的TextToSQL的生成工作 🔄 。

Vanna 是 MIT 许可的开源 Python RAG（检索增强生成）框架，用于 SQL 生成和相关功能。

用户接口

这些是使用 Vanna 构建的一些用户接口。您可以按原样使用这些接口，也可以将其用作自己的自定义接口的起点。

• Jupyter Notebook
• vanna-ai/vanna-streamlit
• vanna-ai/vanna-flask
• vanna-ai/vanna-slack

开始

有关所需数据库、LLM 等的详细信息，请参阅文档。

如果您想在训练后了解它的工作原理，可以尝试这款 Colab 笔记本。

安装

pip install vanna

根据您的用例，您可能需要也可能不需要运行这些命令。有关详细信息，请参阅文档。

可以先使用如下代码创建默认的Vanna对象：

import vanna
from vanna.remote import VannaDefault
vn = VannaDefault(model='model_name', api_key='api_key')
vn.connect_to_sqlite('https://vanna.ai/Chinook.sqlite')
vn.ask("What are the top 10 albums by sales?")

导入

如果要自定义 LLM 或向量数据库，请参阅文档。

# The import statement will vary depending on your LLM and vector database. This is an example for OpenAI + ChromaDB

from vanna.openai.openai_chat import OpenAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore

class MyVanna(ChromaDB_VectorStore, OpenAI_Chat):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        OpenAI_Chat.__init__(self, config=config)

vn = MyVanna(config={'api_key': 'sk-...', 'model': 'gpt-4-...'})

# See the documentation for other options

样例

import vanna
from vanna.openai.openai_chat import OpenAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore
 
class MyVanna(ChromaDB_VectorStore, OpenAI_Chat):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        OpenAI_Chat.__init__(self, config=config)
 
vn = MyVanna(config={'api_key': 'sk-*************', 'model': 'gpt-3.5-turbo'})
 
# 训练向量库
vn.train(ddl="""
    SELECT customer_name, SUM(sales_amount) as total_sales
FROM sales_wn
GROUP BY customer_name
ORDER BY total_sales DESC
LIMIT 10;
""")
 
vn.train(sql="SELECT name, age FROM my-table WHERE name = 'John Doe'")
 
 
# 询问问题
vn.ask("What are the top 10 customers by sales?")

训练

根据您的用例，您可能需要也可能不需要运行这些命令。有关详细信息，请参阅文档。

显示这些陈述是为了让您了解它的工作原理。

使用 DDL 语句进行训练

DDL 语句包含有关数据库中的表名、列、数据类型和关系的信息。

vn.train(ddl="""
    CREATE TABLE IF NOT EXISTS my-table (
        id INT PRIMARY KEY,
        name VARCHAR(100),
        age INT
    )
""")

带文档的训练

有时，您可能希望添加有关业务术语或定义的文档。

vn.train(documentation="Our business defines XYZ as ...")

使用 SQL 进行训练

您还可以将 SQL 查询添加到训练数据中。如果您已经有一些查询，这将很有用。您只需从编辑器中复制并粘贴这些内容即可开始生成新的 SQL。

vn.train(sql="SELECT name, age FROM my-table WHERE name = 'John Doe'")

提问

vn.ask("What are the top 10 customers by sales?")

你会得到 SQL

SELECT c.c_name as customer_name,
        sum(l.l_extendedprice * (1 - l.l_discount)) as total_sales
FROM   snowflake_sample_data.tpch_sf1.lineitem l join snowflake_sample_data.tpch_sf1.orders o
        ON l.l_orderkey = o.o_orderkey join snowflake_sample_data.tpch_sf1.customer c
        ON o.o_custkey = c.c_custkey
GROUP BY customer_name
ORDER BY total_sales desc limit 10;

如果已连接到数据库，则将获得以下表：

您还将获得一个自动的 Plotly 图表：