Use custom user dir

This commit is contained in:
Heiner Lohaus 2023-10-04 03:15:17 +02:00
parent 5b2efa4aaf
commit 136770e05b
2 changed files with 16 additions and 26 deletions

View File

@ -1,7 +1,6 @@
from __future__ import annotations
import uuid
import json
import uuid, json, time
from .base_provider import AsyncGeneratorProvider
from .helper import get_browser, get_cookies, format_prompt
@ -56,8 +55,13 @@ class OpenaiChat(AsyncGeneratorProvider):
line = line[6:]
if line == b"[DONE]":
break
try:
line = json.loads(line)
if "message" in line and not line["message"]["end_turn"]:
except:
continue
if "message" not in line or "message_type" not in line["message"]["metadata"]:
continue
if line["message"]["metadata"]["message_type"] == "next":
new_message = line["message"]["content"]["parts"][0]
yield new_message[len(last_message):]
last_message = new_message
@ -68,11 +72,9 @@ class OpenaiChat(AsyncGeneratorProvider):
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
except ImportError:
return
driver = get_browser()
if not driver:
except ImportError:
return
driver.get(f"{cls.url}/")
@ -83,6 +85,7 @@ class OpenaiChat(AsyncGeneratorProvider):
javascript = "return (await (await fetch('/api/auth/session')).json())['accessToken']"
return driver.execute_script(javascript)
finally:
time.sleep(1)
driver.quit()
@classmethod

View File

@ -2,7 +2,7 @@ from __future__ import annotations
import asyncio, sys
from asyncio import AbstractEventLoop
from os import path
import browser_cookie3
# Change event loop policy on windows
@ -60,24 +60,11 @@ def format_prompt(messages: list[dict[str, str]], add_special_tokens=False):
def get_browser(user_data_dir: str = None):
try:
from undetected_chromedriver import Chrome
except ImportError:
return None
from platformdirs import user_config_dir
def get_user_data_dir():
dirs = [
'~/.config/google-chrome/Default',
'~/.var/app/com.google.Chrome/config/google-chrome/Default',
'%LOCALAPPDATA%\\Google\\Chrome\\User Data\\Default',
'~/Library/Application Support/Google/Chrome/Default',
]
from os import path
for dir in dirs:
dir = path.expandvars(dir)
if path.exists(dir):
return dir
if not user_data_dir:
user_data_dir = get_user_data_dir()
user_data_dir = user_config_dir("g4f")
user_data_dir = path.join(user_data_dir, "Default")
return Chrome(user_data_dir=user_data_dir)