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

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

任务感知知识压缩

设想一家私募股权公司正在评估以 5 亿美元收购一家制造企业。尽职调查团队必须分析覆盖 12 家子公司、历时 5 年的财务报表；同时还要面对超过 200 份供应商合同、来自 8 个设施的环境合规报告，以及 50 多起法律案件。当分析师询问在当前供应商条款和未决诉讼条件下的合并财务风险时，RAG 的相似度搜索无法找出能够回答该问题的内容。相关信息分散在数百份文档中，而且其中的关联并不具备词汇相似性。

TAKC 通过使用 LLM 生成更短、更聚焦于任务的文档摘要来解决这类问题；针对不同任务，同一份文档会产生不同的摘要。

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

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

TAKC 提供的是以压缩形式访问_整个_知识库的能力，而不只是访问相似度搜索返回的 top-k 片段。由于压缩过程会同时查看文档，系统能够保留文档之间的关联。它还会基于相同源材料，为不同任务生成不同的压缩输出。一份 10-K 文件的财务压缩版本，与该文件的法律风险压缩版本会截然不同。该压缩过程可将 Token 数量降低 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 Token 的片段，片段之间保留 50 Token 重叠，以防止边界处的信息丢失。随后，它会为每个片段异步调用压缩 Lambda 函数，从而实现并行处理。对于大规模摄取，可为压缩函数配置预留并发，并在切分与压缩步骤之间放置 Amazon Simple Queue Service（Amazon SQS）队列，以平稳处理限流。

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

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

模型之所以知道应保留什么，是因为提示词明确规定了对该任务重要的信息。这一区别使其成为任务感知压缩，而非通用压缩。系统将压缩输出以 takc:financial:medium 这样的键存储在 Amazon ElastiCache Serverless 中，并备份至 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 值配置，无需修改代码。AWS Lambda（Python 3.12+）负责数据处理和查询逻辑。Amazon ElastiCache Serverless 存储压缩缓存，而 Amazon S3 保存原始数据、分片和缓存备份（使用 KMS 加密）。Amazon API Gateway 暴露 REST 端点，Amazon Cognito 提供基于 JWT 的身份验证。AWS WAF 通过速率限制和托管安全规则保护 API 前端。Amazon CloudWatch 提供监控和指标，AWS Key Management Service（AWS KMS） 管理加密密钥。

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

成本比较

对于一个包含 100,000 Token 的知识库、每天被查询 1,000 次的场景，其输入 Token 用量如下（不包括输出 Token，因为响应长度与上下文大小无关）：

方法	每次查询的输入 Token	每日输入 Token	相对输入成本
完整上下文	100,000	100,000,000	100%
RAG（top-10 片段）	~10,000	10,000,000	10%
TAKC 轻度（8×）	~12,500	12,500,000	12.5%
TAKC 中等（16×）	~6,250	6,250,000	6.25%
TAKC 高度（32×）	~3,125	3,125,000	3.1%
TAKC 超高（64×）	~1,563	1,563,000	1.6%

压缩带来的 Token 节省直接源自压缩比。实际节省幅度取决于所使用的 Bedrock 模型定价及查询模式。

TAKC 在摄取阶段需要通过一次性 Bedrock 调用承担前期压缩成本。对于变更频率较低、但会被重复查询的知识库，这一成本能够被摊销。对于每小时都发生变化的知识库，RAG 的按查询检索模式可能更实用。

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

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

因素	更适合 TAKC	更适合 RAG
查询类型	跨文档推理、综合分析	范围狭窄的事实查询
知识库稳定性	每天或更低频率变更	每小时或更高频率变更
任务可预测性	任务类型定义明确	查询模式不可预测
覆盖范围要求	必须考虑整个语料库	只有少量文档相关
来源归因	不要求	要求（用户需要查看来源）
Token 预算	紧张	灵活

