@llumiverse/drivers 0.15.0 → 0.17.0

package/src/vertexai/models/imagen.ts

@@ -0,0 +1,401 @@
+import { AIModel, Completion, ExecutionOptions, ImageGeneration, Modalities, ModelType, PromptRole, PromptSegment, readStreamAsBase64 } from "@llumiverse/core";
+import { VertexAIDriver } from "../index.js";
+
+const projectId = process.env.GOOGLE_PROJECT_ID;
+const location = 'us-central1';
+
+import aiplatform, { protos } from '@google-cloud/aiplatform';
+
+// Imports the Google Cloud Prediction Service Client library
+const { PredictionServiceClient } = aiplatform.v1;
+
+// Import the helper module for converting arbitrary protobuf.Value objects
+import { helpers } from '@google-cloud/aiplatform';
+import { ImagenOptions } from "../../../../core/src/options/vertexai.js";
+
+interface ImagenBaseReference {
+    referenceType: "REFERENCE_TYPE_RAW" | "REFERENCE_TYPE_MASK" | "REFERENCE_TYPE_SUBJECT" |
+    "REFERENCE_TYPE_CONTROL" | "REFERENCE_TYPE_STYLE";
+    referenceId: number;
+    referenceImage: {
+        bytesBase64Encoded: string; //10MB max
+    }
+}
+
+export enum ImagenTaskType {
+    TEXT_IMAGE = "TEXT_IMAGE",
+    EDIT_MODE_INPAINT_REMOVAL = "EDIT_MODE_INPAINT_REMOVAL",
+    EDIT_MODE_INPAINT_INSERTION = "EDIT_MODE_INPAINT_INSERTION",
+    EDIT_MODE_BGSWAP = "EDIT_MODE_BGSWAP",
+    EDIT_MODE_OUTPAINT = "EDIT_MODE_OUTPAINT",
+    CUSTOMIZATION_SUBJECT = "CUSTOMIZATION_SUBJECT",
+    CUSTOMIZATION_STYLE = "CUSTOMIZATION_STYLE",
+    CUSTOMIZATION_CONTROLLED = "CUSTOMIZATION_CONTROLLED",
+    CUSTOMIZATION_INSTRUCT = "CUSTOMIZATION_INSTRUCT",
+}
+
+export enum ImagenMaskMode {
+    MASK_MODE_USER_PROVIDED = "MASK_MODE_USER_PROVIDED",
+    MASK_MODE_BACKGROUND = "MASK_MODE_BACKGROUND",
+    MASK_MODE_FOREGROUND = "MASK_MODE_FOREGROUND",
+    MASK_MODE_SEMANTIC = "MASK_MODE_SEMANTIC",
+}
+
+interface ImagenReferenceRaw extends ImagenBaseReference {
+    referenceType: "REFERENCE_TYPE_RAW";
+}
+
+interface ImagenReferenceMask extends Omit<ImagenBaseReference, "referenceImage"> {
+    referenceType: "REFERENCE_TYPE_MASK";
+    maskImageConfig: {
+        maskMode?: ImagenMaskMode;
+        maskClasses?: number[]; //Used for MASK_MODE_SEMANTIC, based on https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/imagen-api-customization#segment-ids
+        dilation?: number; //Recommendation depends on mode: Inpaint: 0.01, BGSwap: 0.0, Outpaint: 0.01-0.03
+    }
+    referenceImage?: {  //Only used for MASK_MODE_USER_PROVIDED
+        bytesBase64Encoded: string; //10MB max
+    }
+}
+
+interface ImagenReferenceSubject extends ImagenBaseReference {
+    referenceType: "REFERENCE_TYPE_SUBJECT";
+    subjectImageConfig: {
+        subjectDescription: string;
+        subjectType: "SUBJECT_TYPE_PERSON" | "SUBJECT_TYPE_ANIMAL" | "SUBJECT_TYPE_PRODUCT" | "SUBJECT_TYPE_DEFAULT";
+    }
+}
+
+interface ImagenReferenceControl extends ImagenBaseReference {
+    referenceType: "REFERENCE_TYPE_CONTROL";
+    controlImageConfig: {
+        controlType: "CONTROL_TYPE_FACE_MESH" | "CONTROL_TYPE_CANNY" | "CONTROL_TYPE_SCRIBBLE";
+        enableControlImageComputation?: boolean; //If true, the model will compute the control image
+    }
+}
+
+interface ImagenReferenceStyle extends ImagenBaseReference {
+    referenceType: "REFERENCE_TYPE_STYLE";
+    styleImageConfig: {
+        styleDescription?: string;
+    }
+}
+
+type ImagenMessage = ImagenReferenceRaw | ImagenReferenceMask | ImagenReferenceSubject | ImagenReferenceControl | ImagenReferenceStyle;
+
+export interface ImagenPrompt {
+    prompt: string;
+    referenceImages?: ImagenMessage[];
+    subjectDescription?: string; //Used for image customization to describe in the reference image
+    negativePrompt?: string; //Used for negative prompts
+}
+
+// Specifies the location of the api endpoint
+const clientOptions = {
+    apiEndpoint: `${location}-aiplatform.googleapis.com`,
+};
+
+// Instantiates a client
+const predictionServiceClient = new PredictionServiceClient(clientOptions);
+
+function getImagenParameters(taskType: string, options: ImagenOptions) {
+    const commonParameters = {
+        sampleCount: options?.number_of_images,
+        seed: options?.seed,
+        safetySetting: options?.safety_setting,
+        personGeneration: options?.person_generation,
+        negativePrompt: taskType ? undefined : "", //Filled in later from the prompt
+        //TODO: Add more safety and prompt rejection information
+        //includeSafetyAttributes: true,
+        //includeRaiReason: true,
+    };
+    switch (taskType) {
+        case ImagenTaskType.EDIT_MODE_INPAINT_REMOVAL:
+            return {
+                ...commonParameters,
+                editMode: "EDIT_MODE_INPAINT_REMOVAL",
+                editConfig: {
+                    baseSteps: options?.edit_steps,
+                },
+            }
+        case ImagenTaskType.EDIT_MODE_INPAINT_INSERTION:
+            return {
+                ...commonParameters,
+                editMode: "EDIT_MODE_INPAINT_INSERTION",
+                editConfig: {
+                    baseSteps: options?.edit_steps,
+                },
+            }
+        case ImagenTaskType.EDIT_MODE_BGSWAP:
+            return {
+                ...commonParameters,
+                editMode: "EDIT_MODE_BGSWAP",
+                editConfig: {
+                    baseSteps: options?.edit_steps,
+                },
+            }
+        case ImagenTaskType.EDIT_MODE_OUTPAINT:
+            return {
+                ...commonParameters,
+                editMode: "EDIT_MODE_OUTPAINT",
+                editConfig: {
+                    baseSteps: options?.edit_steps,
+                },
+            }
+        case ImagenTaskType.TEXT_IMAGE:
+            return {
+                ...commonParameters,
+                // You can't use a seed value and watermark at the same time.
+                addWatermark: options?.add_watermark,
+                aspectRatio: options?.aspect_ratio,
+                enhancePrompt: options?.enhance_prompt,
+            };
+        case ImagenTaskType.CUSTOMIZATION_SUBJECT:
+        case ImagenTaskType.CUSTOMIZATION_CONTROLLED:
+        case ImagenTaskType.CUSTOMIZATION_INSTRUCT:
+        case ImagenTaskType.CUSTOMIZATION_STYLE:
+            return {
+                ...commonParameters,
+            }
+        default:
+            throw new Error("Task type not supported");
+    }
+}
+
+export class ImagenModelDefinition {
+
+    model: AIModel
+
+    constructor(modelId: string) {
+        this.model = {
+            id: modelId,
+            name: modelId,
+            provider: 'vertexai',
+            type: ModelType.Image,
+            can_stream: false,
+        };
+    }
+
+    async createPrompt(_driver: VertexAIDriver, segments: PromptSegment[], options: ExecutionOptions): Promise<ImagenPrompt> {
+        const splits = options.model.split("/");
+        const modelName = splits[splits.length - 1];
+        options = { ...options, model: modelName };
+
+        const prompt: ImagenPrompt = {
+            prompt: "",
+        }
+
+        //Collect text prompts, Imagen does not support roles, so everything gets merged together 
+        // however we still respect our typical pattern. System First, Safety Last.
+        const system: string[] = [];
+        const user: string[] = [];
+        const safety: string[] = [];
+        const negative: string[] = [];
+
+        const mask_mode = (options.model_options as ImagenOptions)?.mask_mode;
+        const imagenOptions = options.model_options as ImagenOptions;
+
+        for (const msg of segments) {
+            if (msg.role === PromptRole.safety) {
+                safety.push(msg.content);
+            } else if (msg.role === PromptRole.system) {
+                system.push(msg.content);
+            } else if (msg.role === PromptRole.negative) {
+                negative.push(msg.content);
+            } else {
+                //Everything else is assumed to be user or user adjacent.
+                user.push(msg.content);
+            }
+            if (msg.files) {
+                //Get images from messages
+                if (!prompt.referenceImages) {
+                    prompt.referenceImages = [];
+                }
+
+                //Always required, but only used by customisation. 
+                //Each ref ID refers to a single "reference", i.e. object. To provide multiple images of a single ref,
+                //include multiple images in one prompt.
+                const refId = prompt.referenceImages.length + 1;
+                for (const img of msg.files) {
+                    if (img.mime_type?.includes("image")) {
+                        if (msg.role !== PromptRole.mask) {
+                            //Editing based mode requires a reference image
+                            if (imagenOptions?.edit_mode?.includes("EDIT_MODE")) {
+                                prompt.referenceImages.push({
+                                    referenceType: "REFERENCE_TYPE_RAW",
+                                    referenceId: refId,
+                                    referenceImage: {
+                                        bytesBase64Encoded: await readStreamAsBase64(await img.getStream()),
+                                    }
+                                });
+                                //If mask is auto-generated, add a mask reference
+                                if (mask_mode !== ImagenMaskMode.MASK_MODE_USER_PROVIDED) {
+                                    prompt.referenceImages.push({
+                                        referenceType: "REFERENCE_TYPE_MASK",
+                                        referenceId: refId,
+                                        maskImageConfig: {
+                                            maskMode: mask_mode,
+                                            dilation: imagenOptions?.mask_dilation,
+                                        }
+                                    });
+                                }
+                            }
+                            else if ((options.model_options as ImagenOptions)?.edit_mode === ImagenTaskType.CUSTOMIZATION_SUBJECT) {
+                                //First image is always the control image
+                                if (refId == 1) {
+                                    //Customization subject mode requires a control image
+                                    prompt.referenceImages.push({
+                                        referenceType: "REFERENCE_TYPE_CONTROL",
+                                        referenceId: refId,
+                                        referenceImage: {
+                                            bytesBase64Encoded: await readStreamAsBase64(await img.getStream()),
+                                        },
+                                        controlImageConfig: {
+                                            controlType: imagenOptions?.controlType === "CONTROL_TYPE_FACE_MESH" ? "CONTROL_TYPE_FACE_MESH" : "CONTROL_TYPE_CANNY",
+                                            enableControlImageComputation: imagenOptions?.controlImageComputation,
+                                        }
+                                    });
+                                } else {
+                                    // Subject images
+                                    prompt.referenceImages.push({
+                                        referenceType: "REFERENCE_TYPE_SUBJECT",
+                                        referenceId: refId,
+                                        referenceImage: {
+                                            bytesBase64Encoded: await readStreamAsBase64(await img.getStream()),
+                                        },
+                                        subjectImageConfig: {
+                                            subjectDescription: prompt.subjectDescription ?? msg.content,
+                                            subjectType: imagenOptions?.subjectType ?? "SUBJECT_TYPE_DEFAULT",
+                                        }
+                                    });
+                                }
+                            } else if ((options.model_options as ImagenOptions)?.edit_mode === ImagenTaskType.CUSTOMIZATION_STYLE) {
+                                // Style images
+                                prompt.referenceImages.push({
+                                    referenceType: "REFERENCE_TYPE_STYLE",
+                                    referenceId: refId,
+                                    referenceImage: {
+                                        bytesBase64Encoded: await readStreamAsBase64(await img.getStream()),
+                                    },
+                                    styleImageConfig: {
+                                        styleDescription: prompt.subjectDescription ?? msg.content,
+                                    }
+                                });
+                            } else if ((options.model_options as ImagenOptions)?.edit_mode === ImagenTaskType.CUSTOMIZATION_CONTROLLED) {
+                                // Control images
+                                prompt.referenceImages.push({
+                                    referenceType: "REFERENCE_TYPE_CONTROL",
+                                    referenceId: refId,
+                                    referenceImage: {
+                                        bytesBase64Encoded: await readStreamAsBase64(await img.getStream()),
+                                    },
+                                    controlImageConfig: {
+                                        controlType: imagenOptions?.controlType === "CONTROL_TYPE_FACE_MESH" ? "CONTROL_TYPE_FACE_MESH" : "CONTROL_TYPE_CANNY",
+                                        enableControlImageComputation: imagenOptions?.controlImageComputation,
+                                    }
+                                });
+                            } else if ((options.model_options as ImagenOptions)?.edit_mode === ImagenTaskType.CUSTOMIZATION_INSTRUCT) {
+                                // Control images
+                                prompt.referenceImages.push({
+                                    referenceType: "REFERENCE_TYPE_RAW",
+                                    referenceId: refId,
+                                    referenceImage: {
+                                        bytesBase64Encoded: await readStreamAsBase64(await img.getStream()),
+                                    },
+                                });
+                            }
+                        }
+                        //If mask is user-provided, add a mask reference
+                        if (msg.role === PromptRole.mask && mask_mode === ImagenMaskMode.MASK_MODE_USER_PROVIDED) {
+                            prompt.referenceImages.push({
+                                referenceType: "REFERENCE_TYPE_MASK",
+                                referenceId: refId,
+                                referenceImage: {
+                                    bytesBase64Encoded: await readStreamAsBase64(await img.getStream()),
+                                },
+                                maskImageConfig: {
+                                    maskMode: mask_mode,
+                                    dilation: imagenOptions?.mask_dilation,
+                                }
+                            });
+                        }
+                    }
+                }
+            }
+        }
+
+        //Extract the text from the segments
+        prompt.prompt += [system.join("\n\n"), user.join("\n\n"), safety.join("\n\n")].join("\n\n");
+
+        //Negative prompt
+        if (negative.length > 0) {
+            prompt.negativePrompt = negative.join(", ");
+        }
+
+        console.log(prompt);
+
+        return prompt
+    }
+
+    async requestImageGeneration(driver: VertexAIDriver, prompt: ImagenPrompt, options: ExecutionOptions): Promise<Completion<ImageGeneration>> {
+        if (options.model_options?._option_id !== "vertexai-imagen") {
+            driver.logger.warn("Invalid model options", {options: options.model_options });
+        }
+        options.model_options = options.model_options as ImagenOptions;
+
+        if (options.output_modality !== Modalities.image) {
+            throw new Error(`Image generation requires image output_modality`);
+        }
+
+        const taskType: string = options.model_options.edit_mode ?? ImagenTaskType.TEXT_IMAGE;
+
+        driver.logger.info("Task type: " + taskType);
+
+        const modelName = options.model.split("/").pop() ?? '';
+
+        // Configure the parent resource
+        const endpoint = `projects/${projectId}/locations/${location}/publishers/google/models/${modelName}`;
+
+        const instanceValue = helpers.toValue(prompt);
+        if (!instanceValue) {
+            throw new Error('No instance value found');
+        }
+        const instances = [instanceValue];
+
+        let parameter: any = getImagenParameters(taskType, options.model_options);
+        parameter.negativePrompt = prompt.negativePrompt ?? undefined;
+
+        const numberOfImages = options.model_options?.number_of_images ?? 1;
+
+        // Remove all undefined values
+        parameter = Object.fromEntries(
+            Object.entries(parameter).filter(([_, v]) => v !== undefined)
+        ) as any;
+
+        const parameters = helpers.toValue(parameter);
+
+        const request: protos.google.cloud.aiplatform.v1.IPredictRequest = {
+            endpoint,
+            instances,
+            parameters,
+        };
+
+        // Predict request
+        const [response] = await predictionServiceClient.predict(request, { timeout: 120000 * numberOfImages }); //Extended timeout for image generation
+        const predictions = response.predictions;
+
+        if (!predictions) {
+            throw new Error('No predictions found');
+        }
+
+        // Extract base64 encoded images from predictions
+        const images: string[] = predictions.map(prediction =>
+            prediction.structValue?.fields?.bytesBase64Encoded?.stringValue ?? ''
+        );
+
+        return {
+            result: {
+                images
+            },
+        };
+    }
+}

