From 136770e05b63fd1482665eac5f90eba5bd10e4c5 Mon Sep 17 00:00:00 2001
From: Heiner Lohaus <heiner@lohaus.eu>
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/Provider')

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