在实践中，生产系统通常会同时受益于两种方法。RAG 能够高效处理快速查找；TAKC 则处理检索式方法难以发现关联的分析型查询。查询复杂度分析器可以在两者之间进行路由。当用户需要将响应追溯至特定源文档时，应使用 RAG。对于既需要跨文档推理又需要可审计性的受监管工作负载，可以将 TAKC 与 RAG 结合：使用 TAKC 生成分析性响应，再使用 RAG 检索支撑审计轨迹的源文档。

快速开始

前提条件

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

部署和测试

克隆仓库 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
pi

If you’re using Retrieval-Augmented Generation (RAG) for complex analytical tasks that span hundreds of documents, such as financial due diligence or regulatory compliance reviews, you’ve likely hit its ceiling. Similarity search surfaces relevant fragments but often misses cross-document connections. This post shows you how to address that gap using task-aware knowledge compression (TAKC), a technique that pre-compresses entire knowledge bases into task-specific representations deployed on AWS. You can deploy a complete open-source implementation in your own account.

Task-aware knowledge compression

Consider a private equity firm evaluating a $500 million acquisition of a manufacturing company. The due diligence team must analyze financial statements spanning 12 subsidiaries and 5 years. They also face 200+ supplier contracts, environmental compliance reports from 8 facilities, and 50+ legal cases. When an analyst asks about consolidated financial risks given current supplier terms and pending litigation, RAG’s similarity search cannot surface that response. Hundreds of documents hold the relevant information, and the connections between them share no lexical similarity.

TAKC addresses this type of problem by using an LLM to produce shorter, task-focused summaries of documents, with different summaries for different tasks.

Different tasks require different information from the same document. An annual report compressed for financial analysis needs revenue figures, margins, and cash flow data. The same report compressed for a compliance review needs regulatory citations and violation histories. Generic summarization attempts to cover everything, which dilutes the information density for any specific use case. TAKC compresses documents through the lens of a specific task, keeping what matters and discarding the rest. The Ingestion pipeline section shows how the compression prompt specifies exactly what information to preserve. For production deployments, store task-type prompts in a versioned configuration (such as AWS Systems Manager Parameter Store or a dedicated Amazon Simple Storage Service (Amazon S3) prefix) so that prompt changes are auditable and recompression can be triggered when prompts are updated.

The system compresses documents offline, once per document per task type. At query time, the system retrieves a pre-compressed representation rather than the original document. The system then answers questions using the compressed version instead of the full document. If the compressed representation lacks sufficient detail, the query complexity analyzer routes the question to a lower compression tier that retains more context.

TAKC provides access to the entire knowledge base in compressed form, not just the top-k chunks that a similarity search returns. The system preserves connections between documents because the compression sees documents together. It also produces different compressed outputs for different tasks from the same source material. The financial compression of a 10-K filing looks nothing like the legal risk compression of that same filing. The compression reduces token count by 8x to 64x while targeting task-relevant information for retention.

Multi-rate compression

Different queries require different levels of fidelity. A question such as “What was Q3 revenue?” needs far less context than a request to analyze relationships between supplier payment terms and quarterly cash flow across subsidiaries.

TAKC addresses this by maintaining four compression tiers for each task type. At the lightest compression (8x), the system retains approximately 87.5% less context. It preserves enough detail for multi-step reasoning and cross-document synthesis. At the medium tier (16x), context reduction reaches approximately 93.8%. This level suits analytical queries with moderate complexity. The high tier (32x) reduces context by approximately 96.9% and serves factual lookups and well-defined questions. At the ultra tier (64x), reduction reaches approximately 98.4%. This level is appropriate for classification tasks and keyword lookups.

A query complexity analyzer routes incoming questions to the appropriate tier based on signals like query length, question type, and presence of analytical language. Straightforward factual questions hit the ultra-compressed cache while complex analytical questions use the lightly compressed cache. This happens transparently to the user.

Most enterprise queries are lookups that can be served from the higher compression tiers at minimal cost. The infrequent complex queries consume larger context budgets only when needed. This tier-based routing complements existing RAG optimizations like metadata filtering and query reformatting, which can narrow the document set before compression is applied. To validate compression quality, compare LLM responses at each tier against responses generated from the full uncompressed document. The reference implementation includes test scripts that run this comparison for your specific task types and documents.

