@llumiverse/drivers 0.16.0 → 0.18.0

package/src/bedrock/index.ts

@@ -1,18 +1,24 @@
-import { Bedrock, CreateModelCustomizationJobCommand, FoundationModelSummary, GetModelCustomizationJobCommand, GetModelCustomizationJobCommandOutput, ModelCustomizationJobStatus, StopModelCustomizationJobCommand } from "@aws-sdk/client-bedrock";
-import { BedrockRuntime, ConverseRequest, ConverseResponse, ConverseStreamOutput, InferenceConfiguration } from "@aws-sdk/client-bedrock-runtime";
+import {
+    Bedrock, CreateModelCustomizationJobCommand, FoundationModelSummary, GetModelCustomizationJobCommand,
+    GetModelCustomizationJobCommandOutput, ModelCustomizationJobStatus, ModelModality, StopModelCustomizationJobCommand
+} from "@aws-sdk/client-bedrock";
+import { BedrockRuntime, ConverseRequest, ConverseResponse, ConverseStreamOutput, InferenceConfiguration, Tool } from "@aws-sdk/client-bedrock-runtime";
 import { S3Client } from "@aws-sdk/client-s3";
-import { AbstractDriver, AIModel, Completion, CompletionChunkObject, DataSource, DriverOptions, EmbeddingsOptions, EmbeddingsResult, ExecutionOptions, ExecutionTokenUsage, ImageGeneration, Modalities, PromptOptions, PromptSegment, TextFallbackOptions, TrainingJob, TrainingJobStatus, TrainingOptions } from "@llumiverse/core";
+import { AwsCredentialIdentity, Provider } from "@aws-sdk/types";
+import {
+    AbstractDriver, AIModel, Completion, CompletionChunkObject, DataSource, DriverOptions, EmbeddingsOptions, EmbeddingsResult,
+    ExecutionOptions, ExecutionTokenUsage, ImageGeneration, Modalities, PromptOptions, PromptSegment,
+    TextFallbackOptions, ToolDefinition, ToolUse, TrainingJob, TrainingJobStatus, TrainingOptions,
+    BedrockClaudeOptions, BedrockPalmyraOptions, getMaxTokensLimit, NovaCanvasOptions,
+    modelModalitiesToArray, getModelCapabilities
+} from "@llumiverse/core";
 import { transformAsyncIterator } from "@llumiverse/core/async";
 import { formatNovaPrompt, NovaMessagesPrompt } from "@llumiverse/core/formatters";
-import { AwsCredentialIdentity, Provider } from "@smithy/types";
-import mnemonist from "mnemonist";
-import { BedrockClaudeOptions, NovaCanvasOptions } from "../../../core/src/options/bedrock.js";
-import { converseConcatMessages, converseRemoveJSONprefill, converseSystemToMessages, fortmatConversePrompt } from "./converse.js";
+import { LRUCache } from "mnemonist";
+import { converseConcatMessages, converseJSONprefill, converseSystemToMessages, formatConversePrompt } from "./converse.js";
 import { formatNovaImageGenerationPayload, NovaImageGenerationTaskType } from "./nova-image-payload.js";
 import { forceUploadFile } from "./s3.js";
 
