Beyond RAG: Task-aware knowledge compression for enterprise AI on AWS

超越 RAG：AWS 上企业 AI 的任务感知知识压缩 | 人工智能

超越 RAG：AWS 上企业 AI 的任务感知知识压缩

如果您正在使用检索增强生成（RAG）处理需要分析数百个文档的复杂分析任务（如财务尽职调查或合规审查），您可能已经遇到了其局限性。相似性搜索可以提取相关片段，但常常忽略文档间的关联。本文将展示如何通过任务感知知识压缩（TAKC）技术填补这一空白，该技术通过 AWS 部署将整个知识库预压缩为特定任务的表示形式。您可以在自己的账户中部署完整的开源实现。

任务感知知识压缩

考虑一家私募股权公司正在评估对一家制造公司的 5 亿美元收购。尽职调查团队需要分析涵盖 12 家子公司和 5 年的财务报表，同时还要处理 200 多份供应商合同、8 个设施的环境合规报告以及 50 多起法律案件。当分析师询问在当前供应商条款和未决诉讼背景下合并财务风险时，RAG 的相似性搜索无法提供答案。数百份文档包含相关信息，但它们之间的关联缺乏词汇相似性。

TAKC 通过使用大型语言模型（LLM）生成更短且任务导向的文档摘要，不同任务对应不同摘要，从而解决此类问题。

相同文档针对不同任务需要不同信息。为财务分析压缩的年度报告需要收入数据、利润率和现金流信息，而为合规审查压缩的同一报告则需要法规引用和违规历史。通用摘要试图覆盖所有内容，这会稀释特定使用场景的信息密度。TAKC 通过特定任务视角压缩文档，保留关键信息并丢弃其余内容。摄入管道部分展示了压缩提示如何精确指定需要保留的信息。在生产部署中，将任务类型提示存储在版本化配置（如 AWS Systems Manager Parameter Store 或专用 Amazon Simple Storage Service (Amazon S3) 前缀）中，使提示变更可审计，并在提示更新时触发重新压缩。

系统按文档和任务类型离线压缩文档，每个文档每个任务类型仅处理一次。查询时，系统检索预压缩表示而非原始文档。然后系统使用压缩版本而非完整文档回答问题。如果压缩表示缺乏足够细节，查询复杂度分析器会将问题路由到保留更多上下文的较低压缩层级。

TAKC 以压缩形式提供整个知识库的访问权限，而不仅仅是相似性搜索返回的前 k 个片段。系统保留文档间关联，因为压缩过程将文档视为整体。它还为相同源材料的不同任务生成不同压缩输出。同一 10-K 文件的财务压缩与法律风险压缩完全不同。压缩将标记数量减少 8 倍至 64 倍，同时针对任务相关信息进行保留。

多速率压缩

不同类型的查询需要不同级别的精确度。像“第三季度收入是多少？”这样的问题所需上下文远少于要求分析子公司之间供应商付款条款与季度现金流关系的请求。

TAKC通过为每种任务类型维护四个压缩层级来解决这一问题。在最轻度压缩（8倍）层级，系统保留约87.5%的上下文信息。它保留了足够的细节以支持多步骤推理和跨文档合成。在中等压缩（16倍）层级，上下文减少量达到约93.8%。该层级适用于中等复杂度的分析查询。在高级压缩（32倍）层级，上下文减少约96.9%，适用于事实检索和明确定义的问题。在超高压缩（64倍）层级，减少量达到约98.4%。该层级适用于分类任务和关键词检索。

查询复杂度分析器根据查询长度、问题类型和是否存在分析性语言等信号，将传入的问题路由到合适的层级。简单的事实性问题会命中超高压缩缓存，而复杂的分析性问题则使用轻度压缩缓存。这一过程对用户完全透明。

大多数企业查询都是可以通过高压缩层级以最低成本服务的检索请求。只有在必要时，少数复杂查询才会消耗更大的上下文预算。这种基于层级的路由机制与现有的RAG优化（如元数据过滤和查询重格式化）相辅相成，这些优化可以在压缩前缩小文档集范围。为验证压缩质量，可将每个层级的LLM响应与完整未压缩文档生成的响应进行对比。参考实现包含测试脚本，可针对您的特定任务类型和文档执行此对比。

AWS上的架构

该实现运行在AWS上，包含两个解耦的无服务器流水线：一个用于数据摄入，一个用于查询处理。图1展示了这两个流水线。

TAKC架构。流水线流程处理数据摄入和压缩。用户流程处理经过身份验证的查询

我们选择AWS Lambda进行计算，因为每个函数调用都是短暂且事件驱动的。在摄入期间工作负载以突发方式处理数据，并在突发之间处理可变的查询负载，这使无服务器架构成为自然选择。

我们选择Amazon API Gateway将查询接口作为REST端点公开。对于缓存，我们选择Amazon ElastiCache Serverless来处理复合键（takc:{task}:{rate}）的读取操作，无需管理分片。Amazon Cognito处理JWT颁发和令牌刷新，无需自定义认证代码，从而减少实现复杂度。

数据摄入流水线

当文档进入Amazon S3并位于任务类型前缀下（例如，raw-data/financial/）时，S3事件通知会触发AWS Lambda函数。该函数将文档分割成256个标记的段，段间重叠50个标记以防止边界信息丢失。然后，该函数异步调用压缩Lambda函数处理每个段，实现并行处理。对于大规模摄入，可在压缩函数上配置保留并发，并在分段和压缩步骤之间放置Amazon Simple Queue Service（Amazon SQS）队列，以优雅处理限流。

