Optimizing RAG Costs with Query-Aware Compression on Amazon Bedrock

Introduction

Retrieval Augmented Generation, or RAG, is usually tuned for recall rather than economy. To make sure the right information is available at answer time, retrieval pulls back a broad set of chunks, often five to twenty, so the primary foundation model has thorough source material to work with. That builds confidence that nothing important gets missed, but it also means the primary model reads a large volume of input tokens on every query. At scale, those input tokens become a meaningful share of the total cost of running RAG in production.

This blog explains a post-retrieval optimization pattern for RAG applications built on Amazon Bedrock: query-aware context compression. After retrieval finishes but before the primary model generates an answer, a smaller, lower-cost model filters the retrieved chunks against the user’s question and passes only the relevant material forward. The result is a meaningful reduction in input tokens and cost, with answer quality largely preserved and, as a secondary benefit, less irrelevant context available to cause hallucination.

Why RAG Costs Add Up?

A typical RAG flow embeds the query, retrieves the top-k chunks from a vector index, optionally reranks them, concatenates everything into the prompt, and sends the bundle to the primary model. Because retrieved context scales with the number of chunks and their sizes, technical documentation and legal RAG workloads commonly land in the range of several thousand input tokens per query, most of which are paid to the most expensive model in the pipeline. The open, composable architecture of Amazon Bedrock makes it possible to insert a custom post-retrieval step that trims that context before it reaches the primary model, without changing how retrieval itself works.

How Query-Aware Compression Works?

The pattern adds exactly one step to a standard RAG pipeline. A smaller model, such as Anthropic’s Claude Haiku on Amazon Bedrock, reads the retrieved chunks alongside the user’s query and outputs only the verbatim spans directly relevant to answering that question. It does not summarize or paraphrase; it extracts. The primary model, such as Anthropic’s Claude Sonnet, then receives this compressed context instead of the full set of chunks and generates the final answer. Both calls run within a single AWS Lambda function using the Amazon Bedrock Converse API, adding one model invocation to the existing flow rather than managing a separate service.

The economics come down to two factors: the price gap between the smaller and primary models on Amazon Bedrock, and how much of the retrieved context the smaller model can safely discard for a given query. Since the smaller model costs far less per token, the savings come from shrinking the context before the expensive answer call, even after the added cost of the compression call itself. The pattern pays off most when the retrieved context is large, the price gap is wide, and a meaningful share of what was retrieved is irrelevant to the question asked.

Implementation on Amazon Bedrock

The compression prompt is the most important piece of the implementation. It instructs the smaller model to extract evidence rather than summarize it, forbids paraphrasing or rewriting, and requires that chunk identifiers be preserved. Hence, citations remain traceable back to their source. The compression call runs at a temperature of zero, keeping extraction deterministic so the model reliably copies spans as they appear.

Inside the AWS Lambda function, the first call sends the query and all retrieved chunks to the smaller model via the Amazon Bedrock Converse API. It receives only the relevant verbatim spans in response. The second call sends the compressed evidence, along with the original query, to the primary model, which answers using only the provided evidence and cites its sources. This two-step structure fits between existing retrieval logic and the final answer call without requiring any change to how chunks are retrieved or ranked upstream.

Results and Considerations for Production

Before recommending the pattern, its authors evaluated it against a corpus of more than 500,000 documents across nine enterprise source types and 500 questions spanning ten categories, comparing a baseline pipeline against compression alone and against compression combined with reranking. Compression alone reduced tokens reaching the primary model by roughly 8.6x and cut query cost by about a third, while rerank-and-compression reached around 10x fewer tokens and roughly 36 percent lower cost. Answer quality stayed within a few points of the uncompressed baseline across correctness, completeness, citation accuracy, and conciseness, and the hallucination rate dropped since the primary model had less irrelevant material to draw incorrect claims from.

The tradeoff is latency: the compression call adds a step to the request path, though some of that time is recovered since the primary model then processes a smaller, more focused context. For strict sub-second latency budgets, this needs to be measured against the specific context sizes involved. Selecting the smaller model also matters: pairing models from the same family, such as Claude Haiku with Claude Sonnet, keeps formatting and instruction handling consistent. The pattern fits best where retrieved context is large, the primary model is the dominant cost, and questions are narrow relative to the volume retrieved, such as compliance assistants over long regulatory bulletins, customer-support copilots over large knowledge bases, and internal engineering assistants over lengthy runbooks and wikis.

Conclusion

Query-aware compression is a small addition with an outsized effect on RAG economics: a single AWS Lambda function inserted between retrieval and the final answer call, using the Amazon Bedrock Converse API to run a smaller model as a filter before the primary model sees the context. It delivered substantial reductions in tokens and cost while keeping answer quality close to baseline and reducing hallucinations. It layers cleanly on top of other Amazon Bedrock capabilities, such as prompt caching, Intelligent Prompt Routing, and the Rerank API, making it one of several compounding levers to bring RAG costs under control without sacrificing answer quality.

Drop a query if you have any questions regarding Query-aware, and we will get back to you quickly.