-const { LRUCache } = mnemonist;
-
 const supportStreamingCache = new LRUCache<string, boolean>(4096);
 
 enum BedrockModelType {
@@ -44,8 +50,8 @@ export interface BedrockDriverOptions extends DriverOptions {
      */
     region: string;
     /**
-     * Tthe bucket name to be used for training.
-     * It will be created oif nto already exixts
+     * The bucket name to be used for training.
+     * It will be created if does not already exist.
      */
     training_bucket?: string;
 
@@ -105,10 +111,10 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
         if (opts.model.includes("canvas")) {
             return await formatNovaPrompt(segments, opts.result_schema);
         }
-        return await fortmatConversePrompt(segments, opts.result_schema);
+        return await formatConversePrompt(segments, opts.result_schema);
     }
 
-    static getExtractedExecuton(result: ConverseResponse, _prompt?: BedrockPrompt): CompletionChunkObject {
+    static getExtractedExecution(result: ConverseResponse, _prompt?: BedrockPrompt): CompletionChunkObject {
         return {
             result: result.output?.message?.content?.map(c => c.text).join("\n") ?? "",
             token_usage: {
@@ -145,18 +151,45 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
     };
 
     async requestTextCompletion(prompt: ConverseRequest, options: ExecutionOptions): Promise<Completion> {
+        let conversation = updateConversation(options.conversation as ConverseRequest, prompt);
 
-        const payload = this.preparePayload(prompt, options);
+        const payload = this.preparePayload(conversation, options);
         const executor = this.getExecutor();
 
         const res = await executor.converse({
             ...payload,
         });
 
+        conversation = updateConversation(conversation, {
+            messages: [res.output?.message ?? { content: [{ text: "" }], role: "assistant" }],
+            modelId: prompt.modelId,
+        });
+
+        let tool_use: ToolUse[] | undefined = undefined;
+        //Get tool requests
+        if (res.stopReason == "tool_use") {
+            tool_use = res.output?.message?.content?.reduce((tools: ToolUse[], c) => {
+                if (c.toolUse) {
+                    tools.push({
+                        tool_name: c.toolUse.name ?? "",
+                        tool_input: c.toolUse.input as any,
+                        id: c.toolUse.toolUseId ?? "",
+                    } satisfies ToolUse);
+                }
+                return tools;
+            }, []);
+            //If no tools were used, set to undefined
+            if (tool_use && tool_use.length == 0) {
+                tool_use = undefined;
+            }
+        }
+
         const completion = {
-            ...BedrockDriver.getExtractedExecuton(res, prompt),
+            ...BedrockDriver.getExtractedExecution(res, prompt),
             original_response: options.include_original_response ? res : undefined,
-        } satisfies Completion;
+            conversation: conversation,
+            tool_use: tool_use,
+        };
 
         return completion;
     }
@@ -167,13 +200,13 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
         if (arnMatch) {
             return arnMatch[1];
         }
-        
+
         // Match common AWS regions directly in string
         const regionMatch = modelString.match(/(?:us|eu|ap|sa|ca|me|af)[-](east|west|central|south|north|southeast|southwest|northeast|northwest)[-][1-9]/);
         if (regionMatch) {
             return regionMatch[0];
         }
-    
+
         return defaultRegion;
     }
 
@@ -267,11 +300,17 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
         let additionalField = {};
 
         if (options.model.includes("amazon")) {
+            if (options.result_schema) {
+                prompt.messages = converseJSONprefill(prompt.messages);
+            }
             //Titan models also exists but does not support any additional options
             if (options.model.includes("nova")) {
                 additionalField = { inferenceConfig: { topK: model_options?.top_k } };
             }
         } else if (options.model.includes("claude")) {
+            if (options.result_schema) {
+                prompt.messages = converseJSONprefill(prompt.messages);
+            }
             if (options.model.includes("claude-3-7")) {
                 const thinking_options = options.model_options as BedrockClaudeOptions;
                 const thinking = thinking_options?.thinking_mode ?? false;
@@ -279,13 +318,13 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
                     model_options.max_tokens = thinking ? 128000 : 8192;
                 }
                 additionalField = {
-                    top_k: model_options?.top_k,
+                    ...additionalField,
                     reasoning_config: {
                         type: thinking ? "enabled" : "disabled",
                         budget_tokens: thinking_options?.thinking_budget_tokens,
                     }
                 };
-                if(thinking && (thinking_options?.thinking_budget_tokens ?? 0) > 64000){
+                if (thinking && (thinking_options?.thinking_budget_tokens ?? 0) > 64000) {
                     additionalField = {
                         ...additionalField,
                         anthorpic_beta: ["output-128k-2025-02-19"]
@@ -294,22 +333,11 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
             }
             //Needs max_tokens to be set
             if (!model_options?.max_tokens) {
-                if (options.model.includes("claude-3-5")) {
-                    model_options.max_tokens = 8192;
-
-                    //Bug with AWS Converse Sonnet 3.5, does not effect Haiku.
-                    //See https://github.com/boto/boto3/issues/4279
-                    if (options.model.includes("claude-3-5-sonnet")) {
-                        model_options.max_tokens = 4096;
-                    }
-                } else {
-                    model_options.max_tokens = 4096;
-                }
+                model_options.max_tokens = getMaxTokensLimit(options.model, model_options);
             }
-            additionalField = { top_k: model_options?.top_k };
+            additionalField = { ...additionalField, top_k: model_options?.top_k };
         } else if (options.model.includes("meta")) {
-            //If last message is "```json", remove it. Model requires the final message to be a user message
-            prompt.messages = converseRemoveJSONprefill(prompt.messages);
+            //LLaMA models support no additional options
         } else if (options.model.includes("mistral")) {
             //7B instruct and 8x7B instruct
             if (options.model.includes("7b")) {
@@ -320,14 +348,14 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
                     prompt.system = undefined;
                     prompt.messages = converseConcatMessages(prompt.messages);
                 }
+                if (options.result_schema) {
+                    prompt.messages = converseJSONprefill(prompt.messages);
+                }
             } else {
                 //Other models such as Mistral Small,Large and Large 2
                 //Support no additional fields.
-                prompt.messages = converseRemoveJSONprefill(prompt.messages);
             }
         } else if (options.model.includes("ai21")) {
-            //If last message is "```json", remove it. Model requires the final message to be a user message
-            prompt.messages = converseRemoveJSONprefill(prompt.messages);
             //Jamba models support no additional options
             //Jurassic 2 models do.
             if (options.model.includes("j2")) {
@@ -344,8 +372,6 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
             }
         } else if (options.model.includes("cohere.command")) {
             // If last message is "```json", remove it.
-            // Model requires the final message to be a user message or does not support assistant messages
-            prompt.messages = converseRemoveJSONprefill(prompt.messages);
             //Command R and R plus
             if (options.model.includes("cohere.command-r")) {
                 additionalField = {
@@ -363,6 +389,16 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
                     prompt.messages = converseConcatMessages(prompt.messages);
                 }
             }
+        } else if (options.model.includes("palmyra")) {
+            const palmyraOptions = options.model_options as BedrockPalmyraOptions;
+            additionalField = {
+                seed: palmyraOptions?.seed,
+                presence_penalty: palmyraOptions?.presence_penalty,
+                frequency_penalty: palmyraOptions?.frequency_penalty,
+                min_tokens: palmyraOptions?.min_tokens,
+            }
+        } else if (options.model.includes("deepseek")) {
+            //DeepSeek models support no additional options
         }
 
         //If last message is "```json", add corresponding ``` as a stop sequence.
@@ -378,7 +414,9 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
             }
         }
 
-        return {
+        const tool_defs = getToolDefinitions(options.tools);
+
+        const request: ConverseRequest = {
             messages: prompt.messages,
             system: prompt.system,
             modelId: options.model,
@@ -390,8 +428,17 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
             } satisfies InferenceConfiguration,
             additionalModelRequestFields: {
                 ...additionalField,
-            },
-        } satisfies ConverseRequest;
+            }
+        };
+
+        //Only add tools if they are defined
+        if (tool_defs) {
+            request.toolConfig = {
+                tools: tool_defs,
+            }
+        }
+
+        return request;
     }
 
 
