@llumiverse/drivers 0.14.0 → 0.16.0

package/src/bedrock/index.ts

@@ -1,18 +1,38 @@
 import { Bedrock, CreateModelCustomizationJobCommand, FoundationModelSummary, GetModelCustomizationJobCommand, GetModelCustomizationJobCommandOutput, ModelCustomizationJobStatus, StopModelCustomizationJobCommand } from "@aws-sdk/client-bedrock";
-import { BedrockRuntime, InvokeModelCommandOutput, ResponseStream } from "@aws-sdk/client-bedrock-runtime";
+import { BedrockRuntime, ConverseRequest, ConverseResponse, ConverseStreamOutput, InferenceConfiguration } from "@aws-sdk/client-bedrock-runtime";
 import { S3Client } from "@aws-sdk/client-s3";
-import { AIModel, AbstractDriver, Completion, DataSource, DriverOptions, EmbeddingsOptions, EmbeddingsResult, ExecutionOptions, PromptOptions, PromptSegment, TrainingJob, TrainingJobStatus, TrainingOptions } from "@llumiverse/core";
+import { AbstractDriver, AIModel, Completion, CompletionChunkObject, DataSource, DriverOptions, EmbeddingsOptions, EmbeddingsResult, ExecutionOptions, ExecutionTokenUsage, ImageGeneration, Modalities, PromptOptions, PromptSegment, TextFallbackOptions, TrainingJob, TrainingJobStatus, TrainingOptions } from "@llumiverse/core";
 import { transformAsyncIterator } from "@llumiverse/core/async";
-import { ClaudeMessagesPrompt, formatClaudePrompt } from "@llumiverse/core/formatters";
+import { formatNovaPrompt, NovaMessagesPrompt } from "@llumiverse/core/formatters";
 import { AwsCredentialIdentity, Provider } from "@smithy/types";
 import mnemonist from "mnemonist";