第二个函数调用 Amazon Bedrock 在所有四个压缩层级对片段进行压缩。每个压缩调用都包含一个任务感知提示，告诉模型需要保留哪些信息：

任务：财务分析。保留收入指标、利润率、现金流、债务义务和财务风险指标。
压缩目标：压缩至原始长度的约 1/16。
指令：
- 聚焦与任务相关的事物和关系
- 保留数值数据和指标
- 保持实体及其属性
- 保留因果关系和依赖关系
- 删除冗余或无关信息

由于提示明确指定了任务中重要的信息，模型知道需要保留哪些内容。这种区分正是使该压缩方式具有任务感知特性而非通用压缩的关键。系统将压缩后的输出存储在 Amazon ElastiCache Serverless 中，使用类似 takc:financial:medium 的键名，并备份到 S3 以确保持久性。Redis OSS 数据模型支持多速率缓存查询所需的层次化键结构。缓存条目使用 24 小时生存时间（TTL），并备份到 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 上下文值进行配置，无需修改代码。AWS Lambda（Python 3.12+）处理数据处理和查询逻辑。Amazon ElastiCache Serverless 存储压缩缓存，而 Amazon S3 存储原始数据、片段和缓存备份（KMS 加密）。Amazon API Gateway 暴露 REST 端点，Amazon Cognito 提供基于 JWT 的认证。AWS WAF 通过限流和托管安全规则保护 API。Amazon CloudWatch 提供监控和指标，AWS 密钥管理服务（AWS KMS）管理加密密钥。

将基础设施定义为单个 AWS 云开发工具包（AWS CDK）堆栈，并通过单个命令部署。AWS CDK 提供可重复的部署，并允许通过上下文值自定义参数，如压缩片段大小、Lambda 内存分配和 ElastiCache 存储限制。对于生产部署，建议将有状态资源（S3、ElastiCache、Cognito）分离到独立的堆栈中，以减少影响范围并实现计算和存储层的独立生命周期管理。

成本对比

对于一个包含 100,000 个标记的知识库，每天查询 1,000 次（不包含输出标记，因为响应长度与上下文大小无关）：

| 方法 | 每次查询输入标记 | 每日输入标记 | 相对输入成本 | |--------------|------------------|--------------|--------------| | 完整上下文 | 100,000 | 100,000,000 | 100% | | RAG（前 10 个片段） | ~10,000 | 10,000,000 | 10% |

10%

TAKC Light (8×)

~12,500

12,500,000

12.5%

TAKC Medium (16×)

~6,250

6,250,000

6.25%

TAKC High (32×)

~3,125

3,125,000

3.1%

TAKC Ultra (64×)

~1,563

1,563,000

1.6%

压缩带来的代币节省直接来源于压缩比例。实际节省情况取决于您使用的具体 Bedrock 模型定价和查询模式。

TAKC 需要通过一次性 Bedrock 调用在数据摄入时支付前期压缩成本。对于不常变更且被频繁查询的知识库，这种成本可以分摊。对于每小时都会变更的知识库，RAG 的每查询检索模型可能更实用。

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

根据工作负载特征，下表总结了每种方法更适合的场景：

因素

更倾向 TAKC

更倾向 RAG

查询类型

跨文档推理、综合分析

狭窄的事实查询

知识库稳定性

每日或更少变更

每小时或更多变更

任务可预测性

任务类型明确

查询模式不可预测

覆盖范围要求

必须考虑完整语料库

仅少数文档相关

来源归属

不需要

需要（用户需要查看来源）

代币预算

紧张

灵活

在实际生产系统中，通常两者结合使用效果最佳。RAG 可高效处理快速查询，TAKC 可处理基于检索的方法难以发现关联的分析型查询。查询复杂度分析器可以在两者之间进行路由。当用户需要追溯响应到具体来源文档时，使用 RAG。对于需要跨文档推理且需要审计的合规性工作负载，将 TAKC 与 RAG 结合使用：使用 TAKC 生成分析响应，使用 RAG 检索用于审计追踪的支持性来源文档。

入门指南

先决条件

要部署此参考实现，请确认您具备以下条件：

• 具有 Amazon Bedrock 模型访问权限的 AWS 账户。

• AWS CDK CLI。

• Python 3.12 或更高版本。

• Node.js 18 或更高版本。

部署和测试

克隆仓库 aws-samples/sample-bedrock-takc-compression 并部署 CDK 堆栈：

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 存储桶：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 端点：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 堆栈。首先清空 S3 存储桶，因为 CDK 无法删除包含对象的存储桶：

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

然后销毁堆栈：

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

这将删除 Lambda 函数、API Gateway、Amazon ElastiCache Serverless 缓存、Amazon Cognito 用户池、WAF 网站 ACL 和 Amazon CloudWatch 警报。KMS 密钥会保留 30 天的待删除窗口。要立即安排删除，请运行以下命令：

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

结论

跨越数百个文档的复杂分析任务需要的不仅仅是片段检索。TAKC 提供了专为此类场景设计的解决方案：通过特定任务的视角离线压缩完整知识库，在多个保真度级别缓存这些表示，并根据每个查询的复杂性匹配合适的压缩层级。

AWS 实现方案使用 Amazon Bedrock 进行压缩和推理，使用 Amazon ElastiCache Serverless 进行缓存，并采用完全无服务器架构以按需扩展。可通过 aws-samples/sample-bedrock-takc-compression 获取参考实现、CDK 基础设施和部署脚本。

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

参考资料

• Amazon Bedrock 文档

• Amazon ElastiCache Serverless 文档

• Amazon Bedrock 提示工程指南

作者简介

'"`