使用案例

前言

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

一、项目背景

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

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

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

二、案例

1.初始化

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

<code>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"),code>

"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方法改造