概述
本文使用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)
小結
本文只是實現了:通過本地大模型和數據庫對話的基本功能。然而,通過大模型和數據庫對話的真正的難點是如何適配不同的業務。讓大模型可以理解業務的邏輯和需求,然后根據業務需要分解計算步驟,并計算結果。
參考文檔
- Text-to-SQL Guide (Query Engine + Retriever)