maubot-gpt-bot/gpt/history.py

import redis.asyncio as redis
import aiohttp
import json

class History():
    r: redis.Redis()
    haste_url: ""

    def __init__(self, redis_host, redis_port, haste_url):
        self.r = redis.Redis(host=redis_host, port=redis_port, db=0)
        self.haste_url = haste_url

    def createCacheKey(self, id):
        return f"gpt-history-user-{id}"

    async def reset(self, userData):
        key = self.createCacheKey(userData["sender"] + userData["room_id"])
        await self.r.mset(key, null)

    async def get(self, userData):
        key = self.createCacheKey(userData["sender"] + userData["room_id"])
        history = await self.r.get(key)

        if history is not None:
            history = json.loads(history)
            return history
        else:
            return []


    async def add(self, userData, userMessage, assistantMessage):
        history = await self.get(userData)
        history.append({ "role": "user", "content": userMessage })
        history.append({ "role": "assistant", "content": assistantMessage })

        key = self.createCacheKey(userData["sender"] + userData["room_id"])
        history = json.dumps(history)
        await self.r.psetex(key, 300_000, history) # 5 mins

    async def dump(self, userData):
        history = await self.get(userData)
        if (len(history) == 0):
            return "You have no ChatGPT history at the moment."

        text = ""
        for i in range(0, len(history), 2):
            text += f"You: {history[i]['content']}\nGPT: {history[i+1]['content']}\n\n"

        async with aiohttp.request("POST", self.haste_url, data=text) as response:
            if response.status != 200:
                return "Could not export the ChatGPT history! Please try again later!"
            else:
                haste = await response.json()
                msg = "Your ChatGPT history: https://haste.snrd.eu/raw/{}".format(haste["key"])
                return msg