feat: ♻️ update API key usage and improve token logging

microsoft · Dec 19, 2024 · ec7734f · ec7734f
1 parent c6d166b
commit ec7734f
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 23 deletions.
diff --git a/docs/src/content/docs/getting-started/configuration.mdx b/docs/src/content/docs/getting-started/configuration.mdx
@@ -736,7 +736,7 @@ Open [Google AI Studio](https://aistudio.google.com/app/apikey) and create a new
 Update the `.env` file with the API key.
 
 ```txt title=".env"
-GOOGLE_API_KEY=...
+GEMINI_API_KEY=...
 ```
 
 </li>

diff --git a/packages/core/src/connection.ts b/packages/core/src/connection.ts
@@ -299,20 +299,21 @@ export async function parseTokenFromEnv(
     }
 
     if (provider === MODEL_PROVIDER_GOOGLE) {
-        const token = env.GOOGLE_API_KEY
+        const token = env.GEMINI_API_KEY || env.GOOGLE_API_KEY
         if (!token) return undefined
         if (token === PLACEHOLDER_API_KEY)
-            throw new Error("GOOGLE_API_KEY not configured")
-        const base = env.GOOGLE_API_BASE || GOOGLE_API_BASE
+            throw new Error("GEMINI_API_KEY/GOOGLE_API_BASE not configured")
+        const base =
+            env.GEMINI_API_BASE || env.GOOGLE_API_BASE || GOOGLE_API_BASE
         if (base === PLACEHOLDER_API_BASE)
-            throw new Error("GOOGLE_API_BASE not configured")
+            throw new Error("GEMINI_API_KEY/GOOGLE_API_BASE not configured")
         return {
             provider,
             model,
             base,
             token,
             type: "openai",
-            source: "env: GOOGLE_API_...",
+            source: "env: GEMINI_API_...",
         } satisfies LanguageModelConfiguration
     }
 

diff --git a/packages/core/src/usage.ts b/packages/core/src/usage.ts
@@ -10,7 +10,7 @@ import {
     CreateChatCompletionRequest,
 } from "./chattypes"
 import { MarkdownTrace } from "./trace"
-import { logVerbose } from "./util"
+import { logVerbose, toStringList } from "./util"
 import { parseModelIdentifier } from "./models"
 import { MODEL_PRICINGS, MODEL_PROVIDER_AICI } from "./constants"
 
@@ -226,7 +226,12 @@ export class GenerationStats {
             try {
                 for (const { messages, usage } of this.chatTurns) {
                     trace.item(
-                        `${messages.length} messages, ${usage.total_tokens} tokens`
+                        toStringList(
+                            `${messages.length} messages`,
+                            usage.total_tokens
+                                ? `${usage.total_tokens} tokens`
+                                : undefined
+                        )
                     )
                 }
             } finally {
@@ -256,16 +261,28 @@ export class GenerationStats {
     private logTokens(indent: string) {
         const unknowns = new Set<string>()
         const c = this.cost()
+        const au = this.accumulatedUsage()
+        if (au?.total_tokens > 0 && (this.resolvedModel || c)) {
+            logVerbose(
+                `${indent}${this.label ? `${this.label} (${this.resolvedModel})` : this.resolvedModel}> ${au.total_tokens} tokens (${au.prompt_tokens} -> ${au.completion_tokens}) ${renderCost(c)}`
+            )
+        }
         if (this.model && isNaN(c) && isCosteable(this.model))
             unknowns.add(this.model)
         if (this.chatTurns.length > 1) {
             const chatTurns = this.chatTurns.slice(0, 10)
-            for (const { messages, usage, model: turnModel } of chatTurns) {
+            for (const {
+                messages,
+                usage,
+                model: turnModel,
+            } of chatTurns.filter(
+                ({ usage }) => usage.total_tokens !== undefined
+            )) {
                 const cost = estimateCost(this.model, usage)
                 if (cost === undefined && isCosteable(turnModel))
                     unknowns.add(this.model)
                 logVerbose(
-                    `${indent}  ${messages.length} messages, ${usage.total_tokens} tokens ${renderCost(cost)}`
+                    `${indent}  ${toStringList(`${messages.length} messages`, usage.total_tokens ? `${usage.total_tokens} tokens` : undefined, renderCost(cost))}`
                 )
             }
             if (this.chatTurns.length > chatTurns.length)
@@ -274,12 +291,6 @@ export class GenerationStats {
         const children = this.children.slice(0, 10)
         for (const child of children) child.logTokens(indent + "  ")
         if (this.children.length > children.length) logVerbose(`${indent}  ...`)
-        const au = this.accumulatedUsage()
-        if (au?.total_tokens > 0 && (this.resolvedModel || c)) {
-            logVerbose(
-                `${indent}${this.label ? `${this.label} (${this.resolvedModel})` : this.resolvedModel}> ${au.total_tokens} tokens (${au.prompt_tokens} -> ${au.completion_tokens}) ${renderCost(c)}`
-            )
-        }
         if (unknowns.size)
             logVerbose(`missing pricing for ${[...unknowns].join(", ")}`)
     }

diff --git a/packages/core/src/util.ts b/packages/core/src/util.ts
@@ -47,7 +47,9 @@ export function toArray<T>(a: ArrayLike<T>): T[] {
 }
 
 export function toStringList(...token: string[]) {
-    const md = token.filter((l) => l !== undefined && l !== null).join(", ")
+    const md = token
+        .filter((l) => l !== undefined && l !== null && l !== "")
+        .join(", ")
     return md
 }
 

diff --git a/packages/sample/genaisrc/tools.genai.mjs b/packages/sample/genaisrc/tools.genai.mjs
@@ -5,22 +5,18 @@ script({
 })
 defTool(
     "get_weather",
-    "Get the current weather in a given location",
+    "Get the real-time weather information at a given location or city",
     {
         type: "object",
         properties: {
             location: {
                 type: "string",
                 description: "The city and state, e.g. Chicago, IL",
             },
-            unit: {
-                type: "string",
-                enum: ["celsius", "fahrenheit"],
-            },
         },
         required: ["location"],
     },
-    ({ location, unit }) => {
+    ({ location }) => {
         return location === "Paris" ? "sunny" : "unknown"
     }
 )