vanna

源代码

在 GitHub 上查看源代码: https://github.com/vanna-ai/vanna

基本用法

获取 API 密钥

import vanna as vn
api_key = vn.get_api_key('my-email@example.com')
vn.set_api_key(api_key)

设置模型

vn.set_model('chinook')

提问

vn.ask(question='What are the top 10 artists by sales?')

vn.ask(...) 是对 vn.generate_sql(...)、vn.run_sql(...)、vn.generate_plotly_code(...)、vn.get_plotly_figure(...) 和 vn.generate_followup_questions(...) 的便捷封装。

要查看可以提问的可运行笔记本，请参见此处

训练

您可以添加到模型中的训练数据主要有 3 种类型：SQL、DDL 和文档。

# DDL Statements
vn.train(ddl='CREATE TABLE employees (id INT, name VARCHAR(255), salary INT)')

# Documentation
vn.train(documentation='Our organization\'s definition of sales is the discount price of an item multiplied by the quantity sold.')

# SQL
vn.train(sql='SELECT AVG(salary) FROM employees')

vn.train(...) 是对 vn.add_sql(...)、vn.add_ddl(...) 和 vn.add_documentation(...) 的便捷封装。

要查看可以训练模型的可运行笔记本，请参见此处

命名规范

前缀	定义	示例
`vn.set_`	设置当前会话的变量	[`vn.set_model(...)`][set_model] [`vn.set_api_key(...)`][set_api_key]
`vn.get_`	执行只读操作	[`vn.get_model()`][get_models]
`vn.add_`	向模型添加内容	[`vn.add_sql(...)`][add_sql] [`vn.add_ddl(...)`][add_ddl]
`vn.generate_`	基于模型中的信息使用 AI 生成内容	[`vn.generate_sql(...)`][generate_sql] [`vn.generate_explanation()`][generate_explanation]
`vn.run_`	运行代码（SQL 或 Plotly）	[`vn.run_sql`][run_sql]
`vn.remove_`	从模型中移除内容	[`vn.remove_training_data`][remove_training_data]
`vn.update_`	更新模型中的内容	[`vn.update_model_visibility(...)`][update_model_visibility]
`vn.connect_`	连接到数据库	[`vn.connect_to_snowflake(...)`][connect_to_snowflake]

权限

默认情况下，创建的模型是私有的。您可以向模型添加成员或管理员，或将其设为公开。

用户角色	公开模型		私有模型
用户角色	使用	训练	使用	训练
非成员	✅	❌	❌	❌
成员	✅	❌	✅	❌
管理员	✅	✅	✅	✅

开源与扩展

Vanna.AI 是开源且可扩展的。如果您想在没有服务器的情况下使用 Vanna，请参见此处的示例。

以下是使用 Vanna 默认“本地”版本时，代码库中各种函数实现位置的示例。vanna.base.VannaBase 是基类，它提供了 vanna.base.VannaBase.ask 和 vanna.base.VannaBase.train 函数。这些函数依赖于在子类 vanna.openai_chat.OpenAI_Chat 和 vanna.chromadb_vector.ChromaDB_VectorStore 中实现的抽象方法。vanna.openai_chat.OpenAI_Chat 使用 OpenAI API 生成 SQL 和 Plotly 代码。vanna.chromadb_vector.ChromaDB_VectorStore 使用 ChromaDB 存储训练数据并生成嵌入。

如果您想将 Vanna 与其他 LLM 或数据库一起使用，可以创建自己的 vanna.base.VannaBase 子类并实现抽象方法。

flowchart
    subgraph VannaBase
        ask
        train
    end

    subgraph OpenAI_Chat
        get_sql_prompt
        submit_prompt
        generate_question
        generate_plotly_code
    end

    subgraph ChromaDB_VectorStore
        generate_embedding
        add_question_sql
        add_ddl
        add_documentation
        get_similar_question_sql
        get_related_ddl
        get_related_documentation
    end

API 参考

api_key: Optional[str] = None

