Support questions in impersonation

src/cogs/ollama.py

@@ -953,7 +953,7 @@ class Ollama(commands.Cog):
                         break
         await msg.edit(embed=embed)
 
-    async def generate_truth(self, ctx: commands.Context, entity: str, limit: int = None):
+    async def generate_truth(self, ctx: commands.Context, entity: str, limit: int = None, question: str = None):
         if entity == "trump":
             system = (
                 "You are the former united states president, Donald Trump, convicted of 34 felonies."
@@ -1104,7 +1104,10 @@ class Ollama(commands.Cog):
                         save=False
                     )
                 )
-            self.history.add_message(thread_id, "user", "Generate a new truth post.")
+            if not question:
+                self.history.add_message(thread_id, "user", "Generate a new truth post.")
+            else:
+                self.history.add_message(thread_id, "user", question)
 
             tried = set()
             for _ in range(10):
@@ -1195,7 +1198,7 @@ class Ollama(commands.Cog):
 
     @commands.command(aliases=["trump"])
     @commands.guild_only()
-    async def donald(self, ctx: commands.Context, latest: int = None):
+    async def donald(self, ctx: commands.Context, latest: typing.Optional[int] = None, *, question: str = None):
         """
         Generates a truth social post from trump!
 
@@ -1204,11 +1207,11 @@ class Ollama(commands.Cog):
         This command may take a long time.
         """
         async with ctx.channel.typing():
-            await self.generate_truth(ctx, "trump", latest)
+            await self.generate_truth(ctx, "trump", latest, question=question)
 
     @commands.command(aliases=["tate"])
     @commands.guild_only()
-    async def andrew(self, ctx: commands.Context, latest: int = None):
+    async def andrew(self, ctx: commands.Context, latest: typing.Optional[int] = None, *, question: str = None):
         """
         Generates a truth social post from Andrew Tate
 
@@ -1217,11 +1220,11 @@ class Ollama(commands.Cog):
         This command may take a long time.
         """
         async with ctx.channel.typing():
-            await self.generate_truth(ctx, "tate", latest)
+            await self.generate_truth(ctx, "tate", latest, question=question)
     
     @commands.command(aliases=["sunak"])
     @commands.guild_only()
-    async def rishi(self, ctx: commands.Context, latest: int = None):
+    async def rishi(self, ctx: commands.Context, latest: typing.Optional[int] = None, *, question: str = None):
         """
         Generates a twitter post from Rishi Sunak
 
@@ -1230,11 +1233,11 @@ class Ollama(commands.Cog):
         This command may take a long time.
         """
         async with ctx.channel.typing():
-            await self.generate_truth(ctx, "Rishi Sunak", latest)
+            await self.generate_truth(ctx, "Rishi Sunak", latest, question=question)
     
     @commands.command(aliases=["robinson"])
     @commands.guild_only()
-    async def tommy(self, ctx: commands.Context, latest: int = None):
+    async def tommy(self, ctx: commands.Context, latest: typing.Optional[int] = None, *, question: str = None):
         """
         Generates a twitter post from Tommy Robinson
 
@@ -1243,11 +1246,11 @@ class Ollama(commands.Cog):
         This command may take a long time.
         """
         async with ctx.channel.typing():
-            await self.generate_truth(ctx, "Tommy Robinson 🇬🇧", latest)
+            await self.generate_truth(ctx, "Tommy Robinson 🇬🇧", latest, question=question)
     
     @commands.command(aliases=["fox"])
     @commands.guild_only()
-    async def laurence(self, ctx: commands.Context, latest: int = None):
+    async def laurence(self, ctx: commands.Context, latest: typing.Optional[int] = None, *, question: str = None):
         """
         Generates a twitter post from Laurence Fox
 
@@ -1256,11 +1259,11 @@ class Ollama(commands.Cog):
         This command may take a long time.
         """
         async with ctx.channel.typing():
-            await self.generate_truth(ctx, "Laurence Fox", latest)
+            await self.generate_truth(ctx, "Laurence Fox", latest, question=question)
     
     @commands.command(aliases=["farage"])
     @commands.guild_only()
-    async def nigel(self, ctx: commands.Context, latest: int = None):
+    async def nigel(self, ctx: commands.Context, latest: typing.Optional[int] = None, *, question: str = None):
         """
         Generates a twitter post from Nigel Farage
 
@@ -1269,11 +1272,11 @@ class Ollama(commands.Cog):
         This command may take a long time.
         """
         async with ctx.channel.typing():
-            await self.generate_truth(ctx, "Nigel Farage", latest)
+            await self.generate_truth(ctx, "Nigel Farage", latest, question=question)
 
     @commands.command(aliases=["starmer"])
     @commands.guild_only()
-    async def keir(self, ctx: commands.Context, latest: int = None):
+    async def keir(self, ctx: commands.Context, latest: typing.Optional[int] = None, *, question: str = None):
         """
         Generates a twitter post from Keir Starmer
 
@@ -1282,11 +1285,11 @@ class Ollama(commands.Cog):
         This command may take a long time.
         """
         async with ctx.channel.typing():
-            await self.generate_truth(ctx, "Keir Starmer", latest)
+            await self.generate_truth(ctx, "Keir Starmer", latest, question=question)
     
     @commands.command(aliases=["johnson"])
     @commands.guild_only()
-    async def boris(self, ctx: commands.Context, latest: int = None):
+    async def boris(self, ctx: commands.Context, latest: typing.Optional[int] = None, *, question: str = None):
         """
         Generates a twitter post from Boris Johnson
 
@@ -1295,7 +1298,7 @@ class Ollama(commands.Cog):
         This command may take a long time.
         """
         async with ctx.channel.typing():
-            await self.generate_truth(ctx, "Boris Johnson", latest)
+            await self.generate_truth(ctx, "Boris Johnson", latest, question=question)
     
     @commands.command()
     @commands.guild_only()