perf: gemini 支持cloudflare ai gateway

TBXark · Oct 14, 2024 · 49d9cb1 · 49d9cb1 · Admin-Author · Oct 28, 2024
1 parent 73336e5
commit 49d9cb1
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 14 deletions.
diff --git a/dist/buildinfo.json b/dist/buildinfo.json
diff --git a/dist/index.js b/dist/index.js
diff --git a/dist/timestamp b/dist/timestamp
diff --git a/src/agent/gemini.ts b/src/agent/gemini.ts
@@ -35,11 +35,8 @@ export class Gemini implements ChatAgent {
         if (onStream !== null) {
             console.warn('Stream mode is not supported');
         }
-        const url = `${context.GOOGLE_COMPLETIONS_API}${context.GOOGLE_COMPLETIONS_MODEL}:${
-            // 暂时不支持stream模式
-            // onStream ? 'streamGenerateContent' : 'generateContent'
-            'generateContent'
-        }?key=${context.GOOGLE_API_KEY}`;
+        const mode = 'generateContent'; // onStream ? 'streamGenerateContent' : 'generateContent'
+        const url = `${context.GOOGLE_COMPLETIONS_API}${context.GOOGLE_COMPLETIONS_MODEL}:${mode}`;
 
         const contentsTemp = [...history || [], { role: 'user', content: message }];
         if (prompt) {
@@ -62,7 +59,8 @@ export class Gemini implements ChatAgent {
             method: 'POST',
             headers: {
                 'Content-Type': 'application/json',
-            },
+                'x-goog-api-key': context.GOOGLE_API_KEY,
+            } as Record<string, string>,
             body: JSON.stringify({ contents }),
         });
         const data = await resp.json() as any;

diff --git a/src/config/config.ts b/src/config/config.ts
@@ -148,7 +148,7 @@ export class WorkersConfig {
 export class GeminiConfig {
     // Google Gemini API Key
     GOOGLE_API_KEY: string | null = null;
-    // Google Gemini API
+    // Google Gemini API: Cloudflare AI gateway: https://gateway.ai.cloudflare.com/v1/{account_id}/{gateway_name}/google-ai-studio/v1/models
     GOOGLE_COMPLETIONS_API = 'https://generativelanguage.googleapis.com/v1beta/models/';
     // Google Gemini Model
     GOOGLE_COMPLETIONS_MODEL = 'gemini-pro';