refactor: [ENG-3156] Determine response format in provider class (#4858)

* add determineResponseFormat to provider classes

* add err msg

* fixg

Author: connortbot

packages/__tests__/cost/providers/vertex.test.ts

@@ -241,7 +241,6 @@ describe("VertexProvider", () => {
 
         const parsed = JSON.parse(body);
         expect(parsed.anthropic_version).toBe("vertex-2023-10-16");
-        expect(parsed.model).toBe("claude-3-haiku");
         expect(parsed.anthropic_content).toBe(true);
       });
 
@@ -260,7 +259,6 @@ describe("VertexProvider", () => {
 
         const parsed = JSON.parse(body);
         expect(parsed.anthropic_version).toBe("vertex-2023-10-16");
-        expect(parsed.model).toBe("claude-3-sonnet");
         expect(parsed.should_not_appear).toBeUndefined();
       });
     });

packages/cost/models/provider-helpers.ts

@@ -7,6 +7,7 @@ import type {
   AuthResult,
   RequestBodyContext,
   RequestParams,
+  ResponseFormat,
 } from "./types";
 import { providers, ModelProviderName } from "./providers";
 import { BaseProvider } from "./providers/base";
@@ -265,3 +266,17 @@ export async function buildErrorMessage(
 
   return ok(await provider.buildErrorMessage(response));
 }
+
+export function determineResponseFormat(endpoint: Endpoint): Result<ResponseFormat> {
+  const providerResult = getProvider(endpoint.provider);
+  if (providerResult.error) {
+    return err(providerResult.error);
+  }
+
+  const provider = providerResult.data;
+  if (!provider) {
+    return err(`Provider data is null for: ${endpoint.provider}`);
+  }
+
+  return ok(provider.determineResponseFormat(endpoint));
+}

packages/cost/models/providers/anthropic.ts

@@ -5,6 +5,7 @@ import type {
   Endpoint,
   RequestBodyContext,
   RequestParams,
+  ResponseFormat,
 } from "../types";
 
 export class AnthropicProvider extends BaseProvider {
@@ -41,4 +42,11 @@ export class AnthropicProvider extends BaseProvider {
     const anthropicBody = context.toAnthropic(context.parsedBody, endpoint.providerModelId);
     return JSON.stringify(anthropicBody);
   }
+
+  determineResponseFormat(endpoint: Endpoint): ResponseFormat {
+    if (endpoint.author === "anthropic" || endpoint.providerModelId.includes("claude-")) {
+      return "ANTHROPIC";
+    }
+    return "OPENAI";
+  }
 }

packages/cost/models/providers/base.ts

@@ -6,6 +6,7 @@ import type {
   Endpoint,
   RequestParams,
   ModelProviderConfig,
+  ResponseFormat,
 } from "../types";
 import { CacheProvider } from "../../../common/cache/provider";
 
@@ -67,4 +68,8 @@ export abstract class BaseProvider {
       return `Request failed with status ${response.status}`;
     }
   }
+
+  determineResponseFormat(endpoint: Endpoint): ResponseFormat {
+    return "OPENAI";
+  }
 }

packages/cost/models/providers/bedrock.ts

@@ -10,6 +10,7 @@ import type {
   RequestParams,
   ModelProviderConfig,
   UserEndpointConfig,
+  ResponseFormat,
 } from "../types";
 
 export class BedrockProvider extends BaseProvider {
@@ -121,4 +122,11 @@ export class BedrockProvider extends BaseProvider {
       stream_options: undefined,
     });
   }
+
+  determineResponseFormat(endpoint: Endpoint): ResponseFormat {
+    if (endpoint.author === "anthropic" || endpoint.providerModelId.includes("claude-")) {
+      return "ANTHROPIC";
+    }
+    return "OPENAI";
+  }
 }

packages/cost/models/providers/vertex.ts

@@ -5,6 +5,7 @@ import type {
   AuthContext,
   AuthResult,
   RequestParams,
+  ResponseFormat,
 } from "../types";
 import { getGoogleAccessToken } from "../../auth/gcpServiceAccountAuth";
 import { CacheProvider } from "../../../common/cache/provider";
@@ -125,4 +126,11 @@ export class VertexProvider extends BaseProvider {
       return `Request failed with status ${response.status}`;
     }
   }
+
+  determineResponseFormat(endpoint: Endpoint): ResponseFormat {
+    if (endpoint.author === "anthropic" || endpoint.providerModelId.includes("claude-")) {
+      return "ANTHROPIC";
+    }
+    return "OPENAI";
+  } 
 }

packages/cost/models/types.ts

@@ -38,6 +38,8 @@ export interface Modality {
   outputs: OutputModality[];
 }
 
+export type ResponseFormat = "ANTHROPIC" | "OPENAI";
+
 export type Tokenizer =
   | "Claude"
   | "GPT"

worker/src/lib/ai-gateway/SimpleAIGateway.ts

@@ -13,9 +13,9 @@ import { AttemptExecutor } from "./AttemptExecutor";
 import { Attempt, DisallowListEntry, EscrowInfo } from "./types";
 import { oai2antResponse } from "../clients/llmmapper/router/oai2ant/nonStream";
 import { oai2antStreamResponse } from "../clients/llmmapper/router/oai2ant/stream";
-import { ModelProviderName } from "@helicone-package/cost/models/providers";
 import { RequestParams } from "@helicone-package/cost/models/types";
 import { SecureCacheProvider } from "../util/cache/secureCache";
+import { determineResponseFormat } from "@helicone-package/cost/models/provider-helpers";
 
 export interface AuthContext {
   orgId: string;
@@ -303,22 +303,6 @@ export class SimpleAIGateway {
     );
   }
 
-  private determineResponseFormat(
-    provider: ModelProviderName,
-    modelId: string
-  ): "ANTHROPIC" | "OPENAI" {
-    // TODO: make enum type when there's more map formats
-    if (
-      provider === "anthropic" ||
-      (provider === "bedrock" && modelId.includes("claude-")) ||
-      (provider === "vertex" && modelId.includes("claude-"))
-    ) {
-      return "ANTHROPIC";
-    }
-
-    return "OPENAI";
-  }
-
   private async mapResponse(
     attempt: Attempt,
     response: Response,
@@ -328,17 +312,18 @@ export class SimpleAIGateway {
       return ok(response); // do not map response
     }
 
-    const mappingType = this.determineResponseFormat(
-      attempt.endpoint.provider,
-      attempt.endpoint.providerModelId
-    ); // finds format of the response that we are mapping to OPENAI
+    const mappingType = determineResponseFormat(attempt.endpoint);
+
+    if (isErr(mappingType)) {
+      return err(`Failed to determine response format: ${mappingType.error}`);
+    }
 
-    if (mappingType === "OPENAI") {
+    if (mappingType.data === "OPENAI") {
       return ok(response); // already in OPENAI format
     }
 
     try {
-      if (mappingType === "ANTHROPIC") {
+      if (mappingType.data === "ANTHROPIC") {
         const contentType = response.headers.get("content-type");
         const isStream = contentType?.includes("text/event-stream");