package/src/vertexai/models.ts

@@ -1,89 +1,27 @@
-import { AIModel, Completion, ExecutionOptions, ModelType, PromptOptions, PromptSegment } from "@llumiverse/core";
-import { VertexAIDriver } from "./index.js";
+import { AIModel, Completion, CompletionChunkObject, PromptOptions, PromptSegment, ExecutionOptions } from "@llumiverse/core";
+import { VertexAIDriver , trimModelName} from "./index.js";
 import { GeminiModelDefinition } from "./models/gemini.js";
-
-
-
+import { ClaudeModelDefinition } from "./models/claude.js";
 
 export interface ModelDefinition<PromptT = any> {
     model: AIModel;
     versions?: string[]; // the versions of the model that are available. ex: ['001', '002']
     createPrompt: (driver: VertexAIDriver, segments: PromptSegment[], options: PromptOptions) => Promise<PromptT>;
-    requestCompletion: (driver: VertexAIDriver, prompt: PromptT, options: ExecutionOptions) => Promise<Completion>;
-    requestCompletionStream: (driver: VertexAIDriver, promp: PromptT, options: ExecutionOptions) => Promise<AsyncIterable<string>>;
-}
-
-export function getModelName(model: string) {
-    const i = model.lastIndexOf('@');
-    return i > -1 ? model.substring(0, i) : model;
+    requestTextCompletion: (driver: VertexAIDriver, prompt: PromptT, options: ExecutionOptions) => Promise<Completion>;
+    requestTextCompletionStream: (driver: VertexAIDriver, promp: PromptT, options: ExecutionOptions) => Promise<AsyncIterable<CompletionChunkObject>>;
 }
 
 export function getModelDefinition(model: string): ModelDefinition {
-    const modelName = getModelName(model);
-    const def = Models[modelName];
-    if (!def) {
-        throw new Error(`Unknown model ${model}`);
+    const splits = model.split("/");
+    const publisher = splits[1];
+    const modelName = trimModelName(splits[splits.length - 1]);
+    
+    if (publisher?.includes("anthropic")) {
+        return new ClaudeModelDefinition(modelName);
+    } else if (publisher?.includes("google")) {
+        return new GeminiModelDefinition(modelName);
     }
-    return def;
-}
-
-export function getAIModels() {
-    return Object.values(Models).map(m => m.model);
-}
-
-// Builtin models. VertexAI doesn't provide an API to list models. so we have to hardcode them here.
-export const BuiltinModels: AIModel<string>[] = [
-    {
-        id: "gemini-1.5-flash",
-        name: "Gemini Pro 1.5 Flash",
-        provider: "vertexai",
-        owner: "google",
-        type: ModelType.MultiModal,
-        can_stream: true,
-        is_multimodal: true
 
-    },
-    {
-        id: "gemini-1.5-pro",
-        name: "Gemini Pro 1.5 Pro",
-        provider: "vertexai",
-        owner: "google",
-        type: ModelType.MultiModal,
-        can_stream: true,
-        is_multimodal: true
-
-    },
-    {
-        id: "gemini-1.0-pro",
-        name: "Gemini Pro 1.0",
-        provider: "vertexai",
-        owner: "google",
-        type: ModelType.Text,
-        can_stream: true,
-    },
-    {
-        id: "tablextembedding-gecko",
-        name: "Gecko Text Embeddings",
-        provider: "vertexai",
-        owner: "google",
-        type: ModelType.Embedding,
-    },
-    {
-        id: "textembedding-gecko-multilingual",
-        name: "Gecko Multilingual Text Embeddings",
-        provider: "vertexai",
-        owner: "google",
-        type: ModelType.Embedding,
-    },
-
-
-
-]
-
-
-
-const Models: Record<string, ModelDefinition> = {
-    "gemini-1.5-flash": new GeminiModelDefinition("gemini-1.5-flash"),
-    "gemini-1.5-pro": new GeminiModelDefinition("gemini-1.5-pro"),
-    "gemini-1.0-pro": new GeminiModelDefinition(),
-}
+    //Fallback, assume it is Gemini.
+    return new GeminiModelDefinition(modelName);
+}

