@llumiverse/drivers 0.20.0 → 0.22.0

package/src/vertexai/models/claude.ts

@@ -72,11 +72,13 @@ function maxToken(option: StatelessExecutionOptions): number {
     if (modelOptions && typeof modelOptions.max_tokens === "number") {
         return modelOptions.max_tokens;
     } else {
+        const thinking_budget = modelOptions?.thinking_budget_tokens ?? 0;
+        let maxSupportedTokens = getMaxTokensLimitVertexAi(option.model);
         // Fallback to the default max tokens limit for the model
-        if (option.model.includes('claude-3-7-sonnet') && (modelOptions?.thinking_budget_tokens ?? 0) < 64000) {
-            return 64000; // Claude 3.7 can go up to 128k with a beta header, but when no max tokens is specified, we default to 64k.
+        if (option.model.includes('claude-3-7-sonnet') && (modelOptions?.thinking_budget_tokens ?? 0) < 48000) {
+            maxSupportedTokens = 64000; // Claude 3.7 can go up to 128k with a beta header, but when no max tokens is specified, we default to 64k.
         }
-        return getMaxTokensLimitVertexAi(option.model);
+        return Math.min(16000 + thinking_budget, maxSupportedTokens); // Cap to 16k, to avoid taking up too much context window and quota.
     }
 }
 

package/src/vertexai/models/gemini.ts

@@ -1,5 +1,6 @@
 import {
     Content, FinishReason, FunctionCallingConfigMode, FunctionDeclaration, GenerateContentParameters,
+    GenerateContentResponseUsageMetadata,
     HarmBlockThreshold, HarmCategory, Part, SafetySetting, Schema, Tool, Type
 } from "@google/genai";
 import {
@@ -465,6 +466,12 @@ export function mergeConsecutiveRole(contents: Content[] | undefined): Content[]
     return result;
 }
 
+const supportedFinishReasons: FinishReason[] = [
+    FinishReason.MAX_TOKENS,
+    FinishReason.STOP,
+    FinishReason.FINISH_REASON_UNSPECIFIED
+]
+
 export class GeminiModelDefinition implements ModelDefinition<GenerateContentPrompt> {
 
     model: AIModel
@@ -608,6 +615,32 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
         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("/");
         const modelName = splits[splits.length - 1];
@@ -621,12 +654,7 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
         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,6 +666,13 @@ 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);
 
@@ -648,6 +683,8 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
             }
         }
 
+        
+
         if (tool_use) {
             finish_reason = "tool_use";
         }
@@ -673,12 +710,7 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
         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,6 +720,11 @@ 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);
                         tool_use = collectToolUseParts(candidate.content);
@@ -707,6 +744,7 @@ export class GeminiModelDefinition implements ModelDefinition<GenerateContentPro
             return {
                 result: item.promptFeedback?.blockReasonMessage ?? "",
                 finish_reason: item.promptFeedback?.blockReason ?? "",
+                token_usage: token_usage,
             };
         });
 

package/src/watsonx/index.ts

@@ -31,10 +31,10 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
 
     async requestTextCompletion(prompt: string, options: ExecutionOptions): Promise<Completion<any>> {
         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",
@@ -66,7 +66,7 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
 
     async requestTextCompletionStream(prompt: string, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunk>> {
         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 = {
@@ -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 = {

package/src/xai/index.ts

@@ -1,5 +1,5 @@
 import { AIModel, Completion, DriverOptions, ExecutionOptions, PromptOptions, PromptSegment } from "@llumiverse/core";
-import { formatOpenAILikeMultimodalPrompt, OpenAIPromptFormatterOptions } from "@llumiverse/core/formatters";
+import { formatOpenAILikeMultimodalPrompt, OpenAIPromptFormatterOptions } from "../openai/openai_format.js";
 import { FetchClient } from "@vertesia/api-fetch-client";
 import OpenAI from "openai";
 import { BaseOpenAIDriver } from "../openai/index.js";
@@ -14,7 +14,6 @@ export interface xAiDriverOptions extends DriverOptions {
 
 export class xAIDriver extends BaseOpenAIDriver {
 
-
     service: OpenAI;
     provider: "xai";
     xai_service: FetchClient;
@@ -33,7 +32,7 @@ export class xAIDriver extends BaseOpenAIDriver {
           });
         this.xai_service = new FetchClient(opts.endpoint ?? this.DEFAULT_ENDPOINT ).withAuthCallback(async () => `Bearer ${opts.apiKey}`);
         this.provider = "xai";
-        this.formatPrompt = this._formatPrompt;
+        //this.formatPrompt = this._formatPrompt; //TODO: fix xai prompt formatting
     }
 
     async _formatPrompt(segments: PromptSegment[], opts: PromptOptions): Promise<OpenAI.Chat.Completions.ChatCompletionMessageParam[]> {

package/lib/cjs/openai/azure.js

@@ -1,31 +0,0 @@
-"use strict";
-Object.defineProperty(exports, "__esModule", { value: true });
-exports.AzureOpenAIDriver = void 0;
-const identity_1 = require("@azure/identity");
-const openai_1 = require("openai");
-const index_js_1 = require("./index.js");
-class AzureOpenAIDriver extends index_js_1.BaseOpenAIDriver {
-    service;
-    provider;
-    constructor(opts) {
-        super(opts);
-        if (!opts.azureADTokenProvider && !opts.apiKey) {
-            opts.azureADTokenProvider = this.getDefaultAuth();
-        }
-        this.service = new openai_1.AzureOpenAI({
-            apiKey: opts.apiKey,
-            azureADTokenProvider: opts.azureADTokenProvider,
-            endpoint: opts.endpoint,
-            apiVersion: opts.apiVersion ?? "2024-10-21",
-            deployment: opts.deployment
-        });
-        this.provider = "azure_openai";
-    }
-    getDefaultAuth() {
-        const scope = "https://cognitiveservices.azure.com/.default";
-        const azureADTokenProvider = (0, identity_1.getBearerTokenProvider)(new identity_1.DefaultAzureCredential(), scope);
-        return azureADTokenProvider;
-    }
-}
-exports.AzureOpenAIDriver = AzureOpenAIDriver;
-//# sourceMappingURL=azure.js.map

package/lib/cjs/openai/azure.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"azure.js","sourceRoot":"","sources":["../../../src/openai/azure.ts"],"names":[],"mappings":";;;AAAA,8CAAiF;AAEjF,mCAAqC;AACrC,yCAA8C;AAmB9C,MAAa,iBAAkB,SAAQ,2BAAgB;IAGnD,OAAO,CAAc;IACrB,QAAQ,CAAiB;IAEzB,YAAY,IAA8B;QACtC,KAAK,CAAC,IAAI,CAAC,CAAC;QAEZ,IAAI,CAAC,IAAI,CAAC,oBAAoB,IAAI,CAAC,IAAI,CAAC,MAAM,EAAE,CAAC;YAC7C,IAAI,CAAC,oBAAoB,GAAG,IAAI,CAAC,cAAc,EAAE,CAAC;QACtD,CAAC;QAED,IAAI,CAAC,OAAO,GAAG,IAAI,oBAAW,CAAC;YAC3B,MAAM,EAAE,IAAI,CAAC,MAAM;YACnB,oBAAoB,EAAE,IAAI,CAAC,oBAAoB;YAC/C,QAAQ,EAAE,IAAI,CAAC,QAAQ;YACvB,UAAU,EAAE,IAAI,CAAC,UAAU,IAAI,YAAY;YAC3C,UAAU,EAAE,IAAI,CAAC,UAAU;SAC9B,CAAC,CAAC;QACH,IAAI,CAAC,QAAQ,GAAG,cAAc,CAAC;IACnC,CAAC;IAGD,cAAc;QACV,MAAM,KAAK,GAAG,8CAA8C,CAAC;QAC7D,MAAM,oBAAoB,GAAG,IAAA,iCAAsB,EAAC,IAAI,iCAAsB,EAAE,EAAE,KAAK,CAAC,CAAC;QACzF,OAAO,oBAAoB,CAAC;IAChC,CAAC;CAGJ;AA/BD,8CA+BC"}

package/lib/esm/openai/azure.js

@@ -1,27 +0,0 @@
-import { DefaultAzureCredential, getBearerTokenProvider } from "@azure/identity";
-import { AzureOpenAI } from "openai";
-import { BaseOpenAIDriver } from "./index.js";
-export class AzureOpenAIDriver extends BaseOpenAIDriver {
-    service;
-    provider;
-    constructor(opts) {
-        super(opts);
-        if (!opts.azureADTokenProvider && !opts.apiKey) {
-            opts.azureADTokenProvider = this.getDefaultAuth();
-        }
-        this.service = new AzureOpenAI({
-            apiKey: opts.apiKey,
-            azureADTokenProvider: opts.azureADTokenProvider,
-            endpoint: opts.endpoint,
-            apiVersion: opts.apiVersion ?? "2024-10-21",
-            deployment: opts.deployment
-        });
-        this.provider = "azure_openai";
-    }
-    getDefaultAuth() {
-        const scope = "https://cognitiveservices.azure.com/.default";
-        const azureADTokenProvider = getBearerTokenProvider(new DefaultAzureCredential(), scope);
-        return azureADTokenProvider;
-    }
-}
-//# sourceMappingURL=azure.js.map

package/lib/esm/openai/azure.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"azure.js","sourceRoot":"","sources":["../../../src/openai/azure.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,sBAAsB,EAAE,sBAAsB,EAAE,MAAM,iBAAiB,CAAC;AAEjF,OAAO,EAAE,WAAW,EAAE,MAAM,QAAQ,CAAC;AACrC,OAAO,EAAE,gBAAgB,EAAE,MAAM,YAAY,CAAC;AAmB9C,MAAM,OAAO,iBAAkB,SAAQ,gBAAgB;IAGnD,OAAO,CAAc;IACrB,QAAQ,CAAiB;IAEzB,YAAY,IAA8B;QACtC,KAAK,CAAC,IAAI,CAAC,CAAC;QAEZ,IAAI,CAAC,IAAI,CAAC,oBAAoB,IAAI,CAAC,IAAI,CAAC,MAAM,EAAE,CAAC;YAC7C,IAAI,CAAC,oBAAoB,GAAG,IAAI,CAAC,cAAc,EAAE,CAAC;QACtD,CAAC;QAED,IAAI,CAAC,OAAO,GAAG,IAAI,WAAW,CAAC;YAC3B,MAAM,EAAE,IAAI,CAAC,MAAM;YACnB,oBAAoB,EAAE,IAAI,CAAC,oBAAoB;YAC/C,QAAQ,EAAE,IAAI,CAAC,QAAQ;YACvB,UAAU,EAAE,IAAI,CAAC,UAAU,IAAI,YAAY;YAC3C,UAAU,EAAE,IAAI,CAAC,UAAU;SAC9B,CAAC,CAAC;QACH,IAAI,CAAC,QAAQ,GAAG,cAAc,CAAC;IACnC,CAAC;IAGD,cAAc;QACV,MAAM,KAAK,GAAG,8CAA8C,CAAC;QAC7D,MAAM,oBAAoB,GAAG,sBAAsB,CAAC,IAAI,sBAAsB,EAAE,EAAE,KAAK,CAAC,CAAC;QACzF,OAAO,oBAAoB,CAAC;IAChC,CAAC;CAGJ"}

package/lib/types/openai/azure.d.ts

@@ -1,20 +0,0 @@
-import { DriverOptions } from "@llumiverse/core";
-import { AzureOpenAI } from "openai";
-import { BaseOpenAIDriver } from "./index.js";
-export interface AzureOpenAIDriverOptions extends DriverOptions {
-    /**
-     * The credentials to use to access Azure OpenAI
-     */
-    azureADTokenProvider?: any;
-    apiKey?: string;
-    endpoint?: string;
-    apiVersion?: string;
-    deployment?: string;
-}
-export declare class AzureOpenAIDriver extends BaseOpenAIDriver {
-    service: AzureOpenAI;
-    provider: "azure_openai";
-    constructor(opts: AzureOpenAIDriverOptions);
-    getDefaultAuth(): () => Promise<string>;
-}
-//# sourceMappingURL=azure.d.ts.map

package/lib/types/openai/azure.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"azure.d.ts","sourceRoot":"","sources":["../../../src/openai/azure.ts"],"names":[],"mappings":"AACA,OAAO,EAAE,aAAa,EAAE,MAAM,kBAAkB,CAAC;AACjD,OAAO,EAAE,WAAW,EAAE,MAAM,QAAQ,CAAC;AACrC,OAAO,EAAE,gBAAgB,EAAE,MAAM,YAAY,CAAC;AAE9C,MAAM,WAAW,wBAAyB,SAAQ,aAAa;IAE3D;;OAEG;IACH,oBAAoB,CAAC,EAAE,GAAG,CAAC;IAE3B,MAAM,CAAC,EAAE,MAAM,CAAC;IAEhB,QAAQ,CAAC,EAAE,MAAM,CAAC;IAElB,UAAU,CAAC,EAAE,MAAM,CAAA;IAEnB,UAAU,CAAC,EAAE,MAAM,CAAC;CAEvB;AAED,qBAAa,iBAAkB,SAAQ,gBAAgB;IAGnD,OAAO,EAAE,WAAW,CAAC;IACrB,QAAQ,EAAE,cAAc,CAAC;gBAEb,IAAI,EAAE,wBAAwB;IAkB1C,cAAc;CAOjB"}

package/src/openai/azure.ts

@@ -1,54 +0,0 @@
-import { DefaultAzureCredential, getBearerTokenProvider } from "@azure/identity";
-import { DriverOptions } from "@llumiverse/core";
-import { AzureOpenAI } from "openai";
-import { BaseOpenAIDriver } from "./index.js";
-
-export interface AzureOpenAIDriverOptions extends DriverOptions {
-
-    /**
-     * The credentials to use to access Azure OpenAI
-     */
-    azureADTokenProvider?: any; //type with azure credentials
-    
-    apiKey?: string;
-
-    endpoint?: string;
-
-    apiVersion?: string
-
-    deployment?: string;
-
-}
-
-export class AzureOpenAIDriver extends BaseOpenAIDriver {
-
-
-    service: AzureOpenAI;
-    provider: "azure_openai";
-
-    constructor(opts: AzureOpenAIDriverOptions) {
-        super(opts);
-
-        if (!opts.azureADTokenProvider && !opts.apiKey) {
-            opts.azureADTokenProvider = this.getDefaultAuth();
-        }
-
-        this.service = new AzureOpenAI({
-            apiKey: opts.apiKey,
-            azureADTokenProvider: opts.azureADTokenProvider,          
-            endpoint: opts.endpoint,
-            apiVersion: opts.apiVersion ?? "2024-10-21",
-            deployment: opts.deployment
-        });
-        this.provider = "azure_openai";
-    }
-
-
-    getDefaultAuth() {
-        const scope = "https://cognitiveservices.azure.com/.default";
-        const azureADTokenProvider = getBearerTokenProvider(new DefaultAzureCredential(), scope);
-        return azureADTokenProvider;
-    }      
-
-
-}