feat: Support Gemini via Vertex AI

Author: dividedmind

packages/cli/src/cmds/index/aiEnvVar.ts

@@ -1,4 +1,9 @@
-export const AI_KEY_ENV_VARS = ['OPENAI_API_KEY', 'AZURE_OPENAI_API_KEY', 'ANTHROPIC_API_KEY'];
+export const AI_KEY_ENV_VARS = [
+  'GOOGLE_WEB_CREDENTIALS',
+  'OPENAI_API_KEY',
+  'AZURE_OPENAI_API_KEY',
+  'ANTHROPIC_API_KEY',
+];
 
 export default function detectAIEnvVar(): string | undefined {
   return Object.keys(process.env).find((key) => AI_KEY_ENV_VARS.includes(key));

packages/cli/src/cmds/index/index.ts

@@ -24,8 +24,7 @@ import LocalNavie from '../../rpc/explain/navie/navie-local';
 import RemoteNavie from '../../rpc/explain/navie/navie-remote';
 import { InteractionEvent } from '@appland/navie/dist/interaction-history';
 import { update } from '../../rpc/file/update';
-
-const AI_KEY_ENV_VARS = ['OPENAI_API_KEY', 'AZURE_OPENAI_API_KEY'];
+import { AI_KEY_ENV_VARS } from './aiEnvVar';
 
 export const command = 'index';
 export const describe =

packages/cli/src/rpc/llmConfiguration.ts

@@ -47,9 +47,14 @@ function openAIBaseURL(): string | undefined {
   return baseUrl;
 }
 
+const DEFAULT_BASE_URLS = {
+  anthropic: 'https://api.anthropic.com/v1/',
+  'vertex-ai': 'https://googleapis.com',
+  openai: undefined,
+} as const;
+
 export function getLLMConfiguration(): LLMConfiguration {
-  const baseUrl =
-    SELECTED_BACKEND === 'anthropic' ? 'https://api.anthropic.com/v1/' : openAIBaseURL();
+  const baseUrl = (SELECTED_BACKEND && DEFAULT_BASE_URLS[SELECTED_BACKEND]) ?? openAIBaseURL();
 
   return {
     baseUrl,

packages/navie/package.json

@@ -41,6 +41,7 @@
   "dependencies": {
     "@langchain/anthropic": "^0.3.1",
     "@langchain/core": "^0.2.27",
+    "@langchain/google-vertexai-web": "^0.1.0",
     "@langchain/openai": "^0.2.7",
     "fast-xml-parser": "^4.4.0",
     "js-yaml": "^4.1.0",

packages/navie/src/services/completion-service-factory.ts

@@ -1,5 +1,6 @@
 import { warn } from 'node:console';
 
+import GoogleVertexAICompletionService from './google-vertexai-completion-service';
 import OpenAICompletionService from './openai-completion-service';
 import AnthropicCompletionService from './anthropic-completion-service';
 import CompletionService from './completion-service';
@@ -11,34 +12,43 @@ interface Options {
   trajectory: Trajectory;
 }
 
-type Backend = 'anthropic' | 'openai';
+const BACKENDS = {
+  anthropic: AnthropicCompletionService,
+  openai: OpenAICompletionService,
+  'vertex-ai': GoogleVertexAICompletionService,
+} as const;
 
-function defaultBackend(): Backend {
-  return 'ANTHROPIC_API_KEY' in process.env ? 'anthropic' : 'openai';
+type Backend = keyof typeof BACKENDS;
+
+function defaultBackend(): Backend | undefined {
+  if ('ANTHROPIC_API_KEY' in process.env) return 'anthropic';
+  if ('GOOGLE_WEB_CREDENTIALS' in process.env) return 'vertex-ai';
+  if ('OPENAI_API_KEY' in process.env) return 'openai';
 }
 
 function environmentBackend(): Backend | undefined {
   switch (process.env.APPMAP_NAVIE_COMPLETION_BACKEND) {
     case 'anthropic':
     case 'openai':
+    case 'vertex-ai':
       return process.env.APPMAP_NAVIE_COMPLETION_BACKEND;
     default:
       return undefined;
   }
 }
 
-export const SELECTED_BACKEND: Backend = environmentBackend() ?? defaultBackend();
+export const SELECTED_BACKEND: Backend | undefined = environmentBackend() ?? defaultBackend();
 
 export default function createCompletionService({
   modelName,
   temperature,
   trajectory,
 }: Options): CompletionService {
   const backend = environmentBackend() ?? defaultBackend();
-  if (backend === 'anthropic') {
-    warn('Using Anthropic AI backend');
-    return new AnthropicCompletionService(modelName, temperature, trajectory);
+  if (backend && backend in BACKENDS) {
+    warn(`Using completion service ${backend}`);
+    return new BACKENDS[backend](modelName, temperature, trajectory);
   }
-  warn('Using OpenAI backend');
+  warn(`No completion service available for backend ${backend}. Falling back to OpenAI.`);
   return new OpenAICompletionService(modelName, temperature, trajectory);
 }

packages/navie/src/services/google-vertexai-completion-service.ts

@@ -0,0 +1,138 @@
+import { warn } from 'node:console';
+import { isNativeError } from 'node:util/types';
+
+import type { ChatVertexAI, ChatVertexAIInput } from '@langchain/google-vertexai-web';
+import { zodResponseFormat } from 'openai/helpers/zod';
+import { z } from 'zod';
+
+import Trajectory from '../lib/trajectory';
+import Message from '../message';
+import CompletionService, {
+  CompleteOptions,
+  Completion,
+  CompletionRetries,
+  CompletionRetryDelay,
+  convertToMessage,
+  mergeSystemMessages,
+  Usage,
+} from './completion-service';
+
+const VertexAI = import('@langchain/google-vertexai-web');
+
+export default class GoogleVertexAICompletionService implements CompletionService {
+  constructor(
+    public readonly modelName: string,
+    public readonly temperature: number,
+    private trajectory: Trajectory
+  ) {}
+
+  // Construct a model with non-default options. There doesn't seem to be a way to configure
+  // the model parameters at invocation time like with OpenAI.
+  private async buildModel(options?: ChatVertexAIInput): Promise<ChatVertexAI> {
+    return new (await VertexAI).ChatVertexAI({
+      model: this.modelName,
+      temperature: this.temperature,
+      streaming: true,
+      ...options,
+    });
+  }
+
+  get miniModelName(): string {
+    const miniModel = process.env.APPMAP_NAVIE_MINI_MODEL;
+    return miniModel ?? 'gemini-1.5-flash-002';
+  }
+
+  // Request a JSON object with a given JSON schema.
+  async json<Schema extends z.ZodType>(
+    messages: Message[],
+    schema: Schema,
+    options?: CompleteOptions
+  ): Promise<z.infer<Schema> | undefined> {
+    const model = await this.buildModel({
+      ...options,
+      streaming: false,
+      responseMimeType: 'application/json',
+    });
+    const sentMessages = mergeSystemMessages([
+      ...messages,
+      {
+        role: 'system',
+        content: `Use the following JSON schema for your response:\n\n${JSON.stringify(
+          zodResponseFormat(schema, 'requestedObject').json_schema.schema,
+          null,
+          2
+        )}`,
+      },
+    ]);
+
+    for (const message of sentMessages) this.trajectory.logSentMessage(message);
+
+    const response = await model.invoke(sentMessages.map(convertToMessage));
+
+    this.trajectory.logReceivedMessage({
+      role: 'assistant',
+      content: JSON.stringify(response),
+    });
+
+    const sanitizedContent = response.content.toString().replace(/^`{3,}[^\s]*?$/gm, '');
+    const parsed = JSON.parse(sanitizedContent) as unknown;
+    schema.parse(parsed);
+    return parsed;
+  }
+
+  async *complete(messages: readonly Message[], options?: { temperature?: number }): Completion {
+    const usage = new Usage();
+    const model = await this.buildModel(options);
+    const sentMessages: Message[] = mergeSystemMessages(messages);
+    const tokens = new Array<string>();
+    for (const message of sentMessages) this.trajectory.logSentMessage(message);
+
+    const maxAttempts = CompletionRetries;
+    for (let attempt = 0; attempt < maxAttempts; attempt += 1) {
+      try {
+        // eslint-disable-next-line no-await-in-loop
+        const response = await model.stream(sentMessages.map(convertToMessage));
+
+        // eslint-disable-next-line @typescript-eslint/naming-convention, no-await-in-loop
+        for await (const { content, usage_metadata } of response) {
+          yield content.toString();
+          tokens.push(content.toString());
+          if (usage_metadata) {
+            usage.promptTokens += usage_metadata.input_tokens;
+            usage.completionTokens += usage_metadata.output_tokens;
+          }
+        }
+
+        this.trajectory.logReceivedMessage({
+          role: 'assistant',
+          content: tokens.join(''),
+        });
+
+        break;
+      } catch (cause) {
+        if (attempt < maxAttempts - 1 && tokens.length === 0) {
+          const nextAttempt = CompletionRetryDelay * 2 ** attempt;
+          warn(`Received ${JSON.stringify(cause)}, retrying in ${nextAttempt}ms`);
+          await new Promise<void>((resolve) => {
+            setTimeout(resolve, nextAttempt);
+          });
+          continue;
+        }
+        throw new Error(
+          `Failed to complete after ${attempt + 1} attempt(s): ${errorMessage(cause)}`,
+          {
+            cause,
+          }
+        );
+      }
+    }
+
+    warn(usage.toString());
+    return usage;
+  }
+}
+
+function errorMessage(err: unknown): string {
+  if (isNativeError(err)) return err.cause ? errorMessage(err.cause) : err.message;
+  return String(err);
+}

yarn.lock

@@ -456,6 +456,7 @@ __metadata:
   dependencies:
     "@langchain/anthropic": ^0.3.1
     "@langchain/core": ^0.2.27
+    "@langchain/google-vertexai-web": ^0.1.0
     "@langchain/openai": ^0.2.7
     "@tsconfig/node-lts": ^20.1.3
     "@types/jest": ^29.4.1
@@ -7038,6 +7039,41 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@langchain/google-common@npm:~0.1.0":
+  version: 0.1.1
+  resolution: "@langchain/google-common@npm:0.1.1"
+  dependencies:
+    uuid: ^10.0.0
+    zod-to-json-schema: ^3.22.4
+  peerDependencies:
+    "@langchain/core": ">=0.2.21 <0.4.0"
+  checksum: e460a08eaf5e6902c3cb7e8deb9edddcdb46c6bc38657ee1050d05ab5f17bf864bf298a9f00cc41e2824f8c072d79c1dca9b84a7ce64ebcf5a5357af14f5b9d9
+  languageName: node
+  linkType: hard
+
+"@langchain/google-vertexai-web@npm:^0.1.0":
+  version: 0.1.0
+  resolution: "@langchain/google-vertexai-web@npm:0.1.0"
+  dependencies:
+    "@langchain/google-webauth": ~0.1.0
+  peerDependencies:
+    "@langchain/core": ">=0.2.21 <0.4.0"
+  checksum: 8c32499e4070ddf28de26e3e4354c60303921e0be84aa68bbcbbeecd5e79e78354fb940708dcfc94efbc67f51893e51039288d78418ec00ec3f64a6cb1e5b20e
+  languageName: node
+  linkType: hard
+
+"@langchain/google-webauth@npm:~0.1.0":
+  version: 0.1.0
+  resolution: "@langchain/google-webauth@npm:0.1.0"
+  dependencies:
+    "@langchain/google-common": ~0.1.0
+    web-auth-library: ^1.0.3
+  peerDependencies:
+    "@langchain/core": ">=0.2.21 <0.4.0"
+  checksum: 90d7c04f95e9950ec5fb39a779352f145efa319d2003564b82a183809ef92d64f8f878999e5cb9c75b1bfda83e38c9b650946c928b1d137dd8bf0bebbaddca74
+  languageName: node
+  linkType: hard
+
 "@langchain/openai@npm:>=0.1.0 <0.3.0, @langchain/openai@npm:^0.2.7":
   version: 0.2.7
   resolution: "@langchain/openai@npm:0.2.7"
@@ -28406,6 +28442,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"jose@npm:>= 4.12.0 < 5.0.0":
+  version: 4.15.9
+  resolution: "jose@npm:4.15.9"
+  checksum: 41abe1c99baa3cf8a78ebbf93da8f8e50e417b7a26754c4afa21865d87527b8ac2baf66de2c5f6accc3f7d7158658dae7364043677236ea1d07895b040097f15
+  languageName: node
+  linkType: hard
+
 "joycon@npm:^3.0.1":
   version: 3.1.1
   resolution: "joycon@npm:3.1.1"
@@ -36925,6 +36968,13 @@ resolve@1.1.7:
   languageName: node
   linkType: hard
 
+"rfc4648@npm:^1.5.2":
+  version: 1.5.3
+  resolution: "rfc4648@npm:1.5.3"
+  checksum: 19c81d502582e377125b00fbd7a5cdb0e351f9a1e40182fa9f608b48e1ab852d211b75facb2f4f3fa17f7c6ebc2ef4acca61ae7eb7fbcfa4768f11d2db678116
+  languageName: node
+  linkType: hard
+
 "rfdc@npm:^1.3.0":
   version: 1.3.0
   resolution: "rfdc@npm:1.3.0"
@@ -41838,6 +41888,16 @@ typescript@~4.4.3:
   languageName: node
   linkType: hard
 
+"web-auth-library@npm:^1.0.3":
+  version: 1.0.3
+  resolution: "web-auth-library@npm:1.0.3"
+  dependencies:
+    jose: ">= 4.12.0 < 5.0.0"
+    rfc4648: ^1.5.2
+  checksum: 9e2b303a43ac040037b952c8130260edd44c34956e34277e28faa23ebf63fc40ba3a22b7a3cab5f36ea6f7aac16573415ccacb6e5a970fe4f4d2ee5751ae01d1
+  languageName: node
+  linkType: hard
+
 "web-streams-polyfill@npm:4.0.0-beta.3":
   version: 4.0.0-beta.3
   resolution: "web-streams-polyfill@npm:4.0.0-beta.3"