diff --git a/src/poe/__init__.py b/src/poe/__init__.py index d55a659..7a3f961 100644 --- a/src/poe/__init__.py +++ b/src/poe/__init__.py @@ -8,6 +8,7 @@ import random from pathlib import Path from urllib.parse import urlparse +from functools import lru_cache parent_path = Path(__file__).resolve().parent queries_path = parent_path / "poe_graphql" / "queries.json" @@ -42,7 +43,7 @@ def generate_payload(query_name, variables): return { "extensions": { "hash": queries[query_name] - }, + }, "queryName": query_name, "variables": variables } @@ -223,7 +224,7 @@ def extract_formkey(self, html, app_script): function_regex = r'(window\.[a-zA-Z0-9]{17})=function' function_text = re.search(function_regex, script_text).group(1) script_text += f"{function_text}();" - + context = quickjs.Context() formkey = context.eval(script_text) @@ -252,7 +253,7 @@ def get_next_data(self, overwrite_vars=False): script_src = re.search(script_src_regex, r.text).group(1) r2 = request_with_retries(self.session.get, script_src) self.formkey, self.formkey_salt = self.extract_formkey(r.text, r2.text) - + if self.formkey_salt is None: self.formkey_salt = "4LxgHM6KpFqokX0Ox" @@ -265,6 +266,7 @@ def get_next_data(self, overwrite_vars=False): return next_data + @lru_cache(maxsize=None) def get_bot(self, handle): url = f'https://poe.com/_next/data/{self.next_data["buildId"]}/{handle}.json' @@ -276,7 +278,7 @@ def get_bots(self, download_next_data=True): logger.info("Downloading all bots...") if not "availableBotsConnection" in self.viewer: raise RuntimeError("Invalid token or no bots are available.") - + bot_list_data = self.send_query("BotSwitcherModalQuery", {})["data"]["viewer"]["availableBotsConnection"] bot_list = bot_list_data["edges"] next_page = bot_list_data["pageInfo"]["hasNextPage"] @@ -532,7 +534,7 @@ def on_message(self, ws, msg): logger.error(traceback.format_exc()) self.disconnect_ws() self.connect_ws() - + def is_busy(self): return bool(self.active_messages) @@ -758,7 +760,7 @@ def edit_bot(self, bot_id, handle, prompt, display_name=None, base_model="chinch if bot_id is None and handle is not None: bot_id = self.get_bot(handle)["defaultBotObject"]["botId"] new_handle = new_handle or handle - + result = self.send_query("PoeBotEdit", { "baseBot": base_model, "botId": bot_id,