diff --git a/examples/princess/princess.py b/examples/princess/princess.py index b350ae6..c3b3a52 100644 --- a/examples/princess/princess.py +++ b/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()