package/src/watsonx/index.ts

@@ -1,4 +1,4 @@
-import { AIModel, AbstractDriver, Completion, DriverOptions, EmbeddingsOptions, EmbeddingsResult, ExecutionOptions } from "@llumiverse/core";
+import { AIModel, AbstractDriver, Completion, CompletionChunk, DriverOptions, EmbeddingsOptions, EmbeddingsResult, ExecutionOptions, TextFallbackOptions } from "@llumiverse/core";
 import { transformSSEStream } from "@llumiverse/core/async";
 import { FetchClient } from "api-fetch-client";
 import { GenerateEmbeddingPayload, GenerateEmbeddingResponse, WatsonAuthToken, WatsonxListModelResponse, WatsonxModelSpec, WatsonxTextGenerationPayload, WatsonxTextGenerationResponse } from "./interfaces.js";
@@ -29,13 +29,21 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
         this.fetchClient = new FetchClient(this.endpoint_url).withAuthCallback(async () => this.getAuthToken().then(token => `Bearer ${token}`));
     }
 
-    async requestCompletion(prompt: string, options: ExecutionOptions): Promise<Completion<any>> {
+    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 });
+        }
+        options.model_options = options.model_options as TextFallbackOptions;
+        
         const payload: WatsonxTextGenerationPayload = {
             model_id: options.model,
             input: prompt + "\n",
             parameters: {
-                max_new_tokens: options.max_tokens,
-                //time_limit: options.time_limit,
+                max_new_tokens: options.model_options.max_tokens,
+                temperature: options.model_options.temperature,
+                top_k: options.model_options.top_k,
+                top_p: options.model_options.top_p,
+                stop_sequences: options.model_options.stop_sequence,
             },
             project_id: this.projectId,
         }
