feat(core): add google vertex ai

.docker/selfhost/schema.json

@@ -634,9 +634,9 @@
         },
         "providers.gemini": {
           "type": "object",
-          "description": "The config for the gemini provider.\n@default {\"apiKey\":\"\"}",
+          "description": "The config for the gemini provider.\n@default {\"privateKey\":\"\"}",
           "default": {
-            "apiKey": ""
+            "privateKey": ""
           }
         },
         "providers.perplexity": {

.github/actions/server-test-env/action.yml

@@ -30,9 +30,9 @@ runs:
     - name: Import config
       shell: bash
       run: |
-        printf '{"copilot":{"enabled":true,"providers.fal":{"apiKey":"%s"},"providers.gemini":{"apiKey":"%s"},"providers.openai":{"apiKey":"%s"},"providers.perplexity":{"apiKey":"%s"},"providers.anthropic":{"apiKey":"%s"},"exa":{"key":"%s"}}}' \
+        printf '{"copilot":{"enabled":true,"providers.fal":{"apiKey":"%s"},"providers.gemini":{"privateKey":%s},"providers.openai":{"apiKey":"%s"},"providers.perplexity":{"apiKey":"%s"},"providers.anthropic":{"apiKey":"%s"},"exa":{"key":"%s"}}}' \
         "$COPILOT_FAL_API_KEY" \
-        "$COPILOT_GOOGLE_API_KEY" \
+        "$(printf '%s' "$COPILOT_GOOGLE_PRIVATE_KEY" | jq -aRs .)" \
         "$COPILOT_OPENAI_API_KEY" \
         "$COPILOT_PERPLEXITY_API_KEY" \
         "$COPILOT_ANTHROPIC_API_KEY" \

.github/workflows/build-test.yml

@@ -1000,7 +1000,7 @@ jobs:
         if: ${{ steps.check-blocksuite-update.outputs.skip != 'true' || steps.apifilter.outputs.changed == 'true' }}
         env:
           COPILOT_OPENAI_API_KEY: ${{ secrets.COPILOT_OPENAI_API_KEY }}
-          COPILOT_GOOGLE_API_KEY: ${{ secrets.COPILOT_GOOGLE_API_KEY }}
+          COPILOT_GOOGLE_PRIVATE_KEY: ${{ secrets.COPILOT_GOOGLE_PRIVATE_KEY }}
           COPILOT_FAL_API_KEY: ${{ secrets.COPILOT_FAL_API_KEY }}
           COPILOT_PERPLEXITY_API_KEY: ${{ secrets.COPILOT_PERPLEXITY_API_KEY }}
           COPILOT_ANTHROPIC_API_KEY: ${{ secrets.COPILOT_ANTHROPIC_API_KEY }}
@@ -1104,7 +1104,7 @@ jobs:
         if: ${{ steps.check-blocksuite-update.outputs.skip != 'true' || steps.e2efilter.outputs.changed == 'true' }}
         env:
           COPILOT_OPENAI_API_KEY: ${{ secrets.COPILOT_OPENAI_API_KEY }}
-          COPILOT_GOOGLE_API_KEY: ${{ secrets.COPILOT_GOOGLE_API_KEY }}
+          COPILOT_GOOGLE_PRIVATE_KEY: ${{ secrets.COPILOT_GOOGLE_PRIVATE_KEY }}
           COPILOT_FAL_API_KEY: ${{ secrets.COPILOT_FAL_API_KEY }}
           COPILOT_PERPLEXITY_API_KEY: ${{ secrets.COPILOT_PERPLEXITY_API_KEY }}
           COPILOT_ANTHROPIC_API_KEY: ${{ secrets.COPILOT_ANTHROPIC_API_KEY }}

.github/workflows/copilot-test.yml

@@ -83,7 +83,7 @@ jobs:
         env:
           COPILOT_OPENAI_API_KEY: ${{ secrets.COPILOT_OPENAI_API_KEY }}
           COPILOT_FAL_API_KEY: ${{ secrets.COPILOT_FAL_API_KEY }}
-          COPILOT_GOOGLE_API_KEY: ${{ secrets.COPILOT_GOOGLE_API_KEY }}
+          COPILOT_GOOGLE_PRIVATE_KEY: ${{ secrets.COPILOT_GOOGLE_PRIVATE_KEY }}
           COPILOT_PERPLEXITY_API_KEY: ${{ secrets.COPILOT_PERPLEXITY_API_KEY }}
           COPILOT_ANTHROPIC_API_KEY: ${{ secrets.COPILOT_ANTHROPIC_API_KEY }}
           COPILOT_EXA_API_KEY: ${{ secrets.COPILOT_EXA_API_KEY }}
@@ -158,7 +158,7 @@ jobs:
         env:
           COPILOT_OPENAI_API_KEY: ${{ secrets.COPILOT_OPENAI_API_KEY }}
           COPILOT_FAL_API_KEY: ${{ secrets.COPILOT_FAL_API_KEY }}
-          COPILOT_GOOGLE_API_KEY: ${{ secrets.COPILOT_GOOGLE_API_KEY }}
+          COPILOT_GOOGLE_PRIVATE_KEY: ${{ secrets.COPILOT_GOOGLE_PRIVATE_KEY }}
           COPILOT_PERPLEXITY_API_KEY: ${{ secrets.COPILOT_PERPLEXITY_API_KEY }}
           COPILOT_ANTHROPIC_API_KEY: ${{ secrets.COPILOT_ANTHROPIC_API_KEY }}
           COPILOT_EXA_API_KEY: ${{ secrets.COPILOT_EXA_API_KEY }}

packages/backend/server/package.json

@@ -29,7 +29,8 @@
     "@affine/reader": "workspace:*",
     "@affine/server-native": "workspace:*",
     "@ai-sdk/anthropic": "^1.2.10",
-    "@ai-sdk/google": "^1.2.10",
+    "@ai-sdk/google": "^1.2.18",
+    "@ai-sdk/google-vertex": "^2.2.22",
     "@ai-sdk/openai": "^1.3.21",
     "@ai-sdk/perplexity": "^1.1.6",
     "@apollo/server": "^4.11.3",

packages/backend/server/src/__tests__/copilot-provider.spec.ts

@@ -328,7 +328,7 @@ const actions = [
     messages: [
       {
         role: 'user' as const,
-        content: '',
+        content: 'transcript the audio',
         attachments: [
           'https://cdn.affine.pro/copilot-test/MP9qDGuYgnY+ILoEAmHpp3h9Npuw2403EAYMEA.mp3',
         ],
@@ -350,7 +350,7 @@ const actions = [
     messages: [
       {
         role: 'user' as const,
-        content: '',
+        content: 'transcript the audio',
         attachments: [
           'https://cdn.affine.pro/copilot-test/2ed05eo1KvZ2tWB_BAjFo67EAPZZY-w4LylUAw.m4a',
         ],
@@ -372,7 +372,7 @@ const actions = [
     messages: [
       {
         role: 'user' as const,
-        content: '',
+        content: 'transcript the audio',
         attachments: [
           'https://cdn.affine.pro/copilot-test/nC9-e7P85PPI2rU29QWwf8slBNRMy92teLIIMw.opus',
         ],

packages/backend/server/src/__tests__/mocks/copilot.mock.ts

@@ -24,7 +24,7 @@ export class MockCopilotProvider extends OpenAIProvider {
     'lcm-sd15-i2i',
     'clarity-upscaler',
     'imageutils/rembg',
-    'gemini-2.5-pro-preview-03-25',
+    'gemini-2.5-pro-preview-05-06',
   ];
 
   override readonly capabilities = [

packages/backend/server/src/plugins/copilot/config.ts

@@ -52,7 +52,7 @@ defineModuleConfig('copilot', {
   'providers.gemini': {
     desc: 'The config for the gemini provider.',
     default: {
-      apiKey: '',
+      privateKey: '',
     },
   },
   'providers.perplexity': {

packages/backend/server/src/plugins/copilot/prompt/prompts.ts

@@ -334,7 +334,7 @@ const actions: Prompt[] = [
   {
     name: 'Transcript audio',
     action: 'Transcript audio',
-    model: 'gemini-2.5-pro-preview-03-25',
+    model: 'gemini-2.5-pro-preview-05-06',
     messages: [
       {
         role: 'system',
@@ -1071,6 +1071,8 @@ const chat: Prompt[] = [
       'o4-mini',
       'claude-3-7-sonnet-20250219',
       'claude-3-5-sonnet-20241022',
+      'gemini-2.5-flash-preview-04-17',
+      'gemini-2.5-pro-preview-05-06',
     ],
     messages: [
       {
@@ -1097,11 +1099,12 @@ When referencing information from the provided documents, files or web search re
 1. Use markdown footnote format for citations
 2. Add citations immediately after the relevant sentence or paragraph
 3. Required format: [^reference_index] where reference_index is an increasing positive integer
-4. You MUST include citations at the end of your response in this exact format:
+4. When a single sentence needs multiple citations, write each marker in its own pair of brackets and place them consecutively. Correct: [^2][^4][^12], Incorrect: [^2, 4, 12].
+5. You MUST include citations at the end of your response in this exact format:
   - For documents: [^reference_index]:{"type":"doc","docId":"document_id"}
   - For files: [^reference_index]:{"type":"attachment","blobId":"blob_id","fileName":"file_name","fileType":"file_type"}
   - For web search results: [^reference_index]:{"type":"url","url":"url_path"}
-5. Ensure citations adhere strictly to the required format. Do not add extra spaces in citations like [^ reference_index] or [ ^reference_index].
+6. Ensure citations adhere strictly to the required format. Do not add extra spaces in citations like [^ reference_index] or [ ^reference_index].
 
 ### Citations Structure
 Your response MUST follow this structure:
@@ -1111,6 +1114,7 @@ Your response MUST follow this structure:
 
 Example Output with Citations:
 This is my response with a document citation[^1]. Here is more content with another file citation[^2]. And here is a web search result citation[^3].
+Here is multiple citations: [^1][^2][^3].
 
 [^1]:{"type":"doc","docId":"abc123"}
 [^2]:{"type":"attachment","blobId":"xyz789","fileName":"example.txt","fileType":"text"}

packages/backend/server/src/plugins/copilot/providers/gemini.ts

@@ -1,7 +1,5 @@
-import {
-  createGoogleGenerativeAI,
-  type GoogleGenerativeAIProvider,
-} from '@ai-sdk/google';
+import { GoogleGenerativeAIProviderOptions } from '@ai-sdk/google';
+import { createVertex, type GoogleVertexProvider } from '@ai-sdk/google-vertex';
 import {
   AISDKError,
   generateObject,
@@ -9,6 +7,7 @@ import {
   JSONParseError,
   streamText,
 } from 'ai';
+import { z } from 'zod';
 
 import {
   CopilotPromptInvalid,
@@ -30,10 +29,21 @@ import { chatToGPTMessage } from './utils';
 export const DEFAULT_DIMENSIONS = 256;
 
 export type GeminiConfig = {
-  apiKey: string;
-  baseUrl?: string;
+  privateKey: string;
 };
 
+const PrivateKeySchema = z.object({
+  type: z.string(),
+  client_email: z.string(),
+  private_key: z.string(),
+  private_key_id: z.string(),
+  project_id: z.string(),
+  client_id: z.string(),
+  universe_domain: z.string().optional(),
+});
+
+type PrivateKey = z.infer<typeof PrivateKeySchema>;
+
 export class GeminiProvider
   extends CopilotProvider<GeminiConfig>
   implements CopilotTextToTextProvider
@@ -43,22 +53,50 @@ export class GeminiProvider
   override readonly models = [
     // text to text
     'gemini-2.0-flash-001',
-    'gemini-2.5-pro-preview-03-25',
+    'gemini-2.5-flash-preview-04-17',
+    'gemini-2.5-pro-preview-05-06',
     // embeddings
     'text-embedding-004',
   ];
 
-  #instance!: GoogleGenerativeAIProvider;
+  private readonly MAX_STEPS = 20;
+
+  private readonly CALLOUT_PREFIX = '\n> [!]\n> ';
+
+  #instance!: GoogleVertexProvider;
 
   override configured(): boolean {
-    return !!this.config.apiKey;
+    return !!this.parsePrivateKey(this.config.privateKey);
   }
 
   protected override setup() {
     super.setup();
-    this.#instance = createGoogleGenerativeAI({
-      apiKey: this.config.apiKey,
-      baseURL: this.config.baseUrl,
+
+    // can not throw error here
+    const {
+      type,
+      client_email,
+      private_key,
+      private_key_id,
+      project_id,
+      client_id,
+      universe_domain,
+    } = this.parsePrivateKey(this.config.privateKey) || {};
+
+    this.#instance = createVertex({
+      project: project_id,
+      location: 'us-central1',
+      googleAuthOptions: {
+        credentials: {
+          type,
+          client_email,
+          private_key,
+          private_key_id,
+          project_id,
+          client_id,
+          universe_domain,
+        },
+      },
     });
   }
 
@@ -191,20 +229,56 @@ export class GeminiProvider
       metrics.ai.counter('chat_text_stream_calls').add(1, { model });
       const [system, msgs] = await chatToGPTMessage(messages);
 
-      const { textStream } = streamText({
-        model: this.#instance(model),
+      const { fullStream } = streamText({
+        model: this.#instance(model, {
+          useSearchGrounding: this.withWebSearch(options),
+        }),
         system,
         messages: msgs,
         abortSignal: options.signal,
+        maxSteps: this.MAX_STEPS,
+        providerOptions: {
+          google: this.getGeminiOptions(options, model),
+        },
       });
 
-      for await (const message of textStream) {
-        if (message) {
-          yield message;
+      let lastType;
+      // reasoning, tool-call, tool-result need to mark as callout
+      let prefix: string | null = this.CALLOUT_PREFIX;
+      for await (const chunk of fullStream) {
+        if (chunk) {
+          switch (chunk.type) {
+            case 'text-delta': {
+              let result = chunk.textDelta;
+              if (lastType !== chunk.type) {
+                result = '\n\n' + result;
+              }
+              yield result;
+              break;
+            }
+            case 'reasoning': {
+              if (prefix) {
+                yield prefix;
+                prefix = null;
+              }
+              let result = chunk.textDelta;
+              if (lastType !== chunk.type) {
+                result = '\n\n' + result;
+              }
+              yield this.markAsCallout(result);
+              break;
+            }
+            case 'error': {
+              const error = chunk.error as { type: string; message: string };
+              throw new Error(error.message);
+            }
+          }
+
           if (options.signal?.aborted) {
-            await textStream.cancel();
+            await fullStream.cancel();
             break;
           }
+          lastType = chunk.type;
         }
       }
     } catch (e: any) {
@@ -212,4 +286,36 @@ export class GeminiProvider
       throw this.handleError(e);
     }
   }
+
+  private getGeminiOptions(options: CopilotChatOptions, model: string) {
+    const result: GoogleGenerativeAIProviderOptions = {};
+    if (options?.reasoning && this.isThinkingModel(model)) {
+      result.thinkingConfig = {
+        thinkingBudget: 12000,
+        includeThoughts: true,
+      };
+    }
+    return result;
+  }
+
+  private markAsCallout(text: string) {
+    return text.replaceAll('\n', '\n> ');
+  }
+
+  private isThinkingModel(model: string) {
+    // TODO gemini-2.5-pro-preview is not supported thinking yet
+    return model.startsWith('gemini-2.5-flash-preview');
+  }
+
+  private withWebSearch(options: CopilotChatOptions) {
+    return options?.tools?.includes('webSearch');
+  }
+
+  private parsePrivateKey(jsonString: string): PrivateKey | null {
+    try {
+      return PrivateKeySchema.parse(JSON.parse(jsonString));
+    } catch {
+      return null;
+    }
+  }
 }

yarn.lock

@@ -909,7 +909,8 @@ __metadata:
     "@affine/reader": "workspace:*"
     "@affine/server-native": "workspace:*"
     "@ai-sdk/anthropic": "npm:^1.2.10"
-    "@ai-sdk/google": "npm:^1.2.10"
+    "@ai-sdk/google": "npm:^1.2.18"
+    "@ai-sdk/google-vertex": "npm:^2.2.22"
     "@ai-sdk/openai": "npm:^1.3.21"
     "@ai-sdk/perplexity": "npm:^1.1.6"
     "@apollo/server": "npm:^4.11.3"
@@ -1073,7 +1074,7 @@ __metadata:
   languageName: unknown
   linkType: soft
 
-"@ai-sdk/anthropic@npm:^1.2.10":
+"@ai-sdk/anthropic@npm:1.2.11, @ai-sdk/anthropic@npm:^1.2.10":
   version: 1.2.11
   resolution: "@ai-sdk/anthropic@npm:1.2.11"
   dependencies:
@@ -1085,15 +1086,30 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@ai-sdk/google@npm:^1.2.10":
-  version: 1.2.17
-  resolution: "@ai-sdk/google@npm:1.2.17"
+"@ai-sdk/google-vertex@npm:^2.2.22":
+  version: 2.2.22
+  resolution: "@ai-sdk/google-vertex@npm:2.2.22"
+  dependencies:
+    "@ai-sdk/anthropic": "npm:1.2.11"
+    "@ai-sdk/google": "npm:1.2.18"
+    "@ai-sdk/provider": "npm:1.1.3"
+    "@ai-sdk/provider-utils": "npm:2.2.8"
+    google-auth-library: "npm:^9.15.0"
+  peerDependencies:
+    zod: ^3.0.0
+  checksum: 10/852928d35797cc14802d8c4a7dba2794197f019507a24124e7f87e456ba19f2b4e50438af626dd30e3a2a0f9f10819d6534a31e23c08f7f3d148d56b68167d0a
+  languageName: node
+  linkType: hard
+
+"@ai-sdk/google@npm:1.2.18, @ai-sdk/google@npm:^1.2.18":
+  version: 1.2.18
+  resolution: "@ai-sdk/google@npm:1.2.18"
   dependencies:
     "@ai-sdk/provider": "npm:1.1.3"
     "@ai-sdk/provider-utils": "npm:2.2.8"
   peerDependencies:
     zod: ^3.0.0
-  checksum: 10/588d1d9c9de7dbfe4ddb628c65c2cd06509024bf44d889eb3f9d1156a16899ebfd56db7a727793070a5c38f8d74c3896997319b09875b92182f084ac17a993d4
+  checksum: 10/e8ff1ea1cae8f6c1c17e5526e3e51a8e584bb60d8e407646594c9b07600e06ef43c85518d08aafd3856aa2d46a1ae88111d6c61532bdf8c917859e0baad23432
   languageName: node
   linkType: hard
 
@@ -22692,7 +22708,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"google-auth-library@npm:^9.0.0, google-auth-library@npm:^9.7.0":
+"google-auth-library@npm:^9.0.0, google-auth-library@npm:^9.15.0, google-auth-library@npm:^9.7.0":
   version: 9.15.1
   resolution: "google-auth-library@npm:9.15.1"
   dependencies: