Merge pull request #83 from matrixgpt/v2-cleanup

Add persistence

.env.example

@@ -7,6 +7,12 @@ CHATGPT_CONTEXT=thread
 # (Optional) Explicitly set the ChatGPT model to be used by the API.
 #CHATGPT_MODEL=text-chat-davinci-002-20221122
 
+# Set data store settings
+KEYV_BACKEND=file
+KEYV_URL=
+KEYV_BOT_ENCRYPTION=false
+KEYV_BOT_STORAGE=true
+
 # Matrix Static Settings (required, see notes)
 # Defaults to "https://matrix.org"
 MATRIX_HOMESERVER_URL=

package.json

@@ -11,18 +11,23 @@
   },
   "scripts": {
     "build": "rm -f tsconfig.tsbuildinfo && npx tsc",
-    "start": "node dist/index.js",
+    "start": "node --enable-source-maps dist/index.js",
     "typecheck": "npx tsc"
   },
   "dependencies": {
+    "@keyv/mongo": "^2.1.8",
+    "@keyv/postgres": "^1.4.1",
+    "@keyv/redis": "^2.5.4",
+    "@keyv/sqlite": "^3.6.4",
     "chatgpt": "^4.1.1",
     "dotenv": "^16.0.3",
+    "hash.js": "^1.1.7",
+    "keyv": "^4.5.2",
+    "keyv-file": "^0.2.0",
     "markdown-it": "^13.0.1",
     "matrix-bot-sdk": "^0.6.3",
     "typescript": "^4.9.4",
     "znv": "^0.3.2",
     "zod": "^3.20.2"
-  },
-  "devDependencies": {
   }
 }

src/env.ts

