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 {