@@ -51,19 +59,25 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
                 result: result.generated_token_count,
                 total: result.input_token_count + result.generated_token_count,
             },
-            finish_reason: result.stop_reason,
+            finish_reason: watsonFinishReason(result.stop_reason),
             original_response: options.include_original_response ? res : undefined,
         }
     }
 
-    async requestCompletionStream(prompt: string, options: ExecutionOptions): Promise<AsyncIterable<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 });
+        }
+        options.model_options = options.model_options as TextFallbackOptions;
         const payload: WatsonxTextGenerationPayload = {
             model_id: options.model,
             input: prompt + "\n",
             parameters: {
-                max_new_tokens: options.max_tokens,
-                //time_limit: options.time_limit,
+                max_new_tokens: options.model_options.max_tokens,
+                temperature: options.model_options.temperature,
+                top_k: options.model_options.top_k,
+                top_p: options.model_options.top_p,
+                stop_sequences: options.model_options.stop_sequence,
             },
             project_id: this.projectId,
         }
@@ -75,7 +89,15 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
 
         return transformSSEStream(stream, (data: string) => {
             const json = JSON.parse(data) as WatsonxTextGenerationResponse;
-            return json.results[0]?.generated_text ?? '';
+            return {
+                result: json.results[0]?.generated_text ?? '',
+                finish_reason: watsonFinishReason(json.results[0]?.stop_reason),
+                token_usage: {
+                    prompt: json.results[0].input_token_count,
+                    result: json.results[0].generated_token_count,
+                    total: json.results[0].input_token_count + json.results[0].generated_token_count,
+                },
+            };
         });
 
     }
