feat: Support Gemini via Vertex AI

.vscode/settings.json

@@ -1,11 +1,17 @@
 {
-  "editor.rulers": [100],
+  "editor.rulers": [
+    100
+  ],
   "editor.formatOnSave": true,
   "editor.defaultFormatter": "esbenp.prettier-vscode",
   "jest.virtualFolders": [
     {
       "name": "cli",
       "rootPath": "packages/cli"
+    },
+    {
+      "name": "navie",
+      "rootPath": "packages/navie"
     }
-  ],
-}
+  ]
+}

package.json

@@ -43,5 +43,8 @@
   "packageManager": "[email protected]",
   "dependencies": {
     "puppeteer": "^19.7.2"
+  },
+  "resolutions": {
+    "web-auth-library": "getappmap/web-auth-library#v1.0.3-cjs"
   }
 }

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/.eslintrc.js

@@ -22,6 +22,7 @@ module.exports = {
   },
   root: true,
   rules: {
+    'eslint-comments/disable-enable-pair': ['error', { allowWholeFile: true }],
     'no-param-reassign': ['error', { props: false }],
     'no-console': process.env.NODE_ENV === 'production' ? 'warn' : 'off',
     'no-debugger': process.env.NODE_ENV === 'production' ? 'warn' : 'off',

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/interaction-history.ts

@@ -1,5 +1,5 @@
+/* eslint-disable @typescript-eslint/no-unsafe-declaration-merging */
 import EventEmitter from 'events';
-import * as path from 'path';
 import InteractionState from './interaction-state';
 import { ContextV2 } from './context';
 import { PROMPTS, PromptType } from './prompt';
@@ -283,11 +283,7 @@ export class TechStackEvent extends InteractionEvent {
   }
 }
 
-export interface InteractionHistoryEvents {
-  on(event: 'event', listener: (event: InteractionEvent) => void): void;
-}
-
-export default class InteractionHistory extends EventEmitter implements InteractionHistoryEvents {
+class InteractionHistory extends EventEmitter {
   public readonly events: InteractionEvent[] = [];
 
   // eslint-disable-next-line class-methods-use-this
@@ -333,3 +329,10 @@ export default class InteractionHistory extends EventEmitter implements Interact
     return state;
   }
 }
+
+interface InteractionHistory {
+  on(event: 'event', listener: (event: InteractionEvent) => void): this;
+  on(event: string, listener: (...args: unknown[]) => void): this;
+}
+
+export default InteractionHistory;

packages/navie/src/lib/trajectory.ts

@@ -1,3 +1,4 @@
+/* eslint-disable @typescript-eslint/no-unsafe-declaration-merging */
 import EventEmitter from 'events';
 import Message from '../message';
 
@@ -7,7 +8,7 @@ export interface TrajectoryEvent {
   timestamp: Date;
 }
 
-export default class Trajectory extends EventEmitter {
+class Trajectory extends EventEmitter {
   logSentMessage(message: Message) {
     const event: TrajectoryEvent = {
       type: 'sent',
@@ -26,3 +27,10 @@ export default class Trajectory extends EventEmitter {
     this.emit('event', event);
   }
 }
+
+interface Trajectory {
+  on(event: 'event', listener: (event: TrajectoryEvent) => void): this;
+  on(event: string, listener: (...args: unknown[]) => void): this;
+}
+
+export default Trajectory;

packages/navie/src/navie.ts

@@ -8,7 +8,6 @@ import InteractionHistory, {
   AgentSelectionEvent,
   ClassificationEvent,
   InteractionEvent,
-  InteractionHistoryEvents,
 } from './interaction-history';
 import { ContextV2 } from './context';
 import ProjectInfoService from './services/project-info-service';
@@ -48,7 +47,7 @@ export interface ClientRequest {
   prompt?: string;
 }
 
-export interface INavie extends InteractionHistoryEvents {
+export interface INavie {
   on(event: 'event', listener: (event: InteractionEvent) => void): void;
 
   on(event: 'agent', listener: (agent: string) => void): void;

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';
@@ -10,47 +11,41 @@ interface Options {
   modelName: string;
   temperature: number;
   trajectory: Trajectory;
+  backend?: Backend;
 }
 
-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 environmentBackend(): Backend | undefined {
+function determineCompletionBackend(): Backend {
   switch (process.env.APPMAP_NAVIE_COMPLETION_BACKEND) {
     case 'anthropic':
     case 'openai':
+    case 'vertex-ai':
       return process.env.APPMAP_NAVIE_COMPLETION_BACKEND;
     default:
-      return undefined;
+    // pass
   }
+  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';
+  return 'openai'; // fallback
 }
 
-export const SELECTED_BACKEND: Backend = environmentBackend() ?? defaultBackend();
+export const SELECTED_BACKEND: Backend = determineCompletionBackend();
 
 export default function createCompletionService({
   modelName,
   temperature,
   trajectory,
+  backend = determineCompletionBackend(),
 }: Options): CompletionService {
-  const backend = environmentBackend() ?? defaultBackend();
   const messageTokenReducerService = new MessageTokenReducerService();
-  if (backend === 'anthropic') {
-    warn('Using Anthropic AI backend');
-    return new AnthropicCompletionService(
-      modelName,
-      temperature,
-      trajectory,
-      messageTokenReducerService
-    );
-  }
-  warn('Using OpenAI backend');
-  return new OpenAICompletionService(
-    modelName,
-    temperature,
-    trajectory,
-    messageTokenReducerService
-  );
+  warn(`Using completion service ${backend}`);
+  return new BACKENDS[backend](modelName, temperature, trajectory, messageTokenReducerService);
 }

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

@@ -0,0 +1,137 @@
+import { warn } from 'node:console';
+import { isNativeError } from 'node:util/types';
+
+import { ChatVertexAI, type 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';
+
+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 buildModel(options?: ChatVertexAIInput): ChatVertexAI {
+    return new ChatVertexAI({
+      model: this.modelName,
+      temperature: this.temperature,
+      streaming: true,
+      maxOutputTokens: 8192,
+      ...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 = 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 = 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);
+}