Add Google AI API (khanon/oai-reverse-proxy!57)

src/shared/tokenization/openai.ts

@@ -2,7 +2,11 @@ import { Tiktoken } from "tiktoken/lite";
 import cl100k_base from "tiktoken/encoders/cl100k_base.json";
 import { logger } from "../../logger";
 import { libSharp } from "../file-storage";
-import type { OpenAIChatMessage } from "../../proxy/middleware/request/preprocessors/transform-outbound-payload";
+import type {
+  GoogleAIChatMessage,
+  OpenAIChatMessage,
+} from "../../proxy/middleware/request/preprocessors/transform-outbound-payload";
+import { z } from "zod";
 
 const log = logger.child({ module: "tokenizer", service: "openai" });
 const GPT4_VISION_SYSTEM_PROMPT_SIZE = 170;
@@ -29,11 +33,11 @@ export async function getTokenCount(
     return getTextTokenCount(prompt);
   }
 
-  const gpt4 = model.startsWith("gpt-4");
+  const oldFormatting = model.startsWith("turbo-0301");
   const vision = model.includes("vision");
 
-  const tokensPerMessage = gpt4 ? 3 : 4;
-  const tokensPerName = gpt4 ? 1 : -1; // turbo omits role if name is present
+  const tokensPerMessage = oldFormatting ? 4 : 3;
+  const tokensPerName = oldFormatting ? -1 : 1; // older formatting replaces role with name if name is present
 
   let numTokens = vision ? GPT4_VISION_SYSTEM_PROMPT_SIZE : 0;
 
@@ -228,3 +232,24 @@ export function getOpenAIImageCost(params: {
     token_count: Math.ceil(tokens),
   };
 }
+
+export function estimateGoogleAITokenCount(prompt: string | GoogleAIChatMessage[]) {
+  if (typeof prompt === "string") {
+    return getTextTokenCount(prompt);
+  }
+
+  const tokensPerMessage = 3;
+
+  let numTokens = 0;
+  for (const message of prompt) {
+    numTokens += tokensPerMessage;
+    numTokens += encoder.encode(message.parts[0].text).length;
+  }
+
+  numTokens += 3;
+
+  return {
+    tokenizer: "tiktoken (google-ai estimate)",
+    token_count: numTokens,
+  };
+}