This is a comprehensive design document for your Python package, llm-context-manager. You can use this directly as your project's README.md or as the architectural blueprint for development.

Project: llm-context-manager

1. Background & Problem Statement

Large Language Models (LLMs) are moving from simple "Prompt Engineering" (crafting a single query) to "Context Engineering" (managing a massive ecosystem of retrieved documents, tools, and history).

The current problem is Context Pollution:

1. Overloading: RAG (Retrieval Augmented Generation) pipelines often dump too much data, exceeding token limits.
2. Noise: Duplicate or irrelevant information confuses the model and increases hallucination rates.
3. Formatting Chaos: Different models (Claude vs. Llama vs. GPT) require different formatting (XML vs. Markdown vs. Plain Text), leading to messy, hard-to-maintain string concatenation code.
4. Black Box: Developers rarely see exactly what "context" was sent to the LLM until after a failure occurs.

The Solution: llm-context-manager acts as a middleware layer for the LLM pipeline. It creates a structured, optimized, and budget-aware "context payload" before it reaches the model.

2. Architecture: Where It Fits

The package sits strictly between the Retrieval/Agent Layer (e.g., LangChain, LlamaIndex) and the Execution Layer (the LLM API).

Diagram: The "Before" (Standard Pipeline)

Without llm-context-manager, retrieval is messy and often truncated arbitrarily.

graph LR
    A[User Query] --> B[LangChain Retriever]
    B --> C{Result: 15 Docs}
    C -->|Raw Dump| D[LLM Context Window]
    D -->|Token Limit Exceeded!| E[Truncated/Error]

Diagram: The "After" (With llm-context-manager)

With your package, the context is curated, prioritized, and formatted.

graph LR
    A[User Query] --> B[LangChain Retriever]
    B --> C[Raw Data: 15 Docs + History]
    C --> D[**llm-context-manager**]
    
    subgraph "Your Middleware"
    D --> E["1. Token Budgeting"]
    E --> F["2. Semantic Pruning"]
    F --> G["3. Formatting (XML/JSON)"]
    end
    
    G --> H[Optimized Prompt]
    H --> I[LLM API]

3. Key Features & Tools

Here is the breakdown of the 4 core modules, the features they provide, and the libraries powering them.

Module A: The Budget Controller (budget)

• Goal: Ensure the context never exceeds the model's limit (e.g., 8192 tokens) while keeping the most important information.
• Feature: PriorityQueue. Users assign a priority (Critical, High, Medium, Low) to every piece of context. If the budget is full, "Low" items are dropped first.
• Supported Providers & Tools:
  • OpenAI (gpt-4, o1, o3, etc.): tiktoken — fast, local token counting.
  • HuggingFace (meta-llama/..., mistralai/...): tokenizers — for open-source models.
  • Google (gemini-2.0-flash, gemma-...): google-genai — API-based count_tokens.
  • Anthropic (claude-sonnet-4-20250514, etc.): anthropic — API-based count_tokens.
• Installation (pick what you need):

  pip install llm-context-manager[openai]       # tiktoken
  pip install llm-context-manager[huggingface]   # tokenizers
  pip install llm-context-manager[google]        # google-genai
  pip install llm-context-manager[anthropic]     # anthropic
  pip install llm-context-manager[all]           # everything
  

Module B: The Semantic Pruner (prune)

• Goal: Remove redundancy. If three retrieved documents say "Python is great," keep only the best one.
• Features:
  • Deduplicator (block-level): Calculates cosine similarity between context blocks and removes duplicate blocks. Among duplicates, the highest-priority block is kept.
  • Deduplicator.deduplicate_chunks() (chunk-level): Splits a single block's content by separator (e.g. \n\n), deduplicates the chunks internally, and reassembles the cleaned content. Ideal for RAG results where multiple retrieved chunks within one block are semantically redundant.
• Tools:
  • FastEmbed: Lightweight embedding generation (CPU-friendly, no heavy PyTorch needed).
  • Numpy: For efficient vector math (dot products).
• Installation:

  pip install llm-context-manager[prune]
  

Module C: Context Distillation (distill)

• Goal: Compress individual blocks by removing non-essential tokens (e.g., reduces a 5000-token document to 2500 tokens) using a small ML model.
• Feature: Compressor. Uses LLMLingua-2 (small BERT-based token classifier) to keep only the most important words.
• Tools:
  • llmlingua: Microsoft's library for prompt compression.
  • onnxruntime / transformers: For running the small BERT model.
• Installation:

  pip install llm-context-manager[distill]
  

Module D: The Formatter (format)

• Goal: Adapt the text structure to the specific LLM being used without changing the data.

• Feature: ModelAdapter.

• Claude Mode: Wraps data in XML tags (<doc id="1">...</doc>).

• Llama Mode: Uses specific Markdown headers or [INST] tags.

• Tools:

• Jinja2: For powerful, logic-based string templates.

• Pydantic: To enforce strict schema validation on the input data.

Module E: Observability (inspect)

• Goal: Let the developer see exactly what is happening.
• Feature: ContextVisualizer and Snapshot. Prints a colored bar chart of token usage to the terminal and saves the final prompt to a JSON file for debugging.
• Tools:
• Rich: For beautiful terminal output and progress bars.

4. Installation & Usage Guide

Installation

pip install llm-context-manager

Example Usage Code

This is how a developer would use your package in their Python script.

from context_manager import ContextEngine, ContextBlock
from context_manager.strategies import PriorityPruning

# 1. Initialize the Engine
engine = ContextEngine(
    model="gpt-4",
    token_limit=4000,
    pruning_strategy=PriorityPruning()
)

# 2. Add Data (Raw from LangChain/User)
# System Prompt (Critical Priority)
engine.add(ContextBlock(
    content="You are a helpful AI assistant for financial analysis.",
    role="system",
    priority="critical"
))

# User History (High Priority)
engine.add(ContextBlock(
    content="User: What is the PE ratio of Apple?",
    role="history",
    priority="high"
))

# Retrieved Documents (Medium Priority - acting as 'Low' if space is tight)
docs = ["Apple PE is 28...", "Microsoft PE is 35...", "Banana prices are up..."]
for doc in docs:
    engine.add(ContextBlock(
        content=doc,
        role="rag_context",
        priority="medium"
    ))

# 3. BUILD THE CONTEXT
# This triggers the budgeting, pruning, and formatting
final_prompt = engine.compile()

# 4. INSPECT (Optional)
engine.visualize() 
# Output: [██ system (10%) ][████ history (20%) ][██████ rag (40%) ][    unused (30%)   ]

# 5. SEND TO LLM
openai.ChatCompletion.create(messages=final_prompt)

5. Roadmap for Development

1. v0.1 (MVP): Implement tiktoken counting and PriorityPruning. Just get the math right.
2. v0.2 (Structure): Add Jinja2 templates for "Claude Style" (XML) vs "OpenAI Style" (JSON).
3. v0.3 (Smarts): Integrate FastEmbed to automatically detect and remove duplicate documents.
4. v0.4 (Vis): Add the Rich terminal visualization.

This design gives you a clear path to building a high-value tool that solves a specific, painful problem for AI engineers.