-import { AI21RequestPayload, AmazonRequestPayload, ClaudeRequestPayload, CohereCommandRPayload, CohereRequestPayload, LLama2RequestPayload, MistralPayload } from "./payloads.js";
+import { BedrockClaudeOptions, NovaCanvasOptions } from "../../../core/src/options/bedrock.js";
+import { converseConcatMessages, converseRemoveJSONprefill, converseSystemToMessages, fortmatConversePrompt } 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 {
+    FoundationModel = "foundation-model",
+    InferenceProfile = "inference-profile",
+    CustomModel = "custom-model",
+    Unknown = "unknown",
+};
+
+function converseFinishReason(reason: string | undefined) {
+    //Possible values:
+    //end_turn | tool_use | max_tokens | stop_sequence | guardrail_intervened | content_filtered
+    if (!reason) return undefined;
+    switch (reason) {
+        case 'end_turn': return "stop";
+        case 'max_tokens': return "length";
+        default: return reason;
+    }
+}
+
 export interface BedrockModelCapabilities {
     name: string;
     canStream: boolean;
@@ -40,7 +60,7 @@ export interface BedrockDriverOptions extends DriverOptions {
     credentials?: AwsCredentialIdentity | Provider<AwsCredentialIdentity>;
 }
 
-export type BedrockPrompt = string | ClaudeMessagesPrompt;
+export type BedrockPrompt = NovaMessagesPrompt | ConverseRequest;
 
 export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockPrompt> {
 
@@ -50,6 +70,7 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
 
     private _executor?: BedrockRuntime;
     private _service?: Bedrock;
+    private _service_region?: string;
 
     constructor(options: BedrockDriverOptions) {
         super(options);
@@ -69,164 +90,170 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
         return this._executor;
     }
 
-    getService() {
-        if (!this._service) {
+    getService(region: string = this.options.region) {
+        if (!this._service || this._service_region != region) {
             this._service = new Bedrock({
-                region: this.options.region,
+                region: region,
                 credentials: this.options.credentials,
             });
+            this._service_region = region;
         }
         return this._service;
     }
 
     protected async formatPrompt(segments: PromptSegment[], opts: PromptOptions): Promise<BedrockPrompt> {
-        //TODO move the anthropic test in abstract driver?
-        if (opts.model.includes('anthropic')) {
-            //TODO: need to type better the types aren't checked properly by TS
-            return await formatClaudePrompt(segments, opts.result_schema);
-        } else {
-            return await super.formatPrompt(segments, opts) as string;
+        if (opts.model.includes("canvas")) {
+            return await formatNovaPrompt(segments, opts.result_schema);
         }
+        return await fortmatConversePrompt(segments, opts.result_schema);
     }
 
-    extractDataFromResponse(prompt: BedrockPrompt, response: InvokeModelCommandOutput): Completion {
-
-        const decoder = new TextDecoder();
-        const body = decoder.decode(response.body);
-        const result = JSON.parse(body);
-
-        const getTextAnsStopReason = (): string[] => {
-            if (result.generation) {
-                // LLAMA2
-                return [result.generation, result.stop_reason]; // comes in coirrect format (stop, length)
-            } else if (result.generations) {
-                // Cohere
-                return [result.generations[0].text, cohereFinishReason(result.generations[0].finish_reason)];
-            } else if (result.chat_history) {
-                //Cohere Command R
-                return [result.text, cohereFinishReason(result.finish_reason)];
-            } else if (result.completions) {
-                //A21
-                return [result.completions[0].data?.text, a21FinishReason(result.completions[0].finishReason?.reason)];
-            } else if (result.content) {
-                // Claude
-                //if last prompt.messages is {, add { to the response
-                const p = prompt as ClaudeMessagesPrompt;
-                const lastMessage = (p as ClaudeMessagesPrompt).messages[p.messages.length - 1];
-                const res = lastMessage.content[0].text === '{' ? '{' + result.content[0]?.text : result.content[0]?.text;
-
-                return [res, claudeFinishReason(result.stop_reason)];
-
-            } else if (result.outputs) {
-                // mistral
-                return [result.outputs[0]?.text, result.outputs[0]?.stop_reason]; // the stop reason is in the expected format ("stop" and "length")
-            } else if (result.results) {
-                // Amazon Titan
-                return [result.results[0]?.outputText ?? '', titanFinishReason(result.results[0]?.completionReason)];
-            } else if (result.completion) { // TODO: who uses this?
-                return [result.completion];
-            } else {
-                return [result.toString()];
-            }
-        };
-
-        const [text, finish_reason] = getTextAnsStopReason();
-
-        const promptLength = typeof prompt === 'string' ? prompt.length :
-            (prompt.system || '').length + prompt.messages.reduce((acc, m) => acc + m.content.length, 0);
+    static getExtractedExecuton(result: ConverseResponse, _prompt?: BedrockPrompt): CompletionChunkObject {
         return {
-            result: text,
+            result: result.output?.message?.content?.map(c => c.text).join("\n") ?? "",
             token_usage: {
-                result: text?.length,
-                prompt: promptLength,
-                total: text?.length + promptLength,
+                prompt: result.usage?.inputTokens,
+                result: result.usage?.outputTokens,
+                total: result.usage?.totalTokens,
             },
-            finish_reason
+            finish_reason: converseFinishReason(result.stopReason),
         }
-    }
+    };
+
+    static getExtractedStream(result: ConverseStreamOutput, _prompt?: BedrockPrompt): CompletionChunkObject {
+        let output: string = "";
+        let stop_reason = "";
+        let token_usage: ExecutionTokenUsage | undefined;
+        if (result.contentBlockDelta) {
+            output = result.contentBlockDelta.delta?.text ?? "";
+        }
+        if (result.messageStop) {
+            stop_reason = result.messageStop.stopReason ?? "";
+        }
+        if (result.metadata) {
+            token_usage = {
+                prompt: result.metadata.usage?.inputTokens,
+                result: result.metadata.usage?.outputTokens,
+                total: result.metadata.usage?.totalTokens,
+            }
+        }
+        return {
+            result: output,
+            token_usage: token_usage,
+            finish_reason: converseFinishReason(stop_reason),
+        }
+    };
 
-    async requestCompletion(prompt: BedrockPrompt, options: ExecutionOptions): Promise<Completion> {
+    async requestTextCompletion(prompt: ConverseRequest, options: ExecutionOptions): Promise<Completion> {
 
         const payload = this.preparePayload(prompt, options);
         const executor = this.getExecutor();
-        const res = await executor.invokeModel({
-            modelId: options.model,
-            contentType: "application/json",
-            body: JSON.stringify(payload),
+
+        const res = await executor.converse({
+            ...payload,
         });
-        const completion = this.extractDataFromResponse(prompt, res);
-        if (options.include_original_response) {
-            completion.original_response = res;
-        }
+
+        const completion = {
+            ...BedrockDriver.getExtractedExecuton(res, prompt),
+            original_response: options.include_original_response ? res : undefined,
+        } satisfies Completion;
+
         return completion;
     }
 
+    extractRegion(modelString: string, defaultRegion: string): string {
+        // Match region in full ARN pattern
+        const arnMatch = modelString.match(/arn:aws[^:]*:bedrock:([^:]+):/);
+        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;
+    }
+
+    private async getCanStream(model: string, type: BedrockModelType): Promise<boolean> {
+        let canStream: boolean = false;
+        let error: any = null;
+        const region = this.extractRegion(model, this.options.region);
+        if (type == BedrockModelType.FoundationModel || type == BedrockModelType.Unknown) {
+            try {
+                const response = await this.getService(region).getFoundationModel({
+                    modelIdentifier: model
+                });
+                canStream = response.modelDetails?.responseStreamingSupported ?? false;
+                return canStream;
+            } catch (e) {
+                error = e;
+            }
+        }
+        if (type == BedrockModelType.InferenceProfile || type == BedrockModelType.Unknown) {
+            try {
+                const response = await this.getService(region).getInferenceProfile({
+                    inferenceProfileIdentifier: model
+                });
+                canStream = await this.getCanStream(response.models?.[0].modelArn ?? "", BedrockModelType.FoundationModel);
+                return canStream;
+            } catch (e) {
+                error = e;
+            }
+        }
+        if (type == BedrockModelType.CustomModel || type == BedrockModelType.Unknown) {
+            try {
+                const response = await this.getService(region).getCustomModel({
+                    modelIdentifier: model
+                });
+                canStream = await this.getCanStream(response.baseModelArn ?? "", BedrockModelType.FoundationModel);
+                return canStream;
+            } catch (e) {
+                error = e;
+            }
+        }
+        if (error) {
+            console.warn("Error on canStream check for model: " + model + " region detected: " + region, error);
+        }
+        return canStream;
+    }
+
     protected async canStream(options: ExecutionOptions): Promise<boolean> {
         let canStream = supportStreamingCache.get(options.model);
         if (canStream == null) {
-            const response = await this.getService().getFoundationModel({
-                modelIdentifier: options.model
-            });
-            canStream = response.modelDetails?.responseStreamingSupported ?? false;
+            let type = BedrockModelType.Unknown;
+            if (options.model.includes("foundation-model")) {
+                type = BedrockModelType.FoundationModel;
+            } else if (options.model.includes("inference-profile")) {
+                type = BedrockModelType.InferenceProfile;
+            } else if (options.model.includes("custom-model")) {
+                type = BedrockModelType.CustomModel;
+            }
+            canStream = await this.getCanStream(options.model, type);
             supportStreamingCache.set(options.model, canStream);
         }
         return canStream;
     }
 
-    async requestCompletionStream(prompt: BedrockPrompt, options: ExecutionOptions): Promise<AsyncIterable<string>> {
+    async requestTextCompletionStream(prompt: ConverseRequest, options: ExecutionOptions): Promise<AsyncIterable<CompletionChunkObject>> {
         const payload = this.preparePayload(prompt, options);
         const executor = this.getExecutor();
-        return executor.invokeModelWithResponseStream({
-            modelId: options.model,
-            contentType: "application/json",
-            body: JSON.stringify(payload),
+        return executor.converseStream({
+            ...payload,
         }).then((res) => {
+            const stream = res.stream;
 
-            if (!res.body) {
-                throw new Error("Body not found");
+            if (!stream) {
+                throw new Error("[Bedrock] Stream not found in response");
             }
-            const decoder = new TextDecoder();
 
-            const addBracket = () => {
-                if (typeof prompt === 'object' && (prompt as ClaudeMessagesPrompt).messages) {
-                    const p = prompt as ClaudeMessagesPrompt;
-                    const lastMessage = p.messages[p.messages.length - 1];
-                    return lastMessage.content[0].text === '{';
-                }
-                return false;
-            };
-
-            return transformAsyncIterator(res.body, (stream: ResponseStream) => {
-                const segment = JSON.parse(decoder.decode(stream.chunk?.bytes));
+            return transformAsyncIterator(stream, (stream: ConverseStreamOutput) => {
+                //const segment = JSON.parse(decoder.decode(stream.chunk?.bytes));
                 //console.log("Debug Segment for model " + options.model, JSON.stringify(segment));
-                if (segment.delta) { // who is this?
-                    return segment.delta.text || '';
-                } else if (segment.completion) { // who is this?
-                    return segment.completion;
-                } else if (segment.text) { //cohere
-                    return segment.text;
-                } else if (segment.completions) {
-                    return segment.completions[0].data?.text;
-                } else if (segment.generation) {
-                    return segment.generation;
-                } else if (segment.generations) {
-                    return segment.generations[0].text;
-                } else if (segment.outputs) {
-                    // mistral.mixtral-8x7b-instruct-v0:1
-                    return segment.outputs[0].text;
-                    //segment.outputs[0].stop_reason;
-                } else if (segment.outputText) {
-                    // Amazon Titan
-                    return segment.outputText;
-                    //completionReason
-                    // token count too
-                } else {
-                    segment.toString();
-                }
-
-            },
-                () => addBracket() ? '{' : ''
-            );
+                return BedrockDriver.getExtractedStream(stream, prompt);
+            });
 
         }).catch((err) => {
             this.logger.error("[Bedrock] Failed to stream", err);
@@ -234,79 +261,180 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
         });
     }
 
+    preparePayload(prompt: ConverseRequest, options: ExecutionOptions) {
+        const model_options = options.model_options as TextFallbackOptions;
 
+        let additionalField = {};
 
-    preparePayload(prompt: BedrockPrompt, options: ExecutionOptions) {
-
-        //split arn on / should give provider
-        //TODO: check if works with custom models
-        //const provider = options.model.split("/")[0];
-        const contains = (str: string, substr: string) => str.indexOf(substr) !== -1;
-
-        if (contains(options.model, "meta")) {
-            return {
-                prompt,
-                temperature: options.temperature,
-                max_gen_len: options.max_tokens,
-            } as LLama2RequestPayload
-        } else if (contains(options.model, "claude")) {
-
-            const maxToken = () => {
-                if (options.max_tokens) {
-                    return options.max_tokens;
-                
-                } else if (contains(options.model, "claude-3-5")) {
-                    return 8192;
+        if (options.model.includes("amazon")) {
+            //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.model.includes("claude-3-7")) {
+                const thinking_options = options.model_options as BedrockClaudeOptions;
+                const thinking = thinking_options?.thinking_mode ?? false;
+                if (!model_options?.max_tokens) {
+                    model_options.max_tokens = thinking ? 128000 : 8192;
+                }
+                additionalField = {
+                    top_k: model_options?.top_k,
+                    reasoning_config: {
+                        type: thinking ? "enabled" : "disabled",
+                        budget_tokens: thinking_options?.thinking_budget_tokens,
+                    }
+                };
+                if(thinking && (thinking_options?.thinking_budget_tokens ?? 0) > 64000){
+                    additionalField = {
+                        ...additionalField,
+                        anthorpic_beta: ["output-128k-2025-02-19"]
+                    };
+                }
+            }
+            //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 {
-                    return 4096
+                    model_options.max_tokens = 4096;
                 }
             }
-            return {
-                anthropic_version: "bedrock-2023-05-31",
-                ...(prompt as ClaudeMessagesPrompt),
-                temperature: options.temperature,
-                max_tokens: maxToken(),
-            } as ClaudeRequestPayload;
-        } else if (contains(options.model, "ai21")) {
-            return {
-                prompt: prompt,
-                temperature: options.temperature,
-                maxTokens: options.max_tokens,
-            } as AI21RequestPayload;
-        } else if (contains(options.model, "command-r-plus")) {
-            return {
-                message: prompt as string,
-                max_tokens: options.max_tokens,
-                temperature: options.temperature,
-            } as CohereCommandRPayload;
+            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);
+        } else if (options.model.includes("mistral")) {
+            //7B instruct and 8x7B instruct
+            if (options.model.includes("7b")) {
+                additionalField = { top_k: model_options?.top_k };
+                //Does not support system messages
+                if (prompt.system && prompt.system?.length != 0) {
+                    prompt.messages?.push(converseSystemToMessages(prompt.system));
+                    prompt.system = undefined;
+                    prompt.messages = converseConcatMessages(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")) {
+                additionalField = {
+                    presencePenalty: { scale: model_options?.presence_penalty },
+                    frequencyPenalty: { scale: model_options?.frequency_penalty },
+                };
+                //Does not support system messages
+                if (prompt.system && prompt.system?.length != 0) {
+                    prompt.messages?.push(converseSystemToMessages(prompt.system));
+                    prompt.system = undefined;
+                    prompt.messages = converseConcatMessages(prompt.messages);
+                }
+            }
+        } 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 = {
+                    k: model_options?.top_k,
+                    frequency_penalty: model_options?.frequency_penalty,
+                    presence_penalty: model_options?.presence_penalty,
+                };
+            } else {
+                // Command non-R
+                additionalField = { k: model_options?.top_k };
+                //Does not support system messages
+                if (prompt.system && prompt.system?.length != 0) {
+                    prompt.messages?.push(converseSystemToMessages(prompt.system));
+                    prompt.system = undefined;
+                    prompt.messages = converseConcatMessages(prompt.messages);
+                }
+            }
+        }
 
+        //If last message is "```json", add corresponding ``` as a stop sequence.
+        if (prompt.messages && prompt.messages.length > 0) {
+            if (prompt.messages[prompt.messages.length - 1].content?.[0].text === "```json") {
+                let stopSeq = model_options?.stop_sequence;
+                if (!stopSeq) {
+                    model_options.stop_sequence = ["```"];
+                } else if (!stopSeq.includes("```")) {
+                    stopSeq.push("```");
+                    model_options.stop_sequence = stopSeq;
+                }
+            }
         }
-        else if (contains(options.model, "cohere")) {
-            return {
-                prompt: prompt,
-                temperature: options.temperature,
-                max_tokens: options.max_tokens,
-            } as CohereRequestPayload;
-        } else if (contains(options.model, "amazon")) {
-            return {
-                inputText: "User: " + (prompt as string) + "\nBot:", // see https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-text.html#model-parameters-titan-request-response
-                textGenerationConfig: {
-                    temperature: options.temperature,
-                    topP: options.top_p,
-                    maxTokenCount: options.max_tokens,
-                    //stopSequences: ["\n"],
-                },
-            } as AmazonRequestPayload;
-        } else if (contains(options.model, "mistral")) {
-            return {
-                prompt: prompt,
-                temperature: options.temperature,
-                max_tokens: options.max_tokens,
-            } as MistralPayload;
-        } else {
-            throw new Error("Cannot prepare payload for unknown provider: " + options.model);
+
+        return {
+            messages: prompt.messages,
+            system: prompt.system,
+            modelId: options.model,
+            inferenceConfig: {
+                maxTokens: model_options?.max_tokens,
+                temperature: model_options?.temperature,
+                topP: model_options?.top_p,
+                stopSequences: model_options?.stop_sequence,
+            } satisfies InferenceConfiguration,
+            additionalModelRequestFields: {
+                ...additionalField,
+            },
+        } satisfies ConverseRequest;
+    }
+
+
+    async requestImageGeneration(prompt: NovaMessagesPrompt, options: ExecutionOptions): Promise<Completion<ImageGeneration>> {
+        if (options.output_modality !== Modalities.image) {
+            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 });
         }
+        const model_options = options.model_options as NovaCanvasOptions;
+
+        const executor = this.getExecutor();
+        const taskType = model_options.taskType ?? NovaImageGenerationTaskType.TEXT_IMAGE;
+
+        this.logger.info("Task type: " + taskType);
+
+        if (typeof prompt === "string") {
+            throw new Error("Bad prompt format");
+        }
+
+        const payload = await formatNovaImageGenerationPayload(taskType, prompt, options);
 
+        const res = await executor.invokeModel({
+            modelId: options.model,
+            contentType: "application/json",
+            accept: "application/json",
+            body: JSON.stringify(payload),
+        },
+            {
+                requestTimeout: 60000 * 5
+            });
+
+        const decoder = new TextDecoder();
+        const body = decoder.decode(res.body);
+        const result = JSON.parse(body);
+
+        return {
+            error: result.error,
+            result: {
+                images: result.images,
+            }
+        }
     }
 
     async startTraining(dataset: DataSource, options: TrainingOptions): Promise<TrainingJob> {
@@ -387,13 +515,14 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
     async listModels(): Promise<AIModel[]> {
         this.logger.debug("[Bedrock] listing models");
         // exclude trainable models since they are not executable
-        const filter = (m: FoundationModelSummary) => m.inferenceTypesSupported?.includes("ON_DEMAND") ?? false;
+        // exclude embedding models, not to be used for typical completions.
+        const filter = (m: FoundationModelSummary) => (m.inferenceTypesSupported?.includes("ON_DEMAND") && !m.outputModalities?.includes("EMBEDDING")) ?? false;
         return this._listModels(filter);
     }
 
     async _listModels(foundationFilter?: (m: FoundationModelSummary) => boolean): Promise<AIModel[]> {
         const service = this.getService();
-        const [foundationals, customs] = await Promise.all([
+        const [foundationals, customs, inferenceProfiles] = 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
@@ -402,6 +531,10 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
                 this.logger.warn("[Bedrock] Can't list custom models. Check if the user has the right permissions.");
                 return undefined
             }),
+            service.listInferenceProfiles({}).catch(() => {
+                this.logger.warn("[Bedrock] Can't list inference profiles. Check if the user has the right permissions.");
+                return undefined
+            }),
         ]);
 
         if (!foundationals?.modelSummaries) {
@@ -454,21 +587,41 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
             });
         }
 
+        //add inference profiles
+        if (inferenceProfiles?.inferenceProfileSummaries) {
+            inferenceProfiles.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,
+                };
+
+                aimodels.push(model);
+            });
+        }
+
         return aimodels;
     }
 
-    async generateEmbeddings({ content, model = "amazon.titan-embed-text-v1" }: EmbeddingsOptions): Promise<EmbeddingsResult> {
+    async generateEmbeddings({ text, image, model }: EmbeddingsOptions): Promise<EmbeddingsResult> {
 
         this.logger.info("[Bedrock] Generating embeddings with model " + model);
+        const defaultModel = image ? "amazon.titan-embed-image-v1" : "amazon.titan-embed-text-v2:0";
+        const modelID = model ?? defaultModel;
 
         const invokeBody = {
-            inputText: content
+            inputText: text,
+            inputImage: image
         }
 
         const executor = this.getExecutor();
         const res = await executor.invokeModel(
             {
-                modelId: model,
+                modelId: modelID,
                 contentType: "application/json",
                 body: JSON.stringify(invokeBody),
             }
@@ -485,16 +638,12 @@ export class BedrockDriver extends AbstractDriver<BedrockDriverOptions, BedrockP
 
         return {
             values: result.embedding,
-            model: model,
+            model: modelID,
             token_count: result.inputTextTokenCount
         };
-
     }
-
 }
 
-
-
 function jobInfo(job: GetModelCustomizationJobCommandOutput, jobId: string): TrainingJob {
     const jobStatus = job.status;
     let status = TrainingJobStatus.running;
@@ -517,41 +666,4 @@ function jobInfo(job: GetModelCustomizationJobCommandOutput, jobId: string): Tra
         status,
         details
     }
-}
-
-
-function claudeFinishReason(reason: string | undefined) {
-    if (!reason) return undefined;
-    switch (reason) {
-        case 'end_turn': return "stop";
-        case 'max_tokens': return "length";
-        default: return reason; //stop_sequence
-    }
-}
-
-function cohereFinishReason(reason: string | undefined) {
-    if (!reason) return undefined;
-    switch (reason) {
-        case 'COMPLETE': return "stop";
-        case 'MAX_TOKENS': return "length";
-        default: return reason;
-    }
-}
-
-function a21FinishReason(reason: string | undefined) {
-    if (!reason) return undefined;
-    switch (reason) {
-        case 'endoftext': return "stop";
-        case 'length': return "length";
-        default: return reason;
-    }
-}
-
-function titanFinishReason(reason: string | undefined) {
-    if (!reason) return undefined;
-    switch (reason) {
-        case 'FINISH': return "stop";
-        case 'LENGTH': return "length";
-        default: return reason;
-    }
-}
+}