当前位置：首页 > news >正文

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

news 来源：原创 2024/9/25 17:16:08

概述

本文使用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；

把以下数据插入到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"},
]

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

申明openai的key

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

export OPENAI_API_KEY="xxxxxx"

注意：这可能是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 Ollamafrom llama_index.core.query_engine import NLSQLTableQueryEnginefrom 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)