fig_as_img: bool = False

run_sql: Optional[Callable[[str], pandas.core.frame.DataFrame]] = None

示例

vn.run_sql = lambda sql: pd.read_sql(sql, engine)

为 Vanna.AI 设置 SQL 转 DataFrame 函数。这在 [vn.ask(...)][ask] 函数中使用。除了直接设置，您也可以使用 [vn.connect_to_snowflake(...)][connect_to_snowflake] 来设置。

def get_api_key(email: str, otp_code: Optional[str] = None) -> str:

示例

vn.get_api_key(email="my-email@example.com")

参数

email (str): 用于登录的电子邮件地址。
otp_code (Union[str, None]): 用于登录的一次性密码 (OTP) 代码。如果为 None，则会将 OTP 代码发送到电子邮件地址。

str: API 密钥。

def set_api_key(key: str) -> None:

设置 Vanna.AI 的 API 密钥。

示例

api_key = vn.get_api_key(email="my-email@example.com")
vn.set_api_key(api_key)

参数

key (str): API 密钥。

def get_models() -> List[str]:

示例

models = vn.get_models()

列出用户是成员的模型。

List[str]: 模型名称列表。

def create_model(model: str, db_type: str) -> bool:

示例

vn.create_model(model="my-model", db_type="postgres")

创建新模型。

参数

model (str): 要创建的模型名称。
db_type (str): 模型使用的数据库类型。可以是 "Snowflake"、"BigQuery"、"Postgres" 或其他任何类型。

bool: 如果模型创建成功则为 True，否则为 False。

def add_user_to_model(model: str, email: str, is_admin: bool) -> bool:

示例

vn.add_user_to_model(model="my-model", email="user@example.com")

将用户添加到模型。

参数

model (str): 要将用户添加到的模型名称。
email (str): 要添加的用户电子邮件地址。
is_admin (bool): 用户是否应成为管理员。

bool: 如果用户添加成功则为 True，否则为 False。

def update_model_visibility(public: bool) -> bool:

示例

vn.update_model_visibility(public=True)

设置当前模型的可见性。如果模型可见，任何人都可以看到它。如果不可见，只有模型的成员才能看到它。

参数

public (bool): 模型是否应公开可见。

bool: 如果模型可见性设置成功则为 True，否则为 False。

def set_model(model: str):

设置 Vanna.AI API 使用的模型。

示例

vn.set_model("my-model")

参数

model (str): 要使用的模型名称。

def add_sql(question: str, sql: str, tag: Optional[str] = 'Manually Trained') -> bool:

将问题及其对应的 SQL 查询添加到模型的训练数据中。推荐的调用方式是使用 [vn.train(sql=...)][train]。

示例

vn.add_sql(
    question="What is the average salary of employees?",
    sql="SELECT AVG(salary) FROM employees"
)

参数

question (str): 要存储的问题。
sql (str): 要存储的 SQL 查询。
tag (Union[str, None]): 与问题和 SQL 查询关联的标签。

bool: 如果问题和 SQL 查询存储成功则为 True，否则为 False。

def add_ddl(ddl: str) -> bool:

将 DDL 语句添加到模型的训练数据中

示例

vn.add_ddl(
    ddl="CREATE TABLE employees (id INT, name VARCHAR(255), salary INT)"
)

参数

ddl (str): 要存储的 DDL 语句。

bool: 如果 DDL 语句存储成功则为 True，否则为 False。

def add_documentation(documentation: str) -> bool:

将文档添加到模型的训练数据中

示例

vn.add_documentation(
    documentation="Our organization's definition of sales is the discount price of an item multiplied by the quantity sold."
)

参数

documentation (str): 要存储的文档字符串。

bool: 如果文档字符串存储成功则为 True，否则为 False。

@dataclass

class TrainingPlanItem

TrainingPlanItem(item_type: str, item_group: str, item_name: str, item_value: str)

item_type: str

item_group: str

item_name: str

item_value: str

ITEM_TYPE_SQL = 'sql'

