zen: gemini

5 files changed, 79 insertions, 1 deletions

diff --git a/packages/console/app/src/routes/zen/util/provider/anthropic.ts b/packages/console/app/src/routes/zen/util/provider/anthropic.ts
index d8d1cd741..887a6e4b5 100644
--- a/packages/console/app/src/routes/zen/util/provider/anthropic.ts
+++ b/packages/console/app/src/routes/zen/util/provider/anthropic.ts
@@ -30,6 +30,7 @@ export const anthropicHelper = {
       service_tier: "standard_only",
     }
   },
+  streamSeparator: "\n\n",
   createUsageParser: () => {
     let usage: Usage
 
diff --git a/packages/console/app/src/routes/zen/util/provider/google.ts b/packages/console/app/src/routes/zen/util/provider/google.ts
new file mode 100644
index 000000000..afde42096
--- /dev/null
+++ b/packages/console/app/src/routes/zen/util/provider/google.ts
@@ -0,0 +1,74 @@
+import { ProviderHelper } from "./provider"
+
+/*
+{
+  promptTokenCount: 11453,
+  candidatesTokenCount: 71,
+  totalTokenCount: 11625,
+  cachedContentTokenCount: 8100,
+  promptTokensDetails: [
+    {modality: "TEXT",tokenCount: 11453}
+  ],
+  cacheTokensDetails: [
+    {modality: "TEXT",tokenCount: 8100}
+  ],
+  thoughtsTokenCount: 101
+}
+*/
+
+type Usage = {
+  promptTokenCount?: number
+  candidatesTokenCount?: number
+  totalTokenCount?: number
+  cachedContentTokenCount?: number
+  promptTokensDetails?: { modality: string; tokenCount: number }[]
+  cacheTokensDetails?: { modality: string; tokenCount: number }[]
+  thoughtsTokenCount?: number
+}
+
+export const googleHelper = {
+  format: "google",
+  modifyUrl: (providerApi: string, model?: string, isStream?: boolean) =>
+    `${providerApi}/models/${model}:${isStream ? "streamGenerateContent?alt=sse" : "generateContent"}`,
+  modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => {
+    headers.set("x-goog-api-key", apiKey)
+  },
+  modifyBody: (body: Record<string, any>) => {
+    return body
+  },
+  streamSeparator: "\r\n\r\n",
+  createUsageParser: () => {
+    let usage: Usage
+
+    return {
+      parse: (chunk: string) => {
+        if (!chunk.startsWith("data: ")) return
+
+        let json
+        try {
+          json = JSON.parse(chunk.slice(6)) as { usageMetadata?: Usage }
+        } catch (e) {
+          return
+        }
+
+        if (!json.usageMetadata) return
+        usage = json.usageMetadata
+      },
+      retrieve: () => usage,
+    }
+  },
+  normalizeUsage: (usage: Usage) => {
+    const inputTokens = usage.promptTokenCount ?? 0
+    const outputTokens = usage.candidatesTokenCount ?? 0
+    const reasoningTokens = usage.thoughtsTokenCount ?? 0
+    const cacheReadTokens = usage.cachedContentTokenCount ?? 0
+    return {
+      inputTokens: inputTokens - cacheReadTokens,
+      outputTokens,
+      reasoningTokens,
+      cacheReadTokens,
+      cacheWrite5mTokens: undefined,
+      cacheWrite1hTokens: undefined,
+    }
+  },
+} satisfies ProviderHelper
diff --git a/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts b/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts
index 8a9170ef1..5771ed4fa 100644
--- a/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts
+++ b/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts
@@ -33,6 +33,7 @@ export const oaCompatHelper = {
       ...(body.stream ? { stream_options: { include_usage: true } } : {}),
     }
   },
+  streamSeparator: "\n\n",
   createUsageParser: () => {
     let usage: Usage
 
diff --git a/packages/console/app/src/routes/zen/util/provider/openai.ts b/packages/console/app/src/routes/zen/util/provider/openai.ts
index e79e83579..dff6e13fb 100644
--- a/packages/console/app/src/routes/zen/util/provider/openai.ts
+++ b/packages/console/app/src/routes/zen/util/provider/openai.ts
@@ -21,6 +21,7 @@ export const openaiHelper = {
   modifyBody: (body: Record<string, any>) => {
     return body
   },
+  streamSeparator: "\n\n",
   createUsageParser: () => {
     let usage: Usage
 
diff --git a/packages/console/app/src/routes/zen/util/provider/provider.ts b/packages/console/app/src/routes/zen/util/provider/provider.ts
index d0f123968..8366f3a63 100644
--- a/packages/console/app/src/routes/zen/util/provider/provider.ts
+++ b/packages/console/app/src/routes/zen/util/provider/provider.ts
@@ -26,9 +26,10 @@ import {
 
 export type ProviderHelper = {
   format: ZenData.Format
-  modifyUrl: (providerApi: string) => string
+  modifyUrl: (providerApi: string, model?: string, isStream?: boolean) => string
   modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => void
   modifyBody: (body: Record<string, any>) => Record<string, any>
+  streamSeparator: string
   createUsageParser: () => {
     parse: (chunk: string) => void
     retrieve: () => any