Make model context window optional (Issue #362)

Author: bhouston

packages/agent/src/core/llm/providers/anthropic.ts

@@ -12,8 +12,18 @@ import {
   ProviderOptions,
 } from '../types.js';
 
-// Cache for model context window sizes
-const modelContextWindowCache: Record<string, number> = {};
+const ANTHROPIC_CONTEXT_WINDOWS: Record<string, number> = {
+  'claude-3-7-sonnet-20250219': 200000,
+  'claude-3-7-sonnet-latest': 200000,
+  'claude-3-5-sonnet-20241022': 200000,
+  'claude-3-5-sonnet-latest': 200000,
+  'claude-3-haiku-20240307': 200000,
+  'claude-3-opus-20240229': 200000,
+  'claude-3-sonnet-20240229': 200000,
+  'claude-2.1': 100000,
+  'claude-2.0': 100000,
+  'claude-instant-1.2': 100000,
+};
 
 /**
  * Anthropic-specific options
@@ -87,7 +97,7 @@ function addCacheControlToMessages(
 function tokenUsageFromMessage(
   message: Anthropic.Message,
   model: string,
-  contextWindow: number,
+  contextWindow: number | undefined,
 ) {
   const usage = new TokenUsage();
   usage.input = message.usage.input_tokens;
@@ -100,7 +110,7 @@ function tokenUsageFromMessage(
   return {
     usage,
     totalTokens,
-    maxTokens: contextWindow,
+    contextWindow,
   };
 }
 
@@ -131,64 +141,12 @@ export class AnthropicProvider implements LLMProvider {
     });
   }
 
-  /**
-   * Fetches the model context window size from the Anthropic API
-   *
-   * @returns The context window size
-   * @throws Error if the context window size cannot be determined
-   */
-  private async getModelContextWindow(): Promise<number> {
-    const cachedContextWindow = modelContextWindowCache[this.model];
-    if (cachedContextWindow !== undefined) {
-      return cachedContextWindow;
-    }
-    const response = await this.client.models.list();
-
-    if (!response?.data || !Array.isArray(response.data)) {
-      throw new Error(`Invalid response from models.list() for ${this.model}`);
-    }
-
-    // Try to find the exact model
-    let model = response.data.find((m) => m.id === this.model);
-
-    // If not found, try to find a model that starts with the same name
-    // This helps with model aliases like 'claude-3-sonnet-latest'
-    if (!model) {
-      // Split by '-latest' or '-20' to get the base model name
-      const parts = this.model.split('-latest');
-      const modelPrefix =
-        parts.length > 1 ? parts[0] : this.model.split('-20')[0];
-
-      if (modelPrefix) {
-        model = response.data.find((m) => m.id.startsWith(modelPrefix));
-
-        if (model) {
-          console.info(
-            `Model ${this.model} not found, using ${model.id} for context window size`,
-          );
-        }
-      }
-    }
-
-    // Using type assertion to access context_window property
-    // The Anthropic API returns context_window but it may not be in the TypeScript definitions
-    if (model && 'context_window' in model) {
-      const contextWindow = (model as any).context_window;
-      // Cache the result for future use
-      modelContextWindowCache[this.model] = contextWindow;
-      return contextWindow;
-    } else {
-      throw new Error(
-        `No context window information found for model: ${this.model}`,
-      );
-    }
-  }
-
   /**
    * Generate text using Anthropic API
    */
   async generateText(options: GenerateOptions): Promise<LLMResponse> {
-    const modelContextWindow = await this.getModelContextWindow();
+    const modelContextWindow = ANTHROPIC_CONTEXT_WINDOWS[this.model];
+
     const { messages, functions, temperature = 0.7, maxTokens, topP } = options;
 
     // Extract system message
@@ -252,7 +210,7 @@ export class AnthropicProvider implements LLMProvider {
       toolCalls: toolCalls,
       tokenUsage: tokenInfo.usage,
       totalTokens: tokenInfo.totalTokens,
-      maxTokens: tokenInfo.maxTokens,
+      contextWindow: tokenInfo.contextWindow,
     };
   }
 

packages/agent/src/core/llm/providers/ollama.ts

@@ -24,8 +24,7 @@ import {
 
 // Define model context window sizes for Ollama models
 // These are approximate and may vary based on specific model configurations
-const OLLAMA_MODEL_LIMITS: Record<string, number> = {
-  default: 4096,
+const OLLAMA_CONTEXT_WINDOWS: Record<string, number> = {
   llama2: 4096,
   'llama2-uncensored': 4096,
   'llama2:13b': 4096,
@@ -136,19 +135,21 @@ export class OllamaProvider implements LLMProvider {
     const totalTokens = tokenUsage.input + tokenUsage.output;
 
     // Extract the base model name without specific parameters
-    const baseModelName = this.model.split(':')[0];
     // Check if model exists in limits, otherwise use base model or default
-    const modelMaxTokens =
-      OLLAMA_MODEL_LIMITS[this.model] ||
-      (baseModelName ? OLLAMA_MODEL_LIMITS[baseModelName] : undefined) ||
-      4096; // Default fallback
+    let contextWindow = OLLAMA_CONTEXT_WINDOWS[this.model];
+    if (!contextWindow) {
+      const baseModelName = this.model.split(':')[0];
+      if (baseModelName) {
+        contextWindow = OLLAMA_CONTEXT_WINDOWS[baseModelName];
+      }
+    }
 
     return {
       text: content,
       toolCalls: toolCalls,
       tokenUsage: tokenUsage,
       totalTokens,
-      maxTokens: modelMaxTokens,
+      contextWindow,
     };
   }
 

packages/agent/src/core/llm/providers/openai.ts

@@ -20,8 +20,7 @@ import type {
 } from 'openai/resources/chat';
 
 // Define model context window sizes for OpenAI models
-const OPENAI_MODEL_LIMITS: Record<string, number> = {
-  default: 128000,
+const OPENA_CONTEXT_WINDOWS: Record<string, number> = {
   'o3-mini': 200000,
   'o1-pro': 200000,
   o1: 200000,
@@ -136,14 +135,14 @@ export class OpenAIProvider implements LLMProvider {
 
       // Calculate total tokens and get max tokens for the model
       const totalTokens = tokenUsage.input + tokenUsage.output;
-      const modelMaxTokens = OPENAI_MODEL_LIMITS[this.model] || 8192; // Default fallback
+      const contextWindow = OPENA_CONTEXT_WINDOWS[this.model];
 
       return {
         text: content,
         toolCalls,
         tokenUsage,
         totalTokens,
-        maxTokens: modelMaxTokens,
+        contextWindow,
       };
     } catch (error) {
       throw new Error(`Error calling OpenAI API: ${(error as Error).message}`);

packages/agent/src/core/llm/types.ts

@@ -82,7 +82,7 @@ export interface LLMResponse {
   tokenUsage: TokenUsage;
   // Add new fields for context window tracking
   totalTokens?: number; // Total tokens used in this request
-  maxTokens?: number; // Maximum allowed tokens for this model
+  contextWindow?: number; // Maximum allowed tokens for this model
 }
 
 /**

packages/agent/src/core/toolAgent/statusUpdates.ts

@@ -14,12 +14,14 @@ import { ToolContext } from '../types.js';
  */
 export function generateStatusUpdate(
   totalTokens: number,
-  maxTokens: number,
+  contextWindow: number | undefined,
   tokenTracker: TokenTracker,
   context: ToolContext,
 ): Message {
   // Calculate token usage percentage
-  const usagePercentage = Math.round((totalTokens / maxTokens) * 100);
+  const usagePercentage = contextWindow
+    ? Math.round((totalTokens / contextWindow) * 100)
+    : undefined;
 
   // Get active sub-agents
   const activeAgents = context.agentTracker ? getActiveAgents(context) : [];
@@ -35,7 +37,9 @@ export function generateStatusUpdate(
   // Format the status message
   const statusContent = [
     `--- STATUS UPDATE ---`,
-    `Token Usage: ${formatNumber(totalTokens)}/${formatNumber(maxTokens)} (${usagePercentage}%)`,
+    contextWindow !== undefined
+      ? `Token Usage: ${formatNumber(totalTokens)}/${formatNumber(contextWindow)} (${usagePercentage}%)`
+      : '',
     `Cost So Far: ${tokenTracker.getTotalCost()}`,
     ``,
     `Active Sub-Agents: ${activeAgents.length}`,
@@ -47,9 +51,10 @@ export function generateStatusUpdate(
     `Active Browser Sessions: ${activeSessions.length}`,
     ...activeSessions.map((s) => `- ${s.id}: ${s.description}`),
     ``,
-    usagePercentage >= 50
-      ? `Your token usage is high (${usagePercentage}%). It is recommended to use the 'compactHistory' tool now to reduce context size.`
-      : `If token usage gets high (>50%), consider using the 'compactHistory' tool to reduce context size.`,
+    usagePercentage !== undefined &&
+      (usagePercentage >= 50
+        ? `Your token usage is high (${usagePercentage}%). It is recommended to use the 'compactHistory' tool now to reduce context size.`
+        : `If token usage gets high (>50%), consider using the 'compactHistory' tool to reduce context size.`),
     `--- END STATUS ---`,
   ].join('\n');
 

packages/agent/src/core/toolAgent/toolAgentCore.ts

@@ -151,34 +151,35 @@ export const toolAgent = async (
       maxTokens: localContext.maxTokens,
     };
 
-    const { text, toolCalls, tokenUsage, totalTokens, maxTokens } =
+    const { text, toolCalls, tokenUsage, totalTokens, contextWindow } =
       await generateText(provider, generateOptions);
 
     tokenTracker.tokenUsage.add(tokenUsage);
 
     // Send status updates based on frequency and token usage threshold
     statusUpdateCounter++;
-    if (totalTokens && maxTokens) {
-      const usagePercentage = Math.round((totalTokens / maxTokens) * 100);
-      const shouldSendByFrequency =
-        statusUpdateCounter >= STATUS_UPDATE_FREQUENCY;
-      const shouldSendByUsage = usagePercentage >= TOKEN_USAGE_THRESHOLD;
+    if (totalTokens) {
+      let statusTriggered = false;
+      statusTriggered ||= statusUpdateCounter >= STATUS_UPDATE_FREQUENCY;
+
+      if (contextWindow) {
+        const usagePercentage = Math.round((totalTokens / contextWindow) * 100);
+        statusTriggered ||= usagePercentage >= TOKEN_USAGE_THRESHOLD;
+      }
 
       // Send status update if either condition is met
-      if (shouldSendByFrequency || shouldSendByUsage) {
+      if (statusTriggered) {
         statusUpdateCounter = 0;
 
         const statusMessage = generateStatusUpdate(
           totalTokens,
-          maxTokens,
+          contextWindow,
           tokenTracker,
           localContext,
         );
 
         messages.push(statusMessage);
-        logger.debug(
-          `Sent status update to agent (token usage: ${usagePercentage}%)`,
-        );
+        logger.debug(`Sent status update to agent`);
       }
     }