text-gen-bot/index.js

import config from './config.json' assert {type: "json"};
import { MatrixClient, SimpleFsStorageProvider, AutojoinRoomsMixin } from "matrix-bot-sdk";
import fs from "fs";
import { PythonShell } from 'python-shell';

const storage = new SimpleFsStorageProvider("storage.json");
const client = new MatrixClient(config.homeserver, config.token, storage);
const pyFile = "textgen.py";

AutojoinRoomsMixin.setupOnClient(client);
client.start().then(() => console.log(`Client has started!\n`));

// let messageCounter = 0;
const messageCounters = new Map(); // room ID, message count
let trainingCounter = 0;

client.on("room.message", (roomId, event) => {
    if (!event["content"] || event["sender"] === config.user) return;

    // ++messageCounter;
    messageCounters.set(roomId, (messageCounters.get(roomId) ?? 0) + 1);

    ++trainingCounter;
    let userMessage = event["content"]["body"].split(" ");

    if (userMessage[0].startsWith(config.prefix)) {
        userMessage[0] = userMessage[0].replace(config.prefix, '').toLowerCase();
    } else {
        fs.appendFile(config.file, userMessage.join(' ') + "\n", function (err) {
            if (err) throw err;
        });
    };

    // ? send message if:
    // ?    - enough messages have been sent
    // ?    - commanded
    if (!(messageCounter.get(roomId) % config.frequency) || userMessage[0] === "speak") {
        console.log("Generating message...");

        userMessage.shift()
        const options = { args: ['generate'] };

        PythonShell.run(pyFile, options, (err, message) => {
            if (err) throw err;
            client.sendText(roomId, message.toString());
            console.log("Message sent!");
        }); // ? send generated message to room
    };

    // ? retrain if:
    // ?    - enough message have been sent
    // ?    - commanded
    if (trainingCounter >= config.retrain || userMessage[0] === "train") {
        console.log("Retraining the AI...");
        client.sendText(roomId, "Retraining the AI...");

        trainingCounter = 0;
        const options = { args: ['train'] };

        PythonShell.run(pyFile, options, (err, message) => {
            if (err) throw err;
            console.log(message.toString());
        });
        console.log("Training finished!");
        client.sendText(roomId, "Training finished!");
    };
});

function lineCount(text) {
    return fs.readFileSync(text).toString().split("\n").length - 1;
};