@@ -5,6 +5,10 @@ dotenv.config();
 
 export const {
   DATA_PATH,
+  KEYV_BACKEND,
+  KEYV_URL,
+  KEYV_BOT_ENCRYPTION,
+  KEYV_BOT_STORAGE,
   /** Matrix Bot Settings */
   MATRIX_HOMESERVER_URL,
   MATRIX_ACCESS_TOKEN,
@@ -29,6 +33,10 @@ export const {
   CHATGPT_MODEL
 } = parseEnv(process.env, {
   DATA_PATH:                   { schema: z.string().default("./storage"),          description: "Set to /storage/ if using docker, ./storage if running without" },
+  KEYV_BACKEND:                { schema: z.enum(["file", "other"]).default("file"),description: "Set the Keyv backend to 'file' or 'other' if other set KEYV_URL" },
+  KEYV_URL:                    { schema: z.string().default(""),                   description: "Set Keyv backend for storage, in-memory if blank, ignored if KEYV_BACKEND set to `file`"},
+  KEYV_BOT_ENCRYPTION:         { schema: z.boolean().default(false),               description: "Set to true to use a Keyv backend to store bot encryption keys. Uses a file if false."},
+  KEYV_BOT_STORAGE:            { schema: z.boolean().default(false),               description: "Set to true to use a Keyv backend to store bot data. Uses a file if false."},
   /** Matrix Bot Settings */
   MATRIX_HOMESERVER_URL:       { schema: z.string().default("https://matrix.org"), description: "Set matrix homeserver with 'https://' prefix" },
   MATRIX_ACCESS_TOKEN:         { schema: z.string().optional(),                    description: "Set MATRIX_BOT_USERNAME & MATRIX_BOT_PASSWORD to print MATRIX_ACCESS_TOKEN or follow https://webapps.stackexchange.com/questions/131056/how-to-get-an-access-token-for-element-riot-matrix" },

src/index.ts

@@ -1,33 +1,48 @@
+import { ChatGPTAPI } from 'chatgpt'
+import Keyv from 'keyv'
+import { KeyvFile } from 'keyv-file';
 import {
-  MatrixAuth, MatrixClient, SimpleFsStorageProvider, AutojoinRoomsMixin,
+  MatrixAuth, MatrixClient, AutojoinRoomsMixin,
-  LogService, LogLevel,
+  LogService, LogLevel, RichConsoleLogger,
-  RichConsoleLogger,
+  ICryptoStorageProvider, RustSdkCryptoStorageProvider, IStorageProvider, SimpleFsStorageProvider,
-  ICryptoStorageProvider,
-  RustSdkCryptoStorageProvider,
 } from "matrix-bot-sdk";
 
 import * as path from "path";
-import { DATA_PATH, OPENAI_API_KEY, MATRIX_HOMESERVER_URL, MATRIX_ACCESS_TOKEN, MATRIX_AUTOJOIN, MATRIX_BOT_PASSWORD, MATRIX_BOT_USERNAME, MATRIX_ENCRYPTION, MATRIX_THREADS, CHATGPT_CONTEXT, CHATGPT_MODEL } from './env.js'
+import { DATA_PATH, KEYV_URL, OPENAI_API_KEY, MATRIX_HOMESERVER_URL, MATRIX_ACCESS_TOKEN, MATRIX_AUTOJOIN, MATRIX_BOT_PASSWORD, MATRIX_BOT_USERNAME, MATRIX_ENCRYPTION, MATRIX_THREADS, CHATGPT_CONTEXT, CHATGPT_MODEL, KEYV_BOT_ENCRYPTION, KEYV_BOT_STORAGE, KEYV_BACKEND } from './env.js'
-import { parseMatrixUsernamePretty } from './utils.js';
 import CommandHandler from "./handlers.js"
-import { ChatGPTAPI } from 'chatgpt'
+import { KeyvCryptoStorageProvider, KeyvStorageProvider } from './storage.js'
+import { parseMatrixUsernamePretty } from './utils.js';
 
 LogService.setLogger(new RichConsoleLogger());
-
 // Shows the Matrix sync loop details - not needed most of the time
 // LogService.setLevel(LogLevel.DEBUG);
-
 LogService.setLevel(LogLevel.INFO);
-
 // LogService.muteModule("Metrics");
 LogService.trace = LogService.debug;
 
-const storage = new SimpleFsStorageProvider(path.join(DATA_PATH, "bot.json")); // /storage/bot.json
+if (KEYV_URL && KEYV_BACKEND === 'file') LogService.warn('config', 'KEYV_URL is ignored when KEYV_BACKEND is set to `file`')
+
+let chatgptStore:Keyv
+if (KEYV_BACKEND === 'file'){
+  chatgptStore = new Keyv({store: new KeyvFile({ filename: path.join(DATA_PATH, `chatgpt-bot-api.json`),})})
+} else {
+  chatgptStore = new Keyv(KEYV_URL, { namespace: 'chatgpt-bot-api' });
+}
+
+let storage: IStorageProvider
+if (KEYV_BOT_STORAGE) {
+  storage = new KeyvStorageProvider('chatgpt-bot-storage');
+} else {
+  storage = new SimpleFsStorageProvider(path.join(DATA_PATH, "bot.json")); // /storage/bot.json
+}
 
-// Prepare a crypto store if we need that
 let cryptoStore: ICryptoStorageProvider;
 if (MATRIX_ENCRYPTION) {
-  cryptoStore = new RustSdkCryptoStorageProvider(path.join(DATA_PATH, "encrypted")); // /storage/encrypted
+  if (KEYV_BOT_ENCRYPTION) {
+    cryptoStore = new KeyvCryptoStorageProvider('chatgpt-bot-encryption');
+  } else {
+    cryptoStore = new RustSdkCryptoStorageProvider(path.join(DATA_PATH, "encrypted")); // /storage/encrypted
+  }
 }
 
 async function main() {
@@ -40,29 +55,14 @@ async function main() {
   }
   if (!MATRIX_THREADS && CHATGPT_CONTEXT !== "room") throw Error("You must set CHATGPT_CONTEXT to 'room' if you set MATRIX_THREADS to false")
   const client: MatrixClient = new MatrixClient(MATRIX_HOMESERVER_URL, MATRIX_ACCESS_TOKEN, storage, cryptoStore);
-
-  // use puppeteer to bypass cloudflare (headful because of captchas)  
   const chatGPT: ChatGPTAPI = new ChatGPTAPI({
     apiKey: OPENAI_API_KEY,
     completionParams: {
-      model: CHATGPT_MODEL
+      model: CHATGPT_MODEL,
-    }
+    },
+    messageStore: chatgptStore
   })
 
-//   // call `api.refreshSession()` every hour to refresh the session
-//   setInterval(() => {
-//     chatGPT.refreshSession().then(() => {
-//       LogService.info('ChatGPT session reset');
-//     });
-//   }, 60 * 60 * 1000);
-
-//   // call `api.resetSession()` every 24 hours to reset the session
-//   setInterval(() => {
-//     chatGPT.resetSession().then(() => {
-//       LogService.info('ChatGPT session reset');
-//     });
-//   }, 24 * 60 * 60 * 1000);
-
   // Automatically join rooms the bot is invited to
   if (MATRIX_AUTOJOIN) {
     AutojoinRoomsMixin.setupOnClient(client);

src/storage.ts

@@ -0,0 +1,191 @@
+import Keyv from 'keyv'
+import { KeyvFile } from 'keyv-file';
+import * as sha512 from "hash.js/lib/hash/sha/512.js";
+import * as sha256 from "hash.js/lib/hash/sha/256.js";
+import * as path from "path";
+import { IAppserviceCryptoStorageProvider, IAppserviceStorageProvider, ICryptoRoomInformation, ICryptoStorageProvider, IFilterInfo, IStorageProvider } from "matrix-bot-sdk";
+import { DATA_PATH, KEYV_BACKEND, KEYV_URL } from './env.js';
+
+/**
+ * A storage provider that uses the disk to store information.
+ * @category Storage providers
+ */
+export class KeyvStorageProvider implements IStorageProvider, IAppserviceStorageProvider {
+
+    private completedTransactions = [];
+    private db: Keyv;
+    
+    /**
+     * Creates a new simple file system storage provider.
+     * @param {Keyv} keyvStore A Keyv instance for storing data.
+     */
+    constructor(namespace: string, private trackTransactionsInMemory = true, private maxInMemoryTransactions = 20) {
+        if (KEYV_BACKEND === 'file'){
+            this.db = new Keyv({store: new KeyvFile({ filename: path.join(DATA_PATH, `${namespace}.json`),})})
+        } else {
+            this.db = new Keyv(KEYV_URL, { namespace: namespace });
+        }
+        this.db.set('syncToken', null)
+        this.db.set('filter', null)
+        this.db.set('appserviceUsers', {}) // userIdHash => { data }
+        this.db.set('appserviceTransactions', {}) // txnIdHash => { data }
+        this.db.set('kvStore', {}) // key => value (str)
+    }
+
+    setSyncToken(token: string | null): Promise<any> | void {
+        this.db.set('syncToken', token);
+    }
+
+    getSyncToken(): string | Promise<string | null> | null{
+        return this.db.get('syncToken');
+    }
+
+    setFilter(filter: IFilterInfo): void {
+        this.db.set('filter', filter);
+    }
+
+    getFilter(): IFilterInfo | Promise<IFilterInfo> {
+        return this.db.get('filter');
+    }
+
+    addRegisteredUser(userId: string): Promise<any> | void {
+        const key = sha512().update(userId).digest('hex');
+        this.db.set(`appserviceUsers.${key}.userId`, userId)
+        this.db.set(`appserviceUsers.${key}.registered`, true)
+    }
+
+    isUserRegistered(userId: string): boolean | Promise<boolean> {
+        const key = sha512().update(userId).digest('hex');
+        return this.db.get(`appserviceUsers.${key}.registered`);
+    }
+
+    isTransactionCompleted(transactionId: string): boolean | Promise<boolean> {
+        if (this.trackTransactionsInMemory) {
+            return this.completedTransactions.indexOf(transactionId) !== -1;
+        }
+
+        const key = sha512().update(transactionId).digest('hex');
+        return this.db.get(`appserviceTransactions.${key}.completed`);
+    }
+
+    setTransactionCompleted(transactionId: string): Promise<any> | void {
+        if (this.trackTransactionsInMemory) {
+            if (this.completedTransactions.indexOf(transactionId) === -1) {
+                this.completedTransactions.push(transactionId);
+            }
+            if (this.completedTransactions.length > this.maxInMemoryTransactions) {
+                this.completedTransactions = this.completedTransactions.reverse().slice(0, this.maxInMemoryTransactions).reverse();
+            }
+            return;
+        }
+
+        const key = sha512().update(transactionId).digest('hex');
+        this.db.set(`appserviceTransactions.${key}.txnId`, transactionId)
+        this.db.set(`appserviceTransactions.${key}.completed`, true)
+    }
+
+    readValue(key: string): string|Promise<string|null|undefined>|null|undefined{
+        return this.db.get(key)
+    }
+
+    storeValue(key: string, value: string): Promise<any> | void {
+        this.db.set(key, value);
+    }
+
+    storageForUser(userId: string): IStorageProvider {
+        return new NamespacedKeyvProvider(userId, this);
+    }
+}
+
+/**
+ * A namespaced storage provider that uses the disk to store information.
+ * @category Storage providers
+ */
+class NamespacedKeyvProvider implements IStorageProvider {
+    constructor(private prefix: string, private parent: KeyvStorageProvider) {
+    }
+
+    setFilter(filter: IFilterInfo): Promise<any> | void {
+        return this.parent.storeValue(`${this.prefix}_int_filter`, JSON.stringify(filter));
+    }
+
+    getFilter(): IFilterInfo | Promise<IFilterInfo> {
+        return Promise.resolve(this.parent.readValue(`${this.prefix}_int_filter`)).then(r => r ? JSON.parse(r) : r);
+    }
+
+    setSyncToken(token: string | null): Promise<any> | void {
+        return this.parent.storeValue(`${this.prefix}_int_syncToken`, token || "");
+    }
+
+    getSyncToken(): string | Promise<string | null> | null {
+        return Promise.resolve(this.parent.readValue(`${this.prefix}_int_syncToken`)).then(r => r ?? null);
+    }
+
+    readValue(key: string): string | Promise<string | null | undefined> | null | undefined {
+        return this.parent.readValue(`${this.prefix}_kv_${key}`);
+    }
+
+    storeValue(key: string, value: string): Promise<any> | void {
+        return this.parent.storeValue(`${this.prefix}_kv_${key}`, value);
+    }
+}
+
+/**
+ * A crypto storage provider for the default rust-sdk store (sled, file-based).
+ * @category Storage providers
+ */
+export class KeyvCryptoStorageProvider implements ICryptoStorageProvider {
+    private db: Keyv;
+
+    /**
+     * Creates a new rust-sdk storage provider.
+     * @param {string} namespace The *directory* to persist database details to.
+     */
+    public constructor(public readonly namespace: string) {
+        if (KEYV_BACKEND === 'file'){
+            this.db = new Keyv({store: new KeyvFile({ filename: path.join(DATA_PATH, `${namespace}.json`),})})
+        } else {
+            this.db = new Keyv(KEYV_URL, { namespace: namespace });
+        }
+        this.db.set('deviceId', null)
+        this.db.set('rooms', {})
+    }
+
+    public async getDeviceId(): Promise<string> {
+        return this.db.get('deviceId');
+    }
+
+    public async setDeviceId(deviceId: string): Promise<void> {
+        this.db.set('deviceId', deviceId);
+    }
+
+    public async getRoom(roomId: string): Promise<ICryptoRoomInformation> {
+        const key = sha512().update(roomId).digest('hex');
+        return this.db.get(`rooms.${key}`);
+    }
+
+    public async storeRoom(roomId: string, config: ICryptoRoomInformation): Promise<void> {
+        const key = sha512().update(roomId).digest('hex');
+        this.db.set(`rooms.${key}`, config);
+    }
+}
+
+/**
+ * An appservice crypto storage provider for the default rust-sdk store (sled, file-based).
+ * @category Storage providers
+ */
+export class KeyvAppserviceCryptoStorageProvider extends KeyvCryptoStorageProvider implements IAppserviceCryptoStorageProvider {
+    /**
+     * Creates a new rust-sdk storage provider.
+     * @param {string} baseNamespace The *directory* to persist database details to.
+     */
+    public constructor(private baseNamespace: string) {
+        super(baseNamespace +  "_default");
+    }
+
+    public storageForUser(userId: string): ICryptoStorageProvider {
+        // sha256 because sha512 is a bit big for some operating systems
+        const key = sha256().update(userId).digest('hex');
+        return new KeyvCryptoStorageProvider(this.baseNamespace + "_" + key);
+    }
+}