From 136770e05b63fd1482665eac5f90eba5bd10e4c5 Mon Sep 17 00:00:00 2001 From: Heiner Lohaus Date: Wed, 4 Oct 2023 03:15:17 +0200 Subject: Use custom user dir --- g4f/Provider/OpenaiChat.py | 19 +++++++++++-------- g4f/Provider/helper.py | 23 +++++------------------ 2 files changed, 16 insertions(+), 26 deletions(-) (limited to 'g4f') diff --git a/g4f/Provider/OpenaiChat.py b/g4f/Provider/OpenaiChat.py index ca148da2..8f183fb9 100644 --- a/g4f/Provider/OpenaiChat.py +++ b/g4f/Provider/OpenaiChat.py @@ -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 - line = json.loads(line) - if "message" in line and not line["message"]["end_turn"]: + try: + line = json.loads(line) + 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: + driver = get_browser() + 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 diff --git a/g4f/Provider/helper.py b/g4f/Provider/helper.py index 8f09239a..c127f241 100644 --- a/g4f/Provider/helper.py +++ b/g4f/Provider/helper.py @@ -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 undetected_chromedriver import Chrome + 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) \ No newline at end of file -- cgit v1.2.3