RAG性能升级：AWS任务感知知识压缩方案

如果你正在使用 Retrieval-Augmented Generation（RAG，检索增强生成）处理横跨数百份文档的复杂分析任务，例如财务尽职调查或监管合规审查，那么很可能已经触及了它的能力上限。相似度搜索能够找出相关片段，却经常遗漏文档之间的关联。本文将介绍如何使用 task-aware knowledge compression（TAKC，任务感知知识压缩）来弥补这一缺陷。这项技术会预先将整个知识库压缩成面向特定任务的表示，并部署在 AWS 上。你可以在自己的账户中部署一套完整的开源实现。

任务感知知识压缩

假设一家私募股权公司正在评估一项价值 5 亿美元的制造企业收购案。尽职调查团队必须分析覆盖 12 家子公司、时间跨度达 5 年的财务报表。此外，他们还需要审查 200 多份供应商合同、来自 8 个生产设施的环境合规报告，以及 50 多起法律案件。当分析师询问「结合当前供应商条款和待决诉讼，公司合并层面的财务风险有哪些」时，RAG 的相似度搜索无法直接找出答案。相关信息分散在数百份文档中，而这些信息之间的联系并不存在词汇层面的相似性。

TAKC 使用 LLM 为文档生成更短、聚焦特定任务的摘要，从而解决这类问题；针对不同任务，同一份文档会生成不同的摘要。

面对不同任务，同一份文档中需要保留的信息也不同。为财务分析而压缩年度报告时，需要保留收入数字、利润率和现金流数据；为合规审查压缩同一份报告时，则需要保留监管条款引用和违规记录。通用摘要试图覆盖所有内容，结果会稀释每个具体使用场景所需的信息密度。TAKC 会从特定任务的视角压缩文档，只保留重要信息，其余内容则予以舍弃。「数据摄取流水线」一节将展示如何通过压缩 prompt 明确指定需要保留的信息。对于生产环境部署，应将任务类型对应的 prompt 存储在带版本管理的配置中，例如 AWS Systems Manager Parameter Store，或专用的 Amazon Simple Storage Service（Amazon S3）前缀。这样便可以审计 prompt 的变更，并在 prompt 更新时触发重新压缩。

系统会离线压缩文档，每份文档针对每种任务类型只压缩一次。查询时，系统检索的是预先压缩的表示，而不是原始文档。随后，系统使用压缩版本而非完整文档回答问题。如果压缩表示缺少足够的细节，查询复杂度分析器就会将问题路由到压缩程度更低、保留上下文更多的层级。

TAKC 能够以压缩形式访问整个知识库，而不只是相似度搜索返回的 top-k 个分块。由于压缩过程会同时查看多份文档，因此系统能够保留文档之间的关联。对于相同的源材料，它还可以针对不同任务生成不同的压缩结果。一份 10-K 文件的财务分析压缩结果，与同一文件的法律风险压缩结果会截然不同。压缩可以将 token 数量减少到原来的 1/8 至 1/64，同时有针对性地保留与任务相关的信息。

多压缩率机制

不同查询对信息保真度的要求不同。「第三季度的收入是多少？」这类问题所需的上下文，远少于分析各子公司的供应商付款条款与季度现金流之间关系所需的上下文。

为此，TAKC 会为每种任务类型维护四个压缩层级。在压缩程度最低的层级（8x），系统会将上下文减少约 87.5%，同时保留足以支持多步推理和跨文档综合分析的细节。在中等层级（16x），上下文减少约 93.8%，适用于复杂度适中的分析型查询。在高压缩层级（32x），上下文减少约 96.9%，适用于事实查找和定义明确的问题。在超高压缩层级（64x），上下文减少约 98.4%，适用于分类任务和关键词查找。

查询复杂度分析器会依据查询长度、问题类型，以及是否包含分析型语言等信号，将传入的问题路由到合适的层级。简单直接的事实型问题会使用超高压缩缓存，复杂的分析型问题则会使用低压缩缓存。整个过程对用户完全透明。

大多数企业查询都是信息查找，可以通过更高的压缩层级以极低成本完成。只有少量复杂查询会在确有需要时消耗更大的上下文预算。这种基于层级的路由机制，可以与 metadata filtering、query reformatting 等现有 RAG 优化手段配合使用，在执行压缩之前先缩小文档集合。为了验证压缩质量，可以将各压缩层级下的 LLM 响应，与根据完整未压缩文档生成的响应进行比较。参考实现中包含测试脚本，可针对你的特定任务类型和文档执行这种比较。

该实现在 AWS 上以两条相互解耦的 serverless 流水线运行：一条负责数据摄取，另一条负责查询。图 1 展示了这两条流水线。

