From d6828b6a907fbf094542ff0a07a96c2c73e1b4ba Mon Sep 17 00:00:00 2001 From: Abhiiishek44 Date: Mon, 25 May 2026 13:51:51 +0530 Subject: [PATCH 1/3] fix: improve follow-up prompt timeout and abort handling --- packages/components/src/followUpPrompts.ts | 387 +++++++++++++++------ 1 file changed, 281 insertions(+), 106 deletions(-) diff --git a/packages/components/src/followUpPrompts.ts b/packages/components/src/followUpPrompts.ts index bc3b2e0f942..3ed57431173 100644 --- a/packages/components/src/followUpPrompts.ts +++ b/packages/components/src/followUpPrompts.ts @@ -10,6 +10,10 @@ import { StructuredOutputParser } from '@langchain/core/output_parsers' import { ChatGroq } from '@langchain/groq' import { Ollama } from 'ollama' +const FOLLOWUP_TIMEOUT_MS = 15000 +const FOLLOWUP_MAX_RETRIES = 2 +const FOLLOWUP_RETRY_BASE_DELAY_MS = 500 + const FollowUpPromptType = z .object({ questions: z.array(z.string()) @@ -20,6 +24,85 @@ export interface FollowUpPromptResult { questions: string[] } +const sleep = (ms: number): Promise => new Promise((resolve) => setTimeout(resolve, ms)) + +const getErrorStatus = (error: any): number | undefined => { + return error?.status ?? error?.statusCode ?? error?.response?.status ?? error?.cause?.status +} + +const isTimeoutError = (error: any): boolean => { + const errorCode = error?.code + if (error?.name === 'TimeoutError' || error?.name === 'AbortError') return true + if (typeof error?.message === 'string' && /timeout|timed out/i.test(error.message)) return true + return ['ETIMEDOUT', 'ECONNRESET', 'EAI_AGAIN', 'ENOTFOUND', 'ECONNREFUSED'].includes(errorCode) +} + +const isRetryableError = (error: any): boolean => { + const status = getErrorStatus(error) + if (status && [429, 500, 502, 503].includes(status)) return true + return isTimeoutError(error) +} + +const executeWithRetry = async ( + action: (signal: AbortSignal) => Promise, + options: { + provider: string + timeoutMs: number + maxRetries?: number + baseDelayMs?: number + logger?: { error?: (message: string) => void } + } +): Promise => { + const maxRetries = options.maxRetries ?? FOLLOWUP_MAX_RETRIES + const baseDelayMs = options.baseDelayMs ?? FOLLOWUP_RETRY_BASE_DELAY_MS + + for (let attempt = 0; attempt <= maxRetries; attempt += 1) { + const abortController = new AbortController() + const timeoutId = setTimeout(() => abortController.abort(), options.timeoutMs) + let abortHandler: (() => void) | undefined + const abortPromise = new Promise((_, reject) => { + abortHandler = () => { + const timeoutError = new Error('Follow-up prompt request timed out') + timeoutError.name = 'TimeoutError' + reject(timeoutError) + } + abortController.signal.addEventListener('abort', abortHandler, { once: true }) + }) + + try { + const result = await Promise.race([action(abortController.signal), abortPromise]) + return result as T + } catch (error: any) { + const retryable = isRetryableError(error) + const retryCount = attempt + const status = getErrorStatus(error) + const message = error?.message ?? 'Unknown error' + options.logger?.error?.( + `[FollowUpPrompt] ${JSON.stringify({ + provider: options.provider, + retryCount, + timeoutMs: options.timeoutMs, + status, + message, + retryable + })}` + ) + + if (!retryable || attempt >= maxRetries) { + throw error + } + + const delayMs = baseDelayMs * Math.pow(2, attempt) + await sleep(delayMs) + } finally { + clearTimeout(timeoutId) + if (abortHandler) abortController.signal.removeEventListener('abort', abortHandler) + } + } + + throw new Error('Follow-up prompt retry attempts exhausted') +} + export const generateFollowUpPrompts = async ( followUpPromptsConfig: FollowUpPromptConfig, apiMessageContent: string, @@ -29,23 +112,36 @@ export const generateFollowUpPrompts = async ( if (!followUpPromptsConfig.status) return undefined const providerConfig = followUpPromptsConfig[followUpPromptsConfig.selectedProvider] if (!providerConfig) return undefined + const logger = options?.logger + const timeoutFromOptions = options?.followUpPromptTimeoutMs ?? process.env.FOLLOW_UP_PROMPT_TIMEOUT_MS + const timeoutMs = Number(timeoutFromOptions) + const resolvedTimeoutMs = Number.isFinite(timeoutMs) && timeoutMs > 0 ? timeoutMs : FOLLOWUP_TIMEOUT_MS const credentialId = providerConfig.credentialId as string const credentialData = await getCredentialData(credentialId ?? '', options) const followUpPromptsPrompt = providerConfig.prompt.replace('{history}', apiMessageContent) switch (followUpPromptsConfig.selectedProvider) { case FollowUpPromptProvider.ANTHROPIC: { - const llm = new ChatAnthropic({ - apiKey: credentialData.anthropicApiKey, - model: providerConfig.modelName, - temperature: parseFloat(`${providerConfig.temperature}`) - }) - // @ts-ignore - const structuredLLM = llm.withStructuredOutput(FollowUpPromptType, { - method: 'functionCalling' - }) - const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) - return structuredResponse as FollowUpPromptResult + return executeWithRetry( + async (signal) => { + const llm = new ChatAnthropic({ + apiKey: credentialData.anthropicApiKey, + model: providerConfig.modelName, + temperature: parseFloat(`${providerConfig.temperature}`) + }) + // @ts-ignore + const structuredLLM = llm.withStructuredOutput(FollowUpPromptType, { + method: 'functionCalling' + }) + const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt, { signal }) + return structuredResponse as FollowUpPromptResult + }, + { + provider: FollowUpPromptProvider.ANTHROPIC, + timeoutMs: resolvedTimeoutMs, + logger + } + ) } case FollowUpPromptProvider.AZURE_OPENAI: { const azureOpenAIApiKey = credentialData['azureOpenAIApiKey'] @@ -53,115 +149,194 @@ export const generateFollowUpPrompts = async ( const azureOpenAIApiDeploymentName = credentialData['azureOpenAIApiDeploymentName'] const azureOpenAIApiVersion = credentialData['azureOpenAIApiVersion'] - const llm = new AzureChatOpenAI({ - azureOpenAIApiKey, - azureOpenAIApiInstanceName, - azureOpenAIApiDeploymentName, - azureOpenAIApiVersion, - model: providerConfig.modelName, - temperature: parseFloat(`${providerConfig.temperature}`) - }) - // use structured output parser because withStructuredOutput is not working - const parser = StructuredOutputParser.fromZodSchema(FollowUpPromptType as any) - const formatInstructions = parser.getFormatInstructions() - const prompt = PromptTemplate.fromTemplate(` - ${providerConfig.prompt} - - {format_instructions} - `) - const chain = prompt.pipe(llm).pipe(parser) - const structuredResponse = await chain.invoke({ - history: apiMessageContent, - format_instructions: formatInstructions - }) - return structuredResponse as FollowUpPromptResult + return executeWithRetry( + async (signal) => { + const llm = new AzureChatOpenAI({ + azureOpenAIApiKey, + azureOpenAIApiInstanceName, + azureOpenAIApiDeploymentName, + azureOpenAIApiVersion, + model: providerConfig.modelName, + temperature: parseFloat(`${providerConfig.temperature}`) + }) + // use structured output parser because withStructuredOutput is not working + const parser = StructuredOutputParser.fromZodSchema(FollowUpPromptType as any) + const formatInstructions = parser.getFormatInstructions() + const prompt = PromptTemplate.fromTemplate(` + ${providerConfig.prompt} + + {format_instructions} + `) + const chain = prompt.pipe(llm).pipe(parser) + const structuredResponse = await chain.invoke( + { + history: apiMessageContent, + format_instructions: formatInstructions + }, + { signal } + ) + return structuredResponse as FollowUpPromptResult + }, + { + provider: FollowUpPromptProvider.AZURE_OPENAI, + timeoutMs: resolvedTimeoutMs, + logger + } + ) } case FollowUpPromptProvider.GOOGLE_GENAI: { - const model = new ChatGoogleGenerativeAI({ - apiKey: credentialData.googleGenerativeAPIKey, - model: providerConfig.modelName, - temperature: parseFloat(`${providerConfig.temperature}`) - }) - const structuredLLM = model.withStructuredOutput(FollowUpPromptType, { - method: 'functionCalling' - }) - const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) - return structuredResponse as FollowUpPromptResult + return executeWithRetry( + async (signal) => { + const model = new ChatGoogleGenerativeAI({ + apiKey: credentialData.googleGenerativeAPIKey, + model: providerConfig.modelName, + temperature: parseFloat(`${providerConfig.temperature}`) + }) + const structuredLLM = model.withStructuredOutput(FollowUpPromptType, { + method: 'functionCalling' + }) + const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt, { signal }) + return structuredResponse as FollowUpPromptResult + }, + { + provider: FollowUpPromptProvider.GOOGLE_GENAI, + timeoutMs: resolvedTimeoutMs, + logger + } + ) } case FollowUpPromptProvider.MISTRALAI: { - const model = new ChatMistralAI({ - apiKey: credentialData.mistralAIAPIKey, - model: providerConfig.modelName, - temperature: parseFloat(`${providerConfig.temperature}`) - }) - // @ts-ignore - const structuredLLM = model.withStructuredOutput(FollowUpPromptType, { - method: 'functionCalling' - }) - const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) - return structuredResponse as FollowUpPromptResult + return executeWithRetry( + async (signal) => { + const model = new ChatMistralAI({ + apiKey: credentialData.mistralAIAPIKey, + model: providerConfig.modelName, + temperature: parseFloat(`${providerConfig.temperature}`) + }) + // @ts-ignore + const structuredLLM = model.withStructuredOutput(FollowUpPromptType, { + method: 'functionCalling' + }) + const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt, { signal }) + return structuredResponse as FollowUpPromptResult + }, + { + provider: FollowUpPromptProvider.MISTRALAI, + timeoutMs: resolvedTimeoutMs, + logger + } + ) } case FollowUpPromptProvider.OPENAI: { - const model = new ChatOpenAI({ - apiKey: credentialData.openAIApiKey, - model: providerConfig.modelName, - temperature: parseFloat(`${providerConfig.temperature}`), - useResponsesApi: true - }) - // @ts-ignore - const structuredLLM = model.withStructuredOutput(FollowUpPromptType, { - method: 'functionCalling' - }) - const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) - return structuredResponse as FollowUpPromptResult + return executeWithRetry( + async (signal) => { + const model = new ChatOpenAI({ + apiKey: credentialData.openAIApiKey, + model: providerConfig.modelName, + temperature: parseFloat(`${providerConfig.temperature}`), + useResponsesApi: true + }) + // @ts-ignore + const structuredLLM = model.withStructuredOutput(FollowUpPromptType, { + method: 'functionCalling' + }) + const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt, { signal }) + return structuredResponse as FollowUpPromptResult + }, + { + provider: FollowUpPromptProvider.OPENAI, + timeoutMs: resolvedTimeoutMs, + logger + } + ) } case FollowUpPromptProvider.GROQ: { - const llm = new ChatGroq({ - apiKey: credentialData.groqApiKey, - model: providerConfig.modelName, - temperature: parseFloat(`${providerConfig.temperature}`) - }) - const structuredLLM = llm.withStructuredOutput(FollowUpPromptType, { - method: 'functionCalling' - }) - const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) - return structuredResponse as FollowUpPromptResult + return executeWithRetry( + async (signal) => { + const llm = new ChatGroq({ + apiKey: credentialData.groqApiKey, + model: providerConfig.modelName, + temperature: parseFloat(`${providerConfig.temperature}`) + }) + const structuredLLM = llm.withStructuredOutput(FollowUpPromptType, { + method: 'functionCalling' + }) + const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt, { signal }) + return structuredResponse as FollowUpPromptResult + }, + { + provider: FollowUpPromptProvider.GROQ, + timeoutMs: resolvedTimeoutMs, + logger + } + ) } case FollowUpPromptProvider.OLLAMA: { - const ollamaClient = new Ollama({ - host: providerConfig.baseUrl || 'http://127.0.0.1:11434' - }) - - const response = await ollamaClient.chat({ - model: providerConfig.modelName, - messages: [ - { - role: 'user', - content: followUpPromptsPrompt - } - ], - format: { - type: 'object', - properties: { - questions: { - type: 'array', - items: { - type: 'string' + return executeWithRetry( + async (signal) => { + const ollamaClient = new Ollama({ + host: providerConfig.baseUrl || 'http://127.0.0.1:11434' + }) + + // Ollama client does not accept AbortSignal directly in this SDK call. + // Wrap the chat promise in a race so the caller-provided signal can + // abort the request and trigger our retry/timeout behavior. + const chatPromise = ollamaClient.chat({ + model: providerConfig.modelName, + messages: [ + { + role: 'user', + content: followUpPromptsPrompt + } + ], + format: { + type: 'object', + properties: { + questions: { + type: 'array', + items: { + type: 'string' + }, + minItems: 3, + maxItems: 3, + description: 'Three follow-up questions based on the conversation history' + } }, - minItems: 3, - maxItems: 3, - description: 'Three follow-up questions based on the conversation history' + required: ['questions'], + additionalProperties: false + }, + options: { + temperature: parseFloat(`${providerConfig.temperature}`) } - }, - required: ['questions'], - additionalProperties: false + }) + + let abortHandler: (() => void) | undefined + const abortPromise = new Promise((_, reject) => { + abortHandler = () => { + const timeoutError = new Error('Follow-up prompt request timed out') + timeoutError.name = 'TimeoutError' + reject(timeoutError) + } + signal.addEventListener('abort', abortHandler, { once: true }) + }) + + try { + const response = await Promise.race([chatPromise, abortPromise]) + const result = FollowUpPromptType.parse(JSON.parse((response as any).message.content)) + if (!result.questions) { + throw new Error('Follow-up prompt response missing questions') + } + return { questions: result.questions } + } finally { + if (abortHandler) signal.removeEventListener('abort', abortHandler) + } }, - options: { - temperature: parseFloat(`${providerConfig.temperature}`) + { + provider: FollowUpPromptProvider.OLLAMA, + timeoutMs: resolvedTimeoutMs, + logger } - }) - const result = FollowUpPromptType.parse(JSON.parse(response.message.content)) - return result + ) } } } else { From 3469c0d871dec6c83c517a1075614151de922ae1 Mon Sep 17 00:00:00 2001 From: Abhiiishek44 Date: Tue, 9 Jun 2026 13:54:44 +0530 Subject: [PATCH 2/3] fix: retry 504 follow-up prompt errors --- .../components/src/followUpPrompts.test.ts | 79 +++++++++++++++++++ packages/components/src/followUpPrompts.ts | 2 +- 2 files changed, 80 insertions(+), 1 deletion(-) create mode 100644 packages/components/src/followUpPrompts.test.ts diff --git a/packages/components/src/followUpPrompts.test.ts b/packages/components/src/followUpPrompts.test.ts new file mode 100644 index 00000000000..9767a49ac75 --- /dev/null +++ b/packages/components/src/followUpPrompts.test.ts @@ -0,0 +1,79 @@ +import { afterEach, beforeEach, describe, expect, it, jest } from '@jest/globals' + +const mockInvoke = jest.fn<() => Promise<{ questions: string[] }>>() +const mockWithStructuredOutput = jest.fn(() => ({ invoke: mockInvoke })) + +jest.mock('@langchain/anthropic', () => ({ + ChatAnthropic: jest.fn() +})) + +jest.mock('@langchain/google-genai', () => ({ + ChatGoogleGenerativeAI: jest.fn() +})) + +jest.mock('@langchain/mistralai', () => ({ + ChatMistralAI: jest.fn() +})) + +jest.mock('@langchain/openai', () => ({ + ChatOpenAI: jest.fn(() => ({ withStructuredOutput: mockWithStructuredOutput })), + AzureChatOpenAI: jest.fn() +})) + +jest.mock('@langchain/groq', () => ({ + ChatGroq: jest.fn() +})) + +jest.mock('ollama', () => ({ + Ollama: jest.fn() +})) + +jest.mock('./utils', () => ({ + getCredentialData: jest.fn() +})) + +import { getCredentialData } from './utils' +import { FollowUpPromptProvider } from './Interface' +import { generateFollowUpPrompts } from './followUpPrompts' + +describe('generateFollowUpPrompts retry classification', () => { + beforeEach(() => { + jest.clearAllMocks() + jest.mocked(getCredentialData).mockResolvedValue({ openAIApiKey: 'openai-key' }) + }) + + afterEach(() => { + jest.useRealTimers() + }) + + it('retries an Axios-style 504 error', async () => { + jest.useFakeTimers() + const gatewayTimeoutError = Object.assign(new Error('Request failed with status code 504'), { + status: 504, + statusCode: 504, + response: { status: 504 } + }) + mockInvoke.mockRejectedValueOnce(gatewayTimeoutError).mockResolvedValueOnce({ questions: ['What happened next?'] }) + + const result = generateFollowUpPrompts( + { + status: true, + selectedProvider: FollowUpPromptProvider.OPENAI, + [FollowUpPromptProvider.OPENAI]: { + credentialId: 'cred-1', + modelName: 'gpt-4o-mini', + prompt: 'Generate follow-ups for {history}', + temperature: '0' + } + } as any, + 'hello', + {} + ) + + await jest.advanceTimersByTimeAsync(500) + + await expect(result).resolves.toEqual({ questions: ['What happened next?'] }) + + expect(mockInvoke).toHaveBeenCalledTimes(2) + }) +}) diff --git a/packages/components/src/followUpPrompts.ts b/packages/components/src/followUpPrompts.ts index 3ed57431173..327b200d7ce 100644 --- a/packages/components/src/followUpPrompts.ts +++ b/packages/components/src/followUpPrompts.ts @@ -39,7 +39,7 @@ const isTimeoutError = (error: any): boolean => { const isRetryableError = (error: any): boolean => { const status = getErrorStatus(error) - if (status && [429, 500, 502, 503].includes(status)) return true + if (status && [408, 429, 500, 502, 503, 504].includes(status)) return true return isTimeoutError(error) } From b5163593cb6c75abff4cee04464ae6e6d5bec31c Mon Sep 17 00:00:00 2001 From: Abhishek Kumbhar Date: Sun, 21 Jun 2026 00:49:24 +0530 Subject: [PATCH 3/3] refactor: remove redundant Ollama timeout wrapper --- packages/components/src/followUpPrompts.ts | 30 +++++----------------- 1 file changed, 6 insertions(+), 24 deletions(-) diff --git a/packages/components/src/followUpPrompts.ts b/packages/components/src/followUpPrompts.ts index 327b200d7ce..16eaf14b414 100644 --- a/packages/components/src/followUpPrompts.ts +++ b/packages/components/src/followUpPrompts.ts @@ -273,15 +273,15 @@ export const generateFollowUpPrompts = async ( } case FollowUpPromptProvider.OLLAMA: { return executeWithRetry( - async (signal) => { + async () => { const ollamaClient = new Ollama({ host: providerConfig.baseUrl || 'http://127.0.0.1:11434' }) // Ollama client does not accept AbortSignal directly in this SDK call. - // Wrap the chat promise in a race so the caller-provided signal can - // abort the request and trigger our retry/timeout behavior. - const chatPromise = ollamaClient.chat({ + // However, executeWithRetry already handles the timeout and abort logic + // via Promise.race at the outer level, making this wrapper redundant. + const response = await ollamaClient.chat({ model: providerConfig.modelName, messages: [ { @@ -310,26 +310,8 @@ export const generateFollowUpPrompts = async ( } }) - let abortHandler: (() => void) | undefined - const abortPromise = new Promise((_, reject) => { - abortHandler = () => { - const timeoutError = new Error('Follow-up prompt request timed out') - timeoutError.name = 'TimeoutError' - reject(timeoutError) - } - signal.addEventListener('abort', abortHandler, { once: true }) - }) - - try { - const response = await Promise.race([chatPromise, abortPromise]) - const result = FollowUpPromptType.parse(JSON.parse((response as any).message.content)) - if (!result.questions) { - throw new Error('Follow-up prompt response missing questions') - } - return { questions: result.questions } - } finally { - if (abortHandler) signal.removeEventListener('abort', abortHandler) - } + const result = FollowUpPromptType.parse(JSON.parse((response as any).message.content)) + return { questions: result.questions } }, { provider: FollowUpPromptProvider.OLLAMA,