@llumiverse/drivers 0.22.0 → 0.23.0-dev-20251118

package/src/vertexai/models/claude.ts

@@ -9,6 +9,17 @@ import { ModelDefinition } from "../models.js";
 import { MessageCreateParamsBase, MessageCreateParamsNonStreaming, RawMessageStreamEvent } from "@anthropic-ai/sdk/resources/messages.js";
 import { MessageStreamParams } from "@anthropic-ai/sdk/resources/index.mjs";
 
+export const ANTHROPIC_REGIONS: Record<string, string> = {
+    us: "us-east5",
+    europe: "europe-west1",
+    global: "global",
+}
+
+export const NON_GLOBAL_ANTHROPIC_MODELS = [
+    "claude-3-5",
+    "claude-3",
+];
+
 interface ClaudePrompt {
     messages: MessageParam[];
     system?: TextBlockParam[];
@@ -87,7 +98,7 @@ async function collectFileBlocks(segment: PromptSegment, restrictedTypes: true):
 async function collectFileBlocks(segment: PromptSegment, restrictedTypes?: false): Promise<ContentBlockParam[]>;
 async function collectFileBlocks(segment: PromptSegment, restrictedTypes: boolean = false): Promise<ContentBlockParam[]> {
     const contentBlocks: ContentBlockParam[] = [];
-    
+
     for (const file of segment.files || []) {
         if (file.mime_type?.startsWith("image/")) {
             const allowedTypes = ["image/png", "image/jpeg", "image/gif", "image/webp"];
@@ -128,7 +139,7 @@ async function collectFileBlocks(segment: PromptSegment, restrictedTypes: boolea
             }
         }
     }
-    
+
     return contentBlocks;
 }
 
@@ -191,7 +202,7 @@ export class ClaudeModelDefinition implements ModelDefinition<ClaudePrompt> {
                         text: segment.content
                     } satisfies TextBlockParam);
                 }
-                
+
                 // Collect file blocks with type safety
                 const fileBlocks = await collectFileBlocks(segment, true);
                 contentBlocks.push(...fileBlocks);
