From de27d42a2bdd32bdbf83fc6d7b5c1392b58cd4a9 Mon Sep 17 00:00:00 2001 From: nexy7574 Date: Sun, 9 Jun 2024 23:09:09 +0100 Subject: [PATCH] Support questions in impersonation --- src/cogs/ollama.py | 39 +++++++++++++++++++++------------------ 1 file changed, 21 insertions(+), 18 deletions(-) diff --git a/src/cogs/ollama.py b/src/cogs/ollama.py index 08074f7..98772fa 100644 --- a/src/cogs/ollama.py +++ b/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()