jonbot/jonbot.py

@@ -1,10 +1,11 @@
 import config
 import openai
+from openai import AsyncOpenAI
 import discord
+from discord.ext import commands
 import asyncio
 import logging
 from logging.handlers import TimedRotatingFileHandler
-from openai import AsyncOpenAI
 
 
 logger = logging.getLogger("garflog")
@@ -37,35 +38,42 @@ model = config.TXT_MODEL
 intents = discord.Intents.default()
 intents.messages = True
 intents.message_content = True
-client = discord.Client(intents=intents)
 
+client = commands.Bot(
+    command_prefix=["Jonbot ", "jonbot ", "Jon", "jon"],
+    case_insensitive=True,
+    intents=intents,
+)
 
 @client.event
 async def on_ready():
     print(f"Logged in as {client.user.name} running {model}.", flush=True)
 
+@client.command(name="chat")
+async def jonchat(ctx, *, prompt):
+    if "is this true" in prompt.lower():
+        messages = [msg async for msg in ctx.channel.history(limit=2)]
+        prompt = messages[1].content
+        prompt = f"Is this true: {prompt}"
+    answer = await generate_chat(prompt)
+    logger.info(
+        f"Chat Request - User: {ctx.author.name}, Server: {ctx.guild.name}, Prompt: {prompt}"
+    )
+    await ctx.reply(answer)
+
 @client.event
 async def on_message(message):
     if message.author == client.user:
         return
-    if message.content.lower().startswith("hey jon") or isinstance(message.channel, discord.DMChannel):
-        question = message.content[7:] if message.content.lower().startswith("hey jon") else message.content
-        try:
-            response = openai.ChatCompletion.create(
-                model=model,
-                messages=[
-                        {"role": "system", "content": "Pretend you are friendly Jon Arbuckle."},
-                        {"role": "user", "content": f"{question}"}
-                    ],
-                max_tokens=400
-            )
-            answer = response['choices'][0]['message']['content']
-            answer = answer.replace("AI language model", "American citizen")
-            answer = answer.replace("language model AI", "citizen of the United States")
-            await message.channel.send(answer)
-        except Exception as e:
-            e = str(e)
-            await message.channel.send(f"`JonBot Error: {e}`")
+
+    content = message.content.strip()
+    lower = content.lower()
+    if lower.startswith("hey jon") or isinstance(
+        message.channel, discord.DMChannel
+    ):
+        ctx = await client.get_context(message)
+        await jonchat(ctx, prompt=content)
+
 
 oai = AsyncOpenAI(
     api_key=config.OPENAI_TOKEN,
@@ -84,7 +92,7 @@ async def generate_chat(self, question: str) -> str:
             temperature=1.2,
         )
         answer = response.choices[0].message.content
-        return answer.replace("an AI language model", "a cartoon animal")
+        return answer
     except openai.BadRequestError as e:
         logger.error(e)
         return f"`JonBot Error: {e}`"