@llumiverse/drivers 0.20.0 → 0.22.0-dev.1

package/src/vertexai/models/gemini.ts

@@ -1,11 +1,13 @@
 import {
-    Content, FinishReason, FunctionCallingConfigMode, FunctionDeclaration, GenerateContentParameters,
+    Content, FinishReason, FunctionCallingConfigMode, FunctionDeclaration, GenerateContentConfig, GenerateContentParameters,
+    GenerateContentResponseUsageMetadata,
     HarmBlockThreshold, HarmCategory, Part, SafetySetting, Schema, Tool, Type
 } from "@google/genai";
 import {
-    AIModel, Completion, CompletionChunkObject, ExecutionOptions,
-    ExecutionTokenUsage, JSONObject, JSONSchema, ModelType, PromptOptions, PromptRole,
-    PromptSegment, readStreamAsBase64, ToolDefinition, ToolUse
+    AIModel, Completion, CompletionChunkObject, CompletionResult, ExecutionOptions,
+    ExecutionTokenUsage, getMaxTokensLimitVertexAi, JSONObject, JSONSchema, ModelType, PromptOptions, PromptRole,
+    PromptSegment, readStreamAsBase64, StatelessExecutionOptions, ToolDefinition, ToolUse,
+    VertexAIGeminiOptions
 } from "@llumiverse/core";
 import { asyncMap } from "@llumiverse/core/async";
 import { VertexAIDriver, GenerateContentPrompt } from "../index.js";
@@ -44,42 +46,52 @@ const geminiSafetySettings: SafetySetting[] = [
 ];
 
 function getGeminiPayload(options: ExecutionOptions, prompt: GenerateContentPrompt): GenerateContentParameters {
-    const model_options = options.model_options as any;
+    const model_options = options.model_options as VertexAIGeminiOptions | undefined;
     const tools = getToolDefinitions(options.tools);
 
     const useStructuredOutput = supportsStructuredOutput(options) && !tools;
 
+    const thinkingConfigNeeded = model_options?.include_thoughts
+        || model_options?.thinking_budget_tokens
+        || options.model.includes("gemini-2.5");
+
+    const configNanoBanana: GenerateContentConfig = {
+        responseModalities: ["TEXT", "IMAGE"]
+    }
+
+    const config: GenerateContentConfig = {
+        systemInstruction: prompt.system,
+        safetySettings: geminiSafetySettings,
+        tools: tools ? [tools] : undefined,
+        toolConfig: tools ? {
+            functionCallingConfig: {
+                mode: FunctionCallingConfigMode.AUTO,
+            }
+        } : undefined,
+        candidateCount: 1,
+        //JSON/Structured output
+        responseMimeType: useStructuredOutput ? "application/json" : undefined,
+        responseSchema: useStructuredOutput ? parseJSONtoSchema(options.result_schema, true) : undefined,
+        //Model options
+        temperature: model_options?.temperature,
+        topP: model_options?.top_p,
+        topK: model_options?.top_k,
+        maxOutputTokens: geminiMaxTokens(options),
+        stopSequences: model_options?.stop_sequence,
+        presencePenalty: model_options?.presence_penalty,
+        frequencyPenalty: model_options?.frequency_penalty,
+        seed: model_options?.seed,
+        thinkingConfig: thinkingConfigNeeded ?
+            {
+                includeThoughts: model_options?.include_thoughts ?? false,
+                thinkingBudget: geminiThinkingBudget(options),
+            } : undefined,
+    }
+
     return {
         model: options.model,
         contents: prompt.contents,
-        config: {
-            systemInstruction: prompt.system,
-            safetySettings: geminiSafetySettings,
-            tools: tools ? [tools] : undefined,
-            toolConfig: tools ? {
-                functionCallingConfig: {
-                    mode: FunctionCallingConfigMode.AUTO,
-                }
-            } : undefined,
-            candidateCount: 1,
-            //JSON/Structured output
-            responseMimeType: useStructuredOutput ? "application/json" : undefined,
-            responseSchema: useStructuredOutput ? parseJSONtoSchema(options.result_schema, true) : undefined,
-            //Model options
-            temperature: model_options?.temperature,
-            topP: model_options?.top_p,
-            topK: model_options?.top_k,
-            maxOutputTokens: model_options?.max_tokens,
-            stopSequences: model_options?.stop_sequence,
-            presencePenalty: model_options?.presence_penalty,
-            frequencyPenalty: model_options?.frequency_penalty,
-            seed: model_options?.seed,
-            thinkingConfig: model_options?.include_thoughts || model_options?.thinking_budget_tokens ?
-                {
-                    includeThoughts: model_options?.include_thoughts,
-                    thinkingBudget: model_options?.thinking_budget_tokens,
-                } : undefined,
-        }
+        config: options.model.toLowerCase().includes("image") ? configNanoBanana : config,
     };
 }
 
