VannaAI 介绍及使用 - 第四篇

使用案例

前言

本篇主要记录实际使用的案例，仅供参考，请大家继续关注指正。

一、项目背景

1.项目需要使用Vanna的text-to-sql能力，但是呢gpt3.5能力稍弱一些，所以我们想要使用更强大的gpt4o。

2.因为gpt4o国内使用受限，且涉及公司数据，所以我们是需要另外部门提供的api，通过将请求参数发送到他们的接口，他们再去请求gpt4o，将结果返给我们，所以需要自定义请求方式，不能使用Vanna源码的请求。

3.我们发现在实际使用过程中，用户的有些问题是没法生成sql回答的，那么我们需要大模型能正常回答问题，不止是生成sql，但是呢又要不影响正常sql的生成，那可能涉及对prompt提示词的改造。

 二、案例

1.初始化

这里因为是样例，没有完整项目的代码那么完善，仅作为举例说明展示。

from vanna.openai import OpenAI_Chat
from vanna.vannadb import VannaDB_VectorStore
from openai import OpenAI
import time
import json

def ask(
        question: str,
        user: str,
        password: str,
        host: str,
        port: int,
        database: str,
):
    vn.connect_to_mysql(user=user, password=password, host=host, port=port, dbname=database)
    # 记录开始时间
    start_time = time.time()
    # 调用新的 ask 方法
    sql, data, fig, llm_response = vn.ask(question=question, )
    print(f"调用 vn.ask 的耗时: {time.time() - start_time:.4f} 秒")
    # 处理返回的结果
    if llm_response is not None:
        result = {
            "sql": None,
            "data": llm_response,
            "chart": None,
        }
    else:
        if data is not None:
            image = fig.to_json() if fig else None
            result = {
                "sql": sql,
                "data": data.to_json(orient="records"),
                "chart": image
            }
        else:
            result = {
                "sql": sql,
                "data": None,
                "chart": None
            }

    print(json.dumps(result, ensure_ascii=False))

if __name__ == "__main__":
    # 使用自己的vanna模型及api key
    MY_VANNA_MODEL = "my_vanna_model"
    MY_VANNA_API_KEY = "my_vanna_api_key"

    # 使用自己的llm模型及所需配置
    my_llm_api_key = 'EMPTY' #因为是第三方提供的接口，暂不需要apikey
    my_llm_base_url = 'http://0.0.0.0:8000/xx/xx/xx'  # 第三方提供的完整接口，后续不需要拼接
    my_llm_name = 'gpt4o' #告知第三方我们需要的是gpt4o

    # 初始化客户端
    client = OpenAI(
        api_key=my_llm_api_key ,
        base_url=my_llm_base_url
    )

    class MyVanna(VannaDB_VectorStore, OpenAI_Chat):
        def __init__(self, client=None, config=None):
            VannaDB_VectorStore.__init__(self, vanna_model=MY_VANNA_MODEL, vanna_api_key=MY_VANNA_API_KEY,config=config)
            OpenAI_Chat.__init__(self, client=client, config=config)

    # 使用自定义的大模型及vanna提供的向量库
    vn = MyVanna(client=client, config={"model": my_llm_name, })

    # 自定义问题，用于测试
    question = "中国有哪些省份"
    
    user = "root"
    password = "1234"
    host = "127.0.0.1"
    port = 3306
    database = "自己的库名"
    
    ask(question, user, password, host, port, database)

这里设置一个main是为了模拟传入不同的参数和问题，这个可以自己根据业务需求去调整，比如从数据库中取出相应参数，或者配置文件读取，从前端请求传入问题之类的。这里只是便于测试，简单举例。

还定义了一个ask方法，这是为了方便处理返回结果，其中也做了数据库初始化连接，这里也可以根据自己的业务情况，调整结果处理方式，我这里vn.ask返回四个参数是因为我对Vanna的ask源码做了改动，以为我需要额外返回自然回答的情况。详情见后续的ask代码。

 2.ask方法改造

ask方法在Vanna的base.py文件中，改造后的ask方法如下，直接上代码：

    def ask(
            self,
            question: Union[str, None] = None,
            # session_id: str = None,  # 添加 session_id 参数
            print_results: bool = True,
            #        auto_train: bool = True,
            auto_train: bool = False,  # 关闭默认训练，随着时间的推移，prompt会越来越多，token消耗会越来越快
            visualize: bool = True,  # if False, will not generate plotly code
    ) -> Union[
        Tuple[
            Union[str, None],
            Union[pd.DataFrame, None],
            Union[plotly.graph_objs.Figure, None],
            Union[str, None]
        ],
        None,
    ]:
        """
        **Example:**
        ```python
        vn.ask("What are the top 10 customers by sales?")
        ```

        Ask Vanna.AI a question and get the SQL query that answers it.

        Args:
            question (str): The question to ask.
            print_results (bool): Whether to print the results of the SQL query.
            auto_train (bool): Whether to automatically train Vanna.AI on the question and SQL query.
            visualize (bool): Whether to generate plotly code and display the plotly figure.

        Returns:
            Tuple[str, pd.DataFrame, plotly.graph_objs.Figure]: The SQL query, the results of the SQL query, and the plotly figure.
        """

        # 判断否text2sql对话，默认为是
        is_text2sql = 1

        if question is None:
            question = input("Enter a question: ")

        try:
            sql = self.generate_sql(question=question)
            if self.is_sql_valid(sql) is False:
                return None, None, None, sql
        except Exception as e:
            print(e)
            is_text2sql = 0