feat: update model configurations and enhance inference handling

AshishKumar4 · AshishKumar4 · commit f7960d5b4e5c · 2025-12-22T17:17:21.000-05:00
diff --git a/worker/agents/inferutils/config.ts b/worker/agents/inferutils/config.ts
@@ -129,10 +129,10 @@ const DEFAULT_AGENT_CONFIG: AgentConfig = {
         temperature: 0.6,
     },
     blueprint: {
-        name: AIModels.GEMINI_3_PRO_PREVIEW,
+        name: AIModels.GEMINI_3_FLASH_PREVIEW,
         reasoning_effort: 'high',
         max_tokens: 64000,
-        fallbackModel: AIModels.GEMINI_2_5_FLASH,
+        fallbackModel: AIModels.GEMINI_2_5_PRO,
         temperature: 1,
     },
     projectSetup: {
@@ -159,10 +159,10 @@ const DEFAULT_AGENT_CONFIG: AgentConfig = {
         fallbackModel: AIModels.GEMINI_2_5_PRO,
     },
     deepDebugger: {
-        name: AIModels.GEMINI_3_PRO_PREVIEW,
+        name: AIModels.GEMINI_3_FLASH_PREVIEW,
         reasoning_effort: 'high',
         max_tokens: 8000,
-        temperature: 0.5,
+        temperature: 1,
         fallbackModel: AIModels.GEMINI_2_5_FLASH,
     },
     fileRegeneration: {
@@ -173,10 +173,10 @@ const DEFAULT_AGENT_CONFIG: AgentConfig = {
         fallbackModel: AIModels.GEMINI_2_5_FLASH,
     },
     agenticProjectBuilder: {
-        name: AIModels.GEMINI_2_5_PRO,
+        name: AIModels.GEMINI_3_FLASH_PREVIEW,
         reasoning_effort: 'high',
         max_tokens: 8000,
-        temperature: 0.5,
+        temperature: 1,
         fallbackModel: AIModels.GEMINI_2_5_FLASH,
     },
 };
diff --git a/worker/agents/inferutils/infer.ts b/worker/agents/inferutils/infer.ts
@@ -19,6 +19,64 @@ The response you provided was either in an incorrect/unparsable format or was in
 Please provide a valid response that matches the expected output format exactly.
 `;
 
+/**
+ * Resolves model configuration with field-by-field merge
+ *
+ * Precedence: userConfig > AGENT_CONFIG defaults
+ * Each field is resolved independently (first defined value wins)
+ */
+function resolveModelConfig(
+    agentActionName: AgentActionKey,
+    userConfig?: ModelConfig,
+): ModelConfig {
+    const defaultConfig = AGENT_CONFIG[agentActionName];
+
+    const merged: ModelConfig = {
+        name: userConfig?.name ?? defaultConfig.name,
+        reasoning_effort: userConfig?.reasoning_effort ?? defaultConfig.reasoning_effort,
+        max_tokens: userConfig?.max_tokens ?? defaultConfig.max_tokens,
+        temperature: userConfig?.temperature ?? defaultConfig.temperature,
+        fallbackModel: userConfig?.fallbackModel ?? defaultConfig.fallbackModel,
+    };
+
+    // Validate model name - try userConfig first, then default
+    const modelCandidates = [userConfig?.name, defaultConfig.name]
+        .filter((n): n is AIModels | string => n !== undefined);
+
+    let validModelName: AIModels | string | undefined;
+    for (const candidate of modelCandidates) {
+        if (!isValidAIModel(candidate)) {
+            logger.warn(`Model ${candidate} not valid, trying next`);
+            continue;
+        }
+        const check = validateAgentConstraints(agentActionName, candidate);
+        if (check.constraintEnabled && !check.valid) {
+            logger.warn(`Model ${candidate} violates constraints for ${agentActionName}`);
+            continue;
+        }
+        validModelName = candidate;
+        break;
+    }
+
+    if (!validModelName) {
+        logger.warn(`No valid model found for ${agentActionName}, using default`);
+        validModelName = defaultConfig.name;
+    }
+    merged.name = validModelName;
+
+    // Validate fallback model
+    if (merged.fallbackModel) {
+        const fallbackCheck = validateAgentConstraints(agentActionName, merged.fallbackModel);
+        if (fallbackCheck.constraintEnabled && !fallbackCheck.valid) {
+            logger.warn(`Fallback ${merged.fallbackModel} violates constraints, using default`);
+            merged.fallbackModel = defaultConfig.fallbackModel;
+        }
+    }
+
+    logger.info(`Resolved config for ${agentActionName}: model=${merged.name}, fallback=${merged.fallbackModel}`);
+    return merged;
+}
+
 /**
  * Helper function to execute AI inference with consistent error handling
  * @param params Parameters for the inference operation
@@ -39,7 +97,6 @@ interface InferenceParamsBase {
         onChunk: (chunk: string) => void;
     };
     reasoning_effort?: ReasoningEffort;
-    modelConfig?: ModelConfig;
     context: InferenceContext;
     onAssistantMessage?: (message: Message) => Promise<void>;
     completionConfig?: CompletionConfig;
@@ -72,98 +129,23 @@ export async function executeInference<T extends z.AnyZodObject>(   {
     agentActionName,
     format,
     modelName,
-    modelConfig,
     context,
     onAssistantMessage,
     completionConfig,
 }: InferenceParamsBase &    {
     schema?: T;
     format?: SchemaFormat;
 }): Promise<InferResponseString | InferResponseObject<T> | null> {
-    let conf: ModelConfig | undefined;
-    
-    if (modelConfig) {
-        // Use explicitly provided model config
-        conf = modelConfig;
-    } else if (context?.metadata.userId && context?.userModelConfigs) {
-        // Try to get user-specific configuration from context cache
-        conf = context.userModelConfigs[agentActionName];
-        if (conf) {
-            logger.info(`Using user configuration for ${agentActionName}: ${JSON.stringify(conf)}`);
-        } else {
-            logger.info(`No user configuration for ${agentActionName}, using AGENT_CONFIG defaults`, context.userModelConfigs, context);
-        }
-
-        // If conf.name is not a valid AIModels enum value, fall back to defaults
-        if (conf && !isValidAIModel(conf.name)) {
-            logger.warn(`User config model ${conf.name} not in defined AIModels, falling back to defaults`);
-            conf = undefined;
-        }
-
-        // If conf.name violates agent constraints, fall back to defaults
-        if (conf && conf.name) {
-            const constraintCheck = validateAgentConstraints(agentActionName, conf.name);
+    // Resolve config with clear precedence: userConfig > defaults
+    const resolvedConfig = resolveModelConfig(
+        agentActionName,
+        context?.userModelConfigs?.[agentActionName],
+    );
 
-            if (constraintCheck.constraintEnabled && !constraintCheck.valid) {
-                logger.warn(
-                    `User config model ${conf.name} violates constraints for ${agentActionName}, falling back to defaults. ` +
-                    `Allowed models: ${constraintCheck.allowedModels?.join(', ')}`
-                );
-                conf = undefined; // Trigger fallback to AGENT_CONFIG[agentActionName]
-            }
-        }
-
-        // Validate fallback model too (if exists in conf)
-        if (conf && conf.fallbackModel) {
-            const fallbackCheck = validateAgentConstraints(agentActionName, conf.fallbackModel);
-
-            if (fallbackCheck.constraintEnabled && !fallbackCheck.valid) {
-                logger.warn(
-                    `User config fallback model ${conf.fallbackModel} violates constraints for ${agentActionName}, removing fallback`
-                );
-                conf.fallbackModel = undefined; // Remove invalid fallback
-            }
-        }
-
-        // If conf.name is not in defined AIModels, fall back to defaults
-        if (conf && !(conf.name in AIModels)) {
-            logger.warn(`User config model ${conf.name} not in defined AIModels, falling back to defaults`);
-            conf = undefined;
-        }
-
-        // If conf.name violates agent constraints, fall back to defaults
-        if (conf && conf.name) {
-            const constraintCheck = validateAgentConstraints(agentActionName, conf.name);
-
-            if (constraintCheck.constraintEnabled && !constraintCheck.valid) {
-                logger.warn(
-                    `User config model ${conf.name} violates constraints for ${agentActionName}, falling back to defaults. ` +
-                    `Allowed models: ${constraintCheck.allowedModels?.join(', ')}`
-                );
-                conf = undefined; // Trigger fallback to AGENT_CONFIG[agentActionName]
-            }
-        }
-
-        // Validate fallback model too (if exists in conf)
-        if (conf && conf.fallbackModel) {
-            const fallbackCheck = validateAgentConstraints(agentActionName, conf.fallbackModel);
-
-            if (fallbackCheck.constraintEnabled && !fallbackCheck.valid) {
-                logger.warn(
-                    `User config fallback model ${conf.fallbackModel} violates constraints for ${agentActionName}, removing fallback`
-                );
-                conf.fallbackModel = undefined; // Remove invalid fallback
-            }
-        }
-    }
-
-    // Use the final config or fall back to AGENT_CONFIG defaults
-    const finalConf = conf || AGENT_CONFIG[agentActionName];
-
-    modelName = modelName || finalConf.name;
-    temperature = temperature || finalConf.temperature || 0.2;
-    maxTokens = maxTokens || finalConf.max_tokens || 16000;
-    reasoning_effort = reasoning_effort || finalConf.reasoning_effort;
+    modelName = modelName || resolvedConfig.name;
+    temperature = temperature ?? resolvedConfig.temperature ?? 0.2;
+    maxTokens = maxTokens || resolvedConfig.max_tokens || 16000;
+    reasoning_effort = reasoning_effort || resolvedConfig.reasoning_effort;
 
     // Exponential backoff for retries
     const backoffMs = (attempt: number) => Math.min(500 * Math.pow(2, attempt), 10000);
@@ -237,20 +219,13 @@ export async function executeInference<T extends z.AnyZodObject>(   {
                     messages.push(createAssistantMessage(error.response));
                     messages.push(createUserMessage(responseRegenerationPrompts));
                     useCheaperModel = true;
-                    
-                    // If this was a repetition error, apply a frequency penalty to the retry
-                    if (error.message.toLowerCase().includes('repetition')) {
-                        logger.info('Applying frequency penalty to retry due to repetition');
-                        // Create a temporary config override for this retry
-                        conf = {
-                            ...finalConf,
-                            frequency_penalty: 0.5 // Apply moderate penalty
-                        };
-                    }
                 }
             } else {
-                // Try using fallback model if available
-                modelName = conf?.fallbackModel || modelName;
+                // Switch to fallback model if available
+                if (resolvedConfig.fallbackModel && resolvedConfig.fallbackModel !== modelName) {
+                    logger.info(`Switching to fallback model: ${resolvedConfig.fallbackModel}`);
+                    modelName = resolvedConfig.fallbackModel;
+                }
             }
 
             if (!isLastAttempt) {
diff --git a/worker/agents/operations/AgenticProjectBuilder.ts b/worker/agents/operations/AgenticProjectBuilder.ts
@@ -4,8 +4,7 @@ import {
     Message,
     ConversationMessage,
 } from '../inferutils/common';
-import { AgentActionKey, ModelConfig } from '../inferutils/config.types';
-import { AGENT_CONFIG } from '../inferutils/config';
+import { AgentActionKey } from '../inferutils/config.types';
 import { withRenderer } from '../tools/customTools';
 import { RenderToolCall } from './UserConversationProcessor';
 import { PROMPT_UTILS } from '../prompts';
@@ -47,7 +46,6 @@ export interface AgenticProjectBuilderInputs {
     toolRenderer: RenderToolCall;
     onToolComplete?: (message: Message) => Promise<void>;
     onAssistantMessage?: (message: Message) => Promise<void>;
-    modelConfigOverride?: ModelConfig;
 }
 
 export interface AgenticProjectBuilderOutputs {
@@ -293,7 +291,6 @@ export class AgenticProjectBuilderOperation extends AgentOperationWithTools<
 
         return {
             agentActionName: 'agenticProjectBuilder' as AgentActionKey,
-            modelConfig: inputs.modelConfigOverride || AGENT_CONFIG.agenticProjectBuilder,
             completionSignalName: 'mark_generation_complete',
             operationalMode: inputs.operationalMode,
             allowWarningInjection: inputs.operationalMode === 'initial',
diff --git a/worker/agents/operations/DeepDebugger.ts b/worker/agents/operations/DeepDebugger.ts
@@ -1,7 +1,6 @@
 import { createSystemMessage, createUserMessage, Message } from '../inferutils/common';
 import { AgentActionKey, ModelConfig } from '../inferutils/config.types';
 import { CompletionConfig, InferError, InferResponseString } from '../inferutils/core';
-import { AGENT_CONFIG } from '../inferutils/config';
 import { buildDebugTools } from '../tools/customTools';
 import { RenderToolCall } from './UserConversationProcessor';
 import { PROMPT_UTILS } from '../prompts';
@@ -82,7 +81,6 @@ export interface DeepDebuggerInputs {
     runtimeErrors?: RuntimeError[];
     streamCb?: (chunk: string) => void;
     toolRenderer?: RenderToolCall;
-    modelConfigOverride?: ModelConfig;
 }
 
 export interface DeepDebuggerOutputs {
@@ -193,7 +191,6 @@ export class DeepDebuggerOperation extends AgentOperationWithTools<
 
         return {
             agentActionName: 'deepDebugger' as AgentActionKey,
-            modelConfig: inputs.modelConfigOverride || AGENT_CONFIG.deepDebugger,
             completionSignalName: 'mark_debugging_complete',
             operationalMode: 'initial' as const,
             allowWarningInjection: true,
diff --git a/worker/agents/operations/common.ts b/worker/agents/operations/common.ts
@@ -1,7 +1,7 @@
 import { StructuredLogger } from "../../logger";
 import { GenerationContext } from "../domain/values/GenerationContext";
 import { Message } from "../inferutils/common";
-import { InferenceContext, AgentActionKey, ModelConfig } from "../inferutils/config.types";
+import { InferenceContext, AgentActionKey } from "../inferutils/config.types";
 import { createUserMessage, createSystemMessage, createAssistantMessage } from "../inferutils/common";
 import { generalSystemPromptBuilder, USER_PROMPT_FORMATTER } from "../prompts";
 import { CodeSerializerType } from "../utils/codeSerializers";
@@ -121,7 +121,6 @@ export abstract class AgentOperationWithTools<
         session: TSession
     ): {
         agentActionName: AgentActionKey;
-        modelConfig: ModelConfig;
         completionSignalName?: string;
         operationalMode?: "initial" | "followup";
         allowWarningInjection?: boolean;
@@ -154,7 +153,6 @@ export abstract class AgentOperationWithTools<
             messages: Message[];
             tools: ToolDefinition<unknown, unknown>[];
             agentActionName: AgentActionKey;
-            modelConfig: ModelConfig;
             streamCb?: (chunk: string) => void;
             onAssistantMessage?: (message: Message) => Promise<void>;
             completionConfig?: CompletionConfig;
@@ -165,7 +163,6 @@ export abstract class AgentOperationWithTools<
             messages,
             tools,
             agentActionName,
-            modelConfig,
             streamCb,
             onAssistantMessage,
             completionConfig,
@@ -235,7 +232,6 @@ export abstract class AgentOperationWithTools<
                 abortSignal: controller.signal
             },
             agentActionName,
-            modelConfig,
             messages,
             tools: wrappedTools,
             stream: wrappedStreamCb
@@ -262,7 +258,6 @@ export abstract class AgentOperationWithTools<
 
         const {
             agentActionName,
-            modelConfig,
             completionSignalName,
             operationalMode,
             allowWarningInjection,
@@ -280,7 +275,6 @@ export abstract class AgentOperationWithTools<
             messages,
             tools: rawTools,
             agentActionName,
-            modelConfig,
             streamCb: callbacks.streamCb,
             onAssistantMessage: callbacks.onAssistantMessage,
             completionConfig,