ITEM_TYPE_DDL = 'ddl'

ITEM_TYPE_IS = 'is'

class TrainingPlan

一个表示训练计划的类。您可以查看其中的内容，并移除不希望用于训练的项目。

示例

plan = vn.get_training_plan()

plan.get_summary()

TrainingPlan(plan: List[TrainingPlanItem])

def get_summary(self) -> List[str]:

示例

plan = vn.get_training_plan()

plan.get_summary()

获取训练计划的摘要。

List[str]: 描述训练计划的字符串列表。

def remove_item(self, item: str):

示例

plan = vn.get_training_plan()

plan.remove_item("Train on SQL: What is the average salary of employees?")

从训练计划中移除一个项目。

参数

item (str): 要移除的项目。

def get_training_plan_postgres( filter_databases: Optional[List[str]] = None, filter_schemas: Optional[List[str]] = None, include_information_schema: bool = False, use_historical_queries: bool = True) -> TrainingPlan:

def get_training_plan_generic(df) -> TrainingPlan:

def get_training_plan_experimental( filter_databases: Optional[List[str]] = None, filter_schemas: Optional[List[str]] = None, include_information_schema: bool = False, use_historical_queries: bool = True) -> TrainingPlan:

实验性：此方法为实验性方法，未来版本中可能会有所更改。

根据数据库中的元数据获取训练计划。目前这仅适用于 Snowflake。

示例

plan = vn.get_training_plan_experimental(filter_databases=["employees"], filter_schemas=["public"])

vn.train(plan=plan)

def train( question: str = None, sql: str = None, ddl: str = None, documentation: str = None, json_file: str = None, sql_file: str = None, plan: TrainingPlan = None) -> bool:

示例

vn.train()

使用问题及其对应的 SQL 查询来训练 Vanna.AI。如果您不带参数调用它，它会检查您是否连接到数据库，并尝试使用该数据库的元数据进行训练。如果您使用 sql 参数调用它，则等同于调用 [add_sql()][add_sql]。如果您使用 ddl 参数调用它，则等同于调用 [add_ddl()][add_ddl]。如果您使用 documentation 参数调用它，则等同于调用 [add_documentation()][add_documentation]。它还可以接受 JSON 文件路径或 SQL 文件路径，分别用于训练一批问题和 SQL 查询，或训练 SQL 查询列表。此外，您可以传递一个 [TrainingPlan][TrainingPlan] 对象。使用 [vn.get_training_plan_experimental()][get_training_plan_experimental] 获取训练计划。

参数

question (str): 用于训练的问题。
sql (str): 用于训练的 SQL 查询。
sql_file (str): SQL 文件路径。
json_file (str): JSON 文件路径。
ddl (str): DDL 语句。
documentation (str): 用于训练的文档。
plan (TrainingPlan): 用于训练的训练计划。

def flag_sql_for_review( question: str, sql: Optional[str] = None, error_msg: Optional[str] = None) -> bool:

示例

vn.flag_sql_for_review(question="What is the average salary of employees?")

将问题及其对应的 SQL 查询标记为待审核。您可以在 [vn.get_all_questions()][get_all_questions] 中看到该标记。

参数

question (str): 要标记的问题。
sql (str): 要标记的 SQL 查询。
error_msg (str): 要标记的错误消息。

bool: 如果问题和 SQL 查询标记成功则为 True，否则为 False。

def remove_sql(question: str) -> bool:

从模型的训练数据中移除问题及其对应的 SQL 查询

示例

vn.remove_sql(question="What is the average salary of employees?")

参数

question (str): 要移除的问题。

def remove_training_data(id: str) -> bool:

从模型中移除训练数据

示例

vn.remove_training_data(id="1-ddl")

参数

id (str): 要移除的训练数据 ID。

def generate_sql(question: str) -> str:

示例

vn.generate_sql(question="What is the average salary of employees?")
# SELECT AVG(salary) FROM employees

使用 Vanna.AI API 生成 SQL 查询。

参数

