Dynamic Context Management: Slash LLM Costs & Boost Performance with Adaptive RAG

Introduction & The Problem

In the rapidly evolving landscape of AI-powered applications, Large Language Models (LLMs) have become indispensable. From conversational agents to complex data analysis tools, LLMs enable capabilities once thought impossible. However, integrating LLMs into production systems introduces significant challenges, chief among them being the efficient management of their 'context window.'

The context window is the limited amount of text (tokens) an LLM can process at any given time. As interactions become more complex or conversational histories grow longer, this window quickly becomes a bottleneck. Developers face a dilemma: truncate valuable information, leading to less accurate or coherent responses, or continuously feed the entire history, drastically increasing API costs and slowing down response times. This isn't merely a technical hurdle; it translates directly into higher operational expenses, degraded user experience, and ultimately, a compromised return on investment for AI initiatives. Without a smart strategy, scaling LLM-powered applications becomes prohibitively expensive and inefficient.

The Solution Concept & Architecture

The solution lies in a more intelligent approach to context management, moving beyond static context windows to a dynamic, adaptive system. This is where Dynamic Retrieval Augmented Generation (RAG) shines. Instead of dumping all available information into the LLM's context, Dynamic RAG intelligently retrieves only the most relevant pieces of information based on the current query and dynamically prunes or summarizes past interactions to fit within cost and performance constraints.

At its core, a Dynamic RAG architecture involves several key components:

• Query Analysis & Rewriting: Understanding user intent and potentially rewriting the query for better retrieval.
• Intelligent Retrieval: Fetching relevant data chunks from a knowledge base (e.g., vector database, document store) based on the refined query.
• Context Pruning & Summarization: Applying algorithms to reduce the historical conversational context to its most salient points, ensuring critical information is retained while less important details are discarded or condensed.
• Dynamic Prompt Construction: Combining the relevant retrieved information and the condensed historical context into an optimized prompt for the LLM.
• LLM Inference: Sending the optimized prompt to the LLM for generating the final response.

This approach ensures that the LLM's context window is always populated with the highest-value information, minimizing token usage and maximizing the quality of responses.

Step-by-Step Implementation

Let's walk through building a simplified Dynamic RAG system using Python, LangChain, and a vector database like ChromaDB. We'll focus on demonstrating the dynamic context management aspect.

First, ensure you have the necessary libraries installed:

pip install langchain langchain-openai chromadb tiktoken

Next, let's set up our knowledge base. For simplicity, we'll use a few dummy documents. In a real application, this would be populated from your internal documentation, FAQs, or other data sources.

import os
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document

# Set your OpenAI API key
os.environ["OPENAI_API_KEY"] = "YOUR_OPENAI_API_KEY"

# 1. Prepare Documents
documents = [
    Document(page_content="The company's Q3 earnings report showed a 15% increase in revenue, primarily driven by new product launches."),
    Document(page_content="Our new AI product, 'CognitoFlow,' automates data analysis tasks, reducing manual effort by 40%."),
    Document(page_content="Customer support wait times have been reduced by 25% due to the implementation of our new chatbot, 'AssistAI'."),
    Document(page_content="The marketing team plans a major campaign for CognitoFlow in Q4, targeting enterprise clients."),
    Document(page_content="Technical specifications for CognitoFlow include integration with major cloud providers and a secure data encryption standard.")
]

# 2. Split Documents into Chunks
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
chunked_documents = text_splitter.split_documents(documents)

# 3. Create Embeddings and Store in Vector DB
embeddings = OpenAIEmbeddings()
vectorstore = Chroma.from_documents(
    documents=chunked_documents,
    embedding=embeddings,
    persist_directory="./chroma_db"
)
vectorstore.persist()
print("Vector database initialized with documents.")

Now, let's implement the dynamic RAG components. We'll simulate a conversational history and use a small LLM (or a prompt-engineered larger LLM) to summarize the history when it exceeds a certain token limit.

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
import tiktoken # For token counting

# Initialize LLM for summarization and main generation
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
summarizer_llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.2) # A smaller LLM for summarization

# Max tokens for the context window (example: ~4k tokens for gpt-3.5-turbo, adjust based on your target model)
MAX_CONTEXT_TOKENS = 3500

def count_tokens(text: str) -> int:
    encoding = tiktoken.encoding_for_model("gpt-4o-mini") # Use your target model's encoding
    return len(encoding.encode(text))

