Merge pull request #19 from AmbireTech/feat/llm-max-timeout

Add max timeout for LLM responses

Author: cmihaylov

lib/index.ts

@@ -7,3 +7,4 @@ export * from './utils/llm/structures'
 export * from './utils/prompts'
 export * from './utils/strategies'
 export * from './utils/errors'
+export * from './utils/timing'

lib/types/types.ts

@@ -47,6 +47,8 @@ export type LlmProcessProps = {
     prompt: string
     model?: string
     llmOptionsOverride?: { [x: string]: any }
+    timeout?: number
+    timeoutMsg?: string
 }
 
 export type LlmProcessOutput = {

lib/utils/llm/gemini.ts

@@ -2,6 +2,7 @@ import { GoogleGenerativeAI } from '@google/generative-ai'
 import { LlmProcessOutput, LlmProcessProps, Strategy } from '../../types'
 import { stringifyError } from '../errors'
 import { StrategiesGoogleSchema } from './structures/google'
+import { timeoutPromise } from '../timing'
 
 export const GEMINI_MODELS = {
     gemini20flashExp: 'gemini-2.0-flash-exp',
@@ -26,7 +27,11 @@ export async function callGemini(llmInput: LlmProcessProps): Promise<LlmProcessO
             },
             ...llmInput.llmOptionsOverride
         })
-        const result = await aiModel.generateContent(llmInput.prompt)
+        const result = await timeoutPromise(
+            aiModel.generateContent(llmInput.prompt),
+            llmInput.timeout || 60,
+            llmInput.timeoutMsg
+        )
         // console.log(JSON.stringify(result))
 
         const content = result.response.text()

lib/utils/llm/grok.ts

@@ -3,6 +3,7 @@ import { zodResponseFormat } from 'openai/helpers/zod'
 import { LlmProcessOutput, LlmProcessProps, Strategy } from '../../types'
 import { StrategiesZodSchema } from './structures/zod'
 import { stringifyError } from '../errors'
+import { timeoutPromise } from '../timing'
 
 export const XAI_MODELS = {
     grok2latest: 'grok-2-latest',
@@ -22,20 +23,24 @@ export async function callGrok(llmInput: LlmProcessProps): Promise<LlmProcessOut
     let outputTokens: number | undefined
 
     try {
-        const completion = await apiClient.chat.completions.create({
-            model,
-            store: true,
-            messages: [
-                {
-                    role: 'system',
-                    content:
-                        'You are an expert in cryptocurrencies, DeFi applications and their use cases. Return output in JSON format.'
-                },
-                { role: 'user', content: llmInput.prompt }
-            ],
-            response_format: zodResponseFormat(StrategiesZodSchema, 'strategies'),
-            ...llmInput.llmOptionsOverride
-        })
+        const completion = await timeoutPromise(
+            apiClient.chat.completions.create({
+                model,
+                store: true,
+                messages: [
+                    {
+                        role: 'system',
+                        content:
+                            'You are an expert in cryptocurrencies, DeFi applications and their use cases. Return output in JSON format.'
+                    },
+                    { role: 'user', content: llmInput.prompt }
+                ],
+                response_format: zodResponseFormat(StrategiesZodSchema, 'strategies'),
+                ...llmInput.llmOptionsOverride
+            }),
+            llmInput.timeout || 60,
+            llmInput.timeoutMsg
+        )
 
         const outputContent = completion.choices[0].message.content || '{}'
         inputTokens = completion.usage?.prompt_tokens || 0

lib/utils/timing.ts

@@ -0,0 +1,16 @@
+export function timeoutPromise(
+    promise: Promise<any>,
+    timeoutSeconds: number,
+    errorMsg?: string
+): Promise<any> {
+    let timer: NodeJS.Timeout
+    const timeout = new Promise(
+        (_, reject) =>
+            (timer = setTimeout(
+                () => reject(new Error(errorMsg || 'Request timeout')),
+                timeoutSeconds * 1000
+            ))
+    )
+
+    return Promise.race([promise.finally(() => clearTimeout(timer)), timeout])
+}