Architecture on AWS

The implementation runs on AWS as two decoupled serverless pipelines: one for ingestion and one for queries. Figure 1 illustrates both pipelines.

TAKC architecture. The pipeline flow handles data ingestion and compression. The user flow handles authenticated queries

We chose AWS Lambda for compute because each function invocation is short-lived and event-driven. The workload processes data in bursts during ingestion and handles variable query loads between bursts, making serverless a natural fit.

We chose Amazon API Gateway to expose the query interface as a REST endpoint. For caching, we chose Amazon ElastiCache Serverless for reads on composite keys (takc:{task}:{rate}) without shard management. Amazon Cognito handles JWT issuance and token refresh without custom auth code, reducing the implementation surface area.

Ingestion pipeline

When a document lands in Amazon S3 under a task-type prefix (for example, raw-data/financial/), an S3 event notification triggers an AWS Lambda function. This function chunks the document into 256-token segments with 50-token overlap to prevent information loss at boundaries. It then invokes a compression Lambda function for each chunk asynchronously, enabling parallel processing. For large-scale ingestion, configure reserved concurrency on the compression function and place an Amazon Simple Queue Service (Amazon SQS) queue between the chunking and compression steps to handle throttling gracefully.

The second function calls Amazon Bedrock to compress the chunks at all four compression tiers. Each compression call includes a task-aware prompt that tells the model what information to preserve:

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

The model knows what to preserve because the prompt specifies what matters for the task. That distinction is what makes this task-aware rather than generic compression. The system stores compressed outputs in Amazon ElastiCache Serverless with keys like takc:financial:medium, and backs them up to S3 for durability. The Redis OSS data model supports the hierarchical key structure needed for multi-rate cache lookups. Cache entries use a 24-hour TTL and are backed up to S3. If a cache entry is evicted or expires, the query function falls back to the S3 backup and re-populates the cache on read.

Query pipeline

A user authenticates through Amazon Cognito, receives a JWT, and sends a query through Amazon API Gateway. AWS WAF fronts the API for rate limiting and threat protection. A Lambda function analyzes query complexity using heuristics (keyword signals and query length), retrieves the appropriate compressed cache from Amazon ElastiCache Serverless, and sends the compressed context plus query to Amazon Bedrock for inference. When routing confidence is low, the system defaults to the medium compression tier as a safe fallback.

The higher-cost Bedrock compression calls happen once during ingestion. The query path is a cache lookup plus inference on compressed context.

The stack uses Amazon Bedrock (Anthropic Claude 3 Haiku, Claude 3 Sonnet, and Amazon Titan Text) for compression and inference. The model selection is configurable through CDK context values without code changes. AWS Lambda (Python 3.12+) handles data processing and query logic. Amazon ElastiCache Serverless stores the compressed cache, while Amazon S3 holds raw data, chunks, and cache backups (KMS-encrypted). Amazon API Gateway exposes the REST endpoints, and Amazon Cognito provides JWT-based authentication. AWS WAF fronts the API with rate limiting and managed security rules. Amazon CloudWatch provides monitoring and metrics, and AWS Key Management Service (AWS KMS) manages encryption keys.

Define the infrastructure as a single AWS Cloud Development Kit (AWS CDK) stack and deploy it with a single command. AWS CDK provides repeatable deployments and lets you customize parameters like compression chunk size, Lambda memory allocation, and ElastiCache storage limits through context values. For production deployments, consider separating stateful resources (S3, ElastiCache, Cognito) into their own stack to reduce blast radius and enable independent lifecycle management of compute and storage layers.

Cost comparison

Input token usage for a 100,000-token knowledge base queried 1,000 times per day (output tokens are excluded because response length is independent of context size):

Approach	Input tokens per query	Daily input tokens	Relative input cost
Full context	100,000	100,000,000	100%
RAG (top-10 chunks)	~10,000	10,000,000	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%

