artemis/cogs/chat.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import logging
+from pathlib import Path
 import re
 from typing import TYPE_CHECKING
 
@@ -11,7 +12,6 @@ from discord.ext import commands
 from huggingface_hub import AsyncInferenceClient
 
 from artemis.utils import config
-from artemis.utils.constants import TEMP_DIR
 
 
 if TYPE_CHECKING:
@@ -107,12 +107,10 @@ class Chat(commands.Cog):
         self.lock = asyncio.Lock()
 
     def read_prompt(self):
-        path = TEMP_DIR / "prompt"
-        return path.read_text() if path.exists() else ""
+        return Path("temp/prompt").read_text() if Path("temp/prompt").exists() else ""
 
     def write_prompt(self, prompt: str):
-        path = TEMP_DIR / "prompt"
-        path.write_text(prompt)
+        Path("temp/prompt").write_text(prompt)
 
     def replace_emojis(self, text: str) -> str:
         return EMOJI_RE.sub(lambda match: emoji_map[match.group(0)], text)
@@ -126,17 +124,18 @@ class Chat(commands.Cog):
         return content
 
     def add_memory(self, role: str, message: str):
-        prompt = self.prompt + "\n\n"
+        prompt = (
+            self.prompt
+            + "The following is a user chat message directed at you, the format will be the same for subsequent messages, respond with only the message content, without specyfing actions."
+            + "\n\n"
+        )
         if len(self.memory) == 0:
             message = prompt + message
-        if len(self.memory) >= 20:
+        if len(self.memory) >= 15:
             del self.memory[0]
             del self.memory[0]
-            self.memory[0] = {
-                "role": "user",
-                "content": (prompt + self.memory[0]["content"]).strip(),
-            }
-        self.memory.append({"role": role, "content": message.strip()})
+            self.memory[0] = {"role": "user", "content": prompt + self.memory[0]["content"]}
+        self.memory.append({"role": role, "content": message})
 
     def add_user_memory(self, message: str):
         self.add_memory("user", message)
@@ -188,6 +187,8 @@ class Chat(commands.Cog):
         if not content:
             return
 
+        content = f"[USERNAME]: {message.author.display_name}\n[MESSAGE]: {content}"
+
         try:
             async with message.channel.typing():
                 async with self.lock:
@@ -212,16 +213,12 @@ class Chat(commands.Cog):
     async def _prompt(self, ctx: commands.Context, *, prompt: str = None):
         """Get or set system prompt and reset chat memory."""
         if not prompt:
-            await ctx.send(repr(self.prompt) if len(self.prompt) > 0 else "No prompt set.")
-        elif prompt == "reset":
-            self.prompt = ""
-            self.write_prompt(self.prompt)
-            await ctx.send("Prompt reset.")
-        else:
-            self.prompt = prompt
-            self.write_prompt(prompt)
-            self.memory = []
-            await ctx.send("Prompt updated.")
+            await ctx.send(self.prompt)
+            return
+        self.prompt = prompt
+        self.write_prompt(prompt)
+        self.memory = []
+        await ctx.send("Prompt updated.")
 
 
 async def setup(bot: Artemis):