fix: latency calculation for claude, gemini, gpt chat (#216)

empirical-run · May 15, 2024 · b379e7f · b379e7f
1 parent a31470d
commit b379e7f
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 10 deletions.
diff --git a/.changeset/short-clouds-double.md b/.changeset/short-clouds-double.md
@@ -0,0 +1,5 @@
+---
+"@empiricalrun/ai": patch
+---
+
+fix: latency calculation for claude, gemini, gpt chat
diff --git a/packages/ai/src/providers/anthropic/index.ts b/packages/ai/src/providers/anthropic/index.ts
@@ -93,7 +93,7 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
   const { contents, systemPrompt } = convertOpenAIToAnthropicAI(messages);
   const { executionDone } = await batchTaskManager.waitForTurn();
   try {
-    const startedAt = Date.now();
+    let startedAt = Date.now();
     const response = await promiseRetry<Anthropic.Beta.Tools.ToolsBetaMessage>(
       (retry, attempt) => {
         return anthropic.beta.tools.messages
@@ -122,6 +122,7 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
               console.warn(
                 `Retrying request for anthropic model: ${body.model}. Retry attempt: ${attempt}`,
               );
+              startedAt = Date.now();
               retry(err);
               throw err;
             }

diff --git a/packages/ai/src/providers/google/index.test.ts b/packages/ai/src/providers/google/index.test.ts
@@ -24,18 +24,15 @@ describe.concurrent("test google ai provider", () => {
       messages: [
         {
           role: "user",
-          content: `List a few popular cookie recipes using this JSON schema: {'type': 'object', 'properties': { 'recipe_name': {'type': 'string'}}}`,
+          content: `List a few popular cookie recipes using this JSON schema: {'type': 'array', 'items': { 'type': 'object', 'properties': { 'recipe_name': { 'type': 'string' }}}}`,
         },
       ],
       response_mime_type: "application/json",
     });
     expect(chatCompletion.choices.length).toBeGreaterThan(0);
-    expect(
-      JSON.parse(chatCompletion.choices[0]?.message.content!),
-    ).toBeDefined();
-    expect(
-      JSON.parse(chatCompletion.choices[0]?.message.content!)[0]["recipe_name"],
-    ).toBeDefined();
+    const contents = chatCompletion.choices[0]?.message.content!;
+    expect(JSON.parse(contents)).toBeDefined();
+    expect(JSON.parse(contents)[0]["recipe_name"]).toBeDefined();
   }, 120000);
 
   test("passthrough model parameters works", async () => {

diff --git a/packages/ai/src/providers/google/index.ts b/packages/ai/src/providers/google/index.ts
@@ -153,7 +153,7 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
   };
   const { executionDone } = await batch.waitForTurn();
   try {
-    const startedAt = Date.now();
+    let startedAt = Date.now();
     const completion = await promiseRetry<GenerateContentResult>(
       (retry, attempt) => {
         return modelInstance
@@ -170,6 +170,7 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
               console.warn(
                 `Retrying request for google model: ${model}. Retry attempt: ${attempt}`,
               );
+              startedAt = Date.now();
               retry(err);
             }
             throw err;

diff --git a/packages/ai/src/providers/openai/index.ts b/packages/ai/src/providers/openai/index.ts
@@ -31,7 +31,7 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
   });
 
   try {
-    const startedAt = Date.now();
+    let startedAt = Date.now();
     const completions = await promiseRetry<IChatCompletion>(
       (retry, attempt) => {
         return openai.chat.completions.create(body).catch((err) => {
@@ -49,6 +49,7 @@ const createChatCompletion: ICreateChatCompletion = async (body) => {
             console.warn(
               `Retrying request for openai model: ${body.model}. Retry attempt: ${attempt}`,
             );
+            startedAt = Date.now();
             retry(err);
             throw err;
           }