Add tate command

src/cogs/ollama.py

@@ -1019,91 +1019,140 @@ class Ollama(commands.Cog):
                         break
         await msg.edit(embed=embed)
 
-    @commands.command()
+    async def generate_truth(self, ctx: commands.Context, entity: str, limit: int = None):
-    @commands.guild_only()
+        if entity == "trump":
-    async def trump(self, ctx: commands.Context):
+            system = (
-        async with ctx.channel.typing():
-            thread_id = self.history.create_thread(
-                ctx.author,
                 "You are the former united states president, Donald Trump, convicted of 34 felonies."
                 " You have been posting approximately 50 times a day on your platform 'Truth Social' spreading your"
-                " words of truth to the american people in order to make america great again. After reading through all"
+                " words of truth to the american people in order to make america great again."
-                " of your previous truth posts, you should generate another one, ready to be posted, "
-                "under 4000 characters. Write only the content to be posted, do not include any pleasantries."
-                " Write using the style of a twitter or facebook post. Do not repeat a previous post."
             )
-            async with httpx.AsyncClient() as client:
+        elif entity == "tate":
-                r = CONFIG["truth"].get("api", "https://bots.nexy7574.co.uk/jimmy/v2")
+            system = (
-                username = CONFIG["truth"].get("username", "1")
+                "You are the 'Top G' Andrew Tate. Andrew Tate is a former professional kickboxer who turned to the "
-                password = CONFIG["truth"].get("password", "2")
+                "world of influencing. His personality is one of arrogance and self-aggrandizement, "
-                response = await client.get(
+                "often referred to as \"King of Toxic Masculinity\" by critics due to his controversial views on "
-                    r + "/api/truths/all",
+                "gender roles, relationships, and other topics. He has been involved in several controversies related "
-                    timeout=60,
+                "to his content online including promoting extremist ideologies and misogynistic views. "
-                    auth=(username, password),
+                "Despite this, he still has a large following and is known for being an entrepreneur who built multiple"
-                )
+                " successful businesses such as Trinity Kickboxing Academy, Hustlers University, "
-                response.raise_for_status()
+                "and Romania's Real Estate Empire. "
-                truths: list[dict] = response.json()
+                "You post advice regarding masculinity and success, often in a controversial manner, on Truth Social. "
-                truths: list[TruthPayload] = list(map(lambda t: TruthPayload.model_validate(t), truths))
+            )
-                for truth in truths:
+        else:
-                    if truth.author == "trump":
+            raise ValueError("Invalid entity; must be one of trump/tate")
-                        await asyncio.to_thread(
+        system += (
-                            functools.partial(
+            "\n\nAfter reading through all"
-                                self.history.add_message,
+            " of your previous truth posts, you should generate another one, ready to be posted, "
-                                thread_id,
+            "under 4000 characters. Write only the content to be posted, do not include any pleasantries."
-                                "assistant",
+            " Write using the style of a twitter or facebook post. Do not repeat a previous post."
-                                truth.content,
+        )
-                                save=False
+        thread_id = self.history.create_thread(
-                            )
+            ctx.author,
-                        )
+            system
-            self.history.add_message(thread_id, "user", "Generate a new truth post.")
+        )
+        async with httpx.AsyncClient() as client:
+            r = CONFIG["truth"].get("api", "https://bots.nexy7574.co.uk/jimmy/v2")
+            username = CONFIG["truth"].get("username", "1")
+            password = CONFIG["truth"].get("password", "2")
+            response = await client.get(
+                r + "/api/truths/all",
+                timeout=60,
+                auth=(username, password),
+            )
+            response.raise_for_status()
+            truths = response.json()
+            truths: list[TruthPayload] = list(map(lambda t: TruthPayload.model_validate(t), truths))
 
-            tried = set()
+            if entity:
-            for _ in range(10):
+                truths = list(filter(lambda t: t.author == entity, truths))
-                server = self.next_server(tried)
+            if limit:
-                if await self.check_server(CONFIG["ollama"][server]["base_url"]):
+                truths.sort(key=lambda t: t.timestamp, reverse=True)  # newest first
-                    break
-                tried.add(server)
-            else:
-                return await ctx.reply("All servers are offline. Please try again later.", delete_after=300)
 
-            client = OllamaClient(CONFIG["ollama"][server]["base_url"])
+            for truth in truths:
-            async with self.servers[server]:
+                await asyncio.to_thread(
-                if not await client.has_model_named("llama2-uncensored", "7b-chat"):
+                    functools.partial(
-                    with client.download_model("llama2-uncensored", "7b-chat") as handler:
+                        self.history.add_message,
-                        await handler.flatten()
+                        thread_id,
-
+                        "assistant",
-                embed = discord.Embed(
+                        truth.content,
-                    title="New Truth!",
+                        save=False
-                    description="",
-                    colour=0x6559FF
-                )
-                msg = await ctx.reply(embed=embed)
-                last_edit = time.time()
-                messages = self.history.get_history(thread_id)
-                with client.new_chat("llama2-uncensored:7b-chat", messages) as handler:
-                    async for ln in handler:
-                        embed.description += ln["message"]["content"]
-                        if len(embed.description) >= 4000:
-                            break
-                        if (time.time() - last_edit) >= 2.5:
-                            await msg.edit(embed=embed)
-                            last_edit = time.time()
-
-                for truth in truths:
-                    if truth.content == embed.description:
-                        embed.add_field(
-                            name="Repeated truth :(",
-                            value="This truth was already truthed. Shit AI."
-                        )
-                        break
-                embed.set_footer(
-                    text="Finished generating truth based off of {:,} messages, using server {!r} | {!s}".format(
-                        len(messages) - 2,
-                        server,
-                        thread_id
                     )
                 )
-                await msg.edit(embed=embed)
+        self.history.add_message(thread_id, "user", "Generate a new truth post.")
+
+        tried = set()
+        for _ in range(10):
+            server = self.next_server(tried)
+            if await self.check_server(CONFIG["ollama"][server]["base_url"]):
+                break
+            tried.add(server)
+        else:
+            return await ctx.reply("All servers are offline. Please try again later.", delete_after=300)
+
+        client = OllamaClient(CONFIG["ollama"][server]["base_url"])
+        async with self.servers[server]:
+            if not await client.has_model_named("llama2-uncensored", "7b-chat"):
+                with client.download_model("llama2-uncensored", "7b-chat") as handler:
+                    await handler.flatten()
+
+            embed = discord.Embed(
+                title="New Truth!",
+                description="",
+                colour=0x6559FF
+            )
+            msg = await ctx.reply(embed=embed)
+            last_edit = time.time()
+            messages = self.history.get_history(thread_id)
+            with client.new_chat("llama2-uncensored:7b-chat", messages) as handler:
+                async for ln in handler:
+                    embed.description += ln["message"]["content"]
+                    if len(embed.description) >= 4000:
+                        break
+                    if (time.time() - last_edit) >= 2.5:
+                        await msg.edit(embed=embed)
+                        last_edit = time.time()
+
+            for truth in truths:
+                if truth.content == embed.description:
+                    embed.add_field(
+                        name="Repeated truth :(",
+                        value="This truth was already truthed. Shit AI."
+                    )
+                    break
+            embed.set_footer(
+                text="Finished generating truth based off of {:,} messages, using server {!r} | {!s}".format(
+                    len(messages) - 2,
+                    server,
+                    thread_id
+                )
+            )
+            await msg.edit(embed=embed)
+
+    @commands.command()
+    @commands.guild_only()
+    async def trump(self, ctx: commands.Context, latest: int = None):
+        """
+        Generates a truth social post from trump!
+
+        <latest> - limit the training history to the latest <latest> truths.
+
+        This command may take a long time.
+        """
+        async with ctx.channel.typing():
+            await self.generate_truth(ctx, "trump", latest)
+
+    @commands.command()
+    @commands.guild_only()
+    async def tate(self, ctx: commands.Context, latest: int = None):
+        """
+        Generates a truth social post from Andrew Tate
+
+        <latest> - limit the training history to the latest <latest> truths.
+
+        This command may take a long time.
+        """
+        async with ctx.channel.typing():
+            await self.generate_truth(ctx, "tate", latest)
 
 
 def setup(bot):