diff --git a/packages/components/src/followUpPrompts.test.ts b/packages/components/src/followUpPrompts.test.ts new file mode 100644 index 00000000000..71199d239f0 --- /dev/null +++ b/packages/components/src/followUpPrompts.test.ts @@ -0,0 +1,100 @@ +import { FollowUpPromptConfig, FollowUpPromptProvider, ICommonObject } from './Interface' +import { generateFollowUpPrompts } from './followUpPrompts' +import { getCredentialData } from './utils' +import { ChatOpenAI } from '@langchain/openai' + +jest.mock('./utils', () => ({ + getCredentialData: jest.fn() +})) + +const invokeMock = jest.fn() +const withStructuredOutputMock = jest.fn(() => ({ + invoke: invokeMock +})) + +jest.mock('@langchain/openai', () => ({ + ChatOpenAI: jest.fn().mockImplementation(() => ({ + withStructuredOutput: withStructuredOutputMock + })), + AzureChatOpenAI: jest.fn() +})) + +const mockedGetCredentialData = getCredentialData as jest.MockedFunction +const mockedChatOpenAI = ChatOpenAI as jest.MockedClass + +const createFollowUpPromptConfig = (status = true): FollowUpPromptConfig => { + const providerConfig = { + credentialId: 'credential-id', + modelName: 'gpt-4.1-mini', + baseUrl: '', + prompt: 'Suggest follow-up questions for {history}', + temperature: '0.2' + } + + return { + status, + selectedProvider: FollowUpPromptProvider.OPENAI, + [FollowUpPromptProvider.ANTHROPIC]: providerConfig, + [FollowUpPromptProvider.AZURE_OPENAI]: providerConfig, + [FollowUpPromptProvider.GOOGLE_GENAI]: providerConfig, + [FollowUpPromptProvider.MISTRALAI]: providerConfig, + [FollowUpPromptProvider.OPENAI]: providerConfig, + [FollowUpPromptProvider.GROQ]: providerConfig, + [FollowUpPromptProvider.OLLAMA]: providerConfig + } +} + +describe('generateFollowUpPrompts', () => { + beforeEach(() => { + jest.useRealTimers() + jest.clearAllMocks() + mockedGetCredentialData.mockResolvedValue({ openAIApiKey: 'openai-key' }) + }) + + it('returns follow-up questions when the provider succeeds', async () => { + const questions = ['What is next?', 'Can you expand?', 'Any constraints?'] + invokeMock.mockResolvedValue({ questions }) + + const result = await generateFollowUpPrompts(createFollowUpPromptConfig(), 'history', {} as ICommonObject) + + expect(result).toEqual({ questions }) + expect(mockedChatOpenAI).toHaveBeenCalledWith({ + apiKey: 'openai-key', + model: 'gpt-4.1-mini', + temperature: 0.2, + useResponsesApi: true + }) + expect(invokeMock).toHaveBeenCalledWith('Suggest follow-up questions for history') + }) + + it('retries once when the first provider call fails', async () => { + const questions = ['Question one?', 'Question two?', 'Question three?'] + invokeMock.mockRejectedValueOnce(new Error('temporary provider failure')).mockResolvedValueOnce({ questions }) + + const result = await generateFollowUpPrompts(createFollowUpPromptConfig(), 'history', {} as ICommonObject) + + expect(result).toEqual({ questions }) + expect(invokeMock).toHaveBeenCalledTimes(2) + }) + + it('returns undefined when provider calls time out', async () => { + jest.useFakeTimers() + invokeMock.mockImplementation(() => new Promise(() => undefined)) + + const resultPromise = generateFollowUpPrompts(createFollowUpPromptConfig(), 'history', {} as ICommonObject) + await Promise.resolve() + + await jest.advanceTimersByTimeAsync(20000) + + await expect(resultPromise).resolves.toBeUndefined() + expect(invokeMock).toHaveBeenCalledTimes(2) + }) + + it('does not read credentials or call a provider when follow-up prompts are disabled', async () => { + const result = await generateFollowUpPrompts(createFollowUpPromptConfig(false), 'history', {} as ICommonObject) + + expect(result).toBeUndefined() + expect(mockedGetCredentialData).not.toHaveBeenCalled() + expect(mockedChatOpenAI).not.toHaveBeenCalled() + }) +}) diff --git a/packages/components/src/followUpPrompts.ts b/packages/components/src/followUpPrompts.ts index bc3b2e0f942..ae73d7822ee 100644 --- a/packages/components/src/followUpPrompts.ts +++ b/packages/components/src/followUpPrompts.ts @@ -20,6 +20,60 @@ export interface FollowUpPromptResult { questions: string[] } +interface FollowUpPromptGenerationOptions { + timeoutMs?: number + maxRetries?: number +} + +type FollowUpPromptTask = () => Promise + +const DEFAULT_FOLLOW_UP_PROMPT_TIMEOUT_MS = 10000 +const DEFAULT_FOLLOW_UP_PROMPT_MAX_RETRIES = 1 + +class FollowUpPromptTimeoutError extends Error { + constructor(timeoutMs: number) { + super(`Follow-up prompt generation timed out after ${timeoutMs}ms`) + this.name = 'FollowUpPromptTimeoutError' + } +} + +const createTimeoutPromise = (timeoutMs: number): { promise: Promise; clear: () => void } => { + let timeout: NodeJS.Timeout + const promise = new Promise((_, reject) => { + timeout = setTimeout(() => reject(new FollowUpPromptTimeoutError(timeoutMs)), timeoutMs) + }) + + return { + promise, + clear: () => clearTimeout(timeout) + } +} + +const invokeWithTimeoutAndRetry = async ( + task: FollowUpPromptTask, + options: FollowUpPromptGenerationOptions = {} +): Promise => { + const timeoutMs = options.timeoutMs ?? DEFAULT_FOLLOW_UP_PROMPT_TIMEOUT_MS + const maxRetries = options.maxRetries ?? DEFAULT_FOLLOW_UP_PROMPT_MAX_RETRIES + + for (let attempt = 0; attempt <= maxRetries; attempt += 1) { + const timeout = createTimeoutPromise(timeoutMs) + + try { + const taskPromise = task() + taskPromise.catch(() => undefined) + + return await Promise.race([taskPromise, timeout.promise]) + } catch { + if (attempt === maxRetries) return undefined + } finally { + timeout.clear() + } + } + + return undefined +} + export const generateFollowUpPrompts = async ( followUpPromptsConfig: FollowUpPromptConfig, apiMessageContent: string, @@ -44,8 +98,10 @@ export const generateFollowUpPrompts = async ( const structuredLLM = llm.withStructuredOutput(FollowUpPromptType, { method: 'functionCalling' }) - const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) - return structuredResponse as FollowUpPromptResult + return await invokeWithTimeoutAndRetry(async () => { + const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) + return structuredResponse as FollowUpPromptResult + }) } case FollowUpPromptProvider.AZURE_OPENAI: { const azureOpenAIApiKey = credentialData['azureOpenAIApiKey'] @@ -70,11 +126,13 @@ export const generateFollowUpPrompts = async ( {format_instructions} `) const chain = prompt.pipe(llm).pipe(parser) - const structuredResponse = await chain.invoke({ - history: apiMessageContent, - format_instructions: formatInstructions + return await invokeWithTimeoutAndRetry(async () => { + const structuredResponse = await chain.invoke({ + history: apiMessageContent, + format_instructions: formatInstructions + }) + return structuredResponse as FollowUpPromptResult }) - return structuredResponse as FollowUpPromptResult } case FollowUpPromptProvider.GOOGLE_GENAI: { const model = new ChatGoogleGenerativeAI({ @@ -85,8 +143,10 @@ export const generateFollowUpPrompts = async ( const structuredLLM = model.withStructuredOutput(FollowUpPromptType, { method: 'functionCalling' }) - const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) - return structuredResponse as FollowUpPromptResult + return await invokeWithTimeoutAndRetry(async () => { + const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) + return structuredResponse as FollowUpPromptResult + }) } case FollowUpPromptProvider.MISTRALAI: { const model = new ChatMistralAI({ @@ -98,8 +158,10 @@ export const generateFollowUpPrompts = async ( const structuredLLM = model.withStructuredOutput(FollowUpPromptType, { method: 'functionCalling' }) - const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) - return structuredResponse as FollowUpPromptResult + return await invokeWithTimeoutAndRetry(async () => { + const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) + return structuredResponse as FollowUpPromptResult + }) } case FollowUpPromptProvider.OPENAI: { const model = new ChatOpenAI({ @@ -112,8 +174,10 @@ export const generateFollowUpPrompts = async ( const structuredLLM = model.withStructuredOutput(FollowUpPromptType, { method: 'functionCalling' }) - const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) - return structuredResponse as FollowUpPromptResult + return await invokeWithTimeoutAndRetry(async () => { + const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) + return structuredResponse as FollowUpPromptResult + }) } case FollowUpPromptProvider.GROQ: { const llm = new ChatGroq({ @@ -124,44 +188,48 @@ export const generateFollowUpPrompts = async ( const structuredLLM = llm.withStructuredOutput(FollowUpPromptType, { method: 'functionCalling' }) - const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) - return structuredResponse as FollowUpPromptResult + return await invokeWithTimeoutAndRetry(async () => { + const structuredResponse = await structuredLLM.invoke(followUpPromptsPrompt) + return structuredResponse as FollowUpPromptResult + }) } 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' - }, - minItems: 3, - maxItems: 3, - description: 'Three follow-up questions based on the conversation history' + return await invokeWithTimeoutAndRetry(async () => { + const response = await 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' + } + }, + required: ['questions'], + additionalProperties: false }, - required: ['questions'], - additionalProperties: false - }, - options: { - temperature: parseFloat(`${providerConfig.temperature}`) - } - }) - const result = FollowUpPromptType.parse(JSON.parse(response.message.content)) - return result + options: { + temperature: parseFloat(`${providerConfig.temperature}`) + } + }) + const result = FollowUpPromptType.parse(JSON.parse(response.message.content)) + return result + }) } } } else {