行为定制

所有 Vanna 函数都继承自 VannaBase 类。这是一个抽象基类，为所有 Vanna 函数提供基本功能。根据您选择的具体配置，实现代码位于继承自该基类的类中。

您可以选择通过创建直接继承自 VannaBase 或继承自其子类的新类来定制 Vanna 的行为。如果您想改变特定函数的行为或添加新功能，这将非常有用。

类实例化

这是一个示例，展示了当配置使用 OpenAI API 和 ChromaDB 向量存储时如何实例化该类。

要定制具体的行为，您可以在实例化 MyVanna 类时覆盖基类中的任何方法。

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-...'})

覆盖特定函数

下面是一个如何覆盖 is_sql_valid 函数的示例。

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)

    def is_sql_valid(self, sql: str) -> bool:
        # Your implementation here

        return False

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

# Example usage
is_valid = vn.is_sql_valid("SELECT user_name, user_email FROM users WHERE user_id = 123")
print(f"Is the SQL valid? {is_valid}")

添加额外的基于 LLM 的函数

如果您想添加一个使用 LLM 的新函数，可以通过向类中添加新方法来实现。假设您想“解释”一个 SQL 查询。您可以向类中添加一个新方法，该方法使用 LLM 为该 SQL 查询生成解释。

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)

    def generate_query_explanation(self, sql: str):
        my_prompt = [
            self.system_message("You are a helpful assistant that will explain a SQL query"),
            self.user_message("Explain this SQL query: " + sql),
        ]

        return self.submit_prompt(prompt=my_prompt)

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

vn.generate_query_explanation("SELECT user_name, user_email FROM users WHERE user_id = 123")

输出： '此 SQL 查询正在从 users 表中选择 user_name 和 user_email 列。它使用 WHERE 子句指定了一个条件，其中 user_id 列必须等于 123。换句话说，它正在检索 user_id 为 123 的用户的 user_name 和 user_email。'