@llumiverse/drivers 0.23.0 → 0.24.0

package/src/vertexai/models/gemini.ts

@@ -5,12 +5,20 @@ import {
 } from "@google/genai";
 import {
     AIModel, Completion, CompletionChunkObject, CompletionResult, ExecutionOptions,
-    ExecutionTokenUsage, getMaxTokensLimitVertexAi, JSONObject, JSONSchema, ModelType, PromptOptions, PromptRole,
-    PromptSegment, readStreamAsBase64, StatelessExecutionOptions, ToolDefinition, ToolUse,
+    ExecutionTokenUsage,
+    getConversationMeta,
+    getMaxTokensLimitVertexAi,
+    incrementConversationTurn,
+    JSONObject, JSONSchema, ModelType, PromptOptions, PromptRole,
+    PromptSegment, readStreamAsBase64, StatelessExecutionOptions,
+    stripBase64ImagesFromConversation,
+    ToolDefinition, ToolUse,
+    truncateLargeTextInConversation,
+    unwrapConversationArray,
     VertexAIGeminiOptions
 } from "@llumiverse/core";
 import { asyncMap } from "@llumiverse/core/async";
-import { VertexAIDriver, GenerateContentPrompt } from "../index.js";
+import { GenerateContentPrompt, VertexAIDriver } from "../index.js";
 import { ModelDefinition } from "../models.js";
 
 function supportsStructuredOutput(options: PromptOptions): boolean {
@@ -467,11 +475,17 @@ function collectToolUseParts(content: Content): ToolUse[] | undefined {
     const parts = content.parts ?? [];
     for (const part of parts) {
         if (part.functionCall) {
-            out.push({
+            const toolUse: ToolUse = {
                 id: part.functionCall.name ?? '',
                 tool_name: part.functionCall.name ?? '',
                 tool_input: part.functionCall.args as JSONObject,
-            });
+            };
+            // Capture thought_signature for Gemini thinking models (2.5+/3.0+)
+            // This must be passed back with the function response
+            if (part.thoughtSignature) {
+                toolUse.thought_signature = part.thoughtSignature;
+            }
+            out.push(toolUse);
         }
     }
     return out.length > 0 ? out : undefined;
@@ -545,7 +559,7 @@ function geminiThinkingConfig(option: StatelessExecutionOptions): ThinkingConfig
     const model_options = option.model_options as VertexAIGeminiOptions | undefined;
     const include_thoughts = model_options?.include_thoughts ?? false;
     if (model_options?.thinking_budget_tokens) {
-        return {includeThoughts: include_thoughts, thinkingBudget: model_options.thinking_budget_tokens};
+        return { includeThoughts: include_thoughts, thinkingBudget: model_options.thinking_budget_tokens };
     }
 
     // Set minimum thinking level by default.
@@ -623,16 +637,18 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
                 if (!msg.tool_use_id) {
                     throw new Error("Tool response missing tool_use_id");
                 }
+                // Build functionResponse part with optional thought_signature for Gemini thinking models
+                const functionResponsePart: Part = {
+                    functionResponse: {
+                        name: msg.tool_use_id,
+                        response: formatFunctionResponse(msg.content || ''),
+                    },
+                    // Include thought_signature if provided (required for Gemini 2.5+/3.0+ thinking models)
+                    thoughtSignature: msg.thought_signature,
+                };
                 contents.push({
                     role: 'user',
-                    parts: [
-                        {
-                            functionResponse: {
-                                name: msg.tool_use_id,
-                                response: formatFunctionResponse(msg.content || ''),
-                            }
-                        }
-                    ]
+                    parts: [functionResponsePart]
                 });
             } else {    // PromptRole.user, PromptRole.assistant, PromptRole.safety
                 const parts: Part[] = [];
@@ -646,14 +662,27 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
                 // File content handling
                 if (msg.files) {
                     for (const f of msg.files) {
-                        const stream = await f.getStream();
-                        const data = await readStreamAsBase64(stream);
-                        parts.push({
-                            inlineData: {
-                                data,
-                                mimeType: f.mime_type
-                            }
-                        });
+                        let fileUrl = await f.getURL();                     
+                        const isGsUrl = fileUrl.startsWith('gs://') || fileUrl.startsWith('https://storage.googleapis.com/');
+
+                        if (isGsUrl) {
+                            parts.push({
+                                fileData: {
+                                    fileUri: fileUrl,
+                                    mimeType: f.mime_type
+                                }
+                            });
+                        } else {
+                            // Inline data handling
+                            const stream = await f.getStream();
+                            const data = await readStreamAsBase64(stream); 
+                            parts.push({
+                                inlineData: {
+                                    data,
+                                    mimeType: f.mime_type
+                                }
+                            });
+                        }
                     }
                 }
 
@@ -742,7 +771,7 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
         const modelName = splits[splits.length - 1];
         options = { ...options, model: modelName };
 
-        let conversation = updateConversation(options.conversation as Content[], prompt.contents);
+        let conversation = updateConversation(options.conversation, prompt.contents);
         prompt.contents = conversation;
 
         // TODO: Remove hack, use global endpoint manually if needed.
@@ -792,12 +821,27 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
             finish_reason = "tool_use";
         }
 
+        // Increment turn counter for deferred stripping
+        conversation = incrementConversationTurn(conversation) as Content[];
+
+        // Strip large base64 image data based on options.stripImagesAfterTurns
+        const currentTurn = getConversationMeta(conversation).turnNumber;
+        const stripOptions = {
+            keepForTurns: options.stripImagesAfterTurns ?? Infinity,
+            currentTurn,
+            textMaxTokens: options.stripTextMaxTokens
+        };
+        let processedConversation = stripBase64ImagesFromConversation(conversation, stripOptions);
+
+        // Truncate large text content if configured
+        processedConversation = truncateLargeTextInConversation(processedConversation, stripOptions);
+
         return {
             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,
-            conversation,
+            conversation: processedConversation,
             tool_use
         } satisfies Completion;
     }
@@ -811,6 +855,10 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
         const modelName = splits[splits.length - 1];
         options = { ...options, model: modelName };
 
+        // Include conversation history in prompt contents (same as non-streaming)
+        const conversation = updateConversation(options.conversation, prompt.contents);
+        prompt.contents = conversation;
+
         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
         }
@@ -897,16 +945,19 @@ function getToolFunction(tool: ToolDefinition): FunctionDeclaration {
     };
 }
 
-
 /**
  * Update the conversation messages
  * @param prompt
  * @param response
  * @returns
  */
-function updateConversation(conversation: Content[], prompt: Content[]): Content[] {
-    return (conversation || [] satisfies Content[]).concat(prompt);
+function updateConversation(conversation: unknown, prompt: Content[]): Content[] {
+    // Unwrap array if wrapped, otherwise treat as array
+    const unwrapped = unwrapConversationArray<Content>(conversation);
+    const convArray = unwrapped ?? (conversation as Content[] || []);
+    return convArray.concat(prompt);
 }
+
 /**
  *
  * Gemini supports JSON output in the response. so we test if the response is a valid JSON object. otherwise we treat the response as a string.

package/src/xai/index.ts

@@ -1,4 +1,4 @@
-import { AIModel, Completion, DriverOptions, ExecutionOptions, PromptOptions, PromptSegment } from "@llumiverse/core";
+import { AIModel, DriverOptions, PromptOptions, PromptSegment, Providers } from "@llumiverse/core";
 import { formatOpenAILikeMultimodalPrompt, OpenAIPromptFormatterOptions } from "../openai/openai_format.js";
 import { FetchClient } from "@vertesia/api-fetch-client";
 import OpenAI from "openai";
@@ -15,7 +15,7 @@ export interface xAiDriverOptions extends DriverOptions {
 export class xAIDriver extends BaseOpenAIDriver {
 
     service: OpenAI;
-    provider: "xai";
+    readonly provider = Providers.xai;
     xai_service: FetchClient;
     DEFAULT_ENDPOINT = "https://api.x.ai/v1";
 
@@ -31,7 +31,6 @@ export class xAIDriver extends BaseOpenAIDriver {
             baseURL: opts.endpoint ?? this.DEFAULT_ENDPOINT,
         });
         this.xai_service = new FetchClient(opts.endpoint ?? this.DEFAULT_ENDPOINT).withAuthCallback(async () => `Bearer ${opts.apiKey}`);
-        this.provider = "xai";
         //this.formatPrompt = this._formatPrompt; //TODO: fix xai prompt formatting
     }
 
@@ -49,17 +48,9 @@ export class xAIDriver extends BaseOpenAIDriver {
 
     }
 
-    extractDataFromResponse(_options: ExecutionOptions, result: OpenAI.Chat.Completions.ChatCompletion): Completion {
-        return {
-            result: result.choices[0].message.content ? [{ type: "text", value: result.choices[0].message.content }] : [],
-            finish_reason: result.choices[0].finish_reason,
-            token_usage: {
-                prompt: result.usage?.prompt_tokens,
-                result: result.usage?.completion_tokens,
-                total: result.usage?.total_tokens,
-            }
-        }
-    }
+    // Note: We intentionally do NOT override extractDataFromResponse here.
+    // The base class implementation properly handles tool_calls extraction.
+    // xAI's API is OpenAI-compatible and returns tool_calls in the same format.
 
     async listModels(): Promise<AIModel[]> {
         const [lm, em] = await Promise.all([
@@ -76,10 +67,12 @@ export class xAIDriver extends BaseOpenAIDriver {
             return {
                 id: model.id,
                 provider: this.provider,
-                name: model.object,
-                description: model.object,
+                name: model.id,
+                description: `${model.id} by ${model.owned_by}`,
                 is_multimodal: model.input_modalities.length > 1,
-                tags: [...model.input_modalities.map(m => `ì:${m}`), ...model.output_modalities.map(m => `ì:${m}`)],
+                input_modalities: model.input_modalities,
+                output_modalities: model.output_modalities,
+                tags: [...model.input_modalities.map(m => `i:${m}`), ...model.output_modalities.map(m => `o:${m}`)],
             } satisfies AIModel;
         });