@llumiverse/drivers 0.18.0 → 0.20.0

package/src/vertexai/index.ts

@@ -1,9 +1,8 @@
-import { GenerateContentRequest, VertexAI } from "@google-cloud/vertexai";
 import {
     AIModel,
     AbstractDriver,
     Completion,
-    CompletionChunkObject,
+    CompletionChunk,
     DriverOptions,
     EmbeddingsResult,
     ExecutionOptions,
@@ -14,7 +13,7 @@ import {
     getModelCapabilities,
     modelModalitiesToArray,
 } from "@llumiverse/core";
-import { FetchClient } from "api-fetch-client";
+import { FetchClient } from "@vertesia/api-fetch-client";
 import { GoogleAuth, GoogleAuthOptions } from "google-auth-library";
 import { JSONClient } from "google-auth-library/build/src/auth/googleauth.js";
 import { TextEmbeddingsOptions, getEmbeddingsForText } from "./embeddings/embeddings-text.js";
@@ -24,6 +23,7 @@ import { getEmbeddingsForImages } from "./embeddings/embeddings-image.js";
 import { v1beta1 } from "@google-cloud/aiplatform";
 import { AnthropicVertex } from "@anthropic-ai/vertex-sdk";
 import { ImagenModelDefinition, ImagenPrompt } from "./models/imagen.js";
+import { GoogleGenAI, Content } from "@google/genai";
 
 export interface VertexAIDriverOptions extends DriverOptions {
     project: string;
@@ -31,8 +31,13 @@ export interface VertexAIDriverOptions extends DriverOptions {
     googleAuthOptions?: GoogleAuthOptions;
 }
 
+export interface GenerateContentPrompt {
+    contents: Content[];
+    system?: Content;
+}
+
 //General Prompt type for VertexAI
-export type VertexAIPrompt = GenerateContentRequest | ImagenPrompt;
+export type VertexAIPrompt = ImagenPrompt | GenerateContentPrompt;
 
 export function trimModelName(model: string) {
     const i = model.lastIndexOf("@");
@@ -46,8 +51,9 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
     aiplatform: v1beta1.ModelServiceClient | undefined;
     anthropicClient: AnthropicVertex | undefined;
     fetchClient: FetchClient | undefined;
+    googleGenAI: GoogleGenAI | undefined;
+    llamaClient: FetchClient & { region?: string } | undefined;
     modelGarden: v1beta1.ModelGardenServiceClient | undefined;
-    vertexai: VertexAI | undefined;
 
     authClient: JSONClient | GoogleAuth<JSONClient>;
 
@@ -57,12 +63,28 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
         this.aiplatform = undefined;
         this.anthropicClient = undefined;
         this.fetchClient = undefined
+        this.googleGenAI = undefined;
         this.modelGarden = undefined;
-        this.vertexai = undefined;
+        this.llamaClient = undefined;
 
         this.authClient = options.googleAuthOptions?.authClient ?? new GoogleAuth(options.googleAuthOptions);
     }
 
+    public getGoogleGenAIClient(): GoogleGenAI {
+        //Lazy initialisation
+        if (!this.googleGenAI) {
+            this.googleGenAI = new GoogleGenAI({
+                project: this.options.project,
+                location: this.options.region,
+                vertexai: true,
+                googleAuthOptions: {
+                    authClient: this.authClient as JSONClient,
+                }
+            });
+        }
+        return this.googleGenAI;
+    }
+
     public getFetchClient(): FetchClient {
         //Lazy initialisation
         if (!this.fetchClient) {
@@ -78,10 +100,29 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
         return this.fetchClient;
     }
 
+    public getLLamaClient(region: string = "us-central1"): FetchClient {
+        //Lazy initialisation
+        if (!this.llamaClient || this.llamaClient["region"] !== region) {
+            this.llamaClient = createFetchClient({
+                region: region,
+                project: this.options.project,
+                apiVersion: "v1beta1",
+            }).withAuthCallback(async () => {
+                const accessTokenResponse = await this.authClient.getAccessToken();
+                const token = typeof accessTokenResponse === 'string' ? accessTokenResponse : accessTokenResponse?.token;
+                return `Bearer ${token}`;
+            });
+            // Store the region for potential client reuse
+            this.llamaClient["region"] = region;
+        }
+        return this.llamaClient;
+    }
+
     public getAnthropicClient(): AnthropicVertex {
         //Lazy initialisation
         if (!this.anthropicClient) {
             this.anthropicClient = new AnthropicVertex({
+                timeout: 20 * 60 * 10000, // Set to 20 minutes, 10 minute default, setting this disables long request error: https://github.com/anthropics/anthropic-sdk-typescript?#long-requests
                 region: "us-east5",
                 projectId: process.env.GOOGLE_PROJECT_ID,
             });
@@ -89,18 +130,6 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
         return this.anthropicClient;
     }
 
-    public getVertexAIClient(): VertexAI {
-        //Lazy initialisation
-        if (!this.vertexai) {
-            this.vertexai = new VertexAI({
-                project: this.options.project,
-                location: this.options.region,
-                googleAuthOptions: this.options.googleAuthOptions,
-            });
-        }
-        return this.vertexai;
-    }
-
     public getAIPlatformClient(): v1beta1.ModelServiceClient {
         //Lazy initialisation
         if (!this.aiplatform) {
@@ -125,6 +154,18 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
         return this.modelGarden;
     }
 
+    validateResult(result: Completion, options: ExecutionOptions) {
+        // Optionally preprocess the result before validation
+        const modelDef = getModelDefinition(options.model);
+        if (typeof modelDef.preValidationProcessing === "function") {
+            const processed = modelDef.preValidationProcessing(result, options);
+            result = processed.result;
+            options = processed.options;
+        }
+
+        super.validateResult(result, options);
+    }
+
     protected canStream(options: ExecutionOptions): Promise<boolean> {
         if (options.output_modality == Modalities.image) {
             return Promise.resolve(false);
@@ -145,7 +186,7 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
     async requestTextCompletionStream(
         prompt: VertexAIPrompt,
         options: ExecutionOptions,
-    ): Promise<AsyncIterable<CompletionChunkObject>> {
+    ): Promise<AsyncIterable<CompletionChunk>> {
         return getModelDefinition(options.model).requestTextCompletionStream(this, prompt, options);
     }
 
@@ -178,14 +219,31 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
         );
 
         //Model Garden Publisher models - Pretrained models
-        const publishers = ["google", "anthropic"];
-        const supportedModels = { google: ["gemini", "imagen"], anthropic: ["claude"] };
+        const publishers = ["google", "anthropic", "meta"];
+        // Meta "maas" models are LLama Models-As-A-Service. Non-maas models are not pre-deployed.
+        const supportedModels = { google: ["gemini", "imagen"], anthropic: ["claude"], meta: ["maas"] };
+        // Additional models not in the listings, but we want to include
+        // TODO: Remove once the models are available in the listing API, or no longer needed
+        const additionalModels = {
+            google: ["imagen-3.0-fast-generate-001"],
+            anthropic: [],
+            meta: [
+                "llama-4-maverick-17b-128e-instruct-maas",
+                "llama-4-scout-17b-16e-instruct-maas",
+                "llama-3.3-70b-instruct-maas",
+                "llama-3.2-90b-vision-instruct-maas",
+                "llama-3.1-405b-instruct-maas",
+                "llama-3.1-70b-instruct-maas",
+                "llama-3.1-8b-instruct-maas",
+            ],
+        }
 
         //Used to exclude retired models that are still in the listing API but not available for use.
         //Or models we do not support yet
         const unsupportedModelsByPublisher = {
             google: ["gemini-pro", "gemini-ultra"],
             anthropic: [],
+            meta: [],
         };
 
         for (const publisher of publishers) {
@@ -228,13 +286,28 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
                     tool_support: modelCapability.tool_support,
                 } satisfies AIModel<string>;
             }));
+
+            // Add additional models that are not in the listing
+            for (const additionalModel of additionalModels[publisher as keyof typeof additionalModels]) {
+                const publisherModelName = `publishers/${publisher}/models/${additionalModel}`;
+                const modelCapability = getModelCapabilities(additionalModel, "vertexai");
+                models.push({
+                    id: publisherModelName,
+                    name: additionalModel,
+                    provider: 'vertexai',
+                    owner: publisher,
+                    input_modalities: modelModalitiesToArray(modelCapability.input),
+                    output_modalities: modelModalitiesToArray(modelCapability.output),
+                    tool_support: modelCapability.tool_support,
+                } satisfies AIModel<string>);
+            }
         }
 
         //Remove duplicates
         const uniqueModels = Array.from(new Set(models.map(a => a.id)))
             .map(id => {
                 return models.find(a => a.id === id) ?? {} as AIModel<string>;
-            });
+            }).sort((a, b) => a.id.localeCompare(b.id));
 
         return uniqueModels;
     }