pleroma-ollama-bot/src/main.ts

import {
  OllamaRequest,
  OllamaResponse,
  NewStatusBody,
  Notification,
  OllamaConfigOptions,
} from "../types.js";
import striptags from "striptags";
import { PrismaClient } from "../generated/prisma/client.js";
import {
  getInstanceEmojis,
  deleteNotification,
  getNotifications,
} from "./api.js";
import { storeUserData, storePromptData } from "./prisma.js";
import {
  isFromWhitelistedDomain,
  alreadyRespondedTo,
  recordPendingResponse,
  trimInputData,
  selectRandomEmoji,
} from "./util.js";

export const prisma = new PrismaClient();

export const envConfig = {
  pleromaInstanceUrl: process.env.PLEROMA_INSTANCE_URL || "",
  pleromaInstanceDomain: process.env.PLEROMA_INSTANCE_DOMAIN || "",
  whitelistOnly: process.env.ONLY_WHITELIST === "true" ? true : false,
  whitelistedDomains: process.env.WHITELISTED_DOMAINS
    ? process.env.WHITELISTED_DOMAINS.split(",")
    : [process.env.PLEROMA_INSTANCE_DOMAIN],
  ollamaUrl: process.env.OLLAMA_URL || "",
  ollamaSystemPrompt:
    process.env.OLLAMA_SYSTEM_PROMPT ||
    "You are a helpful AI assistant. Answer all questions concisely.",
  ollamaModel: process.env.OLLAMA_MODEL || "",
  fetchInterval: process.env.FETCH_INTERVAL
    ? parseInt(process.env.FETCH_INTERVAL)
    : 15000,
  bearerToken: process.env.INSTANCE_BEARER_TOKEN || "",
};

const ollamaConfig: OllamaConfigOptions = {
  temperature: 1.4,
  top_k: 100,
  top_p: 0.8,
};

// this could be helpful
// https://replicate.com/blog/how-to-prompt-llama

const generateOllamaRequest = async (
  notification: Notification
): Promise<OllamaResponse | undefined> => {
  const { whitelistOnly, ollamaModel, ollamaSystemPrompt, ollamaUrl } =
    envConfig;
  try {
    if (
      striptags(notification.status.content).includes("!prompt") &&
      !notification.status.account.bot && // sanity check, sort of
      notification.type === "mention" &&
      notification.status.visibility !== "private" // for safety, let's only respond to public messages
    ) {
      if (whitelistOnly && !isFromWhitelistedDomain(notification)) {
        await deleteNotification(notification);
        return;
      }
      if (await alreadyRespondedTo(notification)) {
        return;
      }
      await recordPendingResponse(notification);
      await storeUserData(notification);
      const ollamaRequestBody: OllamaRequest = {
        model: ollamaModel,
        system: ollamaSystemPrompt,
        prompt: `[INST] @${
          notification.status.account.fqn
        } says: ${trimInputData(notification.status.content)} [/INST]`,
        stream: false,
        options: ollamaConfig,
      };
      const response = await fetch(`${ollamaUrl}/api/generate`, {
        method: "POST",
        body: JSON.stringify(ollamaRequestBody),
      });
      const ollamaResponse: OllamaResponse = await response.json();
      await storePromptData(notification, ollamaResponse);
      return ollamaResponse;
    }
  } catch (error: any) {
    throw new Error(error.message);
  }
};

const postReplyToStatus = async (
  notification: Notification,
  ollamaResponseBody: OllamaResponse
) => {
  const { pleromaInstanceUrl, bearerToken } = envConfig;
  const emojiList = await getInstanceEmojis();
  let randomEmoji;
  if (emojiList) {
    randomEmoji = selectRandomEmoji(emojiList);
  }
  try {
    let mentions: string[];
    const statusBody: NewStatusBody = {
      content_type: "text/markdown",
      status: `${ollamaResponseBody.response} :${randomEmoji}:`,
      in_reply_to_id: notification.status.id,
    };
    if (
      notification.status.mentions &&
      notification.status.mentions.length > 0
    ) {
      mentions = notification.status.mentions.map((mention) => {
        return mention.acct;
      });
      statusBody.to = mentions;
    }

    const response = await fetch(`${pleromaInstanceUrl}/api/v1/statuses`, {
      method: "POST",
      headers: {
        Authorization: `Bearer ${bearerToken}`,
        "Content-Type": "application/json",
      },
      body: JSON.stringify(statusBody),
    });

    if (!response.ok) {
      throw new Error(`New status request failed: ${response.statusText}`);
    }

    await deleteNotification(notification);
  } catch (error: any) {
    throw new Error(error.message);
  }
};

let notifications = [];
const beginFetchCycle = async () => {
  setInterval(async () => {
    notifications = await getNotifications();
    if (notifications.length > 0) {
      await Promise.all(
        notifications.map(async (notification) => {
          try {
            const ollamaResponse = await generateOllamaRequest(notification);
            if (ollamaResponse) {
              postReplyToStatus(notification, ollamaResponse);
            }
          } catch (error: any) {
            throw new Error(error.message);
          }
        })
      );
    }
  }, envConfig.fetchInterval); // lower intervals may cause the bot to respond multiple times to the same message, but we try to mitigate this with the deleteNotification function
};

console.log(
  `Fetching notifications from ${envConfig.pleromaInstanceDomain}, every ${
    envConfig.fetchInterval / 1000
  } seconds.`
);
console.log(
  `Accepting prompts from: ${envConfig.whitelistedDomains.join(", ")}`
);
console.log(
  `Using model: ${envConfig.ollamaModel}\nConfig: ${JSON.stringify(
    ollamaConfig
  )}`
);
await beginFetchCycle();