1
0
Fork 0

refactored the code

This commit is contained in:
Raju Komati 2023-04-29 15:08:10 +05:30
parent afcaf7b70c
commit c7f9f44f99
No known key found for this signature in database
GPG key ID: A581A5D67A8EB090
3 changed files with 7 additions and 7 deletions

View file

@ -23,4 +23,4 @@ class ForeFrontResponse(BaseModel):
model: str model: str
choices: List[Choice] choices: List[Choice]
usage: Usage usage: Usage
text: str text: str

View file

@ -377,10 +377,10 @@ class Poe:
raise RuntimeError('Sorry, the model you provided does not exist. Please check and try again.') raise RuntimeError('Sorry, the model you provided does not exist. Please check and try again.')
self.model = MODELS[model] self.model = MODELS[model]
self.cookie_path = cookie_path self.cookie_path = cookie_path
self.cookie = self.__load_cookie(driver, download_driver, driver_path=driver_path) self.cookie = self.__load_cookie(driver, driver_path=driver_path)
self.client = PoeClient(self.cookie) self.client = PoeClient(self.cookie)
def __load_cookie(self, driver: str, download_driver: bool, driver_path: Optional[str] = None) -> str: def __load_cookie(self, driver: str, driver_path: Optional[str] = None) -> str:
if (cookie_file := Path(self.cookie_path)).exists(): if (cookie_file := Path(self.cookie_path)).exists():
with cookie_file.open() as fp: with cookie_file.open() as fp:
cookie = json.load(fp) cookie = json.load(fp)
@ -468,4 +468,4 @@ class Poe:
print(f'Successfully created bot with name: {response["bot"]["displayName"]}') print(f'Successfully created bot with name: {response["bot"]["displayName"]}')
def list_bots(self) -> list: def list_bots(self) -> list:
return list(self.client.bot_names.values()) return list(self.client.bot_names.values())

View file

@ -225,7 +225,7 @@ class Client:
r = request_with_retries(self.session.post, self.gql_url, data=payload, headers=headers) r = request_with_retries(self.session.post, self.gql_url, data=payload, headers=headers)
data = r.json() data = r.json()
if data["data"] == None: if data["data"] is None:
logger.warn(f'{query_name} returned an error: {data["errors"][0]["message"]} | Retrying ({i + 1}/20)') logger.warn(f'{query_name} returned an error: {data["errors"][0]["message"]} | Retrying ({i + 1}/20)')
time.sleep(2) time.sleep(2)
continue continue
@ -316,7 +316,7 @@ class Client:
return return
# indicate that the response id is tied to the human message id # indicate that the response id is tied to the human message id
elif key != "pending" and value == None and message["state"] != "complete": elif key != "pending" and value is None and message["state"] != "complete":
self.active_messages[key] = message["messageId"] self.active_messages[key] = message["messageId"]
self.message_queues[key].put(message) self.message_queues[key].put(message)
return return
@ -402,7 +402,7 @@ class Client:
logger.info(f"Downloading {count} messages from {chatbot}") logger.info(f"Downloading {count} messages from {chatbot}")
messages = [] messages = []
if cursor == None: if cursor is None:
chat_data = self.get_bot(self.bot_names[chatbot]) chat_data = self.get_bot(self.bot_names[chatbot])
if not chat_data["messagesConnection"]["edges"]: if not chat_data["messagesConnection"]["edges"]:
return [] return []