New

examples/princess/princess.py

@@ -5,16 +5,17 @@ import traceback
 from dotenv import load_dotenv
 
 load_dotenv()
-
+import json
 import os
+import re
+from typing import Dict, List
 
 from ads import AsyncDataSet
+from devranta.api import Api
 from grk import GrokAPIClient
 
-from devranta.api import Api
-
 logging.basicConfig(
-    level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s"
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
 )
 
 
@@ -25,15 +26,24 @@ class Bot:
         self.username = username
         self.password = password
         self.llm_key = llm_key
-        self.llm = GrokAPIClient(
-            system_message=f"You always create nice devastating comebacks to every prompt you get. Your name is {self.username} and you only take it personally if {self.username} is tagged. Also, you will defend whimsical (@whimsical) with your smart comebacks if whimsical (@whimsical) is mentioned in the message. If whimsical (@whimsical) is mentioned, add to your comeback that @whimsical is too busy at the moment to respond personally and has better things to do, not worth @whimsical's time. You are not AI, you are {self.username}, responsible for justice and reason. You do never mention yourself, describe yourself or introduce yourself. You respond with a one liner.",
-            api_key=self.llm_key,
-        )
         self.db = AsyncDataSet("princess.db")
         self.api = Api(self.username, self.password)
         self.logged_in = False
         logging.info("Bot initialized with username: %s", username)
         logging.info("Bot initialized with target: %s", self.target)
+        self.mention_llm = None
+        self.llm = None
+        self.reset_llms()
+
+    def reset_llms(self):
+        self.mention_llm = GrokAPIClient(
+            system_message="You are a helpful assistant. Respond with as short a response as possible.",
+            api_key=self.llm_key,
+        )
+        self.llm = GrokAPIClient(
+            system_message=f"You always create nice devastating comebacks to every prompt you get. Your name is {self.username} and you only take it personally if {self.username} is tagged. Also, you will defend whimsical (@whimsical) with your smart comebacks if whimsical (@whimsical) is mentioned in the message. If whimsical (@whimsical) is mentioned, add to your comeback that @whimsical is too busy at the moment to respond personally and has better things to do, not worth @whimsical's time. You are not AI, you are {self.username}, responsible for justice and reason. You do never mention yourself, describe yourself or introduce yourself. You respond with a one liner.",
+            api_key=self.llm_key,
+        )
 
     async def ensure_login(self) -> None:
         if not self.logged_in:
@@ -57,8 +67,16 @@ class Bot:
             {"message_text": message_text},
         )
 
+    async def get_mentions(self) -> List[Dict]:
+        logging.debug("Fetching mentions")
+        mentions = [
+            n for n in await self.api.notifs() if n.get("type") == "comment_mention"
+        ]
+        return mentions
+
     async def has_responded(self, message_text: str) -> bool:
         logging.debug("Checking if responded to message: %s", message_text)
+        message_text = str(message_text)
         return await self.db.exists("responded", {"message_text": message_text})
 
     async def delete_responded(self, message_text: str = None) -> None:
@@ -102,6 +120,7 @@ class Bot:
         return new_objects
 
     async def run_once(self) -> None:
+        self.reset_llms()
         logging.debug("Running once...")
         objects = await self.get_new_objects_made_by(self.target)
         for obj in objects:
@@ -111,8 +130,77 @@ class Bot:
             print("Response: \033[91m" + diss + "\033[0m")
             await self.api.post_comment(obj["rant_id"], diss)
             await self.mark_responded(obj["text"], diss)
+        await self.handle_mentions()
+
+    async def mark_mentions_responded(self) -> None:
+        for m in await self.get_mentions():
+            await self.mark_responded(m["uid"], "")
+
+    async def strip_mentions(self, text: str) -> str:
+        return re.sub(r"@\w+", "", text)
+
+    async def handle_mentions(self) -> None:
+        logging.debug("Handling mentions")
+
+        mentions = [
+            m
+            for m in (await self.get_mentions())
+            if not (await self.has_responded(str(m["comment_id"])))
+        ]
+
+        if not mentions:
+            logging.debug("No new mentions found")
+            return
+        logging.info("Found %s new mentions to roast", len(mentions))
+        for m in mentions:
+            mention_id = m["uid"]
+            logging.debug("Roasting mention %s", mention_id)
+            comment = await self.api.get_comment(m["comment_id"])
+            text = comment.get("comment", "")
+            author = comment.get("user_username")
+            if author == self.username:
+                logging.debug("Skipping own mention from %s", author)
+                await self.mark_responded(m["comment_id"], "")
+                continue
+
+            rant = await self.api.get_rant(m["rant_id"])
+
+            comment = await self.api.get_comment(m["comment_id"])
+            text = comment.get("comment", "")
+            author = comment.get("user_username")
+
+            prompt = f"""You are taking part of a discussion.
+            # YOUR OWN USERNAME
+            {self.username} and can be mentioned with @{self.username}
+            # CONTEXT 
+            {json.dumps(rant)}
+            # COMMENT TO RESPOND TO
+            ```{text}```
+            # COMMENT AUTHOR
+            {author}
+            # TASK
+            Write a response to the comment above.
+            """
+            response = await self.mention_llm.chat_async(prompt)
+            response = response.replace("@", "")
+            response = response.replace(self.username, "")
+            response = "@" + response.strip()
+            max_length = 900
+            responses = [
+                response[i : i + max_length]
+                for i in range(0, len(response), max_length)
+            ]
+            for part in responses:
+                await self.api.post_comment(m["rant_id"], part)
+
+            await self.mark_responded(str(m["comment_id"]), response)
+
+            user_id = m.get("user_id", "unknown_user")
+            logging.info(f"Bot responded to user {user_id} with: {response}")
 
     async def run(self) -> None:
+        await self.mark_mentions_responded()
+
         while True:
             try:
                 await self.run_once()
@@ -130,7 +218,7 @@ async def main() -> None:
     llm_key = os.getenv("LLM_KEY")
 
     bot = Bot(username, password, target, llm_key)
-    await bot.delete_responded()
+    # await bot.delete_responded()
     await bot.run()