使用亚马逊 Bedrock 的知识库构建问答应用程序 - RetrieveAndGenerate API

背景

通过知识库，我们可以安全地将亚马逊 Bedrock 中的基础模型(FM)连接到我们公司的数据，用于检索增强生成(RAG)。获取额外数据有助于模型生成更相关、更具上下文针对性和更准确的响应，而无需不断重新训练 FM。从知识库检索的所有信息都附有来源归属，以提高透明度并最大限度地减少幻觉。有关使用控制台创建知识库的更多信息，请参阅此帖子 。

在本笔记本中，我们将深入探讨使用亚马逊 Bedrock 的知识库提供的 RetrieveAndGenerate API 构建问答应用程序。此 API 将查询知识库以获取基于相似性搜索的所需数量的文档块，并将其与大型语言模型(LLM)集成以回答问题。

模式

我们可以使用检索增强生成(RAG)模式来实现解决方案。RAG从语言模型(非参数)外部检索数据，并通过添加相关检索数据来增强提示。在这里，我们正在有效地对前一个笔记本或使用控制台创建的知识库执行 RAG。

先决条件

在回答问题之前，必须处理并存储文档在知识库中。

  1. 通过连接我们的 s3 存储桶(数据源)将文档加载到知识库中。
  2. 摄取 - 知识库将根据所选策略将它们拆分成更小的块，生成嵌入并将其存储在相关的矢量存储中，笔记本 0_create_ingest_documents_test_kb.ipynb 会为我们处理这一点。

data_ingestion.png

笔记本演练

对于我们的笔记本，我们将使用亚马逊 Bedrock 的知识库提供的 RetreiveAndGenerate API，该 API 将用户查询转换为嵌入，搜索知识库，获取相关结果，增强提示，然后调用 LLM 生成响应。

我们将使用以下工作流程来完成这个笔记本。

retrieveAndGenerate.png

用例:

数据集

在这个例子中，我们将使用多年来亚马逊的股东信作为文本语料库来执行问答。这些数据已经被摄取到知识库中。我们需要 knowledge base id 和 model ARN 来运行这个例子。我们使用 Anthropic Claude 3 Haiku 模型来生成对用户问题的响应。

Python 3.10

⚠ 对于这个实验室，我们需要基于 Python 3.10 运行时运行笔记本。⚠

设置

安装以下软件包。

%pip install --force-reinstall -q -r ./requirements.txt

# 重启内核
from IPython.core.display import HTML
HTML("<script>Jupyter.notebook.kernel.restart()</script>")

%store -r

import boto3
import pprint
from botocore.client import Config

pp = pprint.PrettyPrinter(indent=2)

bedrock_config = Config(connect_timeout=120, read_timeout=120, retries={'max_attempts': 0})
bedrock_client = boto3.client('bedrock-runtime')
bedrock_agent_client = boto3.client("bedrock-agent-runtime",
                              config=bedrock_config)
boto3_session = boto3.session.Session()
region_name = boto3_session.region_name

model_id = "anthropic.claude-3-haiku-20240307-v1:0" # 尝试使用 claude 3 Haiku 以及 claude 3 Sonnet。对于 claude 3 Sonnet - "anthropic.claude-3-sonnet-20240229-v1:0"
region_id = region_name # 替换为我们运行 sagemaker 笔记本的区域

RetreiveAndGenerate API

在幕后，RetrieveAndGenerate API 将查询转换为嵌入，搜索知识库，然后使用搜索结果作为上下文信息来增强基础模型提示，并返回 FM 生成的响应。对于多轮对话，知识库管理对话的短期记忆，以提供更多上下文结果。

RetrieveAndGenerate API 的输出包括 生成的响应、来源归属以及 检索到的文本块。

def retrieveAndGenerate(input, kbId, sessionId=None, model_id = "anthropic.claude-3-haiku-20240307-v1:0", region_id = "us-east-1"):
    model_arn = f'arn:aws:bedrock:{region_id}::foundation-model/{model_id}'
    if sessionId:
        return bedrock_agent_client.retrieve_and_generate(
            input={
                'text': input
            },
            retrieveAndGenerateConfiguration={
                'type': 'KNOWLEDGE_BASE',
                'knowledgeBaseConfiguration': {
                    'knowledgeBaseId': kbId,
                    'modelArn': model_arn
                }
            },
            sessionId=sessionId
        )
    else:
        return bedrock_agent_client.retrieve_and_generate(
            input={
                'text': input
            },
            retrieveAndGenerateConfiguration={
                'type': 'KNOWLEDGE_BASE',
                'knowledgeBaseConfiguration': {
                    'knowledgeBaseId': kbId,
                    'modelArn': model_arn
                }
            }
        )

query = "亚马逊在生成式 AI 领域做了什么?"
response = retrieveAndGenerate(query, kb_id,model_id=model_id,region_id=region_id)
generated_text = response['output']['text']
pp.pprint(generated_text)

citations = response["citations"]
contexts = []
for citation in citations:
    retrievedReferences = citation["retrievedReferences"]
    for reference in retrievedReferences:
         contexts.append(reference["content"]["text"])

pp.pprint(contexts)

下一步

如果我们想要更定制的体验，可以使用 Retrieve API。此 API 将用户查询转换为嵌入，搜索知识库，并返回相关结果，为我们提供更多控制权，以在语义搜索结果的基础上构建自定义工作流程。 有关示例代码，请尝试以下笔记本:

下一步: 继续学习下一个实验室，了解如何将 Bedrock 知识库与 Langchain 和 Claude 一起使用。记得在会话结束时进行清理。