Skip to content

Feat: update model configurations and enhance inference handling #289

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
AshishKumar4 merged 1 commit into from
Dec 22, 2025
Merged

Feat: update model configurations and enhance inference handling #289

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions worker/agents/inferutils/config.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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: {
Expand All @@ -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: {
Expand All @@ -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,
},
};
Expand Down
169 changes: 72 additions & 97 deletions worker/agents/inferutils/infer.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -39,7 +97,6 @@ interface InferenceParamsBase {
onChunk: (chunk: string) => void;
};
reasoning_effort?: ReasoningEffort;
modelConfig?: ModelConfig;
context: InferenceContext;
onAssistantMessage?: (message: Message) => Promise<void>;
completionConfig?: CompletionConfig;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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) {
Expand Down
5 changes: 1 addition & 4 deletions worker/agents/operations/AgenticProjectBuilder.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -47,7 +46,6 @@ export interface AgenticProjectBuilderInputs {
toolRenderer: RenderToolCall;
onToolComplete?: (message: Message) => Promise<void>;
onAssistantMessage?: (message: Message) => Promise<void>;
modelConfigOverride?: ModelConfig;
}

export interface AgenticProjectBuilderOutputs {
Expand Down Expand Up @@ -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',
Expand Down
3 changes: 0 additions & 3 deletions worker/agents/operations/DeepDebugger.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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';
Expand Down Expand Up @@ -82,7 +81,6 @@ export interface DeepDebuggerInputs {
runtimeErrors?: RuntimeError[];
streamCb?: (chunk: string) => void;
toolRenderer?: RenderToolCall;
modelConfigOverride?: ModelConfig;
}

export interface DeepDebuggerOutputs {
Expand Down Expand Up @@ -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,
Expand Down
8 changes: 1 addition & 7 deletions worker/agents/operations/common.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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";
Expand Down Expand Up @@ -121,7 +121,6 @@ export abstract class AgentOperationWithTools<
session: TSession
): {
agentActionName: AgentActionKey;
modelConfig: ModelConfig;
completionSignalName?: string;
operationalMode?: "initial" | "followup";
allowWarningInjection?: boolean;
Expand Down Expand Up @@ -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;
Expand All @@ -165,7 +163,6 @@ export abstract class AgentOperationWithTools<
messages,
tools,
agentActionName,
modelConfig,
streamCb,
onAssistantMessage,
completionConfig,
Expand Down Expand Up @@ -235,7 +232,6 @@ export abstract class AgentOperationWithTools<
abortSignal: controller.signal
},
agentActionName,
modelConfig,
messages,
tools: wrappedTools,
stream: wrappedStreamCb
Expand All @@ -262,7 +258,6 @@ export abstract class AgentOperationWithTools<

const {
agentActionName,
modelConfig,
completionSignalName,
operationalMode,
allowWarningInjection,
Expand All @@ -280,7 +275,6 @@ export abstract class AgentOperationWithTools<
messages,
tools: rawTools,
agentActionName,
modelConfig,
streamCb: callbacks.streamCb,
onAssistantMessage: callbacks.onAssistantMessage,
completionConfig,
Expand Down
Loading