diff --git a/packages/@n8n/ai-workflow-builder.ee/evaluations/__tests__/cli.test.ts b/packages/@n8n/ai-workflow-builder.ee/evaluations/__tests__/cli.test.ts index e28577586a9..549eac2c2f9 100644 --- a/packages/@n8n/ai-workflow-builder.ee/evaluations/__tests__/cli.test.ts +++ b/packages/@n8n/ai-workflow-builder.ee/evaluations/__tests__/cli.test.ts @@ -14,6 +14,7 @@ import type { SimpleWorkflow } from '@/types/workflow'; // Store mocks for dependencies const mockParseEvaluationArgs = jest.fn(); +const mockArgsToStageModels = jest.fn(); const mockSetupTestEnvironment = jest.fn(); const mockCreateAgent = jest.fn(); const mockGenerateRunId = jest.fn(); @@ -30,6 +31,7 @@ const mockCreatePairwiseEvaluator = jest.fn(); // Mock all external modules jest.mock('../cli/argument-parser', () => ({ parseEvaluationArgs: (): unknown => mockParseEvaluationArgs(), + argsToStageModels: (...args: unknown[]): unknown => mockArgsToStageModels(...args), getDefaultDatasetName: (suite: unknown): unknown => suite === 'pairwise' ? 'notion-pairwise-workflows' : 'workflow-builder-canvas-prompts', getDefaultExperimentName: (suite: unknown): unknown => @@ -98,9 +100,19 @@ function createMockArgs(overrides: Record = {}) { /** Helper to create mock environment */ function createMockEnvironment() { + const mockLlm = mock(); return { parsedNodeTypes: [] as INodeTypeDescription[], - llm: mock(), + llms: { + default: mockLlm, + supervisor: mockLlm, + responder: mockLlm, + discovery: mockLlm, + builder: mockLlm, + configurator: mockLlm, + parameterUpdater: mockLlm, + judge: mockLlm, + }, lsClient: mock(), }; } @@ -147,6 +159,7 @@ describe('CLI', () => { // Setup default mocks mockParseEvaluationArgs.mockReturnValue(createMockArgs()); + mockArgsToStageModels.mockReturnValue({ default: 'claude-sonnet-4.5' }); mockSetupTestEnvironment.mockResolvedValue(createMockEnvironment()); mockCreateAgent.mockReturnValue(createMockAgentInstance()); mockGenerateRunId.mockReturnValue('test-run-id'); diff --git a/packages/@n8n/ai-workflow-builder.ee/evaluations/__tests__/runner-langsmith.test.ts b/packages/@n8n/ai-workflow-builder.ee/evaluations/__tests__/runner-langsmith.test.ts index 270288ae44c..064daa1c41a 100644 --- a/packages/@n8n/ai-workflow-builder.ee/evaluations/__tests__/runner-langsmith.test.ts +++ b/packages/@n8n/ai-workflow-builder.ee/evaluations/__tests__/runner-langsmith.test.ts @@ -174,7 +174,8 @@ describe('Runner - LangSmith Mode', () => { expect(isLangsmithTargetOutput(result)).toBe(true); if (!isLangsmithTargetOutput(result)) throw new Error('Expected LangSmith target output'); - expect(generateWorkflow).toHaveBeenCalledWith('Create a workflow'); + // Callbacks are passed explicitly from the traceable wrapper (undefined in tests without traceable context) + expect(generateWorkflow).toHaveBeenCalledWith('Create a workflow', undefined); expect(evaluator.evaluate).toHaveBeenCalledWith( workflow, expect.objectContaining({ prompt: 'Create a workflow' }), diff --git a/packages/@n8n/ai-workflow-builder.ee/evaluations/cli/argument-parser.ts b/packages/@n8n/ai-workflow-builder.ee/evaluations/cli/argument-parser.ts index e2fc2d39a19..0bb5c32d358 100644 --- a/packages/@n8n/ai-workflow-builder.ee/evaluations/cli/argument-parser.ts +++ b/packages/@n8n/ai-workflow-builder.ee/evaluations/cli/argument-parser.ts @@ -1,9 +1,12 @@ /* eslint-disable @typescript-eslint/naming-convention */ import { z } from 'zod'; -import type { BuilderFeatureFlags } from '../../src/workflow-builder-agent.js'; +import { AVAILABLE_MODELS, DEFAULT_MODEL, type ModelId } from '@/llm-config'; +import type { BuilderFeatureFlags } from '@/workflow-builder-agent'; + import type { LangsmithExampleFilters } from '../harness/harness-types'; import { DEFAULTS } from '../support/constants'; +import type { StageModels } from '../support/environment.js'; export type EvaluationSuite = 'llm-judge' | 'pairwise' | 'programmatic' | 'similarity'; export type EvaluationBackend = 'local' | 'langsmith'; @@ -32,10 +35,39 @@ export interface EvaluationArgs { numJudges: number; featureFlags?: BuilderFeatureFlags; + + // Model configuration + /** Default model for all stages */ + model: ModelId; + /** Model for LLM judge evaluation */ + judgeModel?: ModelId; + /** Model for supervisor stage */ + supervisorModel?: ModelId; + /** Model for responder stage */ + responderModel?: ModelId; + /** Model for discovery stage */ + discoveryModel?: ModelId; + /** Model for builder stage */ + builderModel?: ModelId; + /** Model for configurator stage */ + configuratorModel?: ModelId; + /** Model for parameter updater (within configurator) */ + parameterUpdaterModel?: ModelId; } type CliValueKind = 'boolean' | 'string'; -type FlagGroup = 'input' | 'eval' | 'pairwise' | 'langsmith' | 'output' | 'feature' | 'advanced'; +type FlagGroup = + | 'input' + | 'eval' + | 'pairwise' + | 'langsmith' + | 'output' + | 'feature' + | 'model' + | 'advanced'; + +// Model ID validation schema +const modelIdSchema = z.enum(AVAILABLE_MODELS as [ModelId, ...ModelId[]]); const cliSchema = z .object({ @@ -65,6 +97,16 @@ const cliSchema = z langsmith: z.boolean().optional(), templateExamples: z.boolean().default(false), + + // Model configuration + model: modelIdSchema.default(DEFAULT_MODEL), + judgeModel: modelIdSchema.optional(), + supervisorModel: modelIdSchema.optional(), + responderModel: modelIdSchema.optional(), + discoveryModel: modelIdSchema.optional(), + builderModel: modelIdSchema.optional(), + configuratorModel: modelIdSchema.optional(), + parameterUpdaterModel: modelIdSchema.optional(), }) .strict(); @@ -185,6 +227,56 @@ const FLAG_DEFS: Record = { desc: 'Enable template examples phase', }, + // Model configuration + '--model': { + key: 'model', + kind: 'string', + group: 'model', + desc: `Default model for all stages (default: ${DEFAULT_MODEL})`, + }, + '--judge-model': { + key: 'judgeModel', + kind: 'string', + group: 'model', + desc: 'Model for LLM judge evaluation', + }, + '--supervisor-model': { + key: 'supervisorModel', + kind: 'string', + group: 'model', + desc: 'Model for supervisor stage', + }, + '--responder-model': { + key: 'responderModel', + kind: 'string', + group: 'model', + desc: 'Model for responder stage', + }, + '--discovery-model': { + key: 'discoveryModel', + kind: 'string', + group: 'model', + desc: 'Model for discovery stage', + }, + '--builder-model': { + key: 'builderModel', + kind: 'string', + group: 'model', + desc: 'Model for builder stage', + }, + '--configurator-model': { + key: 'configuratorModel', + kind: 'string', + group: 'model', + desc: 'Model for configurator stage', + }, + '--parameter-updater-model': { + key: 'parameterUpdaterModel', + kind: 'string', + group: 'model', + desc: 'Model for parameter updater', + }, + // Advanced '--judges': { key: 'numJudges', kind: 'string', group: 'advanced', desc: 'Number of LLM judges' }, }; @@ -217,6 +309,7 @@ const GROUP_TITLES: Record = { langsmith: 'LangSmith Options', output: 'Output', feature: 'Feature Flags', + model: 'Model Configuration', advanced: 'Advanced', }; @@ -235,6 +328,7 @@ function formatHelp(): string { 'langsmith', 'output', 'feature', + 'model', 'advanced', ]; @@ -431,6 +525,31 @@ export function parseEvaluationArgs(argv: string[] = process.argv.slice(2)): Eva donts: parsed.donts, numJudges: parsed.numJudges, featureFlags, + // Model configuration + model: parsed.model, + judgeModel: parsed.judgeModel, + supervisorModel: parsed.supervisorModel, + responderModel: parsed.responderModel, + discoveryModel: parsed.discoveryModel, + builderModel: parsed.builderModel, + configuratorModel: parsed.configuratorModel, + parameterUpdaterModel: parsed.parameterUpdaterModel, + }; +} + +/** + * Converts EvaluationArgs to StageModels for use with environment setup. + */ +export function argsToStageModels(args: EvaluationArgs): StageModels { + return { + default: args.model, + supervisor: args.supervisorModel, + responder: args.responderModel, + discovery: args.discoveryModel, + builder: args.builderModel, + configurator: args.configuratorModel, + parameterUpdater: args.parameterUpdaterModel, + judge: args.judgeModel, }; } diff --git a/packages/@n8n/ai-workflow-builder.ee/evaluations/cli/index.ts b/packages/@n8n/ai-workflow-builder.ee/evaluations/cli/index.ts index 69945e16d7a..45f2a1befe1 100644 --- a/packages/@n8n/ai-workflow-builder.ee/evaluations/cli/index.ts +++ b/packages/@n8n/ai-workflow-builder.ee/evaluations/cli/index.ts @@ -5,7 +5,7 @@ * Can be run directly or used as a reference for custom setups. */ -import type { BaseChatModel } from '@langchain/core/language_models/chat_models'; +import type { Callbacks } from '@langchain/core/callbacks/manager'; import type { INodeTypeDescription } from 'n8n-workflow'; import pLimit from 'p-limit'; @@ -13,6 +13,7 @@ import type { SimpleWorkflow } from '@/types/workflow'; import type { BuilderFeatureFlags } from '@/workflow-builder-agent'; import { + argsToStageModels, getDefaultDatasetName, getDefaultExperimentName, parseEvaluationArgs, @@ -34,33 +35,29 @@ import { loadDefaultTestCases, getDefaultTestCaseIds, } from './csv-prompt-loader'; -import { - consumeGenerator, - getChatPayload, - getTracingCallbacks, -} from '../harness/evaluation-helpers'; +import { consumeGenerator, getChatPayload } from '../harness/evaluation-helpers'; import { createLogger } from '../harness/logger'; import { generateRunId, isWorkflowStateValues } from '../langsmith/types'; import { EVAL_TYPES, EVAL_USERS } from '../support/constants'; -import { setupTestEnvironment, createAgent } from '../support/environment'; +import { setupTestEnvironment, createAgent, type ResolvedStageLLMs } from '../support/environment'; /** * Create a workflow generator function. * LangSmith tracing is handled via traceable() in the runner. - * We bridge the trace context to LangChain via getTracingCallbacks(). + * Callbacks are passed explicitly from the runner to ensure correct trace context + * under high concurrency (avoids AsyncLocalStorage race conditions). */ function createWorkflowGenerator( parsedNodeTypes: INodeTypeDescription[], - llm: BaseChatModel, + llms: ResolvedStageLLMs, featureFlags?: BuilderFeatureFlags, -): (prompt: string) => Promise { - return async (prompt: string): Promise => { +): (prompt: string, callbacks?: Callbacks) => Promise { + return async (prompt: string, callbacks?: Callbacks): Promise => { const runId = generateRunId(); - const callbacks = await getTracingCallbacks(); const agent = createAgent({ parsedNodeTypes, - llm, + llms, featureFlags, }); @@ -149,30 +146,35 @@ export async function runV2Evaluation(): Promise { ); } - // Setup environment + // Setup environment with per-stage model configuration const logger = createLogger(args.verbose); const lifecycle = createConsoleLifecycle({ verbose: args.verbose, logger }); - const env = await setupTestEnvironment(logger); + const stageModels = argsToStageModels(args); + const env = await setupTestEnvironment(stageModels, logger); // Validate LangSmith client early if langsmith backend is requested if (args.backend === 'langsmith' && !env.lsClient) { throw new Error('LangSmith client not initialized - check LANGSMITH_API_KEY'); } - // Create workflow generator (tracing handled via traceable() in runner) - const generateWorkflow = createWorkflowGenerator(env.parsedNodeTypes, env.llm, args.featureFlags); + // Create workflow generator with per-stage LLMs + const generateWorkflow = createWorkflowGenerator( + env.parsedNodeTypes, + env.llms, + args.featureFlags, + ); - // Create evaluators based on mode + // Create evaluators based on mode (using judge LLM for evaluation) const evaluators: Array> = []; switch (args.suite) { case 'llm-judge': - evaluators.push(createLLMJudgeEvaluator(env.llm, env.parsedNodeTypes)); + evaluators.push(createLLMJudgeEvaluator(env.llms.judge, env.parsedNodeTypes)); evaluators.push(createProgrammaticEvaluator(env.parsedNodeTypes)); break; case 'pairwise': evaluators.push( - createPairwiseEvaluator(env.llm, { + createPairwiseEvaluator(env.llms.judge, { numJudges: args.numJudges, }), ); diff --git a/packages/@n8n/ai-workflow-builder.ee/evaluations/harness/runner.ts b/packages/@n8n/ai-workflow-builder.ee/evaluations/harness/runner.ts index d056a7201b4..9cf78ad297b 100644 --- a/packages/@n8n/ai-workflow-builder.ee/evaluations/harness/runner.ts +++ b/packages/@n8n/ai-workflow-builder.ee/evaluations/harness/runner.ts @@ -1,10 +1,11 @@ +import type { Callbacks } from '@langchain/core/callbacks/manager'; import type { BaseMessage } from '@langchain/core/messages'; import { evaluate } from 'langsmith/evaluation'; import type { Run, Example } from 'langsmith/schemas'; import { traceable } from 'langsmith/traceable'; import pLimit from 'p-limit'; -import { runWithOptionalLimiter, withTimeout } from './evaluation-helpers'; +import { getTracingCallbacks, runWithOptionalLimiter, withTimeout } from './evaluation-helpers'; import { toLangsmithEvaluationResult } from './feedback'; import type { Evaluator, @@ -778,16 +779,22 @@ async function runLangsmith(config: LangsmithRunConfig): Promise { // Create traceable wrapper ONCE outside target function to avoid context leaking // when running concurrent evaluations. Pass all parameters explicitly (no closures). + // IMPORTANT: Get callbacks INSIDE the traceable wrapper where AsyncLocalStorage context + // is correctly set, then pass them explicitly to genFn to avoid race conditions. const traceableGenerateWorkflow = traceable( async (args: { prompt: string; - genFn: (prompt: string) => Promise; + genFn: (prompt: string, callbacks?: Callbacks) => Promise; limiter?: LlmCallLimiter; genTimeoutMs?: number; }): Promise => { + // Get callbacks inside traceable where context is correct + // Returns undefined if not in a traceable context (e.g., unit tests) + const callbacks = await getTracingCallbacks(); + return await runWithOptionalLimiter(async () => { return await withTimeout({ - promise: args.genFn(args.prompt), + promise: args.genFn(args.prompt, callbacks), timeoutMs: args.genTimeoutMs, label: 'workflow_generation', }); diff --git a/packages/@n8n/ai-workflow-builder.ee/evaluations/support/environment.ts b/packages/@n8n/ai-workflow-builder.ee/evaluations/support/environment.ts index 718b9d7cfd8..3082ce1e05f 100644 --- a/packages/@n8n/ai-workflow-builder.ee/evaluations/support/environment.ts +++ b/packages/@n8n/ai-workflow-builder.ee/evaluations/support/environment.ts @@ -4,10 +4,11 @@ import { MemorySaver } from '@langchain/langgraph'; import { Client } from 'langsmith/client'; import type { INodeTypeDescription } from 'n8n-workflow'; +import { DEFAULT_MODEL, getApiKeyEnvVar, MODEL_FACTORIES, type ModelId } from '@/llm-config'; +import type { BuilderFeatureFlags } from '@/workflow-builder-agent'; +import { WorkflowBuilderAgent } from '@/workflow-builder-agent'; + import { loadNodesFromFile } from './load-nodes.js'; -import { anthropicClaudeSonnet45 } from '../../src/llm-config.js'; -import type { BuilderFeatureFlags } from '../../src/workflow-builder-agent.js'; -import { WorkflowBuilderAgent } from '../../src/workflow-builder-agent.js'; import type { EvalLogger } from '../harness/logger.js'; import { createTraceFilters, @@ -18,9 +19,52 @@ import { /** Maximum memory for trace queue (3GB) */ const MAX_INGEST_MEMORY_BYTES = 3 * 1024 * 1024 * 1024; +// ============================================================================ +// Stage Models Configuration +// ============================================================================ + +/** + * Configuration for per-stage model selection. + * All fields except 'default' are optional - unspecified stages use the default model. + */ +export interface StageModels { + /** Default model for all stages */ + default: ModelId; + /** Model for supervisor stage (routing decisions) */ + supervisor?: ModelId; + /** Model for responder stage (final user responses) */ + responder?: ModelId; + /** Model for discovery stage (node discovery) */ + discovery?: ModelId; + /** Model for builder stage (workflow structure) */ + builder?: ModelId; + /** Model for configurator stage (node configuration) */ + configurator?: ModelId; + /** Model for parameter updater (within configurator) */ + parameterUpdater?: ModelId; + /** Model for LLM judge evaluation */ + judge?: ModelId; +} + +/** + * Resolved LLM instances for each stage. + * All fields are populated (using default model as fallback). + */ +export interface ResolvedStageLLMs { + default: BaseChatModel; + supervisor: BaseChatModel; + responder: BaseChatModel; + discovery: BaseChatModel; + builder: BaseChatModel; + configurator: BaseChatModel; + parameterUpdater: BaseChatModel; + judge: BaseChatModel; +} + export interface TestEnvironment { parsedNodeTypes: INodeTypeDescription[]; - llm: BaseChatModel; + /** Resolved LLM instances for each stage */ + llms: ResolvedStageLLMs; tracer?: LangChainTracer; lsClient?: Client; /** Trace filtering utilities (only present when minimal tracing is enabled) */ @@ -28,16 +72,48 @@ export interface TestEnvironment { } /** - * Sets up the LLM with proper configuration + * Sets up an LLM with proper configuration + * @param modelId - Model identifier (defaults to DEFAULT_MODEL) * @returns Configured LLM instance - * @throws Error if N8N_AI_ANTHROPIC_KEY environment variable is not set + * @throws Error if the required API key environment variable is not set */ -export async function setupLLM(): Promise { - const apiKey = process.env.N8N_AI_ANTHROPIC_KEY; +export async function setupLLM(modelId: ModelId = DEFAULT_MODEL): Promise { + const envVar = getApiKeyEnvVar(modelId); + const apiKey = process.env[envVar]; if (!apiKey) { - throw new Error('N8N_AI_ANTHROPIC_KEY environment variable is required'); + throw new Error(`${envVar} environment variable is required for model ${modelId}`); } - return await anthropicClaudeSonnet45({ apiKey }); + const factory = MODEL_FACTORIES[modelId]; + return await factory({ apiKey }); +} + +/** + * Resolves all stage models to LLM instances. + * Unspecified stages fall back to the default model. + * @param stageModels - Per-stage model configuration + * @returns Resolved LLM instances for each stage + */ +export async function resolveStageModels(stageModels: StageModels): Promise { + const defaultLLM = await setupLLM(stageModels.default); + + // For stages without specific model, use default + // For parameter updater, fall back to configurator if not specified + const configuratorLLM = stageModels.configurator + ? await setupLLM(stageModels.configurator) + : defaultLLM; + + return { + default: defaultLLM, + supervisor: stageModels.supervisor ? await setupLLM(stageModels.supervisor) : defaultLLM, + responder: stageModels.responder ? await setupLLM(stageModels.responder) : defaultLLM, + discovery: stageModels.discovery ? await setupLLM(stageModels.discovery) : defaultLLM, + builder: stageModels.builder ? await setupLLM(stageModels.builder) : defaultLLM, + configurator: configuratorLLM, + parameterUpdater: stageModels.parameterUpdater + ? await setupLLM(stageModels.parameterUpdater) + : configuratorLLM, + judge: stageModels.judge ? await setupLLM(stageModels.judge) : defaultLLM, + }; } /** @@ -97,24 +173,39 @@ export function createLangsmithClient(logger?: EvalLogger): LangsmithClientResul /** * Sets up the test environment with LLM, nodes, and tracing + * @param stageModels - Per-stage model configuration (optional, uses default model if not provided) * @param logger - Optional logger for trace filter output * @returns Test environment configuration */ -export async function setupTestEnvironment(logger?: EvalLogger): Promise { +export async function setupTestEnvironment( + stageModels?: StageModels, + logger?: EvalLogger, +): Promise { const parsedNodeTypes = loadNodesFromFile(); - const llm = await setupLLM(); + + // Use provided stage models or default configuration + const models: StageModels = stageModels ?? { default: DEFAULT_MODEL }; + const llms = await resolveStageModels(models); + const lsClientResult = createLangsmithClient(logger); const lsClient = lsClientResult?.client; const traceFilters = lsClientResult?.traceFilters; const tracer = lsClient ? createTracer(lsClient, 'workflow-builder-evaluation') : undefined; - return { parsedNodeTypes, llm, tracer, lsClient, traceFilters }; + return { + parsedNodeTypes, + llms, + tracer, + lsClient, + traceFilters, + }; } export interface CreateAgentOptions { parsedNodeTypes: INodeTypeDescription[]; - llm: BaseChatModel; + /** Per-stage LLMs resolved from model configuration */ + llms: ResolvedStageLLMs; tracer?: LangChainTracer; featureFlags?: BuilderFeatureFlags; experimentName?: string; @@ -126,12 +217,18 @@ export interface CreateAgentOptions { * @returns Configured WorkflowBuilderAgent */ export function createAgent(options: CreateAgentOptions): WorkflowBuilderAgent { - const { parsedNodeTypes, llm, tracer, featureFlags, experimentName } = options; + const { parsedNodeTypes, llms, tracer, featureFlags, experimentName } = options; return new WorkflowBuilderAgent({ parsedNodeTypes, - llmSimpleTask: llm, - llmComplexTask: llm, + stageLLMs: { + supervisor: llms.supervisor, + responder: llms.responder, + discovery: llms.discovery, + builder: llms.builder, + configurator: llms.configurator, + parameterUpdater: llms.parameterUpdater, + }, checkpointer: new MemorySaver(), tracer, featureFlags, diff --git a/packages/@n8n/ai-workflow-builder.ee/src/agents/responder.agent.ts b/packages/@n8n/ai-workflow-builder.ee/src/agents/responder.agent.ts index 357d4617e73..962304b59bd 100644 --- a/packages/@n8n/ai-workflow-builder.ee/src/agents/responder.agent.ts +++ b/packages/@n8n/ai-workflow-builder.ee/src/agents/responder.agent.ts @@ -1,7 +1,8 @@ import type { BaseChatModel } from '@langchain/core/language_models/chat_models'; -import type { AIMessage, BaseMessage } from '@langchain/core/messages'; -import { HumanMessage } from '@langchain/core/messages'; +import type { BaseMessage } from '@langchain/core/messages'; +import { AIMessage, HumanMessage } from '@langchain/core/messages'; import { ChatPromptTemplate } from '@langchain/core/prompts'; +import type { RunnableConfig } from '@langchain/core/runnables'; import { buildResponderPrompt, @@ -12,6 +13,7 @@ import { import type { CoordinationLogEntry } from '../types/coordination'; import type { DiscoveryContext } from '../types/discovery-types'; +import { isAIMessage } from '../types/langchain'; import type { SimpleWorkflow } from '../types/workflow'; import { getErrorEntry, @@ -150,8 +152,10 @@ export class ResponderAgent { /** * Invoke the responder agent with the given context + * @param context - Responder context with messages and workflow state + * @param config - Optional RunnableConfig for tracing callbacks */ - async invoke(context: ResponderContext): Promise { + async invoke(context: ResponderContext, config?: RunnableConfig): Promise { const agent = systemPrompt.pipe(this.llm); const contextMessage = this.buildContextMessage(context); @@ -159,6 +163,12 @@ export class ResponderAgent { ? [...context.messages, contextMessage] : context.messages; - return (await agent.invoke({ messages: messagesToSend })) as AIMessage; + const result = await agent.invoke({ messages: messagesToSend }, config); + if (!isAIMessage(result)) { + return new AIMessage({ + content: 'I encountered an issue generating a response. Please try again.', + }); + } + return result; } } diff --git a/packages/@n8n/ai-workflow-builder.ee/src/agents/supervisor.agent.ts b/packages/@n8n/ai-workflow-builder.ee/src/agents/supervisor.agent.ts index cf4a4f9f8a8..dcd5e8d41ee 100644 --- a/packages/@n8n/ai-workflow-builder.ee/src/agents/supervisor.agent.ts +++ b/packages/@n8n/ai-workflow-builder.ee/src/agents/supervisor.agent.ts @@ -2,6 +2,7 @@ import type { BaseChatModel } from '@langchain/core/language_models/chat_models' import type { BaseMessage } from '@langchain/core/messages'; import { HumanMessage } from '@langchain/core/messages'; import { ChatPromptTemplate } from '@langchain/core/prompts'; +import type { RunnableConfig } from '@langchain/core/runnables'; import { z } from 'zod'; import { buildSupervisorPrompt } from '@/prompts/agents/supervisor.prompt'; @@ -104,8 +105,10 @@ export class SupervisorAgent { /** * Invoke the supervisor to get routing decision + * @param context - Supervisor context with messages and workflow state + * @param config - Optional RunnableConfig for tracing callbacks */ - async invoke(context: SupervisorContext): Promise { + async invoke(context: SupervisorContext, config?: RunnableConfig): Promise { const agent = systemPrompt.pipe( this.llm.withStructuredOutput(supervisorRoutingSchema, { name: 'routing_decision', @@ -117,6 +120,6 @@ export class SupervisorAgent { ? [...context.messages, contextMessage] : context.messages; - return await agent.invoke({ messages: messagesToSend }); + return await agent.invoke({ messages: messagesToSend }, config); } } diff --git a/packages/@n8n/ai-workflow-builder.ee/src/ai-workflow-builder-agent.service.ts b/packages/@n8n/ai-workflow-builder.ee/src/ai-workflow-builder-agent.service.ts index 79afca4a982..960958fbca4 100644 --- a/packages/@n8n/ai-workflow-builder.ee/src/ai-workflow-builder-agent.service.ts +++ b/packages/@n8n/ai-workflow-builder.ee/src/ai-workflow-builder-agent.service.ts @@ -169,9 +169,15 @@ export class AiWorkflowBuilderService { const agent = new WorkflowBuilderAgent({ parsedNodeTypes: this.parsedNodeTypes, - // We use Sonnet both for simple and complex tasks - llmSimpleTask: anthropicClaude, - llmComplexTask: anthropicClaude, + // Use the same model for all stages in production + stageLLMs: { + supervisor: anthropicClaude, + responder: anthropicClaude, + discovery: anthropicClaude, + builder: anthropicClaude, + configurator: anthropicClaude, + parameterUpdater: anthropicClaude, + }, logger: this.logger, checkpointer: this.sessionManager.getCheckpointer(), tracer: tracingClient diff --git a/packages/@n8n/ai-workflow-builder.ee/src/chains/conversation-compact.ts b/packages/@n8n/ai-workflow-builder.ee/src/chains/conversation-compact.ts index 52a21814cd2..607a524071b 100644 --- a/packages/@n8n/ai-workflow-builder.ee/src/chains/conversation-compact.ts +++ b/packages/@n8n/ai-workflow-builder.ee/src/chains/conversation-compact.ts @@ -1,6 +1,7 @@ import type { BaseChatModel } from '@langchain/core/language_models/chat_models'; import type { BaseMessage } from '@langchain/core/messages'; import { AIMessage, HumanMessage } from '@langchain/core/messages'; +import type { RunnableConfig } from '@langchain/core/runnables'; import z from 'zod'; import { compactPromptTemplate } from '@/prompts/chains/compact.prompt'; @@ -9,6 +10,7 @@ export async function conversationCompactChain( llm: BaseChatModel, messages: BaseMessage[], previousSummary: string = '', + config?: RunnableConfig, ) { // Use structured output for consistent summary format const CompactedSession = z.object({ @@ -44,7 +46,7 @@ export async function conversationCompactChain( conversationText, }); - const structuredOutput = await modelWithStructure.invoke(compactPrompt); + const structuredOutput = await modelWithStructure.invoke(compactPrompt, config); const formattedSummary = `## Previous Conversation Summary diff --git a/packages/@n8n/ai-workflow-builder.ee/src/chains/workflow-name.ts b/packages/@n8n/ai-workflow-builder.ee/src/chains/workflow-name.ts index 48352cc9f06..1a0526bae8c 100644 --- a/packages/@n8n/ai-workflow-builder.ee/src/chains/workflow-name.ts +++ b/packages/@n8n/ai-workflow-builder.ee/src/chains/workflow-name.ts @@ -1,9 +1,14 @@ import type { BaseChatModel } from '@langchain/core/language_models/chat_models'; +import type { RunnableConfig } from '@langchain/core/runnables'; import z from 'zod'; import { workflowNamingPromptTemplate } from '@/prompts/chains/workflow-name.prompt'; -export async function workflowNameChain(llm: BaseChatModel, initialPrompt: string) { +export async function workflowNameChain( + llm: BaseChatModel, + initialPrompt: string, + config?: RunnableConfig, +) { // Use structured output for the workflow name to ensure it meets the required format and length const nameSchema = z.object({ name: z.string().min(10).max(128).describe('Name of the workflow based on the prompt'), @@ -15,7 +20,8 @@ export async function workflowNameChain(llm: BaseChatModel, initialPrompt: strin initialPrompt, }); - const structuredOutput = (await modelWithStructure.invoke(prompt)) as z.infer; + const rawOutput = await modelWithStructure.invoke(prompt, config); + const structuredOutput = nameSchema.parse(rawOutput); return { name: structuredOutput.name, diff --git a/packages/@n8n/ai-workflow-builder.ee/src/llm-config.ts b/packages/@n8n/ai-workflow-builder.ee/src/llm-config.ts index 2a8eb0a99c0..f95082012f1 100644 --- a/packages/@n8n/ai-workflow-builder.ee/src/llm-config.ts +++ b/packages/@n8n/ai-workflow-builder.ee/src/llm-config.ts @@ -1,33 +1,23 @@ // Different LLMConfig type for this file - specific to LLM providers +import type { BaseChatModel } from '@langchain/core/language_models/chat_models'; + import { MAX_OUTPUT_TOKENS } from '@/constants'; import { getProxyAgent } from './utils/http-proxy-agent'; -interface LLMProviderConfig { +/** + * Configuration for LLM provider initialization. + */ +export interface LLMProviderConfig { apiKey: string; baseUrl?: string; headers?: Record; } -export const o4mini = async (config: LLMProviderConfig) => { +export const gpt52 = async (config: LLMProviderConfig) => { const { ChatOpenAI } = await import('@langchain/openai'); return new ChatOpenAI({ - model: 'o4-mini-2025-04-16', - apiKey: config.apiKey, - configuration: { - baseURL: config.baseUrl, - defaultHeaders: config.headers, - fetchOptions: { - dispatcher: getProxyAgent(config.baseUrl ?? 'https://api.openai.com/v1'), - }, - }, - }); -}; - -export const gpt41mini = async (config: LLMProviderConfig) => { - const { ChatOpenAI } = await import('@langchain/openai'); - return new ChatOpenAI({ - model: 'gpt-4.1-mini-2025-04-14', + model: 'gpt-5.2-2025-12-11', apiKey: config.apiKey, temperature: 0, maxTokens: -1, @@ -41,27 +31,10 @@ export const gpt41mini = async (config: LLMProviderConfig) => { }); }; -export const gpt41 = async (config: LLMProviderConfig) => { - const { ChatOpenAI } = await import('@langchain/openai'); - return new ChatOpenAI({ - model: 'gpt-4.1-2025-04-14', - apiKey: config.apiKey, - temperature: 0.3, - maxTokens: -1, - configuration: { - baseURL: config.baseUrl, - defaultHeaders: config.headers, - fetchOptions: { - dispatcher: getProxyAgent(config.baseUrl ?? 'https://api.openai.com/v1'), - }, - }, - }); -}; - export const anthropicClaudeSonnet45 = async (config: LLMProviderConfig) => { const { ChatAnthropic } = await import('@langchain/anthropic'); const model = new ChatAnthropic({ - model: 'claude-sonnet-4-5', + model: 'claude-sonnet-4-5-20250929', apiKey: config.apiKey, temperature: 0, maxTokens: MAX_OUTPUT_TOKENS, @@ -101,3 +74,165 @@ export const anthropicHaiku45 = async (config: LLMProviderConfig) => { return model; }; + +export const anthropicClaudeOpus45 = async (config: LLMProviderConfig) => { + const { ChatAnthropic } = await import('@langchain/anthropic'); + const model = new ChatAnthropic({ + model: 'claude-opus-4-5-20251101', + apiKey: config.apiKey, + temperature: 0, + maxTokens: MAX_OUTPUT_TOKENS, + anthropicApiUrl: config.baseUrl, + clientOptions: { + defaultHeaders: config.headers, + fetchOptions: { + dispatcher: getProxyAgent(config.baseUrl), + }, + }, + }); + + // Remove Langchain default topP parameter since Opus 4.5 doesn't allow setting both temperature and topP + delete model.topP; + + return model; +}; + +// ============================================================================ +// OpenRouter Models +// ============================================================================ + +const OPENROUTER_BASE_URL = 'https://openrouter.ai/api/v1'; + +/** + * Creates an OpenRouter model factory for a given model name. + * Uses OpenAI-compatible API with OpenRouter base URL. + */ +function createOpenRouterModel(modelName: string) { + return async (config: LLMProviderConfig) => { + const { ChatOpenAI } = await import('@langchain/openai'); + return new ChatOpenAI({ + model: modelName, + apiKey: config.apiKey, + temperature: 0, + maxTokens: -1, + configuration: { + baseURL: OPENROUTER_BASE_URL, + defaultHeaders: { + ...config.headers, + 'HTTP-Referer': 'https://n8n.io', + 'X-Title': 'n8n AI Workflow Builder', + }, + fetchOptions: { + dispatcher: getProxyAgent(OPENROUTER_BASE_URL), + }, + }, + }); + }; +} + +// OpenRouter model factories +export const glm47 = createOpenRouterModel('thudm/glm-4-plus'); +export const gemini3Flash = createOpenRouterModel('google/gemini-3-flash-preview'); +export const deepseekV32 = createOpenRouterModel('deepseek/deepseek-chat-v3-0324'); +export const gemini3Pro = createOpenRouterModel('google/gemini-3-pro-preview'); +export const devstral = createOpenRouterModel('mistralai/devstral-small'); + +// ============================================================================ +// Model Registry +// ============================================================================ + +/** + * IMPORTANT: Generation stages currently only support Anthropic models. + * + * Non-Anthropic models (OpenAI, OpenRouter) are available for evaluation/judging + * purposes only. Using them for generation stages (supervisor, discovery, builder, + * configurator, responder, parameterUpdater) will likely fail due to: + * + * 1. Prompt caching: Our prompts use Anthropic's cache_control for efficiency + * 2. Tool schemas: add_nodes and update_parameters tools use passthrough() schemas + * which only Anthropic models handle correctly + * + * TODO: Add provider-agnostic prompt/tool support to enable non-Anthropic generation. + */ + +/** + * Available model identifiers for the eval CLI. + * These can be used with --model, --judge-model, and per-stage model flags. + */ +export type ModelId = + // Native models + | 'claude-opus-4.5' + | 'claude-sonnet-4.5' + | 'claude-haiku-4.5' + | 'gpt-5.2' + // OpenRouter models + | 'glm-4.7' + | 'gemini-3-flash' + | 'deepseek-v3.2' + | 'gemini-3-pro' + | 'devstral'; + +/** + * Model factory functions mapped by model ID. + */ +export const MODEL_FACTORIES: Record< + ModelId, + (config: LLMProviderConfig) => Promise +> = { + // Native models + 'claude-opus-4.5': anthropicClaudeOpus45, + 'claude-sonnet-4.5': anthropicClaudeSonnet45, + 'claude-haiku-4.5': anthropicHaiku45, + 'gpt-5.2': gpt52, + // OpenRouter models + 'glm-4.7': glm47, + 'gemini-3-flash': gemini3Flash, + 'deepseek-v3.2': deepseekV32, + 'gemini-3-pro': gemini3Pro, + devstral, +}; + +/** OpenRouter model IDs for API key resolution */ +const OPENROUTER_MODELS: ModelId[] = [ + 'glm-4.7', + 'gemini-3-flash', + 'deepseek-v3.2', + 'gemini-3-pro', + 'devstral', +]; + +/** + * Get the required API key environment variable for a model. + */ +export function getApiKeyEnvVar(modelId: ModelId): string { + if (OPENROUTER_MODELS.includes(modelId)) { + return 'OPENROUTER_API_KEY'; + } + if (modelId.startsWith('gpt')) { + return 'N8N_AI_OPENAI_KEY'; + } + return 'N8N_AI_ANTHROPIC_KEY'; +} + +/** + * List of available model IDs for CLI help text. + * Explicitly defined to avoid type casting Object.keys(). + */ +export const AVAILABLE_MODELS: readonly ModelId[] = [ + // Native models + 'claude-opus-4.5', + 'claude-sonnet-4.5', + 'claude-haiku-4.5', + 'gpt-5.2', + // OpenRouter models + 'glm-4.7', + 'gemini-3-flash', + 'deepseek-v3.2', + 'gemini-3-pro', + 'devstral', +] as const; + +/** + * Default model used when no model is specified. + */ +export const DEFAULT_MODEL: ModelId = 'claude-sonnet-4.5'; diff --git a/packages/@n8n/ai-workflow-builder.ee/src/multi-agent-workflow-subgraphs.ts b/packages/@n8n/ai-workflow-builder.ee/src/multi-agent-workflow-subgraphs.ts index 1c612449f7e..0ec1e65a40f 100644 --- a/packages/@n8n/ai-workflow-builder.ee/src/multi-agent-workflow-subgraphs.ts +++ b/packages/@n8n/ai-workflow-builder.ee/src/multi-agent-workflow-subgraphs.ts @@ -1,5 +1,5 @@ -import type { BaseChatModel } from '@langchain/core/language_models/chat_models'; import { HumanMessage } from '@langchain/core/messages'; +import type { RunnableConfig } from '@langchain/core/runnables'; import { StateGraph, END, START, type MemorySaver } from '@langchain/langgraph'; import type { Logger } from '@n8n/backend-common'; import type { INodeTypeDescription } from 'n8n-workflow'; @@ -29,7 +29,7 @@ import { handleCreateWorkflowName, handleDeleteMessages, } from './utils/state-modifier'; -import type { BuilderFeatureFlags } from './workflow-builder-agent'; +import type { BuilderFeatureFlags, StageLLMs } from './workflow-builder-agent'; /** * Maps routing decisions to graph node names. @@ -47,8 +47,8 @@ function routeToNode(next: string): string { export interface MultiAgentSubgraphConfig { parsedNodeTypes: INodeTypeDescription[]; - llmSimpleTask: BaseChatModel; - llmComplexTask: BaseChatModel; + /** Per-stage LLM configuration */ + stageLLMs: StageLLMs; logger?: Logger; instanceUrl?: string; checkpointer?: MemorySaver; @@ -60,7 +60,8 @@ export interface MultiAgentSubgraphConfig { } /** - * Creates a subgraph node handler with standardized error handling + * Creates a subgraph node handler with standardized error handling. + * Accepts RunnableConfig as second parameter to propagate callbacks for tracing. */ function createSubgraphNodeHandler< TSubgraph extends BaseSubgraph, Record>, @@ -71,10 +72,15 @@ function createSubgraphNodeHandler< logger?: Logger, recursionLimit?: number, ) { - return async (state: typeof ParentGraphState.State) => { + return async (state: typeof ParentGraphState.State, config?: RunnableConfig) => { try { const input = subgraph.transformInput(state); - const result = await compiledGraph.invoke(input, { recursionLimit }); + // Merge parent config (callbacks, metadata) with recursionLimit + const invokeConfig: RunnableConfig = { + ...config, + recursionLimit, + }; + const result = await compiledGraph.invoke(input, invokeConfig); const output = subgraph.transformOutput(result, state); return output; @@ -122,7 +128,7 @@ function createSubgraphNodeHandler< export function createMultiAgentWorkflowWithSubgraphs(config: MultiAgentSubgraphConfig) { const { parsedNodeTypes, - llmComplexTask, + stageLLMs, logger, instanceUrl, checkpointer, @@ -131,30 +137,31 @@ export function createMultiAgentWorkflowWithSubgraphs(config: MultiAgentSubgraph onGenerationSuccess, } = config; - const supervisorAgent = new SupervisorAgent({ llm: llmComplexTask }); - const responderAgent = new ResponderAgent({ llm: llmComplexTask }); + const supervisorAgent = new SupervisorAgent({ llm: stageLLMs.supervisor }); + const responderAgent = new ResponderAgent({ llm: stageLLMs.responder }); // Create subgraph instances const discoverySubgraph = new DiscoverySubgraph(); const builderSubgraph = new BuilderSubgraph(); const configuratorSubgraph = new ConfiguratorSubgraph(); - // Compile subgraphs + // Compile subgraphs with per-stage LLMs const compiledDiscovery = discoverySubgraph.create({ parsedNodeTypes, - llm: llmComplexTask, + llm: stageLLMs.discovery, logger, featureFlags, }); const compiledBuilder = builderSubgraph.create({ parsedNodeTypes, - llm: llmComplexTask, + llm: stageLLMs.builder, logger, featureFlags, }); const compiledConfigurator = configuratorSubgraph.create({ parsedNodeTypes, - llm: llmComplexTask, + llm: stageLLMs.configurator, + llmParameterUpdater: stageLLMs.parameterUpdater, logger, instanceUrl, featureFlags, @@ -164,27 +171,35 @@ export function createMultiAgentWorkflowWithSubgraphs(config: MultiAgentSubgraph return ( new StateGraph(ParentGraphState) // Add Supervisor Node (only used for initial routing) - .addNode('supervisor', async (state) => { - const routing = await supervisorAgent.invoke({ - messages: state.messages, - workflowJSON: state.workflowJSON, - coordinationLog: state.coordinationLog, - previousSummary: state.previousSummary, - }); + // Accepts config as second param to propagate callbacks for tracing + .addNode('supervisor', async (state, config) => { + const routing = await supervisorAgent.invoke( + { + messages: state.messages, + workflowJSON: state.workflowJSON, + coordinationLog: state.coordinationLog, + previousSummary: state.previousSummary, + }, + config, + ); return { nextPhase: routing.next, }; }) // Add Responder Node (synthesizes final user-facing response) - .addNode('responder', async (state) => { - const response = await responderAgent.invoke({ - messages: state.messages, - coordinationLog: state.coordinationLog, - discoveryContext: state.discoveryContext, - workflowJSON: state.workflowJSON, - previousSummary: state.previousSummary, - }); + // Accepts config as second param to propagate callbacks for tracing + .addNode('responder', async (state, config) => { + const response = await responderAgent.invoke( + { + messages: state.messages, + coordinationLog: state.coordinationLog, + discoveryContext: state.discoveryContext, + workflowJSON: state.workflowJSON, + previousSummary: state.previousSummary, + }, + config, + ); // Call success callback only when generation completed without errors if (onGenerationSuccess && !hasErrorInLog(state.coordinationLog)) { @@ -212,25 +227,27 @@ export function createMultiAgentWorkflowWithSubgraphs(config: MultiAgentSubgraph nextPhase: determineStateAction(state, autoCompactThresholdTokens), })) .addNode('cleanup_dangling', (state) => handleCleanupDangling(state.messages, logger)) - .addNode('compact_messages', async (state) => { + .addNode('compact_messages', async (state, config) => { const isAutoCompact = state.messages[state.messages.length - 1]?.content !== '/compact'; return await handleCompactMessages( state.messages, state.previousSummary ?? '', - llmComplexTask, + stageLLMs.responder, isAutoCompact, + config, ); }) .addNode('delete_messages', (state) => handleDeleteMessages(state.messages)) .addNode('clear_error_state', (state) => handleClearErrorState(state.coordinationLog, logger)) .addNode( 'create_workflow_name', - async (state) => + async (state, config) => await handleCreateWorkflowName( state.messages, state.workflowJSON, - llmComplexTask, + stageLLMs.responder, logger, + config, ), ) // Add Subgraph Nodes (using helper to reduce duplication) diff --git a/packages/@n8n/ai-workflow-builder.ee/src/subgraphs/configurator.subgraph.ts b/packages/@n8n/ai-workflow-builder.ee/src/subgraphs/configurator.subgraph.ts index 607370e53b3..2b4048ec5cd 100644 --- a/packages/@n8n/ai-workflow-builder.ee/src/subgraphs/configurator.subgraph.ts +++ b/packages/@n8n/ai-workflow-builder.ee/src/subgraphs/configurator.subgraph.ts @@ -101,6 +101,8 @@ export const ConfiguratorSubgraphState = Annotation.Root({ export interface ConfiguratorSubgraphConfig { parsedNodeTypes: INodeTypeDescription[]; llm: BaseChatModel; + /** Separate LLM for parameter updater chain (defaults to llm if not provided) */ + llmParameterUpdater?: BaseChatModel; logger?: Logger; instanceUrl?: string; featureFlags?: BuilderFeatureFlags; @@ -124,11 +126,14 @@ export class ConfiguratorSubgraph extends BaseSubgraph< // Check if template examples are enabled const includeExamples = config.featureFlags?.templateExamples === true; + // Use separate LLM for parameter updater if provided + const parameterUpdaterLLM = config.llmParameterUpdater ?? config.llm; + // Create base tools const baseTools = [ createUpdateNodeParametersTool( config.parsedNodeTypes, - config.llm, // Uses same LLM for parameter updater chain + parameterUpdaterLLM, // Uses separate LLM for parameter updater chain config.logger, config.instanceUrl, ), diff --git a/packages/@n8n/ai-workflow-builder.ee/src/subgraphs/test/integration/multi-agent-error-handling.integration.test.ts b/packages/@n8n/ai-workflow-builder.ee/src/subgraphs/test/integration/multi-agent-error-handling.integration.test.ts index 6538b9f994a..34bb6944a1c 100644 --- a/packages/@n8n/ai-workflow-builder.ee/src/subgraphs/test/integration/multi-agent-error-handling.integration.test.ts +++ b/packages/@n8n/ai-workflow-builder.ee/src/subgraphs/test/integration/multi-agent-error-handling.integration.test.ts @@ -110,8 +110,14 @@ describe('Multi-Agent Error Handling - Integration Tests (AI-1812)', () => { // Don't use checkpointer for this test - we're just testing error message const graph = createMultiAgentWorkflowWithSubgraphs({ parsedNodeTypes, - llmSimpleTask: llm, - llmComplexTask: llm, + stageLLMs: { + supervisor: llm, + responder: llm, + discovery: llm, + builder: llm, + configurator: llm, + parameterUpdater: llm, + }, logger: mockLogger, }); diff --git a/packages/@n8n/ai-workflow-builder.ee/src/test/workflow-builder-agent.test.ts b/packages/@n8n/ai-workflow-builder.ee/src/test/workflow-builder-agent.test.ts index 9b95628c5fa..c4e333de424 100644 --- a/packages/@n8n/ai-workflow-builder.ee/src/test/workflow-builder-agent.test.ts +++ b/packages/@n8n/ai-workflow-builder.ee/src/test/workflow-builder-agent.test.ts @@ -56,8 +56,7 @@ import { describe('WorkflowBuilderAgent', () => { let agent: WorkflowBuilderAgent; - let mockLlmSimple: BaseChatModel; - let mockLlmComplex: BaseChatModel; + let mockLlm: BaseChatModel; let mockLogger: Logger; let mockCheckpointer: MemorySaver; let parsedNodeTypes: INodeTypeDescription[]; @@ -68,18 +67,12 @@ describe('WorkflowBuilderAgent', () => { >; beforeEach(() => { - mockLlmSimple = mock({ + mockLlm = mock({ _llmType: jest.fn().mockReturnValue('test-llm'), bindTools: jest.fn().mockReturnThis(), invoke: jest.fn(), }); - mockLlmComplex = mock({ - _llmType: jest.fn().mockReturnValue('test-llm-complex'), - bindTools: jest.fn().mockReturnThis(), - invoke: jest.fn(), - }); - mockLogger = mock({ debug: jest.fn(), info: jest.fn(), @@ -108,8 +101,14 @@ describe('WorkflowBuilderAgent', () => { config = { parsedNodeTypes, - llmSimpleTask: mockLlmSimple, - llmComplexTask: mockLlmComplex, + stageLLMs: { + supervisor: mockLlm, + responder: mockLlm, + discovery: mockLlm, + builder: mockLlm, + configurator: mockLlm, + parameterUpdater: mockLlm, + }, logger: mockLogger, checkpointer: mockCheckpointer, }; @@ -172,7 +171,7 @@ describe('WorkflowBuilderAgent', () => { mockCreateStreamProcessor.mockReturnValue(mockAsyncGenerator); // Mock the LLM to return a simple response - (mockLlmSimple.invoke as jest.Mock).mockResolvedValue({ + (mockLlm.invoke as jest.Mock).mockResolvedValue({ content: 'Mocked response', tool_calls: [], }); diff --git a/packages/@n8n/ai-workflow-builder.ee/src/utils/state-modifier.ts b/packages/@n8n/ai-workflow-builder.ee/src/utils/state-modifier.ts index bd2d2a0d028..233d6080726 100644 --- a/packages/@n8n/ai-workflow-builder.ee/src/utils/state-modifier.ts +++ b/packages/@n8n/ai-workflow-builder.ee/src/utils/state-modifier.ts @@ -1,6 +1,7 @@ import type { BaseChatModel } from '@langchain/core/language_models/chat_models'; import type { BaseMessage } from '@langchain/core/messages'; import { HumanMessage, RemoveMessage } from '@langchain/core/messages'; +import type { RunnableConfig } from '@langchain/core/runnables'; import type { Logger } from '@n8n/backend-common'; import { cleanupDanglingToolCallMessages } from './cleanup-dangling-tool-call-messages'; @@ -142,12 +143,19 @@ export function handleCleanupDangling( * * For manual /compact: Removes all messages, routes to responder for acknowledgment. * For auto-compact: Removes old messages, preserves last user message to continue processing. + * + * @param messages - Conversation messages to compact + * @param previousSummary - Previous conversation summary + * @param llm - Language model for summarization + * @param isAutoCompact - Whether this is auto-compaction (preserve last message) or manual + * @param config - Optional RunnableConfig for tracing callbacks */ export async function handleCompactMessages( messages: BaseMessage[], previousSummary: string, llm: BaseChatModel, isAutoCompact: boolean, + config?: RunnableConfig, ): Promise<{ previousSummary: string; messages: BaseMessage[]; @@ -158,7 +166,7 @@ export async function handleCompactMessages( throw new Error('Cannot compact messages: no HumanMessage found'); } - const compactedMessages = await conversationCompactChain(llm, messages, previousSummary); + const compactedMessages = await conversationCompactChain(llm, messages, previousSummary, config); // For manual /compact: just remove messages, responder will generate acknowledgment // For auto-compact: remove messages but preserve the last user message to continue processing @@ -252,12 +260,19 @@ export function handleClearErrorState( /** * Generates a workflow name from the initial user message. + * + * @param messages - Conversation messages + * @param workflowJSON - Current workflow state + * @param llm - Language model for name generation + * @param logger - Optional logger + * @param config - Optional RunnableConfig for tracing callbacks */ export async function handleCreateWorkflowName( messages: BaseMessage[], workflowJSON: SimpleWorkflow, llm: BaseChatModel, logger?: Logger, + config?: RunnableConfig, ): Promise<{ workflowJSON: SimpleWorkflow }> { if (messages.length === 1 && messages[0] instanceof HumanMessage) { const initialMessage = messages[0]; @@ -267,7 +282,7 @@ export async function handleCreateWorkflowName( } logger?.debug('Generating workflow name'); - const { name } = await workflowNameChain(llm, initialMessage.content); + const { name } = await workflowNameChain(llm, initialMessage.content, config); return { workflowJSON: { ...workflowJSON, name }, diff --git a/packages/@n8n/ai-workflow-builder.ee/src/workflow-builder-agent.ts b/packages/@n8n/ai-workflow-builder.ee/src/workflow-builder-agent.ts index 346e1899f7d..f00ba0695e5 100644 --- a/packages/@n8n/ai-workflow-builder.ee/src/workflow-builder-agent.ts +++ b/packages/@n8n/ai-workflow-builder.ee/src/workflow-builder-agent.ts @@ -37,10 +37,23 @@ export type TypedStateSnapshot = Omit & { values: typeof WorkflowState.State; }; +/** + * Per-stage LLM configuration for the workflow builder. + * All stages must be configured with an LLM instance. + */ +export interface StageLLMs { + supervisor: BaseChatModel; + responder: BaseChatModel; + discovery: BaseChatModel; + builder: BaseChatModel; + configurator: BaseChatModel; + parameterUpdater: BaseChatModel; +} + export interface WorkflowBuilderAgentConfig { parsedNodeTypes: INodeTypeDescription[]; - llmSimpleTask: BaseChatModel; - llmComplexTask: BaseChatModel; + /** Per-stage LLM configuration */ + stageLLMs: StageLLMs; logger?: Logger; checkpointer: MemorySaver; tracer?: LangChainTracer; @@ -80,8 +93,7 @@ export interface ChatPayload { export class WorkflowBuilderAgent { private checkpointer: MemorySaver; private parsedNodeTypes: INodeTypeDescription[]; - private llmSimpleTask: BaseChatModel; - private llmComplexTask: BaseChatModel; + private stageLLMs: StageLLMs; private logger?: Logger; private tracer?: LangChainTracer; private instanceUrl?: string; @@ -90,8 +102,7 @@ export class WorkflowBuilderAgent { constructor(config: WorkflowBuilderAgentConfig) { this.parsedNodeTypes = config.parsedNodeTypes; - this.llmSimpleTask = config.llmSimpleTask; - this.llmComplexTask = config.llmComplexTask; + this.stageLLMs = config.stageLLMs; this.logger = config.logger; this.checkpointer = config.checkpointer; this.tracer = config.tracer; @@ -107,8 +118,7 @@ export class WorkflowBuilderAgent { private createWorkflow(featureFlags?: BuilderFeatureFlags) { return createMultiAgentWorkflowWithSubgraphs({ parsedNodeTypes: this.parsedNodeTypes, - llmSimpleTask: this.llmSimpleTask, - llmComplexTask: this.llmComplexTask, + stageLLMs: this.stageLLMs, logger: this.logger, instanceUrl: this.instanceUrl, checkpointer: this.checkpointer,