added create bot method in Poe class

This commit is contained in:
Raju Komati 2023-04-24 13:13:41 +05:30
parent d97cbceb0d
commit 95fa113a60
No known key found for this signature in database
GPG Key ID: A581A5D67A8EB090
1 changed files with 24 additions and 5 deletions

View File

@ -377,7 +377,7 @@ class Completion:
class Poe:
def __init__(self, model: str = "sage"):
def __init__(self, model: str = "gpt-3.5-turbo"):
self.cookie = self.__load_cookie()
self.model = MODELS[model]
self.client = PoeClient(self.cookie)
@ -402,7 +402,7 @@ class Poe:
print(mail_address)
options = webdriver.FirefoxOptions()
options.add_argument("-headless")
# options.add_argument("-headless")
driver = webdriver.Firefox(options=options)
driver.get("https://www.poe.com")
@ -440,9 +440,28 @@ class Poe:
driver.close()
return cookie
def chat(self, message: str):
def chat(self, message: str, model: Optional[str] = None) -> str:
model = model or self.model
response = None
for chunk in self.client.send_message(self.model, message):
for chunk in self.client.send_message(model, message):
response = chunk["text"]
return response
def create_bot(
self,
name: str,
/,
prompt: str = "",
base_model: str = "gpt-3.5-turbo",
description: str = "",
) -> None:
if base_model not in MODELS:
raise RuntimeError('Sorry, the base_model you provided does not exist. Please check and try again.')
response = self.client.create_bot(
handle=name,
prompt=prompt,
base_model=MODELS[base_model],
description=description,
)
print(f'Successfully created bot with name: {response["bot"]["displayName"]}')