Katkat
Back

Manage Context Windows in an Applied AI System

RogoMachine Learning Engineer · Onsite
MediumMachine Learning EngineerOnsiteML System Design

Last updated

This question tests the ability to design a robust context-management layer for an LLM-based application that balances limited token capacity against competing sources of information.

Problem Statement

You are designing a context-management middleware for an enterprise RAG application that assists analysts in processing financial documents. The system pulls from four sources: the current user prompt, a rolling conversation history (last 5 turns), retrieved document snippets from a vector DB, and tool execution outputs (e.g., Python code execution results).

The model has a hard limit of 32k tokens. Your system must select and structure the optimal context while ensuring the model remains responsive and accurate.

Requirements:

  1. Context Prioritization: Define a strategy for truncating or prioritizing input when the total token count exceeds the limit. How do you handle conflicts between history and fresh search results?
  2. Provenance: Ensure the system tracks which part of the context originated from which source to allow for UI-level citations.
  3. Prompt Injection Defense: Propose a mechanism to isolate user-provided text from system instructions and retrieved evidence.
  4. Graceful Degradation: What happens when the context is too large even after initial truncation? Propose a fallback strategy.

Constraints:

  • Assume an average tokenization overhead of 1.3x for English text.
  • The system must provide a 'System Prompt' (static) and an 'Analysis Prompt' (dynamic).
  • The solution should account for different token costs between input and output.

Hints

Hint 1

Consider a hierarchical prioritization scheme: System Instructions > Current Query > Tool Results > Retrieved Evidence > Conversation History.

Hint 2

For massive context, evaluate the trade-off between semantic summarization of old history versus simple truncation.

Hint 3

Use XML tags or delimiter patterns to logically encapsulate different data sources and define clear boundaries for the model.

Follow-up questions

  • How would you measure the performance impact of your context-selection logic on retrieval-augmented generation accuracy?
  • If the user asks a question that spans a document retrieved in the first turn but not the current turn, how does your system handle memory?
  • How would you implement 'streaming' for the context-building process to minimize Time to First Token (TTFT)?