refactored the code
This commit is contained in:
parent
afcaf7b70c
commit
c7f9f44f99
3 changed files with 7 additions and 7 deletions
|
@ -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
|
||||||
|
|
|
@ -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())
|
||||||
|
|
|
@ -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 []
|
||||||
|
|
Loading…
Reference in a new issue