@@ -335,7 +347,7 @@ function cleanEmptyFieldsContent(content: Content, result_schema?: JSONSchema):
                 const jsonText = JSON.parse(part.text);
                 // Skip cleaning if not an object
                 if (typeof jsonText === 'object' && jsonText !== null && !Array.isArray(jsonText)) {
-                    const cleanedJson = removeEmptyFields(jsonText, result_schema);       
+                    const cleanedJson = removeEmptyFields(jsonText, result_schema);
                     newPart.text = JSON.stringify(cleanedJson);
                 } else {
                     // Keep original if not an object (string, number, array, etc.)
@@ -370,14 +382,14 @@ function removeEmptyFields(object: JSONObject | any[], schema: JSONSchema): JSON
     if (typeof object == 'object' || object === null) {
         return removeEmptyJSONObject(object, schema);
     }
-    
+
     return object;
 }
 
 function removeEmptyJSONObject(object: JSONObject, schema: JSONSchema): JSONObject {
     // Get the original required properties from schema
     const requiredProps = schema.required || [];
-    const cleanedResult: JSONObject = {...object};
+    const cleanedResult: JSONObject = { ...object };
 
     // Process each property
     for (const [key, value] of Object.entries(object)) {
@@ -401,24 +413,46 @@ function removeEmptyJSONObject(object: JSONObject, schema: JSONSchema): JSONObje
 
 function removeEmptyJSONArray(array: any[], schema: JSONSchema): any[] {
     const cleanedArray = array.map(item => {
-        return removeEmptyFields(item, schema); 
+        return removeEmptyFields(item, schema);
     });
 
     // Filter out empty objects from the array
     return cleanedArray.filter(item => !isEmpty(item));
 }
 
-function collectTextParts(content: Content) {
-    const out = [];
+function collectTextParts(content: Content): CompletionResult[] {
+    const results: CompletionResult[] = [];
     const parts = content.parts;
     if (parts) {
         for (const part of parts) {
             if (part.text) {
-                out.push(part.text);
+                results.push({
+                    type: "text",
+                    value: part.text
+                });
             }
         }
     }
-    return out.join('\n');
+    return results;
+}
+
+function collectInlineDataParts(content: Content): CompletionResult[] {
+    const results: CompletionResult[] = [];
+    const parts = content.parts;
+    if (parts) {
+        for (const part of parts) {
+            if (part.inlineData) {
+                const base64ImageBytes: string = part.inlineData.data ?? "";
+                const mimeType = part.inlineData.mimeType ?? "image/png";
+                const imageUrl = `data:${mimeType};base64,${base64ImageBytes}`;
+                results.push({
+                    type: "image",
+                    value: imageUrl
+                });
+            }
+        }
+    }
+    return results;
 }
 
 function collectToolUseParts(content: Content): ToolUse[] | undefined {
@@ -439,7 +473,7 @@ function collectToolUseParts(content: Content): ToolUse[] | undefined {
 export function mergeConsecutiveRole(contents: Content[] | undefined): Content[] {
     if (!contents || contents.length === 0) return [];
 
-    const needsMerging = contents.some((content, i) => 
+    const needsMerging = contents.some((content, i) =>
         i < contents.length - 1 && content.role === contents[i + 1].role
     );
     // If no merging needed, return original array
@@ -465,6 +499,41 @@ export function mergeConsecutiveRole(contents: Content[] | undefined): Content[]
     return result;
 }
 
+const supportedFinishReasons: FinishReason[] = [
+    FinishReason.MAX_TOKENS,
+    FinishReason.STOP,
+    FinishReason.FINISH_REASON_UNSPECIFIED
+]
+
+function geminiMaxTokens(option: StatelessExecutionOptions) {
+    const model_options = option.model_options as VertexAIGeminiOptions | undefined;
+    if (model_options?.max_tokens) {
+        return model_options.max_tokens;
+    }
+    if (option.model.includes("gemini-2.5")) {
+        const maxSupportedTokens = getMaxTokensLimitVertexAi(option.model);
+        const thinkingBudget = geminiThinkingBudget(option) ?? 0;
+        return Math.min(maxSupportedTokens, 16000 + thinkingBudget);
+    }
+    return undefined;
+}
+
+function geminiThinkingBudget(option: StatelessExecutionOptions) {
+    const model_options = option.model_options as VertexAIGeminiOptions | undefined;
+    if (model_options?.thinking_budget_tokens) {
+        return model_options.thinking_budget_tokens;
+    }
+    // Set minimum thinking level by default.
+    // Docs: https://ai.google.dev/gemini-api/docs/thinking#set-budget
+    if (option.model.includes("gemini-2.5")) {
+        if (option.model.includes("pro")) {
+            return 128;
+        }
+        return 0;
+    }
+    return undefined;
+}
+
 export class GeminiModelDefinition implements ModelDefinition<GenerateContentPrompt> {
 
     model: AIModel
@@ -485,8 +554,16 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
             return { result, options };
         }
         try {
-            const jsonResult = JSON.parse(result.result);
-            result.result = JSON.stringify(removeEmptyFields(jsonResult, options.result_schema));
+            // Extract text content for JSON processing - only process first text result
+            const textResult = result.result.find(r => r.type === 'text')?.value;
+            if (textResult) {
+                const jsonResult = JSON.parse(textResult);
+                const cleanedJson = JSON.stringify(removeEmptyFields(jsonResult, options.result_schema));
+                // Replace the text result with cleaned version
+                result.result = result.result.map(r =>
+                    r.type === 'text' ? { ...r, value: cleanedJson } : r
+                );
+            }
             return { result, options };
         } catch (error) {
             // Log error during processing but don't fail the completion
@@ -504,7 +581,7 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
         const schema = options.result_schema;
         let contents: Content[] = [];
         let system: Content | undefined = { role: "user", parts: [] }; // Single content block for system messages
-        
+
         const safety: Content[] = [];
 
         for (const msg of segments) {
@@ -517,7 +594,7 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
 
                 if (msg.content) {
                     system.parts?.push({
-                      text: msg.content
+                        text: msg.content
                     });
                 }
             } else if (msg.role === PromptRole.tool) {
@@ -591,7 +668,7 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
                 }
             }
         }
-        
+
         // If no system messages, set system to undefined.
         if (!system.parts || system.parts.length === 0) {
             system = undefined;
@@ -604,29 +681,58 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
 
         // Merge consecutive messages with the same role. Note: this may not be necessary, works without it, keeping to match previous behavior.
         contents = mergeConsecutiveRole(contents);
-        
+
         return { contents, system };
     }
 
+    usageMetadataToTokenUsage(usageMetadata: GenerateContentResponseUsageMetadata | undefined): ExecutionTokenUsage {
+        if (!usageMetadata || !usageMetadata.totalTokenCount) {
+            return {};
+        }
+        const tokenUsage: ExecutionTokenUsage = { total: usageMetadata.totalTokenCount, prompt: usageMetadata.promptTokenCount };
+
+        //Output/Response side
+        tokenUsage.result = (usageMetadata.candidatesTokenCount ?? 0)
+            + (usageMetadata.thoughtsTokenCount ?? 0)
+            + (usageMetadata.toolUsePromptTokenCount ?? 0);
+
+        if ((tokenUsage.total ?? 0) != (tokenUsage.prompt ?? 0) + tokenUsage.result) {
+            console.warn("[VertexAI] Gemini token usage mismatch: total does not equal prompt + result", {
+                total: tokenUsage.total,
+                prompt: tokenUsage.prompt,
+                result: tokenUsage.result
+            });
+        }
+
+        if (!tokenUsage.result) {
+            tokenUsage.result = undefined; // If no result, mark as undefined
+        }
+
+        return tokenUsage;
+    }
+
     async requestTextCompletion(driver: VertexAIDriver, prompt: GenerateContentPrompt, options: ExecutionOptions): Promise<Completion> {
         const splits = options.model.split("/");
+        let region: string | undefined = undefined;
+        if (splits[0] === "locations" && splits.length >= 2) {
+            region = splits[1];
+        } 
         const modelName = splits[splits.length - 1];
         options = { ...options, model: modelName };
 
         let conversation = updateConversation(options.conversation as Content[], prompt.contents);
         prompt.contents = conversation;
 
-        const client = driver.getGoogleGenAIClient();
+        if (options.model.includes("gemini-2.5-flash-image")) {
+            region = "global"; // Gemini Flash Image only available in global region, this is for nano-banana model
+        }
+
+        const client = driver.getGoogleGenAIClient(region);
 
         const payload = getGeminiPayload(options, prompt);
         const response = await client.models.generateContent(payload);
 
-        const usage = response.usageMetadata;
-        const token_usage: ExecutionTokenUsage = {
-            prompt: usage?.promptTokenCount,
-            result: usage?.candidatesTokenCount,
-            total: usage?.totalTokenCount,
-        }
+        const token_usage: ExecutionTokenUsage = this.usageMetadataToTokenUsage(response.usageMetadata);
 
         let tool_use: ToolUse[] | undefined;
         let finish_reason: string | undefined, result: any;
@@ -638,22 +744,33 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
                 default: finish_reason = candidate.finishReason;
             }
             const content = candidate.content;
+
+            if (candidate.finishReason && !supportedFinishReasons.includes(candidate.finishReason)) {
+                throw new Error(`Unsupported finish reason: ${candidate.finishReason}, `
+                    + `finish message: ${candidate.finishMessage}, `
+                    + `content: ${JSON.stringify(content, null, 2)}, safety: ${JSON.stringify(candidate.safetyRatings, null, 2)}`);
+            }
+
             if (content) {
                 tool_use = collectToolUseParts(content);
 
                 // We clean the content before validation, so we can update the conversation.
                 const cleanedContent = cleanEmptyFieldsContent(content, options.result_schema);
-                result = collectTextParts(cleanedContent);
+                const textResults = collectTextParts(cleanedContent);
+                const imageResults = collectInlineDataParts(cleanedContent);
+                result = [...textResults, ...imageResults];
                 conversation = updateConversation(conversation, [cleanedContent]);
             }
         }
 
+
+
         if (tool_use) {
             finish_reason = "tool_use";
         }
 
         return {
-            result: result ?? '',
+            result: result && result.length > 0 ? result : [{ type: "text" as const, value: '' }],
             token_usage: token_usage,
             finish_reason: finish_reason,
             original_response: options.include_original_response ? response : undefined,
@@ -664,21 +781,24 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
 
     async requestTextCompletionStream(driver: VertexAIDriver, prompt: GenerateContentPrompt, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunkObject>> {
         const splits = options.model.split("/");
+        let region: string | undefined = undefined;
+        if (splits[0] === "locations" && splits.length >= 2) {
+            region = splits[1];
+        } 
         const modelName = splits[splits.length - 1];
         options = { ...options, model: modelName };
 
-        const client = driver.getGoogleGenAIClient();
+        if (options.model.includes("gemini-2.5-flash-image")) {
+            region = "global"; // Gemini Flash Image only available in global region, this is for nano-banana model
+        }
+
+        const client = driver.getGoogleGenAIClient(region);
 
         const payload = getGeminiPayload(options, prompt);
         const response = await client.models.generateContentStream(payload);
 
         const stream = asyncMap(response, async (item) => {
-            const usage = item.usageMetadata;
-            const token_usage: ExecutionTokenUsage = {
-                prompt: usage?.promptTokenCount,
-                result: usage?.candidatesTokenCount,
-                total: usage?.totalTokenCount,
-            }
+            const token_usage: ExecutionTokenUsage = this.usageMetadataToTokenUsage(item.usageMetadata);
             if (item.candidates && item.candidates.length > 0) {
                 for (const candidate of item.candidates) {
                     let tool_use: ToolUse[] | undefined;
@@ -688,14 +808,21 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
                         case FinishReason.STOP: finish_reason = "stop"; break;
                         default: finish_reason = candidate.finishReason;
                     }
+                    if (candidate.finishReason && !supportedFinishReasons.includes(candidate.finishReason)) {
+                        throw new Error(`Unsupported finish reason: ${candidate.finishReason}, `
+                            + `finish message: ${candidate.finishMessage}, `
+                            + `content: ${JSON.stringify(candidate.content, null, 2)}, safety: ${JSON.stringify(candidate.safetyRatings, null, 2)}`);
+                    }
                     if (candidate.content?.role === 'model') {
-                        const text = collectTextParts(candidate.content);
+                        const textResults = collectTextParts(candidate.content);
+                        const imageResults = collectInlineDataParts(candidate.content);
+                        const combinedResults = [...textResults, ...imageResults];
                         tool_use = collectToolUseParts(candidate.content);
                         if (tool_use) {
                             finish_reason = "tool_use";
                         }
                         return {
-                            result: text,
+                            result: combinedResults.length > 0 ? combinedResults : [],
                             token_usage: token_usage,
                             finish_reason: finish_reason,
                             tool_use,
@@ -705,8 +832,9 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
             }
             //No normal output, returning block reason if it exists.
             return {
-                result: item.promptFeedback?.blockReasonMessage ?? "",
+                result: item.promptFeedback?.blockReasonMessage ? [{ type: "text" as const, value: item.promptFeedback.blockReasonMessage }] : [],
                 finish_reason: item.promptFeedback?.blockReason ?? "",
+                token_usage: token_usage,
             };
         });
 

package/src/vertexai/models/imagen.ts

@@ -1,19 +1,12 @@
 import {
-    AIModel, Completion, ExecutionOptions, ImageGeneration, Modalities,
+    AIModel, Completion, ExecutionOptions, Modalities,
     ModelType, PromptRole, PromptSegment, readStreamAsBase64, ImagenOptions
 } from "@llumiverse/core";
 import { VertexAIDriver } from "../index.js";
 
-const projectId = process.env.GOOGLE_PROJECT_ID;
-const location = 'us-central1';
-
-import aiplatform, { protos } from '@google-cloud/aiplatform';
-
-// Imports the Google Cloud Prediction Service Client library
-const { PredictionServiceClient } = aiplatform.v1;
-
 // Import the helper module for converting arbitrary protobuf.Value objects
-import { helpers } from '@google-cloud/aiplatform';
+import { protos, helpers } from '@google-cloud/aiplatform';
+
 interface ImagenBaseReference {
     referenceType: "REFERENCE_TYPE_RAW" | "REFERENCE_TYPE_MASK" | "REFERENCE_TYPE_SUBJECT" |
     "REFERENCE_TYPE_CONTROL" | "REFERENCE_TYPE_STYLE";
@@ -90,14 +83,6 @@ export interface ImagenPrompt {
     negativePrompt?: string; //Used for negative prompts
 }
 
-// Specifies the location of the api endpoint
-const clientOptions = {
-    apiEndpoint: `${location}-aiplatform.googleapis.com`,
-};
-
-// Instantiates a client
-const predictionServiceClient = new PredictionServiceClient(clientOptions);
-
 function getImagenParameters(taskType: string, options: ImagenOptions) {
     const commonParameters = {
         sampleCount: options?.number_of_images,
@@ -337,9 +322,9 @@ export class ImagenModelDefinition {
         return prompt
     }
 
-    async requestImageGeneration(driver: VertexAIDriver, prompt: ImagenPrompt, options: ExecutionOptions): Promise<Completion<ImageGeneration>> {
+    async requestImageGeneration(driver: VertexAIDriver, prompt: ImagenPrompt, options: ExecutionOptions): Promise<Completion> {
         if (options.model_options?._option_id !== "vertexai-imagen") {
-            driver.logger.warn("Invalid model options", {options: options.model_options });
+            driver.logger.warn("Invalid model options", { options: options.model_options });
         }
         options.model_options = options.model_options as ImagenOptions | undefined;
 
@@ -351,10 +336,11 @@ export class ImagenModelDefinition {
 
         driver.logger.info("Task type: " + taskType);
 
-        const modelName = options.model.split("/").pop() ?? '';
+            const modelName = options.model.split("/").pop() ?? '';
 
         // Configure the parent resource
-        const endpoint = `projects/${projectId}/locations/${location}/publishers/google/models/${modelName}`;
+        // TODO: make location configurable, fixed to us-central1 for now
+        const endpoint = `projects/${driver.options.project}/locations/us-central1/publishers/google/models/${modelName}`;
 
         const instanceValue = helpers.toValue(prompt);
         if (!instanceValue) {
@@ -380,8 +366,10 @@ export class ImagenModelDefinition {
             parameters,
         };
 
+        const client = driver.getImagenClient();
+
         // Predict request
-        const [response] = await predictionServiceClient.predict(request, { timeout: 120000 * numberOfImages }); //Extended timeout for image generation
+        const [response] = await client.predict(request, { timeout: 120000 * numberOfImages }); //Extended timeout for image generation
         const predictions = response.predictions;
 
         if (!predictions) {
@@ -394,9 +382,10 @@ export class ImagenModelDefinition {
         );
 
         return {
-            result: {
-                images
-            },
+            result: images.map(image => ({
+                type: "image" as const,
+                value: image
+            })),
         };
     }
 }

package/src/vertexai/models/llama.ts

@@ -1,5 +1,5 @@
 import {
-    AIModel, Completion, CompletionChunk, CompletionChunkObject, ExecutionOptions, ModelType,
+    AIModel, Completion, CompletionChunkObject, ExecutionOptions, ModelType,
     PromptOptions, PromptRole, PromptSegment,
     TextFallbackOptions
 } from "@llumiverse/core";
@@ -195,7 +195,7 @@ export class LLamaModelDefinition implements ModelDefinition<LLamaPrompt> {
         });
 
         return {
-            result: text,
+            result: [{ type: "text", value: text }],
             token_usage: {
                 prompt: result.usage.prompt_tokens,
                 result: result.usage.completion_tokens,
@@ -206,11 +206,11 @@ export class LLamaModelDefinition implements ModelDefinition<LLamaPrompt> {
         };
     }
 
-    async requestTextCompletionStream(driver: VertexAIDriver, prompt: LLamaPrompt, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunk>> {
+    async requestTextCompletionStream(driver: VertexAIDriver, prompt: LLamaPrompt, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunkObject>> {
         const splits = options.model.split("/");
         const modelName = splits[splits.length - 1];
 
-        let conversation = updateConversation(options.conversation as LLamaPrompt, prompt);
+        const conversation = updateConversation(options.conversation as LLamaPrompt, prompt);
 
         const modelOptions = options.model_options as TextFallbackOptions;
 
@@ -247,8 +247,9 @@ export class LLamaModelDefinition implements ModelDefinition<LLamaPrompt> {
         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: choice?.delta?.content ?? '',
+                result: content ? [{ type: "text", value: content }] : [],
                 finish_reason: choice?.finish_reason,
                 token_usage: json.usage ? {
                     prompt: json.usage.prompt_tokens,

package/src/vertexai/models.ts

@@ -1,5 +1,5 @@
-import { AIModel, Completion, PromptSegment, ExecutionOptions, CompletionChunk } from "@llumiverse/core";
-import { VertexAIDriver , trimModelName} from "./index.js";
+import { AIModel, Completion, PromptSegment, ExecutionOptions, CompletionChunkObject } from "@llumiverse/core";
+import { VertexAIDriver, trimModelName } from "./index.js";
 import { GeminiModelDefinition } from "./models/gemini.js";
 import { ClaudeModelDefinition } from "./models/claude.js";
 import { LLamaModelDefinition } from "./models/llama.js";
@@ -9,7 +9,7 @@ export interface ModelDefinition<PromptT = any> {
     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<CompletionChunk>>;
+    requestTextCompletionStream: (driver: VertexAIDriver, prompt: PromptT, options: ExecutionOptions) => Promise<AsyncIterable<CompletionChunkObject>>;
     preValidationProcessing?(result: Completion, options: ExecutionOptions): { result: Completion, options: ExecutionOptions };
 }
 
@@ -17,7 +17,7 @@ export function getModelDefinition(model: string): ModelDefinition {
     const splits = model.split("/");
     const publisher = splits[1];
     const modelName = trimModelName(splits[splits.length - 1]);
-    
+
     if (publisher?.includes("anthropic")) {
         return new ClaudeModelDefinition(modelName);
     } else if (publisher?.includes("google")) {

package/src/watsonx/index.ts

@@ -1,4 +1,4 @@
-import { AIModel, AbstractDriver, Completion, CompletionChunk, DriverOptions, EmbeddingsOptions, EmbeddingsResult, ExecutionOptions, TextFallbackOptions } from "@llumiverse/core";
+import { AIModel, AbstractDriver, 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";
@@ -29,12 +29,12 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
         this.fetchClient = new FetchClient(this.endpoint_url).withAuthCallback(async () => this.getAuthToken().then(token => `Bearer ${token}`));
     }
 
-    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("Invalid model options", { options: options.model_options });
         }
         options.model_options = options.model_options as TextFallbackOptions | undefined;
-        
+
         const payload: WatsonxTextGenerationPayload = {
             model_id: options.model,
             input: prompt + "\n",
@@ -53,7 +53,7 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
         const result = res.results[0];
 
         return {
-            result: result.generated_text,
+            result: [{ type: "text", value: result.generated_text }],
             token_usage: {
                 prompt: result.input_token_count,
                 result: result.generated_token_count,
@@ -64,9 +64,9 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, 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("Invalid model options", { options: options.model_options });
         }
         options.model_options = options.model_options as TextFallbackOptions | undefined;
         const payload: WatsonxTextGenerationPayload = {
@@ -90,7 +90,7 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
         return transformSSEStream(stream, (data: string) => {
             const json = JSON.parse(data) as WatsonxTextGenerationResponse;
             return {
-                result: 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,
@@ -132,7 +132,7 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
             if (now < this.authToken.expiration) {
                 return this.authToken.access_token;
             } else {
-                this.logger.debug("Token expired, refetching", this.authToken, now)
+                this.logger.debug("Token expired, refetching")
             }
         }
         const authToken = await fetch('https://iam.cloud.ibm.com/identity/token', {
@@ -163,7 +163,7 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
         }
 
         if (!options.text) {
-            throw new Error ("No text provided");
+            throw new Error("No text provided");
         }
 
         const payload: GenerateEmbeddingPayload = {