ChatwatsonxAI

import { WatsonXAI } from "@ibm-cloud/watsonx-ai";
import { Stream } from "@ibm-cloud/watsonx-ai/dist/lib/common";
import WatsonxAiMlVml_v1, {
  ObjectStreamed,
  TextChatStreamResponse,
} from "@ibm-cloud/watsonx-ai/dist/watsonx-ai-ml/vml_v1";
import { UserOptions } from "ibm-cloud-sdk-core";
import fs from "fs";
 
interface WatsonxAIConfig extends UserOptions {
  modelId: string;
  projectId: string;
}
 
interface ChatResponse {
  output: string;
}
 
interface WatsonxAICallOptions {
  prompt: string;
  stream?: boolean;
  context?: string;
  outputFormat?: string;
}
 
interface WatsonxAICallWithChatOptions extends WatsonxAICallOptions {
  chatHistory?: Record<"assistant" | "user", string>[];
  image?: { path: string; type: "local" | "remote" }[];
  systemInstruction?: string;
  options?: Omit<
    WatsonxAiMlVml_v1.TextChatParams,
    "modelId" | "messages" | "projectId"
  >;
}
 
export class ChatWatsonxAI {
  private client: WatsonXAI;
  private modelId: string;
  private projectId: string;
 
  constructor(config: WatsonxAIConfig) {
    const { modelId, projectId, serviceUrl, version } = config;
 
    this.client = WatsonXAI.newInstance({
      version: version,
      serviceUrl,
      ...config,
    });
 
    this.modelId = modelId;
    this.projectId = projectId;
  }
 
  async chat(
    options: WatsonxAICallOptions
  ): Promise<
    | ChatResponse
    | Stream<
        WatsonxAiMlVml_v1.ObjectStreamed<WatsonxAiMlVml_v1.TextGenResponse>
      >
  > {
    const { prompt, context, stream, outputFormat } = options;
 
    const content = await this.buildMessage(prompt, context, outputFormat);
 
    if (stream) {
      const response = await this.client.generateTextStream({
        modelId: this.modelId,
        projectId: this.projectId,
        input: content,
        returnObject: true,
      });
      return response;
    }
 
    const response = await this.client.generateText({
      modelId: this.modelId,
      projectId: this.projectId,
      input: content,
    });
 
    return {
      output:
        response.result?.results?.[0]?.generated_text?.replace(
          /^```json\n|\n```$/g,
          ""
        ) ?? "",
    };
  }
 
  async chatWithHistory(
    options: WatsonxAICallWithChatOptions
  ): Promise<ChatResponse | Stream<ObjectStreamed<TextChatStreamResponse>>> {
    const {
      prompt,
      context,
      stream,
      outputFormat,
      image,
      chatHistory,
      systemInstruction,
    } = options;
 
    let userMessages: { role: string; content: string }[] = [];
 
    if (systemInstruction) {
      userMessages.push({ role: "system", content: systemInstruction });
    }
 
    if (chatHistory && chatHistory.length > 0) {
      chatHistory.forEach((chat) => {
        userMessages.push({ role: "user", content: chat.user });
        userMessages.push({ role: "assistant", content: chat.assistant });
      });
    }
 
    const content = await this.buildMessage(
      prompt,
      context,
      outputFormat,
      image
    );
    userMessages.push({
      role: "user",
      content: image && image.length > 0 ? JSON.parse(content) : content,
    });
 
    if (stream) {
      const response = await this.client.textChatStream({
        modelId: this.modelId,
        projectId: this.projectId,
        returnObject: true,
        messages: userMessages,
        ...options.options,
      });
      return response;
    }
 
    const response = await this.client.textChat({
      modelId: this.modelId,
      projectId: this.projectId,
      messages: userMessages,
      ...options.options,
    });
 
    return {
      output:
        response.result?.choices?.[0]?.message?.content?.replace(
          /^```json\n|\n```$/g,
          ""
        ) ?? "",
    };
  }
 
  private async buildMessage(
    prompt: string,
    context?: string,
    outputFormat?: string,
    file?: { path: string; type: "local" | "remote" }[]
  ): Promise<string> {
    let content = prompt;
 
    if (context) {
      content += `\ncontext:\n${context}`;
    }
 
    if (outputFormat) {
      content += `\noutput format:\n${outputFormat}`;
    }
 
    if (file && file.length > 0) {
      const filePromises = file.map(async (f) => ({
        type: "image_url",
        image_url: {
          url: await this.encodeImage(f.path, f.type),
        },
      }));
 
      const resolvedFiles = await Promise.all(filePromises);
 
      const finalContent = [{ type: "text", text: content }, ...resolvedFiles];
      return JSON.stringify(finalContent);
    }
 
    return content;
  }
 
  private async encodeImage(
    path: string,
    type: "local" | "remote"
  ): Promise<string> {
    if (type === "remote") {
      const response = await fetch(path);
      const buffer = Buffer.from(await response.arrayBuffer()).toString(
        "base64"
      );
      return `data:image/jpeg;base64,${buffer}`;
    }
 
    return `data:image/jpeg;base64,${fs.readFileSync(path).toString("base64")}`;
  }
}