The token savings from compression are derived directly from the compression ratios. Actual savings depend on your specific Bedrock model pricing and query patterns.

TAKC requires upfront compression cost through one-time Bedrock calls during ingestion. This cost amortizes for knowledge bases that change infrequently and get queried repeatedly. For knowledge bases that change hourly, RAG’s per-query retrieval model may be more practical.

When to choose TAKC, RAG, or both

This table summarizes when each approach is the better fit, based on your workload characteristics:

Factor	Favors TAKC	Favors RAG
Query type	Cross-document reasoning, synthesis	Narrow factual lookups
Knowledge base stability	Changes daily or less	Changes hourly or more
Task predictability	Well-defined task types	Unpredictable query patterns
Coverage requirement	Must consider full corpus	Only a few documents are relevant
Source attribution	Not required	Required (user needs to see source)
Token budget	Tight	Flexible

In practice, a production system often benefits from both. RAG handles quick lookups efficiently. TAKC handles the analytical queries where retrieval-based approaches miss connections. A query complexity analyzer can route between them. Use RAG when users need to trace responses back to specific source documents. For regulated workloads that require both cross-document reasoning and auditability, combine TAKC with RAG: use TAKC for the analytical response and RAG to retrieve the supporting source documents for the audit trail.

Getting started

Prerequisites

To deploy this reference implementation, check that you have the following:

Deploy and test

Clone the repository aws-samples/sample-bedrock-takc-compression and deploy the 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

After deployment completes, test the pipeline:

1. Upload a document to the 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/
   

2. Wait 2-3 minutes for the ingestion pipeline to chunk and compress the document at all four compression tiers.

3. Query the 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?"}'
   

The system handles chunking, multi-rate compression, caching, and query routing without additional configuration.

Clean up

To avoid ongoing charges, destroy the CDK stack when you’re done testing. First, empty the S3 bucket because CDK cannot delete buckets that contain objects:

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

Then destroy the stack:

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

This removes the Lambda functions, API Gateway, Amazon ElastiCache Serverless cache, Amazon Cognito user pool, WAF web ACL, and Amazon CloudWatch alarms. The KMS key is retained with a 30-day pending deletion window. To schedule its deletion immediately, run the following command:

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

Conclusion

Complex analytical tasks that span hundreds of documents require more than fragment retrieval. TAKC offers an approach built for this. Compress the full knowledge base offline through the lens of specific tasks, cache those representations at multiple fidelity levels, and match the right compression tier to each query’s complexity.

The AWS implementation uses Amazon Bedrock for compression and inference, Amazon ElastiCache Serverless for caching, and a fully serverless architecture that scales with demand. Access the reference implementation, CDK infrastructure, and deployment scripts at aws-samples/sample-bedrock-takc-compression.

To get started, deploy the CDK stack with your own documents and see how the system responds to different query types. If your workload involves cross-document reasoning over stable knowledge bases, TAKC can reduce token costs while improving response quality.

References

---

About the authors

Dhananjay Karanjkar

Dhananjay is a Senior Lead Consultant at AWS Professional Services, specializing in agentic AI systems, multi-agent orchestration, and generative AI security. He holds two US patents and serves as a Responsible AI Champion, with a background spanning financial services, enterprise consulting, and enterprise-scale AI delivery. When not architecting AI solutions, he trains for triathlons, paints oil portraits, and reads voraciously.

Sanjay Chaudhari

Sanjay is a Lead Consultant in AWS Professional Services, where he helps customers migrate and modernize their .NET workloads on AWS. He is deeply interested in the evolving landscape of Agentic AI and its potential to transform enterprise workflows. Outside of work, he enjoys travelling and exploring different cuisines.

Rajeshkumar Kagathara

Rajeshkumar is a Lead Consultant with AWS Professional Services. He helps customers architect, modernize, and deliver scalable cloud-native solutions using AWS services, Java, and Spring Boot. Rajesh is passionate about cloud transformation, enterprise application modernization, and engineering excellence through secure, high-performing architectures. Outside of work, he enjoys playing indoor games and traveling.