question (str): 要为其生成 SQL 查询的问题。

str 或 None: SQL 查询，如果发生错误则为 None。

def get_related_training_data(question: str) -> vanna.types.TrainingData:

示例

training_data = vn.get_related_training_data(question="What is the average salary of employees?")

获取与问题相关的训练数据。

参数

question (str): 要获取相关训练数据的问题。

TrainingData 或 None: 相关训练数据，如果发生错误则为 None。

def generate_meta(question: str) -> str:

示例

vn.generate_meta(question="What tables are in the database?")
# Information about the tables in the database

使用 Vanna.AI API 生成关于数据库元数据的答案。

参数

question (str): 要为其生成答案的问题。

str 或 None: 答案，如果发生错误则为 None。

def generate_followup_questions(question: str, df: pandas.core.frame.DataFrame) -> List[str]:

示例

vn.generate_followup_questions(question="What is the average salary of employees?", df=df)
# ['What is the average salary of employees in the Sales department?', 'What is the average salary of employees in the Engineering department?', ...]

使用 Vanna.AI API 生成后续问题。

参数

question (str): 要为其生成后续问题的问题。
df (pd.DataFrame): 要为其生成后续问题的 DataFrame。

List[str] 或 None: 后续问题列表，如果发生错误则为 None。

def generate_questions() -> List[str]:

示例

vn.generate_questions()
# ['What is the average salary of employees?', 'What is the total salary of employees?', ...]

使用 Vanna.AI API 生成问题。

List[str] 或 None: 问题列表，如果发生错误则为 None。

def ask( question: Optional[str] = None, print_results: bool = True, auto_train: bool = True, generate_followups: bool = True) -> Optional[Tuple[Optional[str], Optional[pandas.core.frame.DataFrame], Optional[plotly.graph_objs._figure.Figure], Optional[List[str]]]]:

示例

# RECOMMENDED IN A NOTEBOOK:
sql, df, fig, followup_questions = vn.ask()


sql, df, fig, followup_questions = vn.ask(question="What is the average salary of employees?")
# SELECT AVG(salary) FROM employees

使用 Vanna.AI API 提问。这将生成一个 SQL 查询并运行它，然后以 dataframe 和 Plotly 图表的形式返回结果。如果您将 print_results 设置为 True，SQL、dataframe 和图表将直接输出到屏幕，而不是作为返回值。

参数

question (str): 要提出的问题。如果为 None，系统将提示您输入问题。
print_results (bool): 是否打印 SQL 查询和结果。
auto_train (bool): 如果 SQL 查询不正确是否自动训练模型。
generate_followups (bool): 是否生成后续问题。

str 或 None: SQL 查询，如果发生错误则为 None。pd.DataFrame 或 None: SQL 查询结果，如果发生错误则为 None。plotly.graph_objs.Figure 或 None: Plotly 图表，如果发生错误则为 None。List[str] 或 None: 后续问题列表，如果发生错误则为 None。

def generate_plotly_code( question: Optional[str], sql: Optional[str], df: pandas.core.frame.DataFrame, chart_instructions: Optional[str] = None) -> str:

示例

vn.generate_plotly_code(
    question="What is the average salary of employees?",
    sql="SELECT AVG(salary) FROM employees",
    df=df
)
# fig = px.bar(df, x="name", y="salary")

使用 Vanna.AI API 生成 Plotly 代码。

参数

question (str): 要为其生成 Plotly 代码的问题。
sql (str): 要为其生成 Plotly 代码的 SQL 查询。
df (pd.DataFrame): 要为其生成 Plotly 代码的 dataframe。
chart_instructions (str): 关于如何绘制图表的可选说明。

str 或 None: Plotly 代码，如果发生错误则为 None。

def get_plotly_figure( plotly_code: str, df: pandas.core.frame.DataFrame, dark_mode: bool = True) -> plotly.graph_objs._figure.Figure:

示例

fig = vn.get_plotly_figure(
    plotly_code="fig = px.bar(df, x='name', y='salary')",
    df=df
)
fig.show()

