feat: support vision ollama

package.json

@@ -1,6 +1,6 @@
 {
   "name": "dcbot",
-  "version": "1.0.3",
+  "version": "1.1.0",
   "description": "",
   "main": "index.js",
   "exports": "./dist/index.js",

src/commands/ai.ts

@@ -4,41 +4,60 @@ import {
   CommandInteraction,
   AutocompleteInteraction,
   MessageContextMenuCommandInteraction,
+  Attachment,
 } from "discord.js";
+import { fetch } from "undici";
 
 import { Ollama } from "ollama";
-const { OLLAMA_URL, OLLAMA_MODEL } = process.env;
+const { OLLAMA_URL, OLLAMA_MODEL, OLLAMA_VISION_MODEL } = process.env;
 
 export default {
   data: new SlashCommandBuilder()
     .setName("ai")
     .setDescription("Ask ai something")
     .addStringOption((option) => option.setName("question").setDescription("Question").setRequired(true))
+    .addAttachmentOption((option) => option.setName("attachment").setDescription("Attachment"))
     .addStringOption((option) => option.setName("model").setDescription("Model Name").setAutocomplete(true)),
   featureId: "ai",
   async execute(interaction: CommandInteraction | MessageContextMenuCommandInteraction) {
     await interaction.deferReply();
     const question =
       interaction instanceof MessageContextMenuCommandInteraction
-        ? interaction.targetMessage
+        ? interaction.targetMessage.content.length
+          ? interaction.targetMessage.content
+          : "Describe this picture."
         : interaction.options.get("question", true).value;
+    const attachments =
+      interaction instanceof MessageContextMenuCommandInteraction
+        ? interaction.targetMessage.attachments.size
+          ? interaction.targetMessage.attachments.toJSON()
+          : []
+        : interaction.options.get("attachment", false)
+        ? [interaction.options.get("attachment")!.attachment as Attachment]
+        : [];
+    const images = await Promise.all(
+      attachments.map((attachment) =>
+        fetch(attachment.url)
+          .then((res) => res.arrayBuffer())
+          .then((res) => Buffer.from(res).toString("base64"))
+      )
+    );
     const modelName =
       interaction instanceof MessageContextMenuCommandInteraction
         ? null
         : (interaction.options.get("model")?.value as string | undefined);
     try {
-      const ollama = new Ollama({ host: OLLAMA_URL || "http://127.0.0.1:11434" });
+      const ollama = new Ollama({
+        host: OLLAMA_URL || "http://127.0.0.1:11434",
+      });
       const response = await ollama.chat({
-        model: modelName || OLLAMA_MODEL || "",
-        messages: [{ role: "user", content: `${question}` }],
-        options: {
-          stop: ["<|eot_id|>"],
-        },
+        model: modelName || (images.length ? OLLAMA_VISION_MODEL : OLLAMA_MODEL) || "",
+        messages: [{ role: "user", content: `${question}`, images: images }],
         keep_alive: 0,
       });
 
       await interaction.followUp({
-        embeds: [new AIEmbed(interaction.client.user, `${question}`, response)],
+        embeds: [new AIEmbed(interaction.client.user, `${question}`, response, attachments)],
       });
     } catch (error) {
       await interaction.followUp({

src/commands/chat.ts

@@ -35,14 +35,14 @@ export default {
           embeds: [
             new InfoEmbed(
               interaction.user,
-              ":calling: Created a thread for you!",
+              `:calling: Created a thread for you!`,
               "Feel free to chat with me. If you have done your conversation, just type `/chat end` to end the chat."
             ),
           ],
           fetchReply: true,
         });
         const thread = await message.startThread({
-          name: `Conversation with ${interaction.user.displayName}`,
+          name: `🦙 ${modelName.split(":")[0]} - Conversation with ${interaction.user.displayName}`,
           autoArchiveDuration: ThreadAutoArchiveDuration.OneHour,
         });
         conversation.set(thread.id, { model: modelName, creater: interaction.user.id, messages: [] });
@@ -85,15 +85,22 @@ export default {
 export const replyConversation = async (message: DiscordMessage<boolean>) => {
   let history = conversation.get(message.channelId);
   if (history && message.content && message.author.id != client.user.id) {
-    history.messages.push({ role: "user", content: message.content });
+    let images: string[] = [];
+    if (message.attachments.size) {
+      images = await Promise.all(
+        message.attachments.toJSON().map((attachment) =>
+          fetch(attachment.url)
+            .then((res) => res.arrayBuffer())
+            .then((res) => Buffer.from(res).toString("base64"))
+        )
+      );
+    }
+    history.messages.push({ role: "user", content: message.content, images: images });
     try {
       const ollama = new Ollama({ host: OLLAMA_URL || "http://127.0.0.1:11434" });
       const response = await ollama.chat({
         model: history.model,
         messages: history.messages,
-        options: {
-          stop: ["<|eot_id|>"],
-        },
         keep_alive: 0,
       });
       history.messages.push({ role: "assistant", content: response.message.content });

src/utils/Embed.ts

@@ -1,4 +1,4 @@
-import { EmbedBuilder, Client, User, APIUser, ClientUser } from "discord.js";
+import { EmbedBuilder, Client, User, APIUser, ClientUser, Attachment } from "discord.js";
 import { TrackMetadata } from "../types/Track";
 import { ChatResponse } from "ollama";
 
@@ -77,12 +77,13 @@ export class InfoEmbed extends EmbedBuilder {
 }
 
 export class AIEmbed extends EmbedBuilder {
-  constructor(user: ClientUser | User | APIUser, question: string, response: ChatResponse) {
+  constructor(user: ClientUser | User | APIUser, question: string, response: ChatResponse, attachments: Attachment[] = []) {
     super();
     this.setColor(0x53fafa)
       .setTitle(":llama: AI answer")
       .setDescription(`Q: ${question}\nA: ${response.message.content}`)
       .setTimestamp()
+      .setImage(attachments.length ? attachments[0].url : null)
       .setFooter({
         text: `${response.model} | ${((response.eval_count * 1000000000) / response.eval_duration).toFixed(2)} tps`,
       });