def summarize_history(chat_history: list, current_summary: str = "") -> str:
    # This function uses an LLM to summarize the chat history.
    # In a real app, you might have more sophisticated logic to choose which parts to summarize.
    history_text = "\n".join([f"{msg.type}: {msg.content}" for msg in chat_history])
    
    summary_prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a helpful assistant. Summarize the following conversation history concisely, focusing on key facts and unresolved questions. Maintain continuity with the provided current summary if any.\n\nCurrent Summary: {current_summary}"),
        ("user", "Conversation History:\n{history_text}\n---\nNew Concise Summary:")
    ])
    
    # Only summarize if there's substantial new history
    if count_tokens(history_text) > 200: # Arbitrary threshold for summarization trigger
        chain = summary_prompt | summarizer_llm | StrOutputParser()
        new_summary = chain.invoke({"history_text": history_text, "current_summary": current_summary})
        print(f"[DEBUG] Summarized history. New summary tokens: {count_tokens(new_summary)}")
        return new_summary
    return current_summary # Return current if not much new to summarize


# Define the RAG chain
def create_rag_chain(vectorstore, llm):
    # 1. Create a history-aware retriever
    contextualize_q_system_prompt = (
        "Given a chat history and the latest user question "
        "which might reference context in the chat history, "
        "formulate a standalone question which can be understood "
        "without the chat history. Do NOT answer the question, "
        "just reformulate it if necessary and otherwise return it as is."
    )
    contextualize_q_prompt = ChatPromptTemplate.from_messages([
        ("system", contextualize_q_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ])
    history_aware_retriever = create_history_aware_retriever(
        llm, vectorstore.as_retriever(), contextualize_q_prompt
    )

    # 2. Create a document combining chain
    qa_system_prompt = (
        "You are an assistant for question-answering tasks. "
        "Use the following retrieved context and chat history to answer the question. "
        "If you don't know the answer, just say that you don't know. "
        "Keep the answer concise and relevant to the provided context."
        "\n\n{context}"
        "\n\nChat History Summary: {summary_context}" # New: include dynamic summary
    )
    qa_prompt = ChatPromptTemplate.from_messages([
        ("system", qa_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ])
    document_chain = create_stuff_documents_chain(llm, qa_prompt)

    # 3. Combine them into a retrieval chain
    # We need to manually manage the history and summary for the main chain
    rag_chain = RunnablePassthrough.assign(
        context=history_aware_retriever,
        summary_context=lambda x: x["summary_context"]
    ) | document_chain

    return rag_chain, history_aware_retriever

# Main conversational loop simulation
chat_history = []
current_summary = ""

rag_chain, retriever = create_rag_chain(vectorstore, llm)

def ask_question(question: str):
    global chat_history, current_summary

    # Prepare current context for token counting
    current_full_context = "\n".join([f"{msg.type}: {msg.content}" for msg in chat_history]) + "\n" + question
    current_context_tokens = count_tokens(current_full_context)

    print(f"[DEBUG] Current full context tokens before processing: {current_context_tokens}")

    # Dynamic context management: summarize if history is too long
    if current_context_tokens > MAX_CONTEXT_TOKENS * 0.75: # Trigger summarization proactively
        print("[DEBUG] Context approaching limit, initiating summarization...")
        current_summary = summarize_history(chat_history, current_summary)
        chat_history = [HumanMessage(content=f"Summary of previous conversation: {current_summary}"), AIMessage(content="Understood.")] # Reset history with summary
        print(f"[DEBUG] Chat history reset with new summary. Summary tokens: {count_tokens(current_summary)}")

    # Retrieve documents (history-aware)
    retrieved_docs = retriever.invoke({"input": question, "chat_history": chat_history})
    retrieved_context_str = "\n".join([doc.page_content for doc in retrieved_docs])

    # Final prompt preparation for the main LLM call
    final_input_for_llm = {
        "input": question,
        "chat_history": chat_history,
        "context": retrieved_docs, # Pass documents explicitly for create_stuff_documents_chain
        "summary_context": current_summary # Pass the summary separately
    }

    # Invoke the RAG chain
    response = rag_chain.invoke(final_input_for_llm)
    ai_response_content = response["answer"]

    # Update chat history
    chat_history.append(HumanMessage(content=question))
    chat_history.append(AIMessage(content=ai_response_content))

    return ai_response_content


print("\n--- Starting Conversation ---")

res1 = ask_question("Tell me about the Q3 earnings report.")
print(f"User: Tell me about the Q3 earnings report.\nAI: {res1}\n")

res2 = ask_question("What were the key drivers for this revenue increase?")
print(f"User: What were the key drivers for this revenue increase?\nAI: {res2}\n")

res3 = ask_question("And how does CognitoFlow contribute to automation?")
print(f"User: And how does CognitoFlow contribute to automation?\nAI: {res3}\n")

# Simulate a long conversation that triggers summarization
for i in range(5):
    ask_question(f"Random question {i+1} to extend history.")

res_long = ask_question("Given our earlier discussion, what's the planned marketing strategy for CognitoFlow?")
print(f"User: Given our earlier discussion, what's the planned marketing strategy for CognitoFlow?\nAI: {res_long}\n")

print("--- Conversation End ---")

In this implementation:

• We initialize a vector store with company documents.
• count_tokens helps us monitor context length.
• summarize_history is a crucial function that leverages a (potentially smaller, cheaper) LLM to condense the chat history when it approaches our defined MAX_CONTEXT_TOKENS. This summary then replaces the verbose history, keeping the overall token count low.
• The main RAG chain dynamically retrieves relevant documents based on the current query and the (potentially summarized) chat history.
• The qa_system_prompt is updated to explicitly include the summary_context, ensuring the LLM has awareness of past discussions without consuming excessive tokens.

Optimization & Best Practices

Implementing Dynamic RAG is a foundational step, but several optimizations can further enhance performance and cost-efficiency:

1. Hybrid Search & Re-ranking: Combine semantic search (vector similarity) with keyword-based search for comprehensive retrieval. Then, use a re-ranking model (e.g., Cohere Rerank) to prioritize the most relevant documents before feeding them to the LLM.
2. Adaptive Chunking: Instead of fixed chunk sizes, use intelligent chunking strategies that consider document structure (headings, paragraphs) to maintain semantic coherence.
3. Multi-stage Summarization: For extremely long histories, consider a hierarchical summarization approach where summaries are periodically summarized themselves.
4. Caching: Cache summarized chat histories and frequently accessed document chunks to reduce redundant computations and vector database lookups.
5. Prompt Engineering for Summarization: Experiment with different summarization prompts to ensure the LLM retains key entities, decisions, and unanswered questions from the conversation history. You might even use a fine-tuned smaller model for this specific task.
6. Token Monitoring & Alerting: Implement robust logging and alerting for token usage. This allows you to proactively identify and address cases where context windows are nearing limits or costs are spiking unexpectedly.
7. User Feedback Loop: Incorporate user feedback on response quality. This data can be used to refine retrieval algorithms, summarization techniques, and prompt engineering, creating a self-improving system.

Business Impact & ROI

The strategic adoption of Dynamic RAG translates directly into tangible business benefits:

• Reduced LLM API Costs: By sending only relevant and condensed information to the LLM, organizations can expect to cut their token usage and associated API costs by 30-60%. This is particularly critical for high-volume applications or those with long user interactions.
• Improved Response Latency: Smaller, more focused prompts mean less data transfer and faster processing times for the LLM, leading to quicker responses and a superior user experience.
• Enhanced Accuracy & Relevance: By prioritizing the most pertinent information, Dynamic RAG minimizes 'hallucinations' and ensures that LLM responses are more accurate, contextually relevant, and directly address user queries.
• Increased Scalability: Efficient context management allows applications to handle more complex queries and longer conversations without hitting prohibitive cost or performance ceilings, enabling seamless scaling as user bases grow.
• Higher User Satisfaction: Users receive faster, more accurate, and more coherent responses, leading to greater engagement and satisfaction with the AI application.

Consider a customer support chatbot that handles thousands of queries daily. A 40% reduction in token usage could save tens of thousands of dollars annually, while simultaneously improving customer satisfaction through faster and more accurate resolutions. This is not just a technical optimization; it's a strategic business advantage.

Conclusion

The era of static, brute-force LLM context management is drawing to a close. Dynamic Retrieval Augmented Generation offers a sophisticated, cost-effective, and highly scalable approach to building intelligent applications. By embracing adaptive context management, developers can build more robust, performant, and economically viable AI solutions. This technique empowers businesses to fully unlock the potential of LLMs, delivering superior user experiences while maintaining a keen eye on operational efficiency and ROI. The future of AI engineering is dynamic, intelligent, and optimized – and Dynamic RAG is a cornerstone of that future.