使用LlamaIndex创建SQL和向量数据库查询引擎

最新推荐文章于 2024-09-09 09:47:53 发布

qq_37836323

最新推荐文章于 2024-09-09 09:47:53 发布

阅读量306

点赞数 4

文章标签：数据库 sql jvm python

本文链接：https://blog.csdn.net/qq_29929123/article/details/140730857

版权

在现代的AI技术中，查询引擎在处理和分析大数据时起着至关重要的作用。本文将介绍如何使用LlamaIndex创建一个自定义的路由查询引擎，该引擎能够在SQL数据库和向量数据库之间进行路由。我们将通过一个简单的例子来演示这个过程。

环境设置

首先，确保安装必要的Python包。如果你在使用Jupyter Notebook，可以使用以下命令来安装LlamaIndex和其他依赖项：

%pip install llama-index-readers-wikipedia
!pip install llama-index

import nest_asyncio
nest_asyncio.apply()

import logging
import sys

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

创建数据库模式和测试数据

我们将创建一个简单的SQL数据库，并插入一些测试数据。

from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, insert

# 创建内存中的SQLite数据库
engine = create_engine("sqlite:///:memory:", future=True)
metadata_obj = MetaData()

# 创建城市统计表
table_name = "city_stats"
city_stats_table = Table(
    table_name,
    metadata_obj,
    Column("city_name", String(16), primary_key=True),
    Column("population", Integer),
    Column("country", String(16), nullable=False),
)

metadata_obj.create_all(engine)

# 插入测试数据
rows = [
    {"city_name": "Toronto", "population": 2930000, "country": "Canada"},
    {"city_name": "Tokyo", "population": 13960000, "country": "Japan"},
    {"city_name": "Berlin", "population": 3645000, "country": "Germany"},
]
for row in rows:
    stmt = insert(city_stats_table).values(**row)
    with engine.begin() as connection:
        cursor = connection.execute(stmt)

# 验证插入的数据
with engine.connect() as connection:
    cursor = connection.exec_driver_sql("SELECT * FROM city_stats")
    print(cursor.fetchall())

加载数据

我们将使用WikipediaReader加载城市的维基百科数据。

!pip install wikipedia

from llama_index.readers.wikipedia import WikipediaReader

cities = ["Toronto", "Berlin", "Tokyo"]
wiki_docs = WikipediaReader().load_data(pages=cities)

构建SQL和向量索引

接下来，我们将分别为SQL数据库和向量数据构建索引。

from llama_index.core import VectorStoreIndex, SQLDatabase
from llama_index.core.query_engine import NLSQLTableQueryEngine

# SQL索引
sql_database = SQLDatabase(engine, include_tables=["city_stats"])
sql_query_engine = NLSQLTableQueryEngine(sql_database=sql_database, tables=["city_stats"])

# 向量索引
vector_indices = []
for wiki_doc in wiki_docs:
    vector_index = VectorStoreIndex.from_documents([wiki_doc])
    vector_indices.append(vector_index)

定义查询引擎工具

我们将定义查询引擎工具，并设置路由查询引擎。

from llama_index.core.tools import QueryEngineTool
from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector

# SQL查询工具
sql_tool = QueryEngineTool.from_defaults(
    query_engine=sql_query_engine,
    description="适用于将自然语言查询转换为SQL查询，该表包含每个城市的人口/国家信息"
)

# 向量查询工具
vector_tools = []
for city, query_engine in zip(cities, vector_query_engines):
    vector_tool = QueryEngineTool.from_defaults(
        query_engine=query_engine,
        description=f"适用于回答关于{city}的语义问题"
    )
    vector_tools.append(vector_tool)

# 路由查询引擎
query_engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),
    query_engine_tools=([sql_tool] + vector_tools),
)

查询示例

我们可以通过路由查询引擎执行一些示例查询。

# 查询哪个城市人口最多
response = query_engine.query("Which city has the highest population?")
print(str(response))

# 查询柏林的历史博物馆
response = query_engine.query("Tell me about the historical museums in Berlin")
print(str(response))

# 查询每个城市所属的国家
response = query_engine.query("Which countries are each city from?")
print(str(response))