Update MyShell.py

This commit is contained in:
H Lohaus 2023-11-16 18:51:24 +01:00 committed by GitHub
parent 0625bd0a06
commit 5c5235671e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -2,15 +2,9 @@ from __future__ import annotations
import time, json import time, json
try:
from selenium.webdriver.remote.webdriver import WebDriver
except ImportError:
class WebDriver():
pass
from ..typing import CreateResult, Messages from ..typing import CreateResult, Messages
from .base_provider import BaseProvider from .base_provider import BaseProvider
from .helper import format_prompt, get_browser from .helper import WebDriver, format_prompt, get_browser
class MyShell(BaseProvider): class MyShell(BaseProvider):
url = "https://app.myshell.ai/chat" url = "https://app.myshell.ai/chat"
@ -27,11 +21,11 @@ class MyShell(BaseProvider):
proxy: str = None, proxy: str = None,
timeout: int = 120, timeout: int = 120,
browser: WebDriver = None, browser: WebDriver = None,
display: bool = True, hidden_display: bool = True,
**kwargs **kwargs
) -> CreateResult: ) -> CreateResult:
if not browser: if not browser:
if display: if hidden_display:
driver, display = get_browser("", True, proxy) driver, display = get_browser("", True, proxy)
else: else:
display = get_browser("", False, proxy) display = get_browser("", False, proxy)
@ -44,9 +38,11 @@ class MyShell(BaseProvider):
driver.get(cls.url) driver.get(cls.url)
try: try:
# Wait for page load
WebDriverWait(driver, timeout).until( WebDriverWait(driver, timeout).until(
EC.presence_of_element_located((By.CSS_SELECTOR, "body:not(.no-js)")) EC.presence_of_element_located((By.CSS_SELECTOR, "body:not(.no-js)"))
) )
# Send message
script = """ script = """
response = await fetch("https://api.myshell.ai/v1/bot/chat/send_message", { response = await fetch("https://api.myshell.ai/v1/bot/chat/send_message", {
"headers": { "headers": {
@ -90,5 +86,5 @@ return content;
if not browser: if not browser:
time.sleep(0.1) time.sleep(0.1)
driver.quit() driver.quit()
if display: if hidden_display:
display.stop() display.stop()