maubot-gpt-bot/gpt/history.py

55 lines
1.8 KiB
Python

import redis.asyncio as redis
import requests
import json
class History():
r: redis.Redis()
def __init__(self, redis_host, redis_port):
self.r = redis.Redis(host=redis_host, port=redis_port, db=0)
def createCacheKey(self, id):
return f"gpt-history-user-{id}"
async def reset(self, userData):
key = self.createCacheKey(userData["sender"])
await self.r.mset(key, null)
async def get(self, userData):
key = self.createCacheKey(userData["sender"])
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"])
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"
response = requests.post("https://haste.snrd.eu/documents", data=text)
if response.status_code != 200:
return "Could not export the ChatGPT history! Please try again later!"
else:
haste = json.loads(response.text)['key']
msg = "Your ChatGPT history: https://haste.snrd.eu/raw/{}".format(haste)
return msg