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有所帮助!