@@ -400,7 +447,7 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
             throw new Error(`Image generation requires image output_modality`);
         }
         if (options.model_options?._option_id !== "bedrock-nova-canvas") {
-            this.logger.warn("Invalid model options", {options: options.model_options });
+            this.logger.warn("Invalid model options", { options: options.model_options });
         }
         const model_options = options.model_options as NovaCanvasOptions;
 
@@ -522,7 +569,7 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
 
     async _listModels(foundationFilter?: (m: FoundationModelSummary) => boolean): Promise<AIModel[]> {
         const service = this.getService();
-        const [foundationals, customs, inferenceProfiles] = await Promise.all([
+        const [foundationModelsList, customModelsList, inferenceProfilesList] = await Promise.all([
             service.listFoundationModels({}).catch(() => {
                 this.logger.warn("[Bedrock] Can't list foundation models. Check if the user has the right permissions.");
                 return undefined
@@ -537,21 +584,65 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
             }),
         ]);
 
-        if (!foundationals?.modelSummaries) {
+        if (!foundationModelsList?.modelSummaries) {
             throw new Error("Foundation models not found");
         }
 
-        let fmodels = foundationals.modelSummaries || [];
+        let foundationModels = foundationModelsList.modelSummaries || [];
         if (foundationFilter) {
-            fmodels = fmodels.filter(foundationFilter);
+            foundationModels = foundationModels.filter(foundationFilter);
         }
 
-        const aimodels: AIModel[] = fmodels.map((m) => {
+        const supportedPublishers = ["amazon", "anthropic", "cohere", "ai21", "mistral", "meta", "deepseek", "writer"];
+        const unsupportedModelsByPublisher = {
+            amazon: ["titan-image-generator", "nova-reel", "nova-sonic", "rerank"],
+            anthropic: [],
+            cohere: ["rerank"],
+            ai21: [],
+            mistral: [],
+            meta: [],
+            deepseek: [],
+            writer: [],
+        };
+
+        // Helper function to check if model should be filtered out
+        const shouldIncludeModel = (modelId?: string, providerName?: string): boolean => {
+            if (!modelId || !providerName) return false;
+
+            const normalizedProvider = providerName.toLowerCase();
+
+            // Check if provider is supported
+            const isProviderSupported = supportedPublishers.some(provider =>
+                normalizedProvider.includes(provider)
+            );
+
+            if (!isProviderSupported) return false;
+
+            // Check if model is in the unsupported list for its provider
+            for (const provider of supportedPublishers) {
+                if (normalizedProvider.includes(provider)) {
+                    const unsupportedModels = unsupportedModelsByPublisher[provider as keyof typeof unsupportedModelsByPublisher] || [];
+                    return !unsupportedModels.some(unsupported =>
+                        modelId.toLowerCase().includes(unsupported)
+                    );
+                }
+            }
+
+            return true;
+        };
+
+        foundationModels = foundationModels.filter(m =>
+            shouldIncludeModel(m.modelId, m.providerName)
+        );
+
+        const aiModels: AIModel[] = foundationModels.map((m) => {
 
             if (!m.modelId) {
                 throw new Error("modelId not found");
             }
 
+            const modelCapability = getModelCapabilities(m.modelArn ?? m.modelId, this.provider);
+
             const model: AIModel = {
                 id: m.modelArn ?? m.modelId,
                 name: `${m.providerName} ${m.modelName}`,
@@ -559,52 +650,78 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
                 //description: ``,
                 owner: m.providerName,
                 can_stream: m.responseStreamingSupported ?? false,
-                is_multimodal: m.inputModalities?.includes("IMAGE") ?? false,
-                tags: m.outputModalities ?? [],
+                input_modalities: m.inputModalities ? formatAmazonModalities(m.inputModalities) : modelModalitiesToArray(modelCapability.input),
+                output_modalities: m.outputModalities ? formatAmazonModalities(m.outputModalities) : modelModalitiesToArray(modelCapability.input),
+                tool_support: modelCapability.tool_support,
             };
 
             return model;
         });
 
         //add custom models
-        if (customs?.modelSummaries) {
-            customs.modelSummaries.forEach((m) => {
+        if (customModelsList?.modelSummaries) {
+            customModelsList.modelSummaries.forEach((m) => {
 
                 if (!m.modelArn) {
                     throw new Error("Model ID not found");
                 }
 
+                const modelCapability = getModelCapabilities(m.modelArn, this.provider);
+
                 const model: AIModel = {
                     id: m.modelArn,
                     name: m.modelName ?? m.modelArn,
                     provider: this.provider,
                     description: `Custom model from ${m.baseModelName}`,
                     is_custom: true,
+                    input_modalities: modelModalitiesToArray(modelCapability.input),
+                    output_modalities: modelModalitiesToArray(modelCapability.output),
+                    tool_support: modelCapability.tool_support,
                 };
 
-                aimodels.push(model);
+                aiModels.push(model);
                 this.validateConnection;
             });
         }
 
         //add inference profiles
-        if (inferenceProfiles?.inferenceProfileSummaries) {
-            inferenceProfiles.inferenceProfileSummaries.forEach((p) => {
+        if (inferenceProfilesList?.inferenceProfileSummaries) {
+            inferenceProfilesList.inferenceProfileSummaries.forEach((p) => {
                 if (!p.inferenceProfileArn) {
                     throw new Error("Profile ARN not found");
                 }
 
-                const model: AIModel = {
-                    id: p.inferenceProfileArn ?? p.inferenceProfileId,
-                    name: p.inferenceProfileName ?? p.inferenceProfileArn,
-                    provider: this.provider,
-                };
+                // Apply the same filtering logic to inference profiles based on their name
+                const profileId = p.inferenceProfileId || "";
+                const profileName = p.inferenceProfileName || "";
+
+                // Extract provider name from profile name or ID
+                let providerName = "";
+                for (const provider of supportedPublishers) {
+                    if (profileName.toLowerCase().includes(provider) || profileId.toLowerCase().includes(provider)) {
+                        providerName = provider;
+                        break;
+                    }
+                }
+
+                const modelCapability = getModelCapabilities(p.inferenceProfileArn ?? p.inferenceProfileId, this.provider);
+
+                if (providerName && shouldIncludeModel(profileId, providerName)) {
+                    const model: AIModel = {
+                        id: p.inferenceProfileArn ?? p.inferenceProfileId,
+                        name: p.inferenceProfileName ?? p.inferenceProfileArn,
+                        provider: this.provider,
+                        input_modalities: modelModalitiesToArray(modelCapability.input),
+                        output_modalities: modelModalitiesToArray(modelCapability.output),
+                        tool_support: modelCapability.tool_support,
+                    };
 
-                aimodels.push(model);
+                    aiModels.push(model);
+                }
             });
         }
 
-        return aimodels;
+        return aiModels;
     }
 
     async generateEmbeddings({ text, image, model }: EmbeddingsOptions): Promise<EmbeddingsResult> {
@@ -666,4 +783,56 @@ function jobInfo(job: GetModelCustomizationJobCommandOutput, jobId: string): Tra
         status,
         details
     }
+}
+
+function getToolDefinitions(tools?: ToolDefinition[]): Tool[] | undefined {
+    return tools ? tools.map(getToolDefinition) : undefined;
+}
+
+function getToolDefinition(tool: ToolDefinition): Tool.ToolSpecMember {
+    return {
+        toolSpec: {
+            name: tool.name,
+            description: tool.description,
+            inputSchema: {
+                json: tool.input_schema as any,
+            }
+        }
+    }
+}
+
+/**
+ * Update the conversation messages
+ * @param prompt
+ * @param response
+ * @returns
+ */
+function updateConversation(conversation: ConverseRequest, prompt: ConverseRequest): ConverseRequest {
+    return {
+        ...conversation,
+        ...prompt,
+        messages: [...(conversation?.messages || []), ...(prompt.messages || [])],
+        system: prompt.system || conversation?.system,
+    };
+}
+
+function formatAmazonModalities(modalities: ModelModality[]): string[] {
+    const standardizedModalities: string[] = [];
+    for (const modality of modalities) {
+        if (modality === ModelModality.TEXT) {
+            standardizedModalities.push("text");
+        } else if (modality === ModelModality.IMAGE) {
+            standardizedModalities.push("image");
+        } else if (modality === ModelModality.EMBEDDING) {
+            standardizedModalities.push("embedding");
+        } else if (modality == "SPEECH") {
+            standardizedModalities.push("audio");
+        } else if (modality == "VIDEO") {
+            standardizedModalities.push("video");
+        } else {
+            // Handle other modalities as needed
+            standardizedModalities.push((modality as string).toString().toLowerCase());
+        }
+    }
+    return standardizedModalities;
 }

package/src/bedrock/nova-image-payload.ts

@@ -1,6 +1,5 @@
-import { ExecutionOptions } from "@llumiverse/core";
+import { ExecutionOptions, NovaCanvasOptions } from "@llumiverse/core";
 import { NovaMessage, NovaMessagesPrompt } from "@llumiverse/core/formatters";
-import { NovaCanvasOptions } from "../../../core/src/options/bedrock.js";
 
 function getFirstImageFromPrompt(prompt: NovaMessage[]) {
 

package/src/bedrock/s3.ts

@@ -59,3 +59,51 @@ export async function forceUploadFile(s3: S3Client, source: ReadableStream, buck
     await tryCreateBucket(s3, bucketName);
     return uploadFile(s3, source, bucketName, file, onProgress);
 }
+
+
+/**
+ * Parse an S3 HTTPS URL into an S3 URI format
+ * s3Url - The S3 HTTPS URL (e.g., https://bucket.s3.region.amazonaws.com/key)
+ * returns The S3 URI (e.g., s3://bucket/key)
+ */
+export function parseS3UrlToUri(s3Url: URL) {
+    try {
+        const url = new URL(s3Url);
+
+        // Extract the hostname which contains the bucket and S3 endpoint
+        const hostname = url.hostname;
+
+        // Parse the hostname to extract the bucket name
+        let bucketName;
+        if (hostname.endsWith('.amazonaws.com')) {
+            if (hostname.includes('.s3.')) {
+                // Format: bucket-name.s3.region.amazonaws.com
+                bucketName = hostname.split('.s3.')[0];
+            } else if (hostname.startsWith('s3.')) {
+                // Format: s3.region.amazonaws.com/bucket-name
+                // In this case, the bucket is actually in the first segment of the pathname
+                bucketName = url.pathname.split('/')[1];
+                // Adjust the pathname to remove the bucket name
+                const pathParts = url.pathname.split('/').slice(2);
+                url.pathname = '/' + pathParts.join('/');
+            } else {
+                throw new Error('Unable to determine bucket name from URL');
+            }
+        } else {
+            throw new Error('Unable to determine bucket name from URL');
+        }
+
+        // The key is the pathname without the leading slash
+        // If we had the bucket name in the path, it's already been removed above
+        let key = url.pathname;
+        if (key.startsWith('/')) {
+            key = key.substring(1);
+        }
+
+        // Construct the S3 URI
+        return `s3://${bucketName}/${key}`;
+    } catch (error) {
+        console.error('Error parsing S3 URL:', error);
+        throw error;
+    }
+}

package/src/groq/index.ts

@@ -28,7 +28,7 @@ export class GroqDriver extends AbstractDriver<GroqDriverOptions, OpenAITextMess
 
     // protected canStream(options: ExecutionOptions): Promise<boolean> {
     //     if (options.result_schema) {
-    //         // not yet streamign json responses
+    //         // not yet streaming json responses
     //         return Promise.resolve(false);
     //     } else {
     //         return Promise.resolve(true);
@@ -59,7 +59,7 @@ export class GroqDriver extends AbstractDriver<GroqDriverOptions, OpenAITextMess
     }
 
     async requestTextCompletion(messages: OpenAITextMessage[], options: ExecutionOptions): Promise<Completion<any>> {
-        if (options.model_options?._option_id !== "text-fallback") {
+        if (options.model_options?._option_id !== "text-fallback" && options.model_options?._option_id !== "groq-deepseek-thinking") {
             this.logger.warn("Invalid model options", {options: options.model_options });
         }
         options.model_options = options.model_options as TextFallbackOptions;
@@ -130,7 +130,7 @@ export class GroqDriver extends AbstractDriver<GroqDriverOptions, OpenAITextMess
             throw new Error("No models found");
         }
 
-        const aimodels = models.data?.map(m => {
+        const aiModels = models.data?.map(m => {
             if (!m.id) {
                 throw new Error("Model id is missing");
             }
@@ -143,7 +143,7 @@ export class GroqDriver extends AbstractDriver<GroqDriverOptions, OpenAITextMess
             }
         });
 
-        return aimodels;
+        return aiModels;
     }
 
     validateConnection(): Promise<boolean> {

package/src/mistral/index.ts

@@ -45,7 +45,7 @@ export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, Open
         // return _options.result_schema ? responseFormatJson : responseFormatText;
 
         //TODO remove this when Mistral properly supports the parameters - it makes an error for now
-        // some models like mixtral mistrall tiny or medium are throwing an error when using the response_format parameter
+        // some models like mixtral mistral tiny or medium are throwing an error when using the response_format parameter
         return undefined
     }
 
@@ -130,7 +130,7 @@ export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, Open
     async listModels(): Promise<AIModel<string>[]> {
         const models: ListModelsResponse = await this.client.get('v1/models');
 
-        const aimodels = models.data.map(m => {
+        const aiModels = models.data.map(m => {
             return {
                 id: m.id,
                 name: m.id,
@@ -140,7 +140,7 @@ export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, Open
             }
         });
 
-        return aimodels;
+        return aiModels;
     }
 
     validateConnection(): Promise<boolean> {
@@ -158,7 +158,7 @@ export class MistralAIDriver extends AbstractDriver<MistralAIDriverOptions, Open
         return {
             values: r.data[0].embedding,
             model,
-            token_count: r.usage.total_tokens
+            token_count: r.usage.total_tokens || r.usage.prompt_tokens + r.usage.completion_tokens,
         }
     }
 

package/src/openai/azure.ts

@@ -8,7 +8,7 @@ export interface AzureOpenAIDriverOptions extends DriverOptions {
     /**
      * The credentials to use to access Azure OpenAI
      */
-    azureADTokenProvider?: any; //type with azure credntials
+    azureADTokenProvider?: any; //type with azure credentials
     
     apiKey?: string;
 
@@ -37,7 +37,7 @@ export class AzureOpenAIDriver extends BaseOpenAIDriver {
             apiKey: opts.apiKey,
             azureADTokenProvider: opts.azureADTokenProvider,          
             endpoint: opts.endpoint,
-            apiVersion: opts.apiVersion ?? "2024-05-01-preview",
+            apiVersion: opts.apiVersion ?? "2024-10-21",
             deployment: opts.deployment
         });
         this.provider = "azure_openai";