@@ -130,15 +152,22 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
         return this.listModels()
             .then(() => true)
             .catch((err) => {
-                this.logger.warn("Failed to connect to WatsonX", err);
+                this.logger.warn("Failed to connect to WatsonX", { error: err });
                 return false
             });
     }
 
     async generateEmbeddings(options: EmbeddingsOptions): Promise<EmbeddingsResult> {
+        if (options.image) {
+            throw new Error("Image embeddings not supported by Watsonx");
+        }
+
+        if (!options.text) {
+            throw new Error ("No text provided");
+        }
 
         const payload: GenerateEmbeddingPayload = {
-            inputs: [options.content],
+            inputs: [options.text],
             model_id: options.model ?? 'ibm/slate-125m-english-rtrvr',
             project_id: this.projectId
         }
@@ -154,6 +183,15 @@ export class WatsonxDriver extends AbstractDriver<WatsonxDriverOptions, string>
 
 }
 
+function watsonFinishReason(reason: string | undefined) {
+    if (!reason) return undefined;
+    switch (reason) {
+        case 'eos_token': return "stop";
+        case 'max_tokens': return "length";
+        default: return reason;
+    }
+}
+
 
 
 /*interface ListModelsParams extends ModelSearchPayload {

package/src/watsonx/interfaces.ts

@@ -6,6 +6,10 @@ export interface WatsonxTextGenerationPayload {
     parameters: {
         max_new_tokens?: number;
         time_limit?: number;
+        stop_sequences?: string[];
+        temperature?: number;
+        top_k?: number;
+        top_p?: number;
     },
     project_id: string;
 }