TAKC 架构。流水线流程负责数据摄取与压缩；用户流程负责处理经过身份认证的查询。

我们选择 AWS Lambda 作为计算服务，是因为每次函数调用持续时间较短，并由事件驱动。该工作负载会在数据摄取期间突发处理大量数据，并在这些高峰之间应对变化不定的查询负载，因此 serverless 是一种自然的选择。

我们选择 Amazon API Gateway，将查询接口公开为 REST endpoint。在缓存方面，我们选择 Amazon ElastiCache Serverless，以便在无须管理分片的情况下读取复合键（takc:{task}:{rate}）。Amazon Cognito 负责 JWT 签发和 token 刷新，无须编写自定义身份认证代码，从而缩小了实现范围。

当文档写入 Amazon S3 中某个任务类型前缀下时，例如 raw-data/financial/，S3 事件通知会触发一个 AWS Lambda 函数。该函数将文档切分为 256-token 的片段，并设置 50-token 的重叠区域，防止边界处的信息丢失。随后，它会为每个分块异步调用压缩 Lambda 函数，从而实现并行处理。对于大规模数据摄取，应为压缩函数配置 reserved concurrency，并在分块与压缩步骤之间放置一个 Amazon Simple Queue Service（Amazon SQS）队列，以便妥善处理限流。

第二个函数调用 Amazon Bedrock，以全部四种压缩层级压缩这些分块。每次压缩调用都会包含一个 task-aware prompt，用于告诉模型需要保留哪些信息：

TASK: Financial analysis. Preserve revenue metrics, margins, cash flow, debt obligations, and financial risk indicators.
COMPRESSION TARGET: Reduce to approximately 1/16 of original length.
INSTRUCTIONS:
- Focus on facts and relationships relevant to the task
- Preserve numerical data and metrics
- Maintain entities and their attributes
- Keep causal relationships and dependencies
- Remove redundant or irrelevant information

模型之所以知道需要保留哪些内容，是因为 prompt 明确规定了当前任务关心的信息。正是这一点，使它成为任务感知压缩，而不是通用压缩。系统会将压缩结果存储在 Amazon ElastiCache Serverless 中，使用类似 takc:financial:medium 的键，并将结果备份至 S3 以保证持久性。Redis OSS 数据模型支持多压缩率缓存查找所需的分层键结构。缓存条目的 TTL 为 24 小时，同时会备份至 S3。如果缓存条目被淘汰或过期，查询函数会回退到 S3 备份，并在读取时重新填充缓存。

用户通过 Amazon Cognito 完成身份认证并获得 JWT，然后通过 Amazon API Gateway 发送查询。AWS WAF 位于 API 前端，负责限流和威胁防护。一个 Lambda 函数使用启发式规则（关键词信号和查询长度）分析查询复杂度，从 Amazon ElastiCache Serverless 中获取适当的压缩缓存，再将压缩后的上下文和查询一起发送至 Amazon Bedrock 进行推理。当路由置信度较低时，系统默认使用中等压缩层级，作为稳妥的回退方案。

成本较高的 Bedrock 压缩调用只会在数据摄取期间执行一次。查询路径只包含一次缓存查找，以及基于压缩上下文的推理。

该技术栈使用 Amazon Bedrock（Anthropic Claude 3 Haiku、Claude 3 Sonnet 和 Amazon Titan Text）完成压缩与推理。无需修改代码，只需通过 CDK context values 即可配置模型选择。AWS Lambda（Python 3.12+）负责数据处理和查询逻辑。Amazon ElastiCache Serverless 存储压缩缓存，Amazon S3 则保存原始数据、分块和缓存备份，并使用 KMS 加密。Amazon API Gateway 公开 REST endpoints，Amazon Cognito 提供基于 JWT 的身份认证。AWS WAF 位于 API 前端，提供限流和托管安全规则。Amazon CloudWatch 提供监控和指标，AWS Key Management Service（AWS KMS）负责管理加密密钥。

将基础设施定义为单个 AWS Cloud Development Kit（AWS CDK）stack，并通过一条命令完成部署。AWS CDK 可以实现可重复部署，并允许你通过 context values 自定义压缩分块大小、Lambda 内存分配和 ElastiCache 存储限制等参数。对于生产环境部署，可以考虑将有状态资源（S3、ElastiCache、Cognito）拆分到单独的 stack 中，以缩小故障影响范围，并支持计算层与存储层的独立生命周期管理。

以下是一个 100,000-token 的知识库每天被查询 1,000 次时的输入 token 使用量（不计输出 token，因为响应长度与上下文大小无关）：

压缩带来的 token 节省量可以直接根据压缩率推导出来。实际节省量取决于你所使用的具体 Bedrock 模型定价和查询模式。

