Vanna项目总结
Vanna是一个开源的Python RAG(检索增强生成)框架,专注于SQL生成和相关功能,主要用于自然语言转SQL查询。
目录结构
核心组件:
src/vanna/base/base.py: 核心抽象基类,定义了整个框架的接口
src/vanna/__init__.py: 公共API和入口点
src/vanna/remote.py: 远程API连接
模块化组织:
LLM模型实现: 支持多种大型语言模型
openai/: OpenAI接口
anthropic/: Anthropic Claude接口
google/: Google Gemini接口
ollama/: Ollama接口
以及其他LLM接口
向量存储实现: 支持多种向量数据库
chromadb/: ChromaDB接口
pgvector/: PostgreSQL向量接口
pinecone/: Pinecone接口
以及其他向量数据库接口
数据库连接器: 支持多种数据库
包括PostgreSQL, MySQL, Snowflake, ClickHouse等
核心逻辑
VannaBase基类:
定义了核心的RAG流程和接口
包含训练模型、生成SQL、运行SQL等基本方法
使用抽象方法要求子类实现特定功能
典型工作流程:
训练阶段:
3. 用户提供SQL示例、DDL定义或业务文档
4. 系统使用向量存储生成和存储嵌入
查询阶段:
3. 用户提出自然语言问题
4. 系统检索相关的训练数据(SQL、DDL、文档)
5. 生成适当的提示并提交给LLM
6. LLM生成目标SQL查询
7. 可选:系统执行SQL并生成可视化
扩展性设计**:
模块化设计允许用户选择不同的LLM和向量存储组合
通过继承适当的实现类来自定义功能
关键类和接口
VannaBase:
核心抽象基类,定义RAG框架的基本结构
实现高级功能如SQL生成、数据库连接等
各种LLM实现:
例如OpenAI_Chat:实现与OpenAI API的交互
提供submit_prompt方法发送提示并获取响应
各种向量存储实现:
例如ChromaDB_VectorStore:实现训练数据的存储和检索
负责嵌入生成、数据添加与检索
远程API接口:
VannaDefault:允许用户连接到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-4-...'})
训练模型:
vn.train(ddl="CREATE TABLE IF NOT EXISTS my-table (...)")
vn.train(documentation="业务定义...")
vn.train(sql="SELECT * FROM table...")
提问并生成SQL:
vn.ask("销售额前10名的客户是谁?")
数据库连接:
vn.connect_to_postgres(host="...", dbname="...", user="...", password="...")
技术亮点
多种LLM支持: 可扩展到多种大语言模型,包括开源和商业模型
灵活的向量存储: 支持多种向量数据库,满足不同需求
广泛的数据库支持: 支持常见的SQL数据库
自动可视化: 生成Plotly图表以可视化查询结果
安全隐私: 数据库内容不会发送到LLM或向量数据库
这个框架的核心功能是使用检索增强生成(RAG)技术,将自然语言问题转换为SQL查询,有助于非技术人员通过自然语言与数据库交互。
1、我的问题是这样的,它这个框架启动阶段可以给vn导入多个DDL的文档定义,我使用的时候是每一张表,一个文件的形式插入给vn的,我想知道,它在生成sql的阶段,怎样使用的这些训练数据?是粗暴的把所有的DDL都插入了上下文还是?
我来解释一下Vanna框架在生成SQL阶段如何使用训练数据的。
当你通过导入多个DDL文件给vn进行训练时,Vanna并不是简单地把所有DDL都粗暴地插入上下文。它使用了更加智能的检索增强生成(RAG)方法。
基于源代码分析,我来解释Vanna框架在生成SQL阶段是如何使用训练数据(尤其是DDL定义)的:
Vanna处理DDL训练数据的工作流程
存储阶段:
当你通过vn.train(ddl="CREATE TABLE...") 或者通过文件形式导入DDL时,每个DDL都会被添加到向量数据库中(如ChromaDB)
DDL被转换为向量嵌入(embeddings)并存储在专门的"ddl_collection"集合中
每个DDL文档都有一个唯一ID,用于后续检索
检索阶段:
当你提出自然语言问题时,Vanna执行get_related_ddl(question)方法
这个方法使用语义相似度搜索,从所有存储的DDL中找出与当前问题最相关的子集
默认情况下,会检索前N个相关DDL(在ChromaDB实现中,默认是10个,由n_results_ddl参数控制)
上下文构建阶段:
在generate_sql方法中,系统会收集这些相关DDL
通过add_ddl_to_prompt方法,系统会将这些相关DDL添加到提示中
重要的是,系统不会将所有DDL都添加到上下文中,而是仅添加与当前问题相关的DDL
同时,系统还会考虑token限制(默认14000),确保不超过LLM的上下文限制
提示构建:
在get_sql_prompt方法中,构建最终发送给LLM的提示
提示包含以下部分:
初始指导(告诉LLM如何生成SQL)
相关DDL表定义(仅包含与问题相关的表)
相关文档(如果有)
相似问题和SQL示例(如果有)
Token控制:
如果所有相关DDL的总长度超过了token限制,系统会优先添加最相关的DDL,直到接近token限制
核心代码解析
在generate_sql方法中:
def generate_sql(self, question: str, allow_llm_to_see_data=False, **kwargs) -> str:
question_sql_list = self.get_similar_question_sql(question, **kwargs)
ddl_list = self.get_related_ddl(question, **kwargs)
doc_list = self.get_related_documentation(question, **kwargs)
prompt = self.get_sql_prompt(
initial_prompt=initial_prompt,
question=question,
question_sql_list=question_sql_list,
ddl_list=ddl_list,
doc_list=doc_list,
**kwargs,
)
# ...后续处理...
DDL检索方法(以ChromaDB实现为例):
def get_related_ddl(self, question: str, **kwargs) -> list:
# 进行相似度搜索,找出与问题最相关的DDL
results = self.ddl_collection.query(
query_embeddings=self.generate_embedding(question),
n_results=self.n_results_ddl,
)
# 返回相关DDL列表
return results["documents"] if "documents" in results else []
提示构建过程中添加DDL:
def add_ddl_to_prompt(self, initial_prompt: str, ddl_list: list[str], max_tokens: int = 14000) -> str:
if len(ddl_list) > 0:
initial_prompt += "\n===Tables \n"
for ddl in ddl_list:
# 检查是否超出token限制
if (self.str_to_approx_token_count(initial_prompt) + self.str_to_approx_token_count(ddl) < max_tokens):
initial_prompt += f"{ddl}\n\n"
return initial_prompt
总结
Vanna不是简单地把所有DDL都丢进上下文,而是通过智能检索方式:
仅选择与当前问题最相关的表结构定义
在token限制内尽可能包含更多相关信息
优先包含最相关的DDL
这种方法有几个明显优势:
即使有上百张表,也只会选择最相关的几张表
大大减少了上下文长度,提高LLM处理效率
降低了干扰信息,提高了SQL生成的准确性
这样设计可以确保即使你导入了大量表结构,系统也能够智能地筛选,只把最相关的表结构送入LLM上下文中进行处理。
2、那么第二个问题是,它怎样把业务定义的相关文档做的召回,又是怎样插入到生成阶段的提示词里面去的?
我来解析Vanna如何处理业务定义文档的召回和插入到生成阶段的提示词中:
业务文档处理流程
基于代码分析,下面是Vanna如何处理业务定义文档进行召回和插入到提示词的详细解析:
业务文档的存储和召回流程
存储阶段:
通过vn.train(documentation="业务定义内容")或类似方法,业务定义文档被添加到系统
每个文档会被转换为向量嵌入(embeddings)并存储在专门的"documentation_collection"集合中
代码实现在add_documentation方法中,它会:
生成一个唯一ID(使用deterministc_uuid)
创建文档的嵌入表示
将文档和嵌入存储到向量数据库(如ChromaDB)
召回阶段:
当用户提出问题时,系统会调用get_related_documentation(question)方法
这个方法使用语义相似度搜索,从所有存储的业务文档中找出与当前问题最相关的子集
在ChromaDB实现中,调用documentation_collection.query方法,默认返回前N个相关文档(由n_results_documentation参数控制,默认为10)
代码实现:
def get_related_documentation(self, question: str, **kwargs) -> list:
return ChromaDB_VectorStore._extract_documents(
self.documentation_collection.query(
query_texts=[question],
n_results=self.n_results_documentation,
)
)
提示构建阶段:
在generate_sql方法中,收集了相关的业务定义文档
通过add_documentation_to_prompt方法将这些文档添加到提示中
与DDL处理类似,系统也会考虑token限制,确保不超过LLM的上下文窗口大小
业务文档以"===Additional Context"的标题添加到提示中,清晰地分隔不同类型的信息
最终提示构建:
在get_sql_prompt中,构建发送给LLM的完整提示
注意这里有一个特殊机制:static_documentation字段
if self.static_documentation != "":
doc_list.append(self.static_documentation)
这允许添加一些始终包含的文档内容,不需要每次都通过相似度检索
关键代码解析
业务文档添加到提示的实现:
def add_documentation_to_prompt(
self,
initial_prompt: str,
documentation_list: list[str],
max_tokens: int = 14000,
) -> str:
if len(documentation_list) > 0:
initial_prompt += "\n===Additional Context \n\n"
for documentation in documentation_list:
# 检查是否超出token限制
if (
self.str_to_approx_token_count(initial_prompt)
+ self.str_to_approx_token_count(documentation)
< max_tokens
):
initial_prompt += f"{documentation}\n\n"
return initial_prompt
整合到SQL生成流程:
def get_sql_prompt(
self,
initial_prompt : str,
question: str,
question_sql_list: list,
ddl_list: list,
doc_list: list,
**kwargs,
):
# 初始提示词
if initial_prompt is None:
initial_prompt = f"You are a {self.dialect} expert. " + \
"Please help to generate a SQL query to answer the question..."
# 添加表结构信息
initial_prompt = self.add_ddl_to_prompt(
initial_prompt, ddl_list, max_tokens=self.max_tokens
)
# 添加静态文档和检索到的相关文档
if self.static_documentation != "":
doc_list.append(self.static_documentation)
initial_prompt = self.add_documentation_to_prompt(
initial_prompt, doc_list, max_tokens=self.max_tokens
)
# 后续添加其他内容...
业务文档召回和使用的特点
智能选择:
与DDL类似,业务文档也只选择与当前问题最相关的内容
使用向量相似度搜索,确保上下文中的文档与问题相关
结构化插入:
业务文档在提示中被明确标记为"Additional Context"
这帮助LLM理解这些内容的性质和如何使用它们
token优化:
实现了智能的token管理,确保不超过LLM的上下文限制
如果所有相关文档无法完全放入上下文,会优先保留最相关的
静态文档支持:
通过static_documentation字段,可以添加始终包含的业务定义
这对于一些关键的、需要在所有查询中考虑的业务规则很有用
实际流程示例
假设你有以下业务文档:
"销售额定义为产品单价乘以数量"
"VIP客户是指购买金额超过10万元的客户"
"区域经理负责特定地理区域的所有销售活动"
当用户提问"哪些VIP客户的销售额超过20万?"时:
系统会计算问题的嵌入向量
查找最相似的业务文档(这里可能是关于VIP客户和销售额的定义)
将这些文档添加到提示中的"Additional Context"部分
同时,系统会检索相关的表结构定义和SQL示例
构建完整提示发送给LLM
LLM根据问题、相关业务文档和表结构生成SQL查询
这种方法确保了生成的SQL不仅语法正确,而且符合业务定义和规则,提高了查询的准确性和业务相关性。
3、第三个问题:它似乎还有将最佳QA对儿,添加给生成阶段的prompts的过程,这个过程又是怎样的?用户问题与什么东西在做相似度搜索?
是的,Vanna确实有将最佳问答对(QA pairs)添加到生成阶段的提示中。我来详细分析这个过程:
问答对(QA pairs)的存储和召回流程
基于代码分析,下面是Vanna如何处理问答对(QA pairs)的召回和添加到提示词的详细解析:
问答对的存储和召回流程
存储阶段:
通过vn.train(question="问题", sql="SQL查询")或者vn.add_question_sql(question, sql)方法,问答对被添加到系统
系统将问题和对应的SQL查询转换为JSON格式,然后生成向量嵌入并存储在"sql_collection"集合中
代码实现在add_question_sql方法中:
def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
question_sql_json = json.dumps(
{
"question": question,
"sql": sql,
},
ensure_ascii=False,
)
id = deterministic_uuid(question_sql_json) + "-sql"
self.sql_collection.add(
documents=question_sql_json,
embeddings=self.generate_embedding(question_sql_json), # 注意这里是对问题和SQL的组合生成嵌入
ids=id,
)
return id
召回阶段:
当用户提出新问题时,系统会调用get_similar_question_sql(question)方法
这个方法使用语义相似度搜索,从所有存储的问答对中找出与当前问题最相似的
检索是基于问题和SQL组合的嵌入进行的,而不仅仅是问题本身
在ChromaDB实现中,默认检索前N个相似问答对(由n_results_sql参数控制,默认为10)
代码实现:
def get_similar_question_sql(self, question: str, **kwargs) -> list:
return ChromaDB_VectorStore._extract_documents(
self.sql_collection.query(
query_texts=[question],
n_results=self.n_results_sql,
)
)
提示构建阶段:
在generate_sql方法中,收集相似的问答对
通过add_sql_to_prompt方法将这些问答对添加到提示中
问答对以"===Question-SQL Pairs"的标题添加到提示中
同样会考虑token限制,确保不超过LLM的上下文窗口大小
代码实现:
def add_sql_to_prompt(self, initial_prompt: str, sql_list: list[str], max_tokens: int = 14000) -> str:
if len(sql_list) > 0:
initial_prompt += "\n===Question-SQL Pairs\n\n"
for question in sql_list:
if (self.str_to_approx_token_count(initial_prompt) + self.str_to_approx_token_count(question["sql"]) < max_tokens):
initial_prompt += f"{question['question']}\n{question['sql']}\n\n"
return initial_prompt
融入对话格式:
在最终构建LLM的prompt时,系统不仅添加了问答对,还将它们转换为对话格式
这是在get_sql_prompt方法中完成的:
message_log = [self.system_message(initial_prompt)]
for example in question_sql_list:
if example is not None and "question" in example and "sql" in example:
message_log.append(self.user_message(example["question"]))
message_log.append(self.assistant_message(example["sql"]))
message_log.append(self.user_message(question))
这种对话格式可以帮助LLM更好地理解问题和SQL之间的对应关系
相似度搜索的关键点
相似度搜索对象:
用户的新问题与存储的问题和SQL组合的嵌入进行相似度比较
在add_question_sql方法中,嵌入是基于JSON格式的问题和SQL的组合生成的
这意味着搜索不仅考虑问题的表述相似性,还考虑整体问答对的相似性
模型行为的影响:
这种方法有助于找到在表述和SQL结构上都相似的例子
对于类似的业务问题,即使表述略有不同,也能找到相关的SQL示例
例子的存储形式:
问答对是作为JSON字符串存储的:{"question": "问题", "sql": "SQL查询"}
检索时会解析这个JSON,提取出问题和SQL部分
融入生成过程的细节
提示中的示例格式:
===Question-SQL Pairs
问题1
SQL查询1
问题2
SQL查询2
对话形式的转换:
系统将这些问答对转换为LLM对话格式:
system_message: "你是SQL专家..."
user_message: "问题1"
assistant_message: "SQL查询1"
user_message: "问题2"
assistant_message: "SQL查询2"
user_message: "当前用户问题"
这种对话形式帮助LLM学习如何回应类似问题
示例的内隐学习:
通过提供相似的问答对,LLM可以从中学习:
特定业务术语到SQL结构的映射
表的连接方式和过滤条件的表达
聚合函数和分组的使用模式
完整流程示例
假设用户之前训练了以下问答对:
问:"谁是销售额最高的客户?" - SQL:"SELECT customer_name, SUM(sales) AS total_sales FROM sales GROUP BY customer_name ORDER BY total_sales DESC LIMIT 1"
问:"2022年总销售额是多少?" - SQL:"SELECT SUM(amount) FROM sales WHERE YEAR(sale_date) = 2022"
当用户提问"销售额前10名的客户是谁?"时:
系统计算这个问题的嵌入向量
搜索最相似的问答对(这里可能是问题1,因为都与"销售额"和"客户"有关)
将这些问答对添加到提示中的"Question-SQL Pairs"部分
转换为对话格式:系统消息 -> 用户问题1 -> 助手SQL1 -> 当前用户问题
同时添加相关的表结构和业务文档
LLM根据这些示例和上下文生成新的SQL,可能是:
SELECT customer_name, SUM(sales) AS total_sales
FROM sales
GROUP BY customer_name
ORDER BY total_sales DESC
LIMIT 10
这种方法特别有效,因为它基于少样本学习(few-shot learning),让LLM通过看到几个相似例子来学习如何将特定风格的问题转换为SQL。这比单纯依赖LLM自身的SQL生成能力要更加精确和一致。
4、第四个问题,以上三个过程其实就是vanna的核心了流程了对吧?那么你能为我模拟构建出来一个最终生成阶段,LLM会接受到的提示词么?包含所有种类的vanna的训练数据的最终的样子,这样,我再构建我自己的TEXT2SQL的agent的时候就可以参考了;
是的,您总结得很准确,这三个过程确实构成了Vanna的核心流程。我将为您模拟一个完整的最终生成阶段LLM会接收到的提示词,包含所有类型的训练数据。
Vanna最终生成阶段的完整提示词模拟
你是一位SQL专家。请帮助生成SQL查询来回答问题。你的回复应该仅基于给定的上下文,并遵循响应指南和格式说明。
===Tables
CREATE TABLE customers (
customer_id INT PRIMARY KEY,
customer_name VARCHAR(100),
email VARCHAR(100),
phone VARCHAR(20),
address VARCHAR(200),
city VARCHAR(50),
state VARCHAR(50),
zip_code VARCHAR(20),
registration_date DATE,
vip_status BOOLEAN
);
CREATE TABLE products (
product_id INT PRIMARY KEY,
product_name VARCHAR(100),
category VARCHAR(50),
price DECIMAL(10, 2),
cost DECIMAL(10, 2),
supplier_id INT,
stock_quantity INT
);
CREATE TABLE orders (
order_id INT PRIMARY KEY,
customer_id INT,
order_date DATE,
total_amount DECIMAL(12, 2),
status VARCHAR(20),
FOREIGN KEY (customer_id) REFERENCES customers(customer_id)
);
CREATE TABLE order_items (
order_item_id INT PRIMARY KEY,
order_id INT,
product_id INT,
quantity INT,
price DECIMAL(10, 2),
discount DECIMAL(10, 2),
FOREIGN KEY (order_id) REFERENCES orders(order_id),
FOREIGN KEY (product_id) REFERENCES products(product_id)
);
===Additional Context
VIP客户定义:VIP客户是指满足以下任一条件的客户:(1)lifetime_value超过10,000元的客户;(2)在过去12个月内购买次数超过24次的客户;(3)被销售经理手动标记为VIP的客户。VIP客户在customer表中通过vip_status字段标记。
销售额计算规则:销售额等于订单项的价格乘以数量再减去折扣金额。即sales = price * quantity - discount。公司总销售额是所有订单项销售额的总和。
产品分类标准:我们的产品分为以下几类:电子产品、家居用品、服装、食品、书籍。每个产品只能属于一个类别,存储在products表的category字段中。
订单状态定义:订单状态包括:已下单(placed)、已付款(paid)、已发货(shipped)、已完成(completed)、已取消(cancelled)。只有状态为"已完成"的订单才计入实际销售业绩。
会计年度规定:公司的会计年度从每年的4月1日开始,到次年的3月31日结束。所有按年统计的报表都应该使用这个会计年度来计算。
===Question-SQL Pairs
查询销售额最高的5个客户以及他们的订单数量
SELECT c.customer_name,
COUNT(DISTINCT o.order_id) AS order_count,
SUM(oi.price * oi.quantity - oi.discount) AS total_sales
FROM customers c
JOIN orders o ON c.customer_id = o.customer_id
JOIN order_items oi ON o.order_id = oi.order_id
WHERE o.status = 'completed'
GROUP BY c.customer_id, c.customer_name
ORDER BY total_sales DESC
LIMIT 5;
获取销售额最高的产品类别
SELECT p.category,
SUM(oi.price * oi.quantity - oi.discount) AS category_sales
FROM products p
JOIN order_items oi ON p.product_id = oi.product_id
JOIN orders o ON oi.order_id = o.order_id
WHERE o.status = 'completed'
GROUP BY p.category
ORDER BY category_sales DESC
LIMIT 1;
计算每个月的销售额和同比增长率
WITH monthly_sales AS (
SELECT
DATE_FORMAT(o.order_date, '%Y-%m') AS month,
SUM(oi.price * oi.quantity - oi.discount) AS sales
FROM orders o
JOIN order_items oi ON o.order_id = oi.order_id
WHERE o.status = 'completed'
GROUP BY DATE_FORMAT(o.order_date, '%Y-%m')
),
sales_with_prev_year AS (
SELECT
month,
sales,
LAG(sales, 12) OVER (ORDER BY month) AS prev_year_sales
FROM monthly_sales
)
SELECT
month,
sales,
prev_year_sales,
CASE
WHEN prev_year_sales IS NULL THEN NULL
ELSE ROUND((sales - prev_year_sales) / prev_year_sales * 100, 2)
END AS growth_rate
FROM sales_with_prev_year
ORDER BY month DESC;
查找从未购买过电子产品的VIP客户
SELECT c.customer_id, c.customer_name, c.email
FROM customers c
WHERE c.vip_status = TRUE
AND c.customer_id NOT IN (
SELECT DISTINCT o.customer_id
FROM orders o
JOIN order_items oi ON o.order_id = oi.order_id
JOIN products p ON oi.product_id = p.product_id
WHERE p.category = '电子产品'
AND o.status = 'completed'
);
===Response Guidelines
1. 如果提供的上下文足够,请为问题生成有效的SQL查询,不要提供任何解释。
2. 如果提供的上下文几乎足够,但需要知道特定列中的具体字符串,请生成中间SQL查询以查找该列中的不同字符串。在查询前加上注释说明intermediate_sql。
3. 如果提供的上下文不足,请解释为什么无法生成查询。
4. 请使用最相关的表。
5. 如果问题之前已经被问过并回答过,请完全按照之前的回答重复。
6. 确保输出的SQL符合SQL标准,可执行,且没有语法错误。
查找2022年购买金额超过5000元的VIP客户,按购买金额降序排列。
这个提示词模拟了Vanna向LLM发送的完整提示,包含了所有三种类型的训练数据。让我为您解释一下各个部分:
提示词结构分析
初始指令:
定义LLM的角色和任务
提供基本指导:基于上下文生成SQL查询
表结构(DDL)部分:
标题:"===Tables"
包含相关表的CREATE TABLE语句
明确表结构、字段和关系
业务定义文档部分:
标题:"===Additional Context"
包含业务规则、术语定义和计算方法
这部分内容帮助LLM理解业务领域知识
问答对部分:
标题:"===Question-SQL Pairs"
包含示例问题和对应的SQL查询
展示了如何将自然语言问题转换为SQL查询
响应指南:
具体指导LLM如何生成响应
处理不同场景(上下文充分/不足)的策略
格式和质量要求
当前问题:
用户实际提出的问题,放在最后
这是LLM需要为其生成SQL的问题
构建自己的TEXT2SQL Agent的参考要点
基于Vanna的这种提示词结构,您在构建自己的TEXT2SQL Agent时可以参考以下几点:
结构化的提示模板:
清晰划分不同类型的上下文信息
使用明确的标题区分各部分内容
智能检索机制:
根据用户问题,动态检索最相关的表结构、业务文档和问答示例
不是简单地包含所有可能的信息,而是选择与当前问题最相关的
多种知识来源融合:
表结构提供技术可能性
业务文档提供领域知识
问答示例提供模式和风格参考
明确的响应指南:
指导LLM如何处理不同情况
定义输出格式和质量标准
Token优化策略:
在有限的上下文窗口中,优先包含最相关的信息
确保不超过LLM的token限制
少样本学习利用:
通过提供几个相关示例,引导LLM学习映射模式
示例应该多样化,覆盖不同类型的SQL查询模式
您可以根据自己的具体需求,调整这个模板的比例和细节。例如,如果您的应用特别依赖复杂的业务规则,可以扩展"Additional Context"部分;如果您有大量高质量的问答对,可以更多地依赖"Question-SQL Pairs"部分。
这种结构化、检索增强的提示词设计,是Vanna成功的关键,也可以作为您构建自己的TEXT2SQL Agent的良好起点。
![[Pasted image 20250324150133.png]]
5、最后一个问题,如果我需要自己阅读这个整个流程的代码,请列出一个文件名列表和核心函数的列表来辅助我的阅读过程
以下是Vanna项目的核心文件和关键函数列表,按照流程顺序排列,帮助您更有效地阅读和理解整个实现过程:
核心文件列表
src/vanna/base/base.py
核心抽象基类,定义了Vanna的基本接口和流程
包含RAG框架的主要逻辑
src/vanna/__init__.py
包的入口点,定义了公共API
包含导出函数和类
src/vanna/{向量存储实现}/
如src/vanna/chromadb/chromadb_vector.py - ChromaDB向量存储实现
如src/vanna/pgvector/pgvector_vector.py - PostgreSQL向量存储实现
负责训练数据的存储和检索
src/vanna/{LLM实现}/
如src/vanna/openai/openai_chat.py - OpenAI模型实现
如src/vanna/anthropic/anthropic_chat.py - Anthropic模型实现
负责与LLM交互生成SQL
src/vanna/remote.py
远程API调用的实现
src/vanna/types.py
定义了项目中使用的数据类型
src/vanna/utils.py
实用函数集合
核心函数列表(按流程顺序)
1. 训练阶段函数
文件
函数名
说明
base.py
train()
统一的训练入口,处理各种类型的训练数据
base.py
generate_question()
从SQL生成自然语言问题(用于反向训练)
chromadb_vector.py
add_ddl()
添加DDL定义到向量存储
chromadb_vector.py
add_documentation()
添加业务文档到向量存储
chromadb_vector.py
add_question_sql()
添加问题-SQL对到向量存储
chromadb_vector.py
generate_embedding()
生成文本的嵌入向量
2. 检索阶段函数
文件
函数名
说明
base.py
generate_sql()
生成SQL的主流程
chromadb_vector.py
get_related_ddl()
检索与问题相关的DDL
chromadb_vector.py
get_related_documentation()
检索与问题相关的业务文档
chromadb_vector.py
get_similar_question_sql()
检索与问题相似的问答对
chromadb_vector.py
_extract_documents()
从查询结果中提取文档
3. 提示词构建函数
文件
函数名
说明
base.py
get_sql_prompt()
构建发送给LLM的完整提示
base.py
add_ddl_to_prompt()
将相关DDL添加到提示中
base.py
add_documentation_to_prompt()
将相关业务文档添加到提示中
base.py
add_sql_to_prompt()
将相似问答对添加到提示中
base.py
str_to_approx_token_count()
估算文本的token数量
4. LLM交互函数
文件
函数名
说明
openai_chat.py
submit_prompt()
将提示发送给LLM并获取响应
openai_chat.py
system_message()
创建系统消息格式
openai_chat.py
user_message()
创建用户消息格式
openai_chat.py
assistant_message()
创建助手消息格式
5. SQL处理和执行函数
文件
函数名
说明
base.py
extract_sql()
从LLM响应中提取SQL
base.py
is_sql_valid()
验证SQL是否有效
base.py
run_sql()
执行SQL查询
base.py
connect_to_*()
连接到各种数据库的方法
base.py
generate_plotly_code()
生成可视化代码
6. 公共API和工具函数
文件
函数名
说明
__init__.py
ask()
统一的问答接口
utils.py
deterministic_uuid()
生成确定性UUID
utils.py
validate_config_path()
验证配置路径
阅读顺序建议
首先查看 base.py 中的 VannaBase 类,了解整体架构和抽象接口
查看训练相关函数:train(), add_ddl(), add_documentation(), add_question_sql()
理解SQL生成流程:generate_sql(), get_sql_prompt(), submit_prompt()
探索向量检索机制:get_related_ddl(), get_related_documentation(), get_similar_question_sql()
学习提示词构建:add_ddl_to_prompt(), add_documentation_to_prompt(), add_sql_to_prompt()
最后了解SQL处理和执行:extract_sql(), run_sql()
关键代码片段
对于核心流程理解,以下代码片段尤为重要:
generate_sql() - SQL生成的主流程
def generate_sql(self, question: str, allow_llm_to_see_data=False, **kwargs) -> str:
question_sql_list = self.get_similar_question_sql(question, **kwargs)
ddl_list = self.get_related_ddl(question, **kwargs)
doc_list = self.get_related_documentation(question, **kwargs)
prompt = self.get_sql_prompt(
initial_prompt=initial_prompt,
question=question,
question_sql_list=question_sql_list,
ddl_list=ddl_list,
doc_list=doc_list,
**kwargs,
)
llm_response = self.submit_prompt(prompt, **kwargs)
return self.extract_sql(llm_response)
get_sql_prompt() - 提示词构建
def get_sql_prompt(self, initial_prompt, question, question_sql_list, ddl_list, doc_list, **kwargs):
if initial_prompt is None:
initial_prompt = f"You are a {self.dialect} expert..."
initial_prompt = self.add_ddl_to_prompt(initial_prompt, ddl_list, max_tokens=self.max_tokens)
if self.static_documentation != "":
doc_list.append(self.static_documentation)
initial_prompt = self.add_documentation_to_prompt(initial_prompt, doc_list, max_tokens=self.max_tokens)
initial_prompt = self.add_sql_to_prompt(initial_prompt, question_sql_list, max_tokens=self.max_tokens)
# 创建消息日志
message_log = [self.system_message(initial_prompt)]
# 添加问答对示例
for example in question_sql_list:
if example is not None and "question" in example and "sql" in example:
message_log.append(self.user_message(example["question"]))
message_log.append(self.assistant_message(example["sql"]))
# 添加当前问题
message_log.append(self.user_message(question))
return message_log
通过这些文件和函数的指引,您应该能够更有效地阅读和理解Vanna的整个代码实现过程。希望这对您构建自己的TEXT2SQL Agent有所帮助!