feat: accept timeout as parameter in empiricalrc.json (#119)

Co-authored-by: Saikat Mitra <saikatmitra91@gmail.com>
empirical-run · Apr 15, 2024 · 9822db6 · 9822db6
1 parent d087119
commit 9822db6
Show file tree

Hide file tree

Showing 11 changed files with 71 additions and 16 deletions.
diff --git a/.changeset/good-spoons-own.md b/.changeset/good-spoons-own.md
@@ -0,0 +1,6 @@
+---
+"@empiricalrun/ai": minor
+"@empiricalrun/types": patch
+---
+
+feat: accept timeout as parameter in empiricalrc.json
diff --git a/docs/models/basics.mdx b/docs/models/basics.mdx
@@ -116,6 +116,24 @@ For example, Mistral models support a `safePrompt` parameter for [guardrailing](
 ]
 ```
 
+#### Configuring request timeout
+
+You can set the timeout duration in milliseconds under model parameters in the `empiricalrc.json` file. This might be required for prompt completions that are expected to take more time, for example while running models like Claude Opus. If no specific value is assigned, the default timeout duration of 30 seconds will be applied.
+
+```json empiricalrc.json
+"runs": [
+  {
+    "type": "model",
+    "provider": "anthropic",
+    "model": "claude-3-opus",
+    "prompt": "Hey I'm {{user_name}}",
+    "parameters": {
+      "timeout": 10000
+    }
+  }
+]
+```
+
 #### Limitations
 
 - These parameters are not supported today: `logit_bias`, `tools`, `tool_choice`, `user`, `stream`

diff --git a/packages/ai/src/constants/index.ts b/packages/ai/src/constants/index.ts
@@ -0,0 +1 @@
+export const DEFAULT_TIMEOUT = 30000;
diff --git a/packages/ai/src/providers/anthropic/index.ts b/packages/ai/src/providers/anthropic/index.ts
@@ -4,6 +4,7 @@ import { ChatCompletionMessageParam } from "openai/resources/chat/completions.mj
 import promiseRetry from "promise-retry";
 import { BatchTaskManager, getPassthroughParams } from "../../utils";
 import { AIError, AIErrorEnum } from "../../error";
+import { DEFAULT_TIMEOUT } from "../../constants";
 
 const batchTaskManager = new BatchTaskManager(5);
 
@@ -54,10 +55,15 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
       "process.env.ANTHROPIC_API_KEY is not set",
     );
   }
+  const { model, messages, ...config } = body;
+  const timeout = config.timeout || DEFAULT_TIMEOUT;
+  if (config.timeout) {
+    delete config.timeout;
+  }
   const anthropic = new Anthropic({
     apiKey: process.env.ANTHROPIC_API_KEY,
+    timeout: timeout,
   });
-  const { model, messages, ...config } = body;
   const { contents, systemPrompt } = convertOpenAIToAnthropicAI(messages);
   const { executionDone } = await batchTaskManager.waitForTurn();
   try {
@@ -130,7 +136,7 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
     executionDone();
     throw new AIError(
       AIErrorEnum.FAILED_CHAT_COMPLETION,
-      `failed chat completion for model ${body.model} with message ${(e as Error).message} `,
+      `Failed to fetch output from model ${body.model}: ${(e as Error).message}`,
     );
   }
 };

diff --git a/packages/ai/src/providers/fireworks/index.ts b/packages/ai/src/providers/fireworks/index.ts
@@ -11,6 +11,9 @@ const batchTaskManager = new BatchTaskManager(10);
 
 const createChatCompletion: ICreateChatCompletion = async (body) => {
   const { model, messages, ...config } = body;
+  if (config.timeout) {
+    delete config.timeout;
+  }
   const payload = JSON.stringify({
     model: `accounts/fireworks/models/${model}`,
     messages,
@@ -35,6 +38,7 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
   const { executionDone } = await batchTaskManager.waitForTurn();
 
   try {
+    const startedAt = Date.now();
     const completion = await promiseRetry<IChatCompletion>(
       (retry) => {
         return fetch("https://api.fireworks.ai/inference/v1/chat/completions", {
@@ -67,12 +71,11 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
       },
       {
         randomize: true,
-        minTimeout: 1000,
       },
     );
-
+    const latency = Date.now() - startedAt;
     executionDone();
-    return completion;
+    return { ...completion, latency };
   } catch (err) {
     throw new AIError(AIErrorEnum.FAILED_CHAT_COMPLETION, "Unknown error");
   }

diff --git a/packages/ai/src/providers/google/index.ts b/packages/ai/src/providers/google/index.ts
@@ -16,6 +16,7 @@ import { BatchTaskManager } from "../../utils";
 import crypto from "crypto";
 import promiseRetry from "promise-retry";
 import { AIError, AIErrorEnum } from "../../error";
+import { DEFAULT_TIMEOUT } from "../../constants";
 
 const batch = new BatchTaskManager(5);
 
@@ -62,7 +63,8 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
   }
   const { model, messages } = body;
   const googleAI = new GoogleGenerativeAI(process.env.GOOGLE_API_KEY!);
-  const modelInstance = googleAI.getGenerativeModel({ model });
+  const timeout = body.timeout || DEFAULT_TIMEOUT;
+  const modelInstance = googleAI.getGenerativeModel({ model }, { timeout });
   const contents = massageOpenAIMessagesToGoogleAI(messages);
   const { executionDone } = await batch.waitForTurn();
   try {
@@ -79,7 +81,6 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
       },
       {
         randomize: true,
-        minTimeout: 2000,
       },
     );
     executionDone();
@@ -131,7 +132,7 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
     executionDone();
     throw new AIError(
       AIErrorEnum.FAILED_CHAT_COMPLETION,
-      `failed chat completion for model ${body.model} with message ${(e as Error).message}`,
+      `Failed to fetch output from model ${body.model} with message ${(e as Error).message}`,
     );
   }
 };

diff --git a/packages/ai/src/providers/mistral/index.ts b/packages/ai/src/providers/mistral/index.ts
@@ -6,6 +6,7 @@ import {
 import { BatchTaskManager, getPassthroughParams } from "../../utils";
 import { ToolCalls, ResponseFormat } from "@mistralai/mistralai";
 import { AIError, AIErrorEnum } from "../../error";
+import { DEFAULT_TIMEOUT } from "../../constants";
 
 type MistralChatMessage = {
   role: string;
@@ -29,9 +30,19 @@ const createChatCompletion: ICreateChatCompletion = async function (body) {
     );
   }
   const MistralClient = await importMistral();
-  const mistralai = new MistralClient(process.env.MISTRAL_API_KEY);
-  const { executionDone } = await batch.waitForTurn();
   const { model, messages, ...config } = body;
+  const mistralai = new MistralClient(
+    process.env.MISTRAL_API_KEY,
+    undefined,
+    // type issue in https://github.com/mistralai/client-js/blob/e33a2f3e5f6fb88fd083e8e7d9c3c081d1c7c0e4/src/client.js#L51, will submit a PR later
+    // @ts-ignore default value for retries
+    5,
+    (config.timeout || DEFAULT_TIMEOUT) / 1000, // Mistral expects values in seconds
+  );
+  if (config.timeout) {
+    delete config.timeout;
+  }
+  const { executionDone } = await batch.waitForTurn();
   try {
     // typecasting as there is a minor difference in role being openai enum vs string
     const mistralMessages = messages as MistralChatMessage[];
@@ -56,7 +67,7 @@ const createChatCompletion: ICreateChatCompletion = async function (body) {
     executionDone();
     throw new AIError(
       AIErrorEnum.FAILED_CHAT_COMPLETION,
-      `failed chat completion for model ${body.model} with message ${(err as Error).message}`,
+      `Failed to fetch output from model ${body.model} with message ${(err as Error).message}`,
     );
   }
 };

diff --git a/packages/ai/src/providers/openai/index.ts b/packages/ai/src/providers/openai/index.ts
@@ -6,6 +6,7 @@ import {
 import OpenAI from "openai";
 import promiseRetry from "promise-retry";
 import { AIError, AIErrorEnum } from "../../error";
+import { DEFAULT_TIMEOUT } from "../../constants";
 
 const createChatCompletion: ICreateChatCompletion = async (body) => {
   const apiKey = process.env.OPENAI_API_KEY;
@@ -15,9 +16,15 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
       "process.env.OPENAI_API_KEY is not set",
     );
   }
+  const timeout = body.timeout || DEFAULT_TIMEOUT;
+  if (body.timeout) {
+    delete body.timeout;
+  }
   const openai = new OpenAI({
     apiKey: process.env.OPENAI_API_KEY,
+    timeout,
   });
+
   try {
     const startedAt = Date.now();
     const completions = await promiseRetry<IChatCompletion>(
@@ -50,7 +57,7 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
   } catch (err) {
     throw new AIError(
       AIErrorEnum.FAILED_CHAT_COMPLETION,
-      `Failed completion for OpenAI ${body.model}: ${(err as any)?.error?.message}`,
+      `Failed to fetch output from model ${body.model}: ${(err as any)?.error?.message}`,
     );
   }
 };

diff --git a/packages/ai/src/utils/index.ts b/packages/ai/src/utils/index.ts
@@ -21,6 +21,7 @@ function isReservedParameter(paramName: string) {
     "seed",
     "stop",
     "top_logprobs",
+    "timeout",
   ];
   return reservedParameters.indexOf(paramName) >= 0;
 }

diff --git a/packages/types/src/index.ts b/packages/types/src/index.ts
@@ -68,6 +68,7 @@ interface ModelParameters {
   seed?: number;
   stop?: string | Array<string>;
   top_logprobs?: number;
+  timeout?: number;
 
   // For other models, we coerce the above known parameters to appropriate slots
   // If users require other parameters, we support passthrough for other key names

diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
-Original file line number
+Diff line change
@@ Expand Up / @@ -21,6 +21,7 @@ function isReservedParameter(paramName: string) { @@
         "seed",
         "stop",
         "top_logprobs",
+        "timeout",
       ];
       return reservedParameters.indexOf(paramName) >= 0;
     }
@@ Expand Down @@