Ai context management system

[jcourson-bg]

Background

Managing conversation history to fit within a Language Model's context window is a common challenge, often leading to complex, manual token counting and pruning logic. This PR addresses this by introducing a robust and developer-friendly context management system. The goal is to provide a clean, composable, and automatic way to ensure prompts remain within token limits, enhancing the AI SDK's "easy AI" philosophy.

Summary

This PR implements a contextWindow middleware for the AI SDK, designed to automatically prune message histories to fit within a specified token limit. Key features include:

• Middleware Integration: Seamlessly wraps any LanguageModelV3 model, making context management transparent to the application logic.
• Configurable Pruning Strategies: Supports sliding-window (default), keep-boundaries, and priority-based strategies to control which messages are retained.
• Message Priority System: Allows users to assign custom priorities (critical, high, normal, low) to messages, enabling fine-grained control over pruning behavior.
• Smart Token Estimation: Provides a fast, dependency-free character-based token estimator, with an option to integrate custom or external tokenizers for higher accuracy.
• Model Presets: Includes pre-configured maxPromptTokens for popular models (e.g., GPT-4o, Claude 3.5 Sonnet, Gemini 1.5 Pro) to simplify setup and ensure safe defaults.
• Pruning Callback: Offers an onPrune callback for observability and debugging when messages are removed.

This system is implemented in a new /packages/ai/src/context/ directory, including core logic, token estimation, and comprehensive tests.

Manual Verification

The following examples demonstrate the functionality:

• examples/ai-core/src/middleware/context-window-basic.ts: Shows basic usage of the contextWindow middleware with a long conversation history, demonstrating automatic pruning.
• examples/ai-core/src/middleware/context-window-preset.ts: Illustrates how to use contextWindowForModel with model-specific presets for simplified configuration.
• examples/ai-core/src/middleware/context-window-priority.ts: Highlights the priority-based strategy with custom assignPriority logic, ensuring important messages are retained.

To verify:

1. Run each example file (pnpm tsx examples/ai-core/src/middleware/context-window-basic.ts, etc.).
2. Observe the console output, specifically the onPrune callback messages, to confirm that messages are being pruned as expected and the final prompt fits within the token limits.

Checklist

• Tests have been added / updated (for bug fixes / features)
• Documentation has been added / updated (for bug fixes / features)
• A patch changeset for relevant packages has been added (for bug fixes / features - run pnpm changeset in the project root)
• I have reviewed this pull request (self-review)

Future Work

• Explore integration with external tokenizers (e.g., tiktoken) directly into presets or as optional dependencies for more precise token counting.
• Consider more advanced pruning strategies, such as summarization-based pruning, though this adds complexity and dependency on LLMs.
• Add more model presets as new models are released or integrated.

---

 

[cursor]

Cursor Agent can help with this pull request. Just @cursor in comments and I'll start working on changes in this branch.
_{Learn more about Cursor Agents}