@llumiverse/drivers 0.22.0 → 0.23.0-dev-20251118

package/src/replicate.ts

@@ -2,7 +2,7 @@ import {
     AIModel,
     AbstractDriver,
     Completion,
-    CompletionChunk,
+    CompletionChunkObject,
     DataSource,
     DriverOptions,
     EmbeddingsResult,
@@ -15,7 +15,7 @@ import {
 } from "@llumiverse/core";
 import { EventStream } from "@llumiverse/core/async";
 import { EventSource } from "eventsource";
-import Replicate, { Prediction } from "replicate";
+import Replicate, { Prediction, Training } from "replicate";
 
 let cachedTrainableModels: AIModel[] | undefined;
 let cachedTrainableModelsTimestamp: number = 0;
@@ -64,12 +64,12 @@ export class ReplicateDriver extends AbstractDriver<DriverOptions, string> {
         };
     }
 
-    async requestTextCompletionStream(prompt: string, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunk>> {
+    async requestTextCompletionStream(prompt: string, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunkObject>> {
         if (options.model_options?._option_id !== "text-fallback") {
-            this.logger.warn("Invalid model options", {options: options.model_options });
+            this.logger.warn({ options: options.model_options }, "Invalid model options");
         }
         options.model_options = options.model_options as TextFallbackOptions;
-        
+
         const model = ReplicateDriver.parseModelId(options.model);
         const predictionData = {
             input: {
@@ -84,11 +84,11 @@ export class ReplicateDriver extends AbstractDriver<DriverOptions, string> {
         const prediction =
             await this.service.predictions.create(predictionData);
 
-        const stream = new EventStream<CompletionChunk>();
+        const stream = new EventStream<CompletionChunkObject>();
 
         const source = new EventSource(prediction.urls.stream!);
         source.addEventListener("output", (e: any) => {
-            stream.push(e.data);
+            stream.push({ result: [{ type: "text", value: e.data }] });
         });
         source.addEventListener("error", (e: any) => {
             let error: any;
@@ -97,7 +97,7 @@ export class ReplicateDriver extends AbstractDriver<DriverOptions, string> {
             } catch (error) {
                 error = JSON.stringify(e);
             }
-            this.logger?.error("Error in SSE stream", {e, error});
+            this.logger.error({ e, error }, "Error in SSE stream");
         });
         source.addEventListener("done", () => {
             try {
@@ -111,7 +111,7 @@ export class ReplicateDriver extends AbstractDriver<DriverOptions, string> {
 
     async requestTextCompletion(prompt: string, options: ExecutionOptions) {
         if (options.model_options?._option_id !== "text-fallback") {
-            this.logger.warn("Invalid model options", {options: options.model_options });
+            this.logger.warn({ options: options.model_options }, "Invalid model options");
         }
         options.model_options = options.model_options as TextFallbackOptions;
         const model = ReplicateDriver.parseModelId(options.model);
@@ -136,9 +136,9 @@ export class ReplicateDriver extends AbstractDriver<DriverOptions, string> {
         //not streaming, wait for the result
         const res = await this.service.wait(prediction, {});
 
-        const text = res.output.join("");
+        const text: string = res.output.join("");
         return {
-            result: text,
+            result: [{ type: "text" as const, value: text }],
             original_response: options.include_original_response ? res : undefined,
         };
     }
@@ -236,7 +236,7 @@ export class ReplicateDriver extends AbstractDriver<DriverOptions, string> {
             this.service.models.versions.list(owner, model),
         ]);
 
-        if (!rModel || !versions || versions.length === 0) {
+        if (!rModel || !versions || (versions as any).results?.length === 0) {
             throw new Error("Model not found or no versions available");
         }
 
@@ -289,7 +289,7 @@ export class ReplicateDriver extends AbstractDriver<DriverOptions, string> {
 
 }
 
-function jobInfo(job: Prediction, modelName?: string): TrainingJob {
+function jobInfo(job: Prediction | Training, modelName?: string): TrainingJob {
     // 'starting' | 'processing' | 'succeeded' | 'failed' | 'canceled'
     const jobStatus = job.status;
     let details: string | undefined;
@@ -298,7 +298,7 @@ function jobInfo(job: Prediction, modelName?: string): TrainingJob {
         status = TrainingJobStatus.succeeded;
     } else if (jobStatus === 'failed') {
         status = TrainingJobStatus.failed;
-        const error = job.error as any; 
+        const error = job.error as any;
         if (typeof error === 'string') {
             details = error;
         } else {

package/src/{test → test-driver}/utils.ts

@@ -8,7 +8,7 @@ export function throwError(message: string, prompt: PromptSegment[]): never {
 
 export function createValidationErrorCompletion(segments: PromptSegment[]) {
     return {
-        result: "An invalid result",
+        result: [{ type: "text", value: "An invalid result" }],
         prompt: segments,
         execution_time: 3000,
         error: {

package/src/togetherai/index.ts

@@ -1,4 +1,4 @@
-import { AIModel, AbstractDriver, Completion, CompletionChunk, DriverOptions, EmbeddingsResult, ExecutionOptions, TextFallbackOptions } from "@llumiverse/core";
+import { AIModel, AbstractDriver, Completion, CompletionChunkObject, DriverOptions, EmbeddingsResult, ExecutionOptions, TextFallbackOptions } from "@llumiverse/core";
 import { transformSSEStream } from "@llumiverse/core/async";
 import { FetchClient } from "@vertesia/api-fetch-client";
 import { TextCompletion, TogetherModelInfo } from "./interfaces.js";
@@ -29,9 +29,9 @@ export class TogetherAIDriver extends AbstractDriver<TogetherAIDriverOptions, st
             } : undefined;
     }
 
-    async requestTextCompletion(prompt: string, options: ExecutionOptions): Promise<Completion<any>> {
+    async requestTextCompletion(prompt: string, options: ExecutionOptions): Promise<Completion> {
         if (options.model_options?._option_id !== "text-fallback") {
-            this.logger.warn("Invalid model options", {options: options.model_options });
+            this.logger.warn({ options: options.model_options }, "Invalid model options");
         }
         options.model_options = options.model_options as TextFallbackOptions;
 
@@ -60,7 +60,7 @@ export class TogetherAIDriver extends AbstractDriver<TogetherAIDriverOptions, st
         const text = choice.text ?? '';
         const usage = res.usage || {};
         return {
-            result: text,
+            result: [{ type: "text", value: text }],
             token_usage: {
                 prompt: usage.prompt_tokens,
                 result: usage.completion_tokens,
@@ -71,9 +71,9 @@ export class TogetherAIDriver extends AbstractDriver<TogetherAIDriverOptions, st
         }
     }
 
-    async requestTextCompletionStream(prompt: string, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunk>> {
+    async requestTextCompletionStream(prompt: string, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunkObject>> {
         if (options.model_options?._option_id !== "text-fallback") {
-            this.logger.warn("Invalid model options", {options: options.model_options });
+            this.logger.warn({ options: options.model_options }, "Invalid model options");
         }
         options.model_options = options.model_options as TextFallbackOptions;
 
@@ -103,7 +103,7 @@ export class TogetherAIDriver extends AbstractDriver<TogetherAIDriverOptions, st
         return transformSSEStream(stream, (data: string) => {
             const json = JSON.parse(data);
             return {
-                result: json.choices[0]?.text ?? '',
+                result: [{ type: "text", value: json.choices[0]?.text ?? '' }],
                 finish_reason: json.choices[0]?.finish_reason,          //Uses expected "stop" , "length" format
                 token_usage: {
                     prompt: json.usage?.prompt_tokens,

package/src/vertexai/embeddings/embeddings-text.ts

@@ -36,9 +36,9 @@ export async function getEmbeddingsForText(driver: VertexAIDriver, options: Text
             title: options.title,
             content: options.content
         }]
-    } as EmbeddingsForTextPrompt;
+    } satisfies EmbeddingsForTextPrompt;
 
-    const model = options.model || "text-embedding-004";
+    const model = options.model || "gemini-embedding-001";
 
     const client = driver.getFetchClient();
 

package/src/vertexai/index.ts

@@ -1,12 +1,15 @@
+import { AnthropicVertex } from "@anthropic-ai/vertex-sdk";
+import { PredictionServiceClient, v1beta1 } from "@google-cloud/aiplatform";
+import { Content, GoogleGenAI, Model } from "@google/genai";
 import {
     AIModel,
     AbstractDriver,
     Completion,
-    CompletionChunk,
+    CompletionChunkObject,
     DriverOptions,
+    EmbeddingsOptions,
     EmbeddingsResult,
     ExecutionOptions,
-    ImageGeneration,
     Modalities,
     ModelSearchPayload,
     PromptSegment,
@@ -14,16 +17,12 @@ import {
     modelModalitiesToArray,
 } from "@llumiverse/core";
 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 { GoogleAuth, GoogleAuthOptions, AuthClient } from "google-auth-library";
+import { getEmbeddingsForImages } from "./embeddings/embeddings-image.js";
 import { TextEmbeddingsOptions, getEmbeddingsForText } from "./embeddings/embeddings-text.js";
 import { getModelDefinition } from "./models.js";
-import { EmbeddingsOptions } from "@llumiverse/core";
-import { getEmbeddingsForImages } from "./embeddings/embeddings-image.js";
-import { v1beta1 } from "@google-cloud/aiplatform";
-import { AnthropicVertex } from "@anthropic-ai/vertex-sdk";
+import { ANTHROPIC_REGIONS, NON_GLOBAL_ANTHROPIC_MODELS } from "./models/claude.js";
 import { ImagenModelDefinition, ImagenPrompt } from "./models/imagen.js";
-import { GoogleGenAI, Content } from "@google/genai";
 
 export interface VertexAIDriverOptions extends DriverOptions {
     project: string;
@@ -54,8 +53,10 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
     googleGenAI: GoogleGenAI | undefined;
     llamaClient: FetchClient & { region?: string } | undefined;
     modelGarden: v1beta1.ModelGardenServiceClient | undefined;
+    imagenClient: PredictionServiceClient | undefined;
 
-    authClient: JSONClient | GoogleAuth<JSONClient>;
+    googleAuth: GoogleAuth<any>;
+    private authClientPromise: Promise<AuthClient> | undefined;
 
     constructor(options: VertexAIDriverOptions) {
         super(options);
@@ -66,19 +67,39 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
         this.googleGenAI = undefined;
         this.modelGarden = undefined;
         this.llamaClient = undefined;
+        this.imagenClient = undefined;
 
-        this.authClient = options.googleAuthOptions?.authClient ?? new GoogleAuth(options.googleAuthOptions);
+        this.googleAuth = new GoogleAuth(options.googleAuthOptions) as GoogleAuth<any>;
+        this.authClientPromise = undefined;
     }
 
-    public getGoogleGenAIClient(): GoogleGenAI {
-        //Lazy initialisation
+    private async getAuthClient(): Promise<AuthClient> {
+        if (!this.authClientPromise) {
+            this.authClientPromise = this.googleAuth.getClient();
+        }
+        return this.authClientPromise;
+    }
+
+    public getGoogleGenAIClient(region: string = this.options.region): GoogleGenAI {
+        //Lazy initialization
+        if (region !== this.options.region) {
+            //Get one off client for different region
+            return new GoogleGenAI({
+                project: this.options.project,
+                location: region,
+                vertexai: true,
+                googleAuthOptions: this.options.googleAuthOptions || {
+                    scopes: ["https://www.googleapis.com/auth/cloud-platform"],
+                }
+            });
+        }
         if (!this.googleGenAI) {
             this.googleGenAI = new GoogleGenAI({
                 project: this.options.project,
-                location: this.options.region,
+                location: region,
                 vertexai: true,
-                googleAuthOptions: {
-                    authClient: this.authClient as JSONClient,
+                googleAuthOptions: this.options.googleAuthOptions || {
+                    scopes: ["https://www.googleapis.com/auth/cloud-platform"],
                 }
             });
         }
@@ -86,14 +107,13 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
     }
 
     public getFetchClient(): FetchClient {
-        //Lazy initialisation
+        //Lazy initialization
         if (!this.fetchClient) {
             this.fetchClient = createFetchClient({
                 region: this.options.region,
                 project: this.options.project,
             }).withAuthCallback(async () => {
-                const accessTokenResponse = await this.authClient.getAccessToken();
-                const token = typeof accessTokenResponse === 'string' ? accessTokenResponse : accessTokenResponse?.token;
+                const token = await this.googleAuth.getAccessToken();
                 return `Bearer ${token}`;
             });
         }
@@ -101,15 +121,14 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
     }
 
     public getLLamaClient(region: string = "us-central1"): FetchClient {
-        //Lazy initialisation
+        //Lazy initialization
         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;
+                const token = await this.googleAuth.getAccessToken();
                 return `Bearer ${token}`;
             });
             // Store the region for potential client reuse
@@ -118,42 +137,80 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
         return this.llamaClient;
     }
 
-    public getAnthropicClient(): AnthropicVertex {
-        //Lazy initialisation
+    public async getAnthropicClient(region: string = this.options.region): Promise<AnthropicVertex> {
+        // Extract region prefix and map if it exists in ANTHROPIC_REGIONS, otherwise use as-is
+        const getRegionPrefix = (r: string) => r.split('-')[0];
+        const regionPrefix = getRegionPrefix(region);
+        const mappedRegion = ANTHROPIC_REGIONS[regionPrefix] || region;
+
+        const defaultRegionPrefix = getRegionPrefix(this.options.region);
+        const defaultMappedRegion = ANTHROPIC_REGIONS[defaultRegionPrefix] || this.options.region;
+
+        // Get auth client to avoid version mismatch with GoogleAuth generic types
+        const authClient = await this.getAuthClient();
+
+        // If mapped region is different from default mapped region, create one-off client
+        if (mappedRegion !== defaultMappedRegion) {
+            return 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: mappedRegion,
+                projectId: this.options.project,
+                authClient,
+            });
+        }
+
+        //Lazy initialization for default region
         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,
+                region: mappedRegion,
+                projectId: this.options.project,
+                authClient,
             });
         }
         return this.anthropicClient;
     }
 
-    public getAIPlatformClient(): v1beta1.ModelServiceClient {
-        //Lazy initialisation
+    public async getAIPlatformClient(): Promise<v1beta1.ModelServiceClient> {
+        //Lazy initialization
         if (!this.aiplatform) {
+            const authClient = await this.getAuthClient();
             this.aiplatform = new v1beta1.ModelServiceClient({
                 projectId: this.options.project,
                 apiEndpoint: `${this.options.region}-${API_BASE_PATH}`,
-                authClient: this.authClient as JSONClient,
+                authClient,
             });
         }
         return this.aiplatform;
     }
 
-    public getModelGardenClient(): v1beta1.ModelGardenServiceClient {
-        //Lazy initialisation
+    public async getModelGardenClient(): Promise<v1beta1.ModelGardenServiceClient> {
+        //Lazy initialization
         if (!this.modelGarden) {
+            const authClient = await this.getAuthClient();
             this.modelGarden = new v1beta1.ModelGardenServiceClient({
                 projectId: this.options.project,
                 apiEndpoint: `${this.options.region}-${API_BASE_PATH}`,
-                authClient: this.authClient as JSONClient,
+                authClient,
             });
         }
         return this.modelGarden;
     }
 
+    public async getImagenClient(): Promise<PredictionServiceClient> {
+        //Lazy initialization
+        if (!this.imagenClient) {
+            // TODO: make location configurable, fixed to us-central1 for now
+            const authClient = await this.getAuthClient();
+            this.imagenClient = new PredictionServiceClient({
+                projectId: this.options.project,
+                apiEndpoint: `us-central1-${API_BASE_PATH}`,
+                authClient,
+            });
+        }
+        return this.imagenClient;
+    }
+
     validateResult(result: Completion, options: ExecutionOptions) {
         // Optionally preprocess the result before validation
         const modelDef = getModelDefinition(options.model);
@@ -180,29 +237,39 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
         return getModelDefinition(options.model).createPrompt(this, segments, options);
     }
 
-    async requestTextCompletion(prompt: VertexAIPrompt, options: ExecutionOptions): Promise<Completion<any>> {
+    async requestTextCompletion(prompt: VertexAIPrompt, options: ExecutionOptions): Promise<Completion> {
         return getModelDefinition(options.model).requestTextCompletion(this, prompt, options);
     }
     async requestTextCompletionStream(
         prompt: VertexAIPrompt,
         options: ExecutionOptions,
-    ): Promise<AsyncIterable<CompletionChunk>> {
+    ): Promise<AsyncIterable<CompletionChunkObject>> {
         return getModelDefinition(options.model).requestTextCompletionStream(this, prompt, options);
     }
 
     async requestImageGeneration(
         _prompt: ImagenPrompt,
         _options: ExecutionOptions,
-    ): Promise<Completion<ImageGeneration>> {
+    ): Promise<Completion> {
         const splits = _options.model.split("/");
         const modelName = trimModelName(splits[splits.length - 1]);
         return new ImagenModelDefinition(modelName).requestImageGeneration(this, _prompt, _options);
     }
 
+    async getGenAIModelsArray(client: GoogleGenAI): Promise<Model[]> {
+        const models: Model[] = [];
+        const pager = await client.models.list();
+        for await (const item of pager) {
+            models.push(item);
+        }
+        return models;
+    }
+
     async listModels(_params?: ModelSearchPayload): Promise<AIModel<string>[]> {
         // Get clients
-        const modelGarden = this.getModelGardenClient();
-        const aiplatform = this.getAIPlatformClient();
+        const modelGarden = await this.getModelGardenClient();
+        const aiplatform = await this.getAIPlatformClient();
+        const globalGenAiClient = this.getGoogleGenAIClient("global");
 
         let models: AIModel<string>[] = [];
 
@@ -241,16 +308,19 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
             parent: `projects/${this.options.project}/locations/${this.options.region}`,
         });
         const publisherPromises = publishers.map(async (publisher) => {
-            let [response] = await modelGarden.listPublisherModels({
+            const [response] = await modelGarden.listPublisherModels({
                 parent: `publishers/${publisher}`,
                 orderBy: "name",
                 listAllVersions: true,
             });
             return { publisher, response };
         });
+
+        const globalGooglePromise = this.getGenAIModelsArray(globalGenAiClient);
         // Await all network requests
-        const [aiplatformResult, ...publisherResults] = await Promise.all([
+        const [aiplatformResult, globalGoogleResult, ...publisherResults] = await Promise.all([
             aiplatformPromise,
+            globalGooglePromise,
             ...publisherPromises,
         ]);
 
@@ -261,7 +331,23 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
                 id: model.name?.split("/").pop() ?? "",
                 name: model.displayName ?? "",
                 provider: "vertexai"
-            })),
+            }))
+        );
+
+        // Process global google models from GenAI
+        models = models.concat(
+            globalGoogleResult.map((model) => {
+                const modelCapability = getModelCapabilities(model.name ?? '', "vertexai");
+                return {
+                    id: "locations/global/" + model.name,
+                    name: "Global " + model.name?.split('/').pop(),
+                    provider: "vertexai",
+                    owner: "google",
+                    input_modalities: modelModalitiesToArray(modelCapability.input),
+                    output_modalities: modelModalitiesToArray(modelCapability.output),
+                    tool_support: modelCapability.tool_support,
+                };
+            })
         );
 
         // Process publisher models
@@ -294,6 +380,36 @@ export class VertexAIDriver extends AbstractDriver<VertexAIDriverOptions, Vertex
                 } satisfies AIModel<string>;
             }));
 
+            // Create global anthropic models for those not in NON_GLOBAL_ANTHROPIC_MODELS
+            if (publisher === 'anthropic') {
+                const globalAnthropicModels = response.filter((model) => {
+                    const modelName = model.name ?? "";
+                    if (retiredModels.some(retiredModel => modelName.includes(retiredModel))) {
+                        return false;
+                    }
+                    if (modelFamily.some(family => modelName.includes(family))) {
+                        if (modelName.includes("claude-3-7")) {
+                            return true;
+                        }
+                        return !NON_GLOBAL_ANTHROPIC_MODELS.some(nonGlobalModel => modelName.includes(nonGlobalModel));
+                    }
+                    return false;
+                }).map(model => {
+                    const modelCapability = getModelCapabilities(model.name ?? '', "vertexai");
+                    return {
+                        id: "locations/global/" + model.name,
+                        name: "Global " + model.name?.split('/').pop(),
+                        provider: 'vertexai',
+                        owner: publisher,
+                        input_modalities: modelModalitiesToArray(modelCapability.input),
+                        output_modalities: modelModalitiesToArray(modelCapability.output),
+                        tool_support: modelCapability.tool_support,
+                    } satisfies AIModel<string>;
+                });
+
+                models = models.concat(globalAnthropicModels);
+            }
+
             // 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}`;