TAKC 需要承担前期压缩成本，即在数据摄取期间执行一次性 Bedrock 调用。对于变更频率较低、但会被反复查询的知识库，这笔成本可以得到摊销。对于每小时都会变化的知识库，RAG 的按查询检索模式可能更加实用。

何时选择 TAKC、RAG，或将两者结合

下表根据工作负载特征，汇总了各类方案更适合的使用场景：

实践中，生产系统往往可以同时受益于这两种方案。RAG 能够高效处理快速查找，TAKC 则负责处理那些基于检索的方法会遗漏信息关联的分析型查询。查询复杂度分析器可以在二者之间进行路由。如果用户需要将响应追溯到具体的源文档，应使用 RAG。对于既要求跨文档推理、又要求可审计性的受监管工作负载，可以结合 TAKC 与 RAG：使用 TAKC 生成分析结果，再使用 RAG 检索支持该结果的源文档，以形成审计追踪记录。

要部署这套参考实现，请确认你具备以下条件：

• 一个已获得 Amazon Bedrock 模型访问权限的 AWS 账户。
• Python 3.12 或更高版本。

克隆 aws-samples/sample-bedrock-takc-compression repository，并部署 CDK stack：

git clone https://github.com/aws-samples/sample-bedrock-takc-compression
cd sample-bedrock-takc-compression/cdk
python3 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
cdk deploy

部署完成后，测试流水线：

将文档上传到 S3 bucket：aws s3 cp your-document.pdf s3://$(aws cloudformation describe-stacks --stack-name TakcStack \ --query 'Stacks[0].Outputs[?OutputKey==DataBucketName].OutputValue' \ --output text)/raw-data/financial/

aws s3 cp your-document.pdf s3://$(aws cloudformation describe-stacks --stack-name TakcStack \
  --query 'Stacks[0].Outputs[?OutputKey==`DataBucketName`].OutputValue' \
  --output text)/raw-data/financial/

等待 2～3 分钟，让数据摄取流水线完成文档分块，并以全部四种压缩层级压缩文档。

查询 API endpoint：curl -X POST $(aws cloudformation describe-stacks --stack-name TakcStack \ --query 'Stacks[0].Outputs[?OutputKey==ApiEndpoint].OutputValue' \ --output text)/query \ -H "Authorization: Bearer $TOKEN" \ -H "Content-Type: application/json" \ -d '{"question": "What are the key financial risks?"}'

curl -X POST $(aws cloudformation describe-stacks --stack-name TakcStack \
  --query 'Stacks[0].Outputs[?OutputKey==`ApiEndpoint`].OutputValue' \
  --output text)/query \
  -H "Authorization: Bearer $TOKEN" \
  -H "Content-Type: application/json" \
  -d '{"question": "What are the key financial risks?"}'

系统会完成分块、多压缩率压缩、缓存和查询路由，无须额外配置。

为避免持续产生费用，测试完成后请销毁 CDK stack。首先清空 S3 bucket，因为 CDK 无法删除其中仍包含对象的 bucket：

aws s3 rm s3://$(aws cloudformation describe-stacks --stack-name TakcStack \
  --query 'Stacks[0].Outputs[?OutputKey==`DataBucketName`].OutputValue' \
  --output text) --recursive

然后销毁 stack：

cd sample-bedrock-takc-compression/cdk
source .venv/bin/activate
cdk destroy

这将删除 Lambda 函数、API Gateway、Amazon ElastiCache Serverless 缓存、Amazon Cognito user pool、WAF Web ACL 和 Amazon CloudWatch alarms。KMS key 会被保留，并进入 30 天的待删除期。要立即安排其删除，请运行以下命令：

aws kms schedule-key-deletion --key-id <key-id> --pending-window-in-days 7

横跨数百份文档的复杂分析任务，需要的不只是片段检索。TAKC 提供了一种专为此类场景设计的方法：从特定任务的视角离线压缩完整知识库，以多种保真度缓存这些表示，再根据每个查询的复杂度选择合适的压缩层级。

这套 AWS 实现使用 Amazon Bedrock 完成压缩和推理，使用 Amazon ElastiCache Serverless 提供缓存，并采用能够随需求扩展的全 serverless 架构。你可以在 aws-samples/sample-bedrock-takc-compression 中获取参考实现、CDK 基础设施和部署脚本。

要开始使用，请使用自己的文档部署 CDK stack，并观察系统如何响应不同类型的查询。如果你的工作负载需要在稳定的知识库上执行跨文档推理，TAKC 可以在提升响应质量的同时降低 token 成本。

Amazon Bedrock 文档

Amazon ElastiCache Serverless 文档

Amazon Bedrock prompt engineering 指南