@@ -208,7 +219,7 @@ export class ClaudeModelDefinition implements ModelDefinition<ClaudePrompt> {
             } else {
                 // Build content blocks for regular messages (all types allowed)
                 const contentBlocks: ContentBlockParam[] = [];
-                
+
                 if (segment.content) {
                     contentBlocks.push({
                         type: 'text',
@@ -250,11 +261,19 @@ export class ClaudeModelDefinition implements ModelDefinition<ClaudePrompt> {
     }
 
     async requestTextCompletion(driver: VertexAIDriver, prompt: ClaudePrompt, options: ExecutionOptions): Promise<Completion> {
-        const client = driver.getAnthropicClient();
+        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 = await driver.getAnthropicClient(region);
         options.model_options = options.model_options as VertexAIClaudeOptions;
 
         if (options.model_options?._option_id !== "vertexai-claude") {
-            driver.logger.warn("Invalid model options", { options: options.model_options });
+            driver.logger.warn({ options: options.model_options }, "Invalid model options");
         }
 
         let conversation = updateConversation(options.conversation as ClaudePrompt, prompt);
@@ -273,7 +292,7 @@ export class ClaudeModelDefinition implements ModelDefinition<ClaudePrompt> {
         conversation = updateConversation(conversation, createPromptFromResponse(result));
 
         return {
-            result: text ?? '',
+            result: text ? [{ type: "text", value: text }] : [{ type: "text", value: '' }],
             tool_use,
             token_usage: {
                 prompt: result.usage.input_tokens,
@@ -287,11 +306,19 @@ export class ClaudeModelDefinition implements ModelDefinition<ClaudePrompt> {
     }
 
     async requestTextCompletionStream(driver: VertexAIDriver, prompt: ClaudePrompt, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunkObject>> {
-        const client = driver.getAnthropicClient();
+        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 = await driver.getAnthropicClient(region);
         const model_options = options.model_options as VertexAIClaudeOptions | undefined;
 
         if (model_options?._option_id !== "vertexai-claude") {
-            driver.logger.warn("Invalid model options", { options: options.model_options });
+            driver.logger.warn({ options: options.model_options }, "Invalid model options");
         }
 
         const { payload, requestOptions } = getClaudePayload(options, prompt);
@@ -303,7 +330,7 @@ export class ClaudeModelDefinition implements ModelDefinition<ClaudePrompt> {
             switch (streamEvent.type) {
                 case "message_start":
                     return {
-                        result: '',
+                        result: [{ type: "text", value: '' }],
                         token_usage: {
                             prompt: streamEvent.message.usage.input_tokens,
                             result: streamEvent.message.usage.output_tokens
@@ -311,7 +338,7 @@ export class ClaudeModelDefinition implements ModelDefinition<ClaudePrompt> {
                     } satisfies CompletionChunkObject;
                 case "message_delta":
                     return {
-                        result: '',
+                        result: [{ type: "text", value: '' }],
                         token_usage: {
                             result: streamEvent.usage.output_tokens
                         },
@@ -321,7 +348,7 @@ export class ClaudeModelDefinition implements ModelDefinition<ClaudePrompt> {
                     // Handle redacted thinking blocks
                     if (streamEvent.content_block.type === "redacted_thinking" && model_options?.include_thoughts) {
                         return {
-                            result: `[Redacted thinking: ${streamEvent.content_block.data}]`
+                            result: [{ type: "text", value: `[Redacted thinking: ${streamEvent.content_block.data}]` }]
                         } satisfies CompletionChunkObject;
                     }
                     break;
@@ -330,12 +357,12 @@ export class ClaudeModelDefinition implements ModelDefinition<ClaudePrompt> {
                     switch (streamEvent.delta.type) {
                         case "text_delta":
                             return {
-                                result: streamEvent.delta.text ?? ''
+                                result: streamEvent.delta.text ? [{ type: "text", value: streamEvent.delta.text }] : []
                             } satisfies CompletionChunkObject;
                         case "thinking_delta":
                             if (model_options?.include_thoughts) {
                                 return {
-                                    result: streamEvent.delta.thinking ?? '',
+                                    result: streamEvent.delta.thinking ? [{ type: "text", value: streamEvent.delta.thinking }] : [],
                                 } satisfies CompletionChunkObject;
                             }
                             break;
@@ -343,7 +370,7 @@ export class ClaudeModelDefinition implements ModelDefinition<ClaudePrompt> {
                             // Signature deltas, signify the end of the thoughts.
                             if (model_options?.include_thoughts) {
                                 return {
-                                    result: '\n\n', // Double newline for more spacing
+                                    result: [{ type: "text", value: '\n\n' }], // Double newline for more spacing
                                 } satisfies CompletionChunkObject;
                             }
                             break;
@@ -353,7 +380,7 @@ export class ClaudeModelDefinition implements ModelDefinition<ClaudePrompt> {
                     // Handle the end of content blocks, for redacted thinking blocks
                     if (model_options?.include_thoughts) {
                         return {
-                            result: '\n\n' // Add double newline for spacing
+                            result: [{ type: "text", value: '\n\n' }] // Add double newline for spacing
                         } satisfies CompletionChunkObject;
                     }
                     break;
@@ -361,7 +388,7 @@ export class ClaudeModelDefinition implements ModelDefinition<ClaudePrompt> {
 
             // Default case for all other event types
             return {
-                result: ''
+                result: []
             } satisfies CompletionChunkObject;
         });
 
@@ -399,8 +426,7 @@ interface RequestOptions {
 }
 
 function getClaudePayload(options: ExecutionOptions, prompt: ClaudePrompt): { payload: MessageCreateParamsBase, requestOptions: RequestOptions | undefined } {
-    const splits = options.model.split("/");
-    const modelName = splits[splits.length - 1];
+    const modelName = options.model; // Model name is already extracted in the calling methods
     const model_options = options.model_options as VertexAIClaudeOptions;
 
     // Add beta header for Claude 3.7 models to enable 128k output tokens

package/src/vertexai/models/gemini.ts

@@ -1,12 +1,13 @@
 import {
-    Content, FinishReason, FunctionCallingConfigMode, FunctionDeclaration, GenerateContentParameters,
+    Content, FinishReason, FunctionCallingConfigMode, FunctionDeclaration, GenerateContentConfig, GenerateContentParameters,
     GenerateContentResponseUsageMetadata,
-    HarmBlockThreshold, HarmCategory, Part, SafetySetting, Schema, Tool, Type
+    HarmBlockThreshold, HarmCategory, Modality, 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";
@@ -45,42 +46,63 @@ 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 = {
+        systemInstruction: prompt.system,
+        safetySettings: geminiSafetySettings,
+        responseModalities: [Modality.TEXT, Modality.IMAGE], // This is an error if only Text, and Only Image just gets blank responses.
+        candidateCount: 1,
+        //Model options
+        temperature: model_options?.temperature,
+        topP: model_options?.top_p,
+        maxOutputTokens: geminiMaxTokens(options),
+        stopSequences: model_options?.stop_sequence,
+        imageConfig: {
+            aspectRatio: model_options?.image_aspect_ratio,
+        }
+    }
+
+    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,
     };
 }
 
@@ -336,7 +358,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.)
@@ -371,14 +393,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)) {
@@ -402,24 +424,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 {
@@ -440,7 +484,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
@@ -469,9 +513,38 @@ export function mergeConsecutiveRole(contents: Content[] | undefined): Content[]
 const supportedFinishReasons: FinishReason[] = [
     FinishReason.MAX_TOKENS,
     FinishReason.STOP,
-    FinishReason.FINISH_REASON_UNSPECIFIED
+    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
@@ -492,8 +565,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
@@ -511,7 +592,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) {
@@ -524,7 +605,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) {
@@ -598,7 +679,7 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
                 }
             }
         }
-        
+
         // If no system messages, set system to undefined.
         if (!system.parts || system.parts.length === 0) {
             system = undefined;
@@ -611,7 +692,7 @@ 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 };
     }
 
@@ -625,7 +706,7 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
         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,
@@ -633,7 +714,7 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
                 result: tokenUsage.result
             });
         }
-        
+
         if (!tokenUsage.result) {
             tokenUsage.result = undefined; // If no result, mark as undefined
         }
@@ -643,13 +724,21 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
 
     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);
@@ -678,19 +767,21 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
 
                 // 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,
@@ -701,10 +792,18 @@ 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);
@@ -726,13 +825,15 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
                             + `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,
@@ -742,7 +843,7 @@ 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({ options: options.model_options }, "Invalid model options");
         }
         options.model_options = options.model_options as ImagenOptions | undefined;
 
@@ -354,7 +339,8 @@ export class ImagenModelDefinition {
         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) {
@@ -362,7 +348,7 @@ export class ImagenModelDefinition {
         }
         const instances = [instanceValue];
 
-        let parameter: any = getImagenParameters(taskType, options.model_options ?? {_option_id: "vertexai-imagen"});
+        let parameter: any = getImagenParameters(taskType, options.model_options ?? { _option_id: "vertexai-imagen" });
         parameter.negativePrompt = prompt.negativePrompt ?? undefined;
 
         const numberOfImages = options.model_options?.number_of_images ?? 1;
@@ -380,8 +366,10 @@ export class ImagenModelDefinition {
             parameters,
         };
 
+        const client = await 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
+            })),
         };
     }
 }