从 dataframe 和 Plotly 代码获取 Plotly 图表。

参数

df (pd.DataFrame): 要使用的 dataframe。
plotly_code (str): 要使用的 Plotly 代码。

plotly.graph_objs.Figure: Plotly 图表。

def get_results(cs, default_database: str, sql: str) -> pandas.core.frame.DataFrame:

已弃用。请改用 vn.run_sql。运行 SQL 查询并将结果作为 pandas dataframe 返回。这只是一个不使用 Vanna.AI API 的辅助函数。

参数

cs: Snowflake 连接游标。
default_database (str): 要使用的默认数据库。
sql (str): 要执行的 SQL 查询。

pd.DataFrame: SQL 查询结果。

def generate_explanation(sql: str) -> str:

示例

vn.generate_explanation(sql="SELECT * FROM students WHERE name = 'John Doe'")
# 'This query selects all columns from the students table where the name is John Doe.'

使用 Vanna.AI API 生成 SQL 查询的解释。

参数

sql (str): 要为其生成解释的 SQL 查询。

str 或 None: 解释，如果发生错误则为 None。

def generate_question(sql: str) -> str:

示例

vn.generate_question(sql="SELECT * FROM students WHERE name = 'John Doe'")
# 'What is the name of the student?'

使用 Vanna.AI API 从 SQL 查询生成问题。

参数

sql (str): 要从中生成问题的问题。

str 或 None: 问题，如果发生错误则为 None。

def get_all_questions() -> pandas.core.frame.DataFrame:

从 Vanna.AI API 获取问题列表。

示例

questions = vn.get_all_questions()

pd.DataFrame 或 None: 问题列表，如果发生错误则为 None。

def get_training_data() -> pandas.core.frame.DataFrame:

获取当前模型的训练数据

示例

training_data = vn.get_training_data()

pd.DataFrame 或 None: 训练数据，如果发生错误则为 None。

def connect_to_sqlite(url: str):

连接到 SQLite 数据库。这只是一个用于设置 [vn.run_sql][run_sql] 的辅助函数。

参数

url (str): 要连接的数据库 URL。

None

def connect_to_snowflake( account: str, username: str, password: str, database: str, schema: Optional[str] = None, role: Optional[str] = None):

使用 Snowflake 连接器连接到 Snowflake。这只是一个用于设置 [vn.run_sql][run_sql] 的辅助函数。

示例

import snowflake.connector

vn.connect_to_snowflake(
    account="myaccount",
    username="myusername",
    password="mypassword",
    database="mydatabase",
    role="myrole",
)

参数

account (str): Snowflake 账户名称。
username (str): Snowflake 用户名。
password (str): Snowflake 密码。
database (str): 要使用的默认数据库。
schema (Union[str, None], optional): 要使用的 schema。默认为 None。
role (Union[str, None], optional): 要使用的角色。默认为 None。

def connect_to_postgres( host: str = None, dbname: str = None, user: str = None, password: str = None, port: int = None):

使用 psycopg2 连接器连接到 postgres。这只是一个用于设置 [vn.run_sql][run_sql] 的辅助函数。示例:

import psycopg2.connect
vn.connect_to_bigquery(
    host="myhost",
    dbname="mydatabase",
    user="myuser",
    password="mypassword",
    port=5432
)

参数

host (str): postgres 主机。
dbname (str): postgres 数据库名称。
user (str): postgres 用户。
password (str): postgres 密码。
port (int): postgres 端口。

def connect_to_bigquery(cred_file_path: str = None, project_id: str = None):

使用 bigquery 连接器连接到 BigQuery。这只是一个用于设置 [vn.run_sql][run_sql] 的辅助函数。示例:

import bigquery.Client
vn.connect_to_bigquery(
    project_id="myprojectid",
    cred_file_path="path/to/credentials.json",
)

参数

project_id (str): GCS 项目 ID。
cred_file_path (str): GCS 凭据文件路径