@llumiverse/drivers 1.3.0 → 1.4.0-dev.20260629.090753Z

package/src/vertexai/models/llama.ts

@@ -1,18 +1,25 @@
 import {
-    AIModel, Completion, CompletionChunkObject, ExecutionOptions, ModelType,
-    PromptOptions, PromptRole, PromptSegment,
-    TextFallbackOptions
-} from "@llumiverse/core";
-import { VertexAIDriver } from "../index.js";
-import { ModelDefinition } from "../models.js";
-import { transformSSEStream } from "@llumiverse/core/async";
+    type AIModel,
+    type Completion,
+    type CompletionChunkObject,
+    type ExecutionOptions,
+    ModelType,
+    type PromptOptions,
+    PromptRole,
+    type PromptSegment,
+    type TextFallbackOptions,
+} from '@llumiverse/core';
+import { transformSSEStream } from '@llumiverse/core/async';
+import type { ServerSentEvent } from '@vertesia/api-fetch-client';
+import type { VertexAIDriver } from '../index.js';
+import type { ModelDefinition } from '../models.js';
 
 interface LLamaMessage {
     role: string;
     content: string;
 }
 
-interface LLamaPrompt {
+export interface LLamaPrompt {
     messages: LLamaMessage[];
 }
 
@@ -61,7 +68,7 @@ interface LLamaStreamResponse {
 /**
  * Convert a stream to a string
  */
-async function streamToString(stream: any): Promise<string> {
+async function streamToString(stream: AsyncIterable<Uint8Array | Buffer | string>): Promise<string> {
     const chunks: Buffer[] = [];
     for await (const chunk of stream) {
         chunks.push(Buffer.from(chunk));
@@ -84,8 +91,7 @@ function updateConversation(conversation: LLamaPrompt | undefined | null, prompt
 }
 
 export class LLamaModelDefinition implements ModelDefinition<LLamaPrompt> {
-
-    model: AIModel
+    model: AIModel;
 
     constructor(modelId: string) {
         this.model = {
@@ -107,7 +113,11 @@ export class LLamaModelDefinition implements ModelDefinition<LLamaPrompt> {
         }
     }
 
-    async createPrompt(_driver: VertexAIDriver, segments: PromptSegment[], options: PromptOptions): Promise<LLamaPrompt> {
+    async createPrompt(
+        _driver: VertexAIDriver,
+        segments: PromptSegment[],
+        options: PromptOptions,
+    ): Promise<LLamaPrompt> {
         const messages: LLamaMessage[] = [];
 
         // Process segments and convert them to the Llama MaaS format
@@ -120,7 +130,7 @@ export class LLamaModelDefinition implements ModelDefinition<LLamaPrompt> {
 
             if (segment.files && segment.files.length > 0) {
                 for (const file of segment.files) {
-                    if (file.mime_type?.startsWith("text/")) {
+                    if (file.mime_type?.startsWith('text/')) {
                         const fileStream = await file.getStream();
                         const fileContent = await streamToString(fileStream);
                         messageContent += `\n\nFile content:\n${fileContent}`;
@@ -130,14 +140,14 @@ export class LLamaModelDefinition implements ModelDefinition<LLamaPrompt> {
 
             messages.push({
                 role: role,
-                content: messageContent
+                content: messageContent,
             });
         }
 
         if (options.result_schema) {
             messages.push({
                 role: 'user',
-                content: "The answer must be a JSON object using the following JSON Schema:\n" + JSON.stringify(options.result_schema)
+                content: `The answer must be a JSON object using the following JSON Schema:\n${JSON.stringify(options.result_schema)}`,
             });
         }
 
@@ -147,15 +157,19 @@ export class LLamaModelDefinition implements ModelDefinition<LLamaPrompt> {
         };
     }
 
-    async requestTextCompletion(driver: VertexAIDriver, prompt: LLamaPrompt, options: ExecutionOptions): Promise<Completion> {
-        const splits = options.model.split("/");
+    async requestTextCompletion(
+        driver: VertexAIDriver,
+        prompt: LLamaPrompt,
+        options: ExecutionOptions,
+    ): Promise<Completion> {
+        const splits = options.model.split('/');
         const modelName = splits[splits.length - 1];
 
         let conversation = updateConversation(options.conversation as LLamaPrompt, prompt);
 
         const modelOptions = options.model_options as TextFallbackOptions;
 
-        const payload: Record<string, any> = {
+        const payload: Record<string, unknown> = {
             model: `meta/${modelName}`,
             messages: conversation.messages,
             stream: false,
@@ -168,19 +182,19 @@ export class LLamaModelDefinition implements ModelDefinition<LLamaPrompt> {
                 google: {
                     model_safety_settings: {
                         enabled: false,
-                        llama_guard_settings: {}
-                    }
-                }
-            }
+                        llama_guard_settings: {},
+                    },
+                },
+            },
         };
 
         // Make POST request to the Llama MaaS API
         const region = this.getLlamaModelRegion(modelName);
         const client = driver.getLLamaClient(region);
         const openaiEndpoint = `endpoints/openapi/chat/completions`;
-        const result = await client.post(openaiEndpoint, {
-            payload
-        }) as LLamaResponse;
+        const result = (await client.post(openaiEndpoint, {
+            payload,
+        })) as LLamaResponse;
 
         // Extract response data
         const assistantMessage = result?.choices[0]?.message;
@@ -188,33 +202,39 @@ export class LLamaModelDefinition implements ModelDefinition<LLamaPrompt> {
 
         // Update conversation with the response
         conversation = updateConversation(conversation, {
-            messages: [{
-                role: assistantMessage?.role,
-                content: text
-            }],
+            messages: [
+                {
+                    role: assistantMessage?.role,
+                    content: text,
+                },
+            ],
         });
 
         return {
-            result: [{ type: "text", value: text }],
+            result: [{ type: 'text', value: text }],
             token_usage: {
                 prompt: result.usage.prompt_tokens,
                 result: result.usage.completion_tokens,
-                total: result.usage.total_tokens
+                total: result.usage.total_tokens,
             },
             finish_reason: result.choices[0].finish_reason,
-            conversation
+            conversation,
         };
     }
 
-    async requestTextCompletionStream(driver: VertexAIDriver, prompt: LLamaPrompt, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunkObject>> {
-        const splits = options.model.split("/");
+    async requestTextCompletionStream(
+        driver: VertexAIDriver,
+        prompt: LLamaPrompt,
+        options: ExecutionOptions,
+    ): Promise<AsyncIterable<CompletionChunkObject>> {
+        const splits = options.model.split('/');
         const modelName = splits[splits.length - 1];
 
         const conversation = updateConversation(options.conversation as LLamaPrompt, prompt);
 
         const modelOptions = options.model_options as TextFallbackOptions;
 
-        const payload: Record<string, any> = {
+        const payload: Record<string, unknown> = {
             model: `meta/${modelName}`,
             messages: conversation.messages,
             stream: true,
@@ -227,10 +247,10 @@ export class LLamaModelDefinition implements ModelDefinition<LLamaPrompt> {
                 google: {
                     model_safety_settings: {
                         enabled: false,
-                        llama_guard_settings: {}
-                    }
-                }
-            }
+                        llama_guard_settings: {},
+                    },
+                },
+            },
         };
 
         // Make POST request to the Llama MaaS API
@@ -239,23 +259,25 @@ export class LLamaModelDefinition implements ModelDefinition<LLamaPrompt> {
         const region = this.getLlamaModelRegion(modelName);
         const client = driver.getLLamaClient(region);
         const openaiEndpoint = `endpoints/openapi/chat/completions`;
-        const stream = await client.post(openaiEndpoint, {
+        const stream = (await client.post(openaiEndpoint, {
             payload,
-            reader: 'sse'
-        });
+            reader: 'sse',
+        })) as ReadableStream<ServerSentEvent>;
 
         return transformSSEStream(stream, (data: string): CompletionChunkObject => {
             const json = JSON.parse(data) as LLamaStreamResponse;
             const choice = json.choices?.[0];
             const content = choice?.delta?.content ?? '';
             return {
-                result: content ? [{ type: "text", value: content }] : [],
+                result: content ? [{ type: 'text', value: content }] : [],
                 finish_reason: choice?.finish_reason,
-                token_usage: json.usage ? {
-                    prompt: json.usage.prompt_tokens,
-                    result: json.usage.completion_tokens,
-                    total: json.usage.total_tokens,
-                } : undefined
+                token_usage: json.usage
+                    ? {
+                          prompt: json.usage.prompt_tokens,
+                          result: json.usage.completion_tokens,
+                          total: json.usage.total_tokens,
+                      }
+                    : undefined,
             };
         });
     }

package/src/vertexai/models.ts

@@ -1,16 +1,36 @@
-import { AIModel, Completion, CompletionChunkObject, ExecutionOptions, LlumiverseError, LlumiverseErrorContext, PromptSegment } from "@llumiverse/core";
-import { VertexAIDriver, trimModelName } from "./index.js";
-import { ClaudeModelDefinition } from "./models/claude.js";
-import { GeminiModelDefinition } from "./models/gemini.js";
-import { LLamaModelDefinition } from "./models/llama.js";
+import type {
+    AIModel,
+    Completion,
+    CompletionChunkObject,
+    ExecutionOptions,
+    LlumiverseError,
+    LlumiverseErrorContext,
+    PromptSegment,
+} from '@llumiverse/core';
+import type { VertexAIDriver, VertexAIPrompt } from './index.js';
+import { ClaudeModelDefinition } from './models/claude.js';
+import { GeminiModelDefinition } from './models/gemini.js';
+import { LLamaModelDefinition } from './models/llama.js';
 
-export interface ModelDefinition<PromptT = any> {
+export function trimModelName(model: string): string {
+    const i = model.lastIndexOf('@');
+    return i > -1 ? model.substring(0, i) : model;
+}
+
+export interface ModelDefinition<PromptT = VertexAIPrompt> {
     model: AIModel;
     versions?: string[]; // the versions of the model that are available. ex: ['001', '002']
-    createPrompt: (driver: VertexAIDriver, segments: PromptSegment[], options: ExecutionOptions) => Promise<PromptT>;
-    requestTextCompletion: (driver: VertexAIDriver, prompt: PromptT, options: ExecutionOptions) => Promise<Completion>;
-    requestTextCompletionStream: (driver: VertexAIDriver, prompt: PromptT, options: ExecutionOptions) => Promise<AsyncIterable<CompletionChunkObject>>;
-    preValidationProcessing?(result: Completion, options: ExecutionOptions): { result: Completion, options: ExecutionOptions };
+    createPrompt(driver: VertexAIDriver, segments: PromptSegment[], options: ExecutionOptions): Promise<PromptT>;
+    requestTextCompletion(driver: VertexAIDriver, prompt: PromptT, options: ExecutionOptions): Promise<Completion>;
+    requestTextCompletionStream(
+        driver: VertexAIDriver,
+        prompt: PromptT,
+        options: ExecutionOptions,
+    ): Promise<AsyncIterable<CompletionChunkObject>>;
+    preValidationProcessing?(
+        result: Completion,
+        options: ExecutionOptions,
+    ): { result: Completion; options: ExecutionOptions };
     /**
      * Format provider-specific errors into standardized LlumiverseError.
      * Optional - if not provided, VertexAIDriver will use default error handling.
@@ -19,13 +39,13 @@ export interface ModelDefinition<PromptT = any> {
 }
 
 export function getModelDefinition(model: string): ModelDefinition {
-    const splits = model.split("/");
+    const splits = model.split('/');
 
     // Handle both formats: "publishers/anthropic/models/..." and "locations/.../publishers/anthropic/models/..."
     let publisher: string | undefined;
     let modelName: string;
 
-    const publisherIndex = splits.indexOf("publishers");
+    const publisherIndex = splits.indexOf('publishers');
     if (publisherIndex !== -1 && publisherIndex + 1 < splits.length) {
         publisher = splits[publisherIndex + 1];
         modelName = trimModelName(splits[splits.length - 1]);
@@ -35,14 +55,14 @@ export function getModelDefinition(model: string): ModelDefinition {
         modelName = trimModelName(splits[splits.length - 1]);
     }
 
-    if (publisher?.includes("anthropic")) {
+    if (publisher?.includes('anthropic')) {
         return new ClaudeModelDefinition(modelName);
-    } else if (publisher?.includes("google")) {
+    } else if (publisher?.includes('google')) {
         return new GeminiModelDefinition(modelName);
-    } else if (publisher?.includes("meta")) {
+    } else if (publisher?.includes('meta')) {
         return new LLamaModelDefinition(modelName);
     }
 
     //Fallback, assume it is Gemini.
     return new GeminiModelDefinition(modelName);
-}
+}

package/src/watsonx/index.ts

@@ -1,7 +1,28 @@
-import { AbstractDriver, AIModel, Completion, CompletionChunkObject, DriverOptions, EmbeddingsOptions, EmbeddingsResult, ExecutionOptions, TextFallbackOptions } from "@llumiverse/core";
-import { transformSSEStream } from "@llumiverse/core/async";
-import { FetchClient } from "@vertesia/api-fetch-client";
-import { GenerateEmbeddingPayload, GenerateEmbeddingResponse, WatsonAuthToken, WatsonxListModelResponse, WatsonxModelSpec, WatsonxTextGenerationPayload, WatsonxTextGenerationResponse } from "./interfaces.js";
+import {
+    type AIModel,
+    type Completion,
+    type CompletionChunkObject,
+    type DriverOptions,
+    type EmbeddingsOptions,
+    type EmbeddingsResult,
+    type ExecutionOptions,
+    LlumiverseError,
+    normalizeEmbeddingsOptions,
+    type TextFallbackOptions,
+    WATSONX_DEFAULT_EMBEDDING_MODEL,
+} from '@llumiverse/core';
+import { transformSSEStream } from '@llumiverse/core/async';
+import { AbstractDriver } from '@llumiverse/core/driver';
+import { FetchClient, type ServerSentEvent } from '@vertesia/api-fetch-client';
+import type {
+    GenerateEmbeddingPayload,
+    GenerateEmbeddingResponse,
+    WatsonAuthToken,
+    WatsonxListModelResponse,
+    WatsonxModelSpec,
+    WatsonxTextGenerationPayload,
+    WatsonxTextGenerationResponse,
+} from './interfaces.js';
 
 interface WatsonxDriverOptions extends DriverOptions {
     apiKey: string;
@@ -9,35 +30,37 @@ interface WatsonxDriverOptions extends DriverOptions {
     endpointUrl: string;
 }
 
-const API_VERSION = "2024-03-14"
+const API_VERSION = '2024-03-14';
 
 export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string> {
-    static PROVIDER = "watsonx";
+    static PROVIDER = 'watsonx';
     provider = WatsonxDriver.PROVIDER;
     apiKey: string;
     endpoint_url: string;
     projectId: string;
     authToken?: WatsonAuthToken;
     fetcher?: FetchClient;
-    fetchClient: FetchClient
+    fetchClient: FetchClient;
 
     constructor(options: WatsonxDriverOptions) {
         super(options);
         this.apiKey = options.apiKey;
         this.projectId = options.projectId;
         this.endpoint_url = options.endpointUrl;
-        this.fetchClient = new FetchClient(this.endpoint_url).withAuthCallback(async () => this.getAuthToken().then(token => `Bearer ${token}`));
+        this.fetchClient = new FetchClient(this.endpoint_url, this.getDriverFetch()).withAuthCallback(async () =>
+            this.getAuthToken().then((token) => `Bearer ${token}`),
+        );
     }
 
     async requestTextCompletion(prompt: string, options: ExecutionOptions): Promise<Completion> {
-        if (options.model_options?._option_id !== undefined && options.model_options?._option_id !== "text-fallback") {
-            this.logger.debug({ options: options.model_options }, "Unexpected option id");
+        if (options.model_options?._option_id !== undefined && options.model_options?._option_id !== 'text-fallback') {
+            this.logger.debug({ options: options.model_options }, 'Unexpected option id');
         }
         options.model_options = options.model_options as TextFallbackOptions | undefined;
 
         const payload: WatsonxTextGenerationPayload = {
             model_id: options.model,
-            input: prompt + "\n",
+            input: `${prompt}\n`,
             parameters: {
                 max_new_tokens: options.model_options?.max_tokens,
                 temperature: options.model_options?.temperature,
@@ -46,14 +69,16 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
                 stop_sequences: options.model_options?.stop_sequence,
             },
             project_id: this.projectId,
-        }
+        };
 
-        const res = await this.fetchClient.post(`/ml/v1/text/generation?version=${API_VERSION}`, { payload }) as WatsonxTextGenerationResponse;
+        const res = (await this.fetchClient.post(`/ml/v1/text/generation?version=${API_VERSION}`, {
+            payload,
+        })) as WatsonxTextGenerationResponse;
 
         const result = res.results[0];
 
         return {
-            result: [{ type: "text", value: result.generated_text }],
+            result: [{ type: 'text', value: result.generated_text }],
             token_usage: {
                 prompt: result.input_token_count,
                 result: result.generated_token_count,
@@ -61,17 +86,20 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
             },
             finish_reason: watsonFinishReason(result.stop_reason),
             original_response: options.include_original_response ? res : undefined,
-        }
+        };
     }
 
-    async requestTextCompletionStream(prompt: string, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunkObject>> {
-        if (options.model_options?._option_id !== undefined && options.model_options?._option_id !== "text-fallback") {
-            this.logger.debug({ options: options.model_options }, "Unexpected option id");
+    async requestTextCompletionStream(
+        prompt: string,
+        options: ExecutionOptions,
+    ): Promise<AsyncIterable<CompletionChunkObject>> {
+        if (options.model_options?._option_id !== undefined && options.model_options?._option_id !== 'text-fallback') {
+            this.logger.debug({ options: options.model_options }, 'Unexpected option id');
         }
         options.model_options = options.model_options as TextFallbackOptions | undefined;
         const payload: WatsonxTextGenerationPayload = {
             model_id: options.model,
-            input: prompt + "\n",
+            input: `${prompt}\n`,
             parameters: {
                 max_new_tokens: options.model_options?.max_tokens,
                 temperature: options.model_options?.temperature,
@@ -80,17 +108,19 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
                 stop_sequences: options.model_options?.stop_sequence,
             },
             project_id: this.projectId,
-        }
+        };
 
-        const stream = await this.fetchClient.post(`/ml/v1/text/generation_stream?version=${API_VERSION}`, {
+        const stream = (await this.fetchClient.post(`/ml/v1/text/generation_stream?version=${API_VERSION}`, {
             payload: payload,
-            reader: 'sse'
-        })
+            reader: 'sse',
+        })) as ReadableStream<ServerSentEvent>;
 
         return transformSSEStream(stream, (data: string) => {
             const json = JSON.parse(data) as WatsonxTextGenerationResponse;
             return {
-                result: json.results[0]?.generated_text ? [{ type: "text", value: json.results[0].generated_text }] : [],
+                result: json.results[0]?.generated_text
+                    ? [{ type: 'text', value: json.results[0].generated_text }]
+                    : [],
                 finish_reason: watsonFinishReason(json.results[0]?.stop_reason),
                 token_usage: {
                     prompt: json.results[0].input_token_count,
@@ -99,17 +129,12 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
                 },
             };
         });
-
     }
 
-
-
     async listModels(): Promise<AIModel<string>[]> {
-
-
-
-        const res = await this.fetchClient.get(`/ml/v1/foundation_model_specs?version=${API_VERSION}`)
-            .catch(err => this.logger.warn("Can't list models on Watsonx: " + err)) as WatsonxListModelResponse;
+        const res = (await this.fetchClient
+            .get(`/ml/v1/foundation_model_specs?version=${API_VERSION}`)
+            .catch((err) => this.logger.warn(`Can't list models on Watsonx: ${err}`))) as WatsonxListModelResponse;
 
         const aiModels = res.resources.map((m: WatsonxModelSpec) => {
             return {
@@ -117,83 +142,94 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
                 name: m.label,
                 description: m.short_description,
                 provider: this.provider,
-            }
+            };
         });
 
         return aiModels;
-
     }
 
     async getAuthToken(): Promise<string> {
-
-
         if (this.authToken) {
             const now = Date.now() / 1000;
             if (now < this.authToken.expiration) {
                 return this.authToken.access_token;
             } else {
-                this.logger.debug("Token expired, refetching")
+                this.logger.debug('Token expired, refetching');
             }
         }
-        const authToken = await fetch('https://iam.cloud.ibm.com/identity/token', {
+        const authToken = (await this.getDriverFetch()('https://iam.cloud.ibm.com/identity/token', {
             method: 'POST',
             headers: {
                 'Content-Type': 'application/x-www-form-urlencoded',
             },
             body: `grant_type=urn:ibm:params:oauth:grant-type:apikey&apikey=${this.apiKey}`,
-        }).then(response => response.json()) as WatsonAuthToken;
+        }).then((response) => response.json())) as WatsonAuthToken;
 
         this.authToken = authToken;
         return this.authToken.access_token;
-
     }
 
     async validateConnection(): Promise<boolean> {
         return this.listModels()
             .then(() => true)
             .catch((err) => {
-                this.logger.warn({ error: err }, "Failed to connect to WatsonX");
-                return false
+                this.logger.warn({ error: err }, 'Failed to connect to WatsonX');
+                return false;
             });
     }
 
     async generateEmbeddings(options: EmbeddingsOptions): Promise<EmbeddingsResult> {
-        if (options.image) {
-            throw new Error("Image embeddings not supported by Watsonx");
-        }
-
-        if (!options.text) {
-            throw new Error("No text provided");
-        }
+        const normalized = normalizeEmbeddingsOptions(options);
+        const texts = normalized.inputs.map((input) => {
+            if (input.type !== 'text') {
+                throw new Error(
+                    `Provider 'watsonx' does not support '${input.type}' embeddings; only 'text' is supported.`,
+                );
+            }
+            return input.text;
+        });
 
         const payload: GenerateEmbeddingPayload = {
-            inputs: [options.text],
-            model_id: options.model ?? 'ibm/slate-125m-english-rtrvr',
-            project_id: this.projectId
-        }
-
-        const res = await this.fetchClient.post(`/ml/v1/text/embeddings?version=${API_VERSION}`, { payload }) as GenerateEmbeddingResponse;
+            inputs: texts,
+            model_id: normalized.model ?? WATSONX_DEFAULT_EMBEDDING_MODEL,
+            project_id: this.projectId,
+        };
 
-        return {
-            values: res.results[0].embedding,
-            model: res.model_id
+        const model = payload.model_id;
+        try {
+            const res = (await this.fetchClient.post(`/ml/v1/text/embeddings?version=${API_VERSION}`, {
+                payload,
+            })) as GenerateEmbeddingResponse;
+            return {
+                model: res.model_id,
+                results: res.results.map((entry) => ({
+                    outputs: [{ values: entry.embedding, modality: 'text' }],
+                })),
+            };
+        } catch (error) {
+            if (LlumiverseError.isLlumiverseError(error)) throw error;
+            if (error instanceof Error && typeof (error as { status?: unknown }).status !== 'number') throw error;
+            throw this.formatLlumiverseError(error, {
+                provider: this.provider,
+                model,
+                operation: 'execute',
+            });
         }
-
     }
-
 }
 
 function watsonFinishReason(reason: string | undefined) {
     if (!reason) return undefined;
     switch (reason) {
-        case 'eos_token': return "stop";
-        case 'max_tokens': return "length";
-        default: return reason;
+        case 'eos_token':
+            return 'stop';
+        case 'max_tokens':
+            return 'length';
+        default:
+            return reason;
     }
 }
 
-
-
 /*interface ListModelsParams extends ModelSearchPayload {
     limit?: number;
 }*/