anthropic mapping on gateway

Author: connortbot

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

@@ -11,6 +11,9 @@ import { gatewayForwarder } from "../../routers/gatewayRouter";
 import { AttemptBuilder } from "./AttemptBuilder";
 import { AttemptExecutor } from "./AttemptExecutor";
 import { Attempt, DisallowListEntry, EscrowInfo } from "./types";
+import { ant2oaiNonStream } from "../clients/llmmapper/router/ant2oai/nonStream";
+import { ant2oaiStreamResponse } from "../clients/llmmapper/router/ant2oai/stream";
+import { ModelProviderName } from "@helicone-package/cost/models/providers";
 
 export interface AuthContext {
   orgId: string;
@@ -138,7 +141,19 @@ export class SimpleAIGateway {
       } else {
         // Success!
         this.requestWrapper.setSuccessfulAttempt(attempt);
-        return result.data;
+        
+        const mappedResponse = await this.mapResponse(
+          attempt,
+          result.data,
+          this.requestWrapper.heliconeHeaders.gatewayConfig.bodyMapping
+        );
+        
+        if (isErr(mappedResponse)) {
+          console.error("Failed to map response:", mappedResponse.error);
+          return result.data;
+        }
+        
+        return mappedResponse.data;
       }
     }
 
@@ -276,6 +291,66 @@ export class SimpleAIGateway {
     );
   }
 
+  private determineResponseFormat(
+    provider: ModelProviderName,
+    modelId: string,
+    bodyMapping?: "OPENAI" | "NO_MAPPING"
+  ): "ANTHROPIC" | "OPENAI" { // TODO: make enum type when there's more map formats
+    if (bodyMapping !== "OPENAI") {
+      return "OPENAI";
+    }
+
+    if (
+      provider === "anthropic" ||
+      (provider === "bedrock" && modelId.includes("claude-"))
+    ) {
+      return "ANTHROPIC";
+    }
+
+    return "OPENAI";
+  }
+
+  private async mapResponse(
+    attempt: Attempt,
+    response: Response,
+    bodyMapping?: "OPENAI" | "NO_MAPPING"
+  ): Promise<Result<Response, string>> {
+    const mappingType = this.determineResponseFormat(
+      attempt.endpoint.provider,
+      attempt.endpoint.providerModelId,
+      bodyMapping
+    );
+    
+    if (mappingType === "OPENAI") {
+      return ok(response);
+    }
+
+    // clone to preserve original for logging
+    const clonedResponse = response.clone();
+    
+    try {
+      if (mappingType === "ANTHROPIC") {
+        const contentType = clonedResponse.headers.get("content-type");
+        const isStream = contentType?.includes("text/event-stream");
+
+        if (isStream) {
+          const mappedResponse = ant2oaiStreamResponse(clonedResponse);
+          return ok(mappedResponse);
+        } else {
+          const mappedResponse = await ant2oaiNonStream(clonedResponse);
+          return ok(mappedResponse);
+        }
+      }
+
+      return ok(response);
+    } catch (error) {
+      console.error("Failed to map response:", error);
+      return err(
+        error instanceof Error ? error.message : "Failed to map response"
+      );
+    }
+  }
+
   private async createErrorResponse(
     errors: Array<{
       attempt: string;

worker/src/lib/clients/llmmapper/providers/anthropic/request/types.ts

@@ -14,8 +14,25 @@ export interface AntRequestBody {
   top_k?: number;
   stop_sequences?: string[];
   stream?: boolean;
+  tools?: AnthropicTool[];
+  tool_choice?: AnthropicToolChoice;
 }
 
+export interface AnthropicTool {
+  name: string;
+  description: string;
+  input_schema: {
+    type: "object";
+    properties: Record<string, any>;
+    required?: string[];
+  };
+}
+
+export type AnthropicToolChoice = 
+  | { type: "auto" }
+  | { type: "any" }
+  | { type: "tool"; name: string };
+
 export interface ContentBlock {
   type: "text" | "image";
   text?: string;

worker/src/lib/clients/llmmapper/providers/openai/request/toAnthropic.ts

@@ -1,4 +1,4 @@
-import { AntRequestBody, ContentBlock } from "../../anthropic/request/types";
+import { AntRequestBody, ContentBlock, AnthropicTool, AnthropicToolChoice } from "../../anthropic/request/types";
 import { OpenAIRequestBody } from "./types";
 
 export function toAnthropic(openAIBody: OpenAIRequestBody): AntRequestBody {
@@ -30,8 +30,19 @@ export function toAnthropic(openAIBody: OpenAIRequestBody): AntRequestBody {
     antBody.metadata = { user_id: openAIBody.user };
   }
 
+  // Map tools from OpenAI format to Anthropic format
+  if (openAIBody.tools) {
+    antBody.tools = mapTools(openAIBody.tools);
+  }
+
+  // Map tool_choice from OpenAI format to Anthropic format
+  if (openAIBody.tool_choice) {
+    antBody.tool_choice = mapToolChoice(openAIBody.tool_choice);
+  }
+
+  // Legacy function_call/functions not supported
   if (openAIBody.function_call || openAIBody.functions) {
-    throw new Error("Function calling and tools are not supported");
+    throw new Error("Legacy function_call and functions are not supported. Use tools instead.");
   }
 
   if (openAIBody.logit_bias) {
@@ -110,3 +121,56 @@ function mapMessages(
     return antMessage;
   });
 }
+
+function mapTools(tools: OpenAIRequestBody["tools"]): AnthropicTool[] {
+  if (!tools) return [];
+  
+  return tools.map((tool) => {
+    if (tool.type !== "function") {
+      throw new Error(`Unsupported tool type: ${tool.type}`);
+    }
+
+    const inputSchema = tool.function.parameters as any || {
+      type: "object",
+      properties: {},
+    };
+
+    return {
+      name: tool.function.name,
+      description: tool.function.description || "",
+      input_schema: {
+        type: "object" as const,
+        properties: inputSchema.properties || {},
+        required: inputSchema.required,
+      },
+    };
+  });
+}
+
+function mapToolChoice(toolChoice: OpenAIRequestBody["tool_choice"]): AnthropicToolChoice {
+  if (!toolChoice) {
+    return { type: "auto" };
+  }
+
+  if (typeof toolChoice === "string") {
+    switch (toolChoice) {
+      case "auto":
+        return { type: "auto" };
+      case "none":
+        // Anthropic doesn't have "none", so we'll omit tools entirely
+        // This should be handled at a higher level
+        return { type: "auto" };
+      default:
+        throw new Error(`Unsupported tool_choice string: ${toolChoice}`);
+    }
+  }
+
+  if (typeof toolChoice === "object" && toolChoice.type === "function") {
+    return {
+      type: "tool",
+      name: toolChoice.function.name,
+    };
+  }
+
+  throw new Error("Unsupported tool_choice format");
+}

worker/src/lib/clients/llmmapper/router/ant2oai/nonStream.ts

@@ -0,0 +1,20 @@
+import { toOpenAI } from "../../providers/anthropic/response/toOpenai";
+import { AntResponseBody } from "../../providers/anthropic/response/types";
+
+export async function ant2oaiNonStream(response: Response): Promise<Response> {
+  try {
+    const anthropicBody = await response.json() as AntResponseBody;
+    const openAIBody = toOpenAI(anthropicBody);
+
+    return new Response(JSON.stringify(openAIBody), {
+      status: response.status,
+      statusText: response.statusText,
+      headers: {
+        "content-type": "application/json",
+      },
+    });
+  } catch (error) {
+    console.error("Failed to transform Anthropic response:", error);
+    return response;
+  }
+}

worker/src/lib/clients/llmmapper/router/ant2oai/stream.ts

@@ -0,0 +1,107 @@
+import { toOpenAI as toOpenAIStreamed } from "../../providers/anthropic/streamedResponse/toOpenai";
+import { AnthropicStreamEvent } from "../../providers/anthropic/streamedResponse/types";
+
+export function ant2oaiStream(
+  stream: ReadableStream<Uint8Array>
+): ReadableStream<Uint8Array> {
+  const decoder = new TextDecoder();
+  const encoder = new TextEncoder();
+  let buffer = "";
+
+  return new ReadableStream({
+    async start(controller) {
+      const reader = stream.getReader();
+
+      try {
+        while (true) {
+          const { done, value } = await reader.read();
+          
+          if (done) {
+            // Process any remaining buffer
+            if (buffer.trim()) {
+              processBuffer(buffer, controller, encoder);
+            }
+            // Send the final [DONE] message
+            controller.enqueue(encoder.encode("data: [DONE]\n\n"));
+            controller.close();
+            break;
+          }
+
+          // Decode the chunk and add to buffer
+          buffer += decoder.decode(value, { stream: true });
+          
+          // Process complete SSE messages (separated by double newlines)
+          const messages = buffer.split("\n\n");
+          
+          // Keep the last incomplete message in the buffer
+          buffer = messages.pop() || "";
+          
+          // Process complete messages
+          for (const message of messages) {
+            if (message.trim()) {
+              processBuffer(message + "\n\n", controller, encoder);
+            }
+          }
+        }
+      } catch (error) {
+        controller.error(error);
+      } finally {
+        reader.releaseLock();
+      }
+    },
+  });
+}
+
+function processBuffer(
+  buffer: string,
+  controller: ReadableStreamDefaultController<Uint8Array>,
+  encoder: TextEncoder
+) {
+  // Split by lines to find data lines
+  const lines = buffer.split("\n");
+  
+  for (const line of lines) {
+    if (line.startsWith("data: ")) {
+      try {
+        const jsonStr = line.slice(6);
+        
+        // Skip the [DONE] message from Anthropic
+        if (jsonStr.trim() === "[DONE]") {
+          continue;
+        }
+        
+        const anthropicEvent: AnthropicStreamEvent = JSON.parse(jsonStr);
+        const openAIEvent = toOpenAIStreamed(anthropicEvent);
+        
+        if (openAIEvent) {
+          const sseMessage = `data: ${JSON.stringify(openAIEvent)}\n\n`;
+          controller.enqueue(encoder.encode(sseMessage));
+        }
+      } catch (error) {
+        // Skip malformed JSON
+        console.error("Failed to parse SSE data:", error);
+      }
+    } else if (line.startsWith("event:") || line.startsWith(":")) {
+      // Skip event type lines and comments
+      continue;
+    }
+  }
+}
+
+export function ant2oaiStreamResponse(response: Response): Response {
+  if (!response.body) {
+    return response;
+  }
+
+  const transformedStream = ant2oaiStream(response.body);
+  
+  return new Response(transformedStream, {
+    status: response.status,
+    statusText: response.statusText,
+    headers: {
+      "content-type": "text/event-stream; charset=utf-8",
+      "cache-control": "no-cache",
+      "connection": "keep-alive",
+    },
+  });
+}