@llumiverse/drivers 0.20.0 → 0.21.0

package/src/openai/azure_openai.ts

@@ -0,0 +1,92 @@
+import { DefaultAzureCredential, getBearerTokenProvider } from "@azure/identity";
+import { AIModel, DriverOptions, getModelCapabilities, modelModalitiesToArray, Providers } from "@llumiverse/core";
+import OpenAI, { 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;
+    readonly provider = Providers.azure_openai;
+
+    //Overload to allow independent instantiation with AzureOpenAI service
+    constructor(serviceOrOpts: AzureOpenAI | AzureOpenAIDriverOptions) {
+        if (serviceOrOpts instanceof AzureOpenAI) {
+            super({});
+            this.service = serviceOrOpts;
+            return;
+        }
+        const opts = serviceOrOpts ?? {};
+        super(opts);
+        if (!opts.azureADTokenProvider && !opts.apiKey) {
+            opts.azureADTokenProvider = this.getDefaultCognitiveServicesAuth();
+        }
+
+        this.service = new AzureOpenAI({
+            apiKey: opts.apiKey,
+            azureADTokenProvider: opts.azureADTokenProvider,          
+            endpoint: opts.endpoint,
+            apiVersion: opts.apiVersion ?? "2024-10-21",
+            deployment: opts.deployment
+        });
+    }
+
+    /**
+     * Get default authentication for Azure Cognitive Services API
+     */
+    getDefaultCognitiveServicesAuth() {
+        const scope = "https://cognitiveservices.azure.com/.default";
+        const azureADTokenProvider = getBearerTokenProvider(new DefaultAzureCredential(), scope);
+        return azureADTokenProvider;
+    }
+    
+    async listModels(): Promise<AIModel[]> {
+        return this._listModels();
+    }
+
+    async _listModels(_filter?: (m: OpenAI.Models.Model) => boolean): Promise<AIModel[]> {
+        if (!this.service.deploymentName) {
+            throw new Error("A specific deployment is not set. Azure OpenAI cannot list deployments. Update your endpoint URL to include the deployment name, e.g., https://your-resource.openai.azure.com/openai/deployments/your-deployment/chat/completions");
+        }
+            
+        //Do a test execution to check if the model works and to get the model ID.
+        let modelID = this.service.deploymentName;
+        try {
+            const testResponse = await this.service.chat.completions.create({
+                model: this.service.deploymentName,
+                messages: [{ role: "user", content: "Hi" }],
+                max_tokens: 1,
+            });
+            modelID = testResponse.model;
+        } catch (error) {
+            this.logger.error("Failed to test model for Azure OpenAI listing :", error);
+        }
+        const modelCapability = getModelCapabilities(modelID, "openai");
+        return [{
+            id: modelID,
+            name: this.service.deploymentName,
+            provider: this.provider,
+            owner: "openai",
+            input_modalities: modelModalitiesToArray(modelCapability.input),
+            output_modalities: modelModalitiesToArray(modelCapability.output),
+            tool_support: modelCapability.tool_support,
+        } satisfies AIModel<string>];
+    }
+}

package/src/openai/index.ts

@@ -11,6 +11,7 @@ import {
     ExecutionTokenUsage,
     JSONSchema,
     ModelType,
+    Providers,
     ToolDefinition,
     ToolUse,
     TrainingJob,
@@ -22,13 +23,11 @@ import {
     supportsToolUse,
 } from "@llumiverse/core";
 import { asyncMap } from "@llumiverse/core/async";
-import { formatOpenAILikeMultimodalPrompt } from "@llumiverse/core/formatters";
+import { formatOpenAILikeMultimodalPrompt } from "./openai_format.js";
 import OpenAI, { AzureOpenAI } from "openai";
+import { ChatCompletionMessageParam } from "openai/resources/chat/completions";
 import { Stream } from "openai/streaming";
 
-//For code readability
-type OpenAIMessageBlock = OpenAI.Chat.Completions.ChatCompletionMessageParam;
-
 //TODO: Do we need a list?, replace with if statements and modernise?
 const supportFineTunning = new Set([
     "gpt-3.5-turbo-1106",
@@ -43,14 +42,15 @@ export interface BaseOpenAIDriverOptions extends DriverOptions {
 
 export abstract class BaseOpenAIDriver extends AbstractDriver<
     BaseOpenAIDriverOptions,
-    OpenAIMessageBlock[]
+    ChatCompletionMessageParam[]
 > {
-    abstract provider: "azure_openai" | "openai" | "xai";
+    //abstract provider: "azure_openai" | "openai" | "xai" | "azure_foundry";
+    abstract provider: Providers.openai | Providers.azure_openai | "xai" | Providers.azure_foundry;
     abstract service: OpenAI | AzureOpenAI;
 
     constructor(opts: BaseOpenAIDriverOptions) {
         super(opts);
-        this.formatPrompt = formatOpenAILikeMultimodalPrompt as any
+        this.formatPrompt = formatOpenAILikeMultimodalPrompt
         //TODO: better type, we send back OpenAI.Chat.Completions.ChatCompletionMessageParam[] but just not compatible with Function call that we don't use here
     }
 
@@ -82,7 +82,7 @@ export abstract class BaseOpenAIDriver extends AbstractDriver<
         };
     }
 
-    async requestTextCompletionStream(prompt: OpenAIMessageBlock[], options: ExecutionOptions): Promise<any> {
+    async requestTextCompletionStream(prompt: ChatCompletionMessageParam[], options: ExecutionOptions): Promise<AsyncIterable<Completion>> {
         if (options.model_options?._option_id !== "openai-text" && options.model_options?._option_id !== "openai-thinking") {
             this.logger.warn("Invalid model options", { options: options.model_options });
         }
@@ -157,7 +157,7 @@ export abstract class BaseOpenAIDriver extends AbstractDriver<
         return asyncMap(stream, mapFn);
     }
 
-    async requestTextCompletion(prompt: OpenAIMessageBlock[], options: ExecutionOptions): Promise<any> {
+    async requestTextCompletion(prompt: ChatCompletionMessageParam[], options: ExecutionOptions): Promise<Completion> {
         if (options.model_options?._option_id !== "openai-text" && options.model_options?._option_id !== "openai-thinking") {
             this.logger.warn("Invalid model options", { options: options.model_options });
         }
@@ -170,7 +170,7 @@ export abstract class BaseOpenAIDriver extends AbstractDriver<
         const toolDefs = getToolDefinitions(options.tools);
         const useTools: boolean = toolDefs ? supportsToolUse(options.model, "openai") : false;
 
-        let conversation = updateConversation(options.conversation as OpenAIMessageBlock[], prompt);
+        let conversation = updateConversation(options.conversation as ChatCompletionMessageParam[], prompt);
 
         let parsedSchema: JSONSchema | undefined = undefined;
         let strictMode = false;
@@ -293,12 +293,6 @@ export abstract class BaseOpenAIDriver extends AbstractDriver<
         const wordBlacklist = ["embed", "whisper", "transcribe", "audio", "moderation", "tts",
             "realtime", "dall-e", "babbage", "davinci", "codex", "o1-pro"];
 
-        if (this.provider === "azure_openai") {
-            //Azure OpenAI has additional information about the models
-            result = result.filter((m) => {
-                return !(m as any)?.capabilities?.embeddings;
-            });
-        }
 
         //OpenAI has very little information, filtering based on name.
         result = result.filter((m) => {
@@ -377,7 +371,7 @@ function jobInfo(job: OpenAI.FineTuning.Jobs.FineTuningJob): TrainingJob {
     }
 }
 
-function insert_image_detail(messages: OpenAIMessageBlock[], detail_level: string): OpenAIMessageBlock[] {
+function insert_image_detail(messages: ChatCompletionMessageParam[], detail_level: string): ChatCompletionMessageParam[] {
     if (detail_level == "auto" || detail_level == "low" || detail_level == "high") {
         for (const message of messages) {
             if (message.role !== 'assistant' && message.content) {
@@ -395,7 +389,7 @@ function insert_image_detail(messages: OpenAIMessageBlock[], detail_level: strin
     return messages;
 }
 
-function convertRoles(messages: OpenAIMessageBlock[], model: string): OpenAIMessageBlock[] {
+function convertRoles(messages: ChatCompletionMessageParam[], model: string): ChatCompletionMessageParam[] {
     //New openai models use developer role instead of system
     if (model.includes("o1") || model.includes("o3")) {
         if (model.includes("o1-mini") || model.includes("o1-preview")) {
@@ -462,7 +456,7 @@ function openAiFinishReason(finish_reason?: string): string | undefined {
     return finish_reason;
 }
 
-function updateConversation(conversation: OpenAIMessageBlock[], message: OpenAIMessageBlock[]): OpenAIMessageBlock[] {
+function updateConversation(conversation: ChatCompletionMessageParam[], message: ChatCompletionMessageParam[]): ChatCompletionMessageParam[] {
     if (!message) {
         return conversation;
     }
@@ -489,12 +483,15 @@ export function collectTools(toolCalls?: OpenAI.Chat.Completions.ChatCompletionM
     return tools.length > 0 ? tools : undefined;
 }
 
-function createPromptFromResponse(response: OpenAI.Chat.Completions.ChatCompletionMessage) : OpenAIMessageBlock[] {
-    const messages: OpenAIMessageBlock[] = [];
+function createPromptFromResponse(response: OpenAI.Chat.Completions.ChatCompletionMessage): ChatCompletionMessageParam[] {
+    const messages: ChatCompletionMessageParam[] = [];
     if (response) {
         messages.push({
             role: response.role,
-            content: response.content,
+            content: [{
+                type: "text",
+                text: response.content ?? ""
+            }],
             tool_calls: response.tool_calls,
         });
     }

package/src/openai/openai.ts

@@ -1,6 +1,6 @@
 
 
-import { DriverOptions } from "@llumiverse/core";
+import { DriverOptions, Providers } from "@llumiverse/core";
 import OpenAI from "openai";
 import { BaseOpenAIDriver } from "./index.js";
 
@@ -19,15 +19,12 @@ export interface OpenAIDriverOptions extends DriverOptions {
 export class OpenAIDriver extends BaseOpenAIDriver {
 
     service: OpenAI;
-    provider: "openai";
+    readonly provider = Providers.openai;
 
     constructor(opts: OpenAIDriverOptions) {
         super(opts);
         this.service = new OpenAI({
             apiKey: opts.apiKey
         });
-        this.provider = "openai";
     }
-
-
 }

package/src/openai/openai_format.ts

@@ -0,0 +1,165 @@
+// This file is used by multiple drivers
+// to format prompts in a way that is compatible with OpenAI's API.
+
+import { PromptRole, PromptOptions, PromptSegment } from "@llumiverse/common";
+import { readStreamAsBase64 } from "@llumiverse/core";
+
+import type {
+    ChatCompletionMessageParam,
+    ChatCompletionContentPartText,
+    ChatCompletionContentPartImage,
+    ChatCompletionUserMessageParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionToolMessageParam
+} from 'openai/resources/chat/completions';
+
+export interface OpenAITextMessage {
+    content: string;
+    role: 'system' | 'user' | 'assistant' | 'developer';
+}
+
+/**
+ * OpenAI text only prompts
+ * @param segments
+ * @returns
+ */
+export function formatOpenAILikeTextPrompt(segments: PromptSegment[]): OpenAITextMessage[] {
+    const system: OpenAITextMessage[] = [];
+    const safety: OpenAITextMessage[] = [];
+    const user: OpenAITextMessage[] = [];
+
+    for (const msg of segments) {
+        if (msg.role === PromptRole.system) {
+            system.push({ content: msg.content, role: "system" });
+        } else if (msg.role === PromptRole.safety) {
+            safety.push({ content: "IMPORTANT: " + msg.content, role: "system" });
+        } else if (msg.role !== PromptRole.negative && msg.role !== PromptRole.mask && msg.role !== PromptRole.tool) {
+            user.push({
+                content: msg.content,
+                role: msg.role || 'user',
+            })
+        }
+    }
+
+    // put system messages first and safety last
+    return system.concat(user).concat(safety);
+}
+
+
+export async function formatOpenAILikeMultimodalPrompt(segments: PromptSegment[], opts: PromptOptions & OpenAIPromptFormatterOptions): Promise<ChatCompletionMessageParam[]> {
+    const system: ChatCompletionMessageParam[] = [];
+    const safety: ChatCompletionMessageParam[] = [];
+    const others: ChatCompletionMessageParam[] = [];
+
+    for (const msg of segments) {
+
+        const parts: (ChatCompletionContentPartImage | ChatCompletionContentPartText)[] = [];
+
+        //generate the parts based on PromptSegment
+        if (msg.files) {
+            for (const file of msg.files) {
+                const stream = await file.getStream();
+                const data = await readStreamAsBase64(stream);
+                parts.push({
+                    type: "image_url",
+                    image_url: {
+                        url: `data:${file.mime_type || "image/jpeg"};base64,${data}`,
+                        //detail: "auto"  //This is modified just before execution to "low" | "high" | "auto"
+                    },
+                })
+            }
+        }
+
+        if (msg.content) {
+            parts.push({
+                text: msg.content,
+                type: "text"
+            })
+        }
+
+
+        if (msg.role === PromptRole.system) {
+            // For system messages, filter to only text parts
+            const textParts = parts.filter((part): part is ChatCompletionContentPartText => part.type === 'text');
+            const systemMsg: ChatCompletionSystemMessageParam = {
+                role: "system",
+                content: textParts.length === 1 && !msg.files ? textParts[0].text : textParts
+            };
+            system.push(systemMsg);
+
+            if (opts.useToolForFormatting && opts.schema) {
+                system.forEach(s => {
+                    if (typeof s.content === 'string') {
+                        s.content = "TOOL: " + s.content;
+                    } else if (Array.isArray(s.content)) {
+                        s.content.forEach((c: any) => {
+                            if (c.type === "text") c.text = "TOOL: " + c.text;
+                        });
+                    }
+                });
+            }
+
+        } else if (msg.role === PromptRole.safety) {
+            const textParts = parts.filter((part): part is ChatCompletionContentPartText => part.type === 'text');
+            const safetyMsg: ChatCompletionSystemMessageParam = {
+                role: "system",
+                content: textParts
+            };
+
+            if (Array.isArray(safetyMsg.content)) {
+                safetyMsg.content.forEach((c: any) => {
+                    if (c.type === "text") c.text = "DO NOT IGNORE - IMPORTANT: " + c.text;
+                });
+            }
+
+            system.push(safetyMsg);
+        } else if (msg.role === PromptRole.tool) {
+            if (!msg.tool_use_id) {
+                throw new Error("Tool use id is required for tool messages")
+            }
+            const toolMsg: ChatCompletionToolMessageParam = {
+                role: "tool",
+                tool_call_id: msg.tool_use_id,
+                content: msg.content || ""
+            };
+            others.push(toolMsg);
+        } else if (msg.role !== PromptRole.negative && msg.role !== PromptRole.mask) {
+            if (msg.role === 'assistant') {
+                const assistantMsg: ChatCompletionAssistantMessageParam = {
+                    role: 'assistant',
+                    content: parts as (ChatCompletionContentPartText)[]
+                };
+                others.push(assistantMsg);
+            } else {
+                const userMsg: ChatCompletionUserMessageParam = {
+                    role: 'user',
+                    content: parts
+                };
+                others.push(userMsg);
+            }
+        }
+
+    }
+
+    if (opts.result_schema && !opts.useToolForFormatting) {
+        const schemaMsg: ChatCompletionSystemMessageParam = {
+            role: "system",
+            content: [{
+                type: "text",
+                text: "IMPORTANT: only answer using JSON, and respecting the schema included below, between the <response_schema> tags. " + `<response_schema>${JSON.stringify(opts.result_schema)}</response_schema>`
+            }]
+        };
+        system.push(schemaMsg);
+    }
+
+    // put system messages first and safety last
+    return ([] as ChatCompletionMessageParam[]).concat(system).concat(others).concat(safety);
+
+}
+
+export interface OpenAIPromptFormatterOptions {
+    multimodal?: boolean
+    useToolForFormatting?: boolean
+    schema?: Object
+}

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/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[]> {