llamaindex实战-使用本地大模型和数据库对话_llamaindex结合本地模型

概述

本文使用NLSQLTableQueryEngine 查询引擎来构建SQL的自然语言处理查询。

请注意，我们需要指定要与该查询引擎一起使用的表。如果我们不这样做，查询引擎将提取所有架构上下文，这可能会溢出 LLM 的上下文窗口。

在以下情况都可以使用NLSQLTableQueryEngine查询引擎：只要您可以预先指定要查询的表，或者所有表结构的总大小加上提示的其余部分，不会超出提示词返回内容的大小。

数据准备

在mysql中创建一张表，并插入一些数据：

 CREATE TABLE `city_stats` (
  `city_name` varchar(16) DEFAULT NULL,
  `population` int(11) DEFAULT NULL,
  `country` varchar(16) DEFAULT NULL
) ENGINE=InnoDB DEFAULT CHARSET=latin1；
1
2
3
4
5

把以下数据插入到city_stats表中：

rows = [ 
    {"city_name": "Toronto", "population": 2930000, "country": "Canada"},
    {"city_name": "Tokyo", "population": 13960000, "country": "Japan"},
    { "city_name": "Chicago", "population": 2679000, "country": "United States" },  
    {"city_name": "Seoul", "population": 9776000, "country": "South Korea"},
]
1
2
3
4
5
6

一旦我们构建了 SQL 数据库，我们就可以使用 NLSQLTableQueryEngine 来构建合成为 SQL 查询的自然语言查询。

申明openai的key

在openai.com官网上申请一个key，并设置环境变量：

export OPENAI_API_KEY="xxxxxx"
1

注意：这可能是llamaindex的一个bug。因为我并没有调用openai的接口，我的机器也无法访问openai网站，但代码提示还是需要让我设置一个OPENAI_API_KEY环境变量。

实现逻辑

（1）通过llama_index的create_engine函数来创建数据库查询引擎

（2）使用SQLDatabase对象来封装查询引擎，通过该对象的封装，就可以把数据库查询引擎带入自然语言处理查询引擎了。

（3）构建大模型对象，可以使用Ollama，或使用其他方式也都可以。

（4）构建自然语言处理查询引擎NLSQLTableQueryEngine，并把数据库查询引擎带入其中。

（5）和NLSQLTableQueryEngine查询引擎进行交互，通过自然语言来和数据库中的数据对话。

完整代码

以下代码通过自然语言的方式和数据对话。

from llama_index.core import SQLDatabase
from llama_index.llms.ollama import Ollama

from llama_index.core.query_engine import NLSQLTableQueryEngine

from sqlalchemy import (
    create_engine,
    select,
)
from sqlalchemy import insert


## prepare database and data 
engine = create_engine("mysql+pymysql://root:@172.16.1.54/llmdb")

# 准备数据
sql_database = SQLDatabase(engine, include_tables=["city_stats"])

# 创建大模型
llm = Ollama(model="llama3", request_timeout=360.0)

# 构建查询引擎
query_engine = NLSQLTableQueryEngine(
    sql_database=sql_database, tables=["city_stats"], llm=llm
)

query_str = "Which city has the highest population?"

response = query_engine.query(query_str)

#display(Markdown(f"<b>{response}</b>"))
print(response)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

小结

本文只是实现了：通过本地大模型和数据库对话的基本功能。然而，通过大模型和数据库对话的真正的难点是如何适配不同的业务。让大模型可以理解业务的逻辑和需求，然后根据业务需要分解计算步骤，并计算结果。

参考文档

• Text-to-SQL Guide (Query Engine + Retriever)