refactor: improve load_yupp_accounts and get_cookies functions for better handling of tokens and cookies

Author: hlohaus

g4f/Provider/Yupp.py

@@ -83,13 +83,13 @@ def create_scraper():
     return scraper
 
 
-def load_yupp_accounts(tokens_str: str):
+def load_yupp_accounts(tokens: str):
     global YUPP_ACCOUNTS, _accounts_loaded
     if _accounts_loaded:
         return
-    if not tokens_str:
+    if not tokens:
         return
-    tokens = [token.strip() for token in tokens_str.split(",") if token.strip()]
+    tokens = [token.strip() for token in (tokens.split(",") if isinstance(tokens, str) else tokens) if token.strip()]
     YUPP_ACCOUNTS = [
         {"token": token, "is_valid": True, "error_count": 0, "last_used": 0.0}
         for token in tokens
@@ -541,7 +541,7 @@ async def create_async_generator(
         if not api_key:
             api_key = AuthManager.load_api_key(cls)
         if not api_key:
-            api_key = get_cookies("yupp.ai", False).get("__Secure-yupp.session-token")
+            api_key = [cookies.get("__Secure-yupp.session-token") for cookies in get_cookies("yupp.ai", False, "all").values() if cookies.get("__Secure-yupp.session-token")]
         if api_key:
             load_yupp_accounts(api_key)
             log_debug(f"Yupp provider initialized with {len(YUPP_ACCOUNTS)} accounts")

g4f/cookies.py

@@ -94,14 +94,16 @@ def get_headers(domain_name: str) -> Dict[str, str]:
 
 
 def get_cookies(domain_name: str, raise_requirements_error: bool = True,
-                single_browser: bool = False, cache_result: bool = True) -> Dict[str, str]:
+                single_browser: Optional[str] = None, cache_result: bool = True) -> Dict[str, str]:
     """Load cookies for a given domain from all supported browsers."""
-    if domain_name in CookiesConfig.cookies:
+    if single_browser != "all" and domain_name in CookiesConfig.cookies:
         return CookiesConfig.cookies[domain_name]
 
     cookies = load_cookies_from_browsers(domain_name, raise_requirements_error, single_browser)
-    if cache_result:
-        CookiesConfig.cookies[domain_name] = cookies
+    if single_browser != "all" and cache_result:
+        if len(cookies) > 0:
+            CookiesConfig.cookies[domain_name] = cookies
+        return CookiesConfig.cookies.get(domain_name, {})
     return cookies
 
 
@@ -115,23 +117,26 @@ def set_cookies(domain_name: str, cookies: Cookies = None) -> None:
 
 def load_cookies_from_browsers(domain_name: str,
                                raise_requirements_error: bool = True,
-                               single_browser: bool = False) -> Cookies:
+                               single_browser: Optional[str] = None) -> Cookies:
     """Helper to load cookies from all supported browsers."""
     if not has_browser_cookie3:
         if raise_requirements_error:
             raise MissingRequirementsError('Install "browser_cookie3" package')
         return {}
 
     cookies = {}
+    all_cookies = {}
     for cookie_fn in BROWSERS:
+        all_cookies[cookie_fn.__name__] = {}
         try:
             cookie_jar = cookie_fn(domain_name=domain_name)
-            if cookie_jar:
-                debug.log(f"Read cookies from {cookie_fn.__name__} for {domain_name}")
             for cookie in cookie_jar:
                 if cookie.name not in cookies and (not cookie.expires or cookie.expires > time.time()):
                     cookies[cookie.name] = cookie.value
-            if single_browser and cookie_jar:
+                    all_cookies[cookie_fn.__name__][cookie.name] = cookie.value
+            if len(all_cookies[cookie_fn.__name__]) > 0:
+                debug.log(f"Total cookies loaded for {domain_name} from {cookie_fn.__name__}: {len(all_cookies[cookie_fn.__name__])}")
+            if single_browser is True and cookie_jar:
                 break
         except BrowserCookieError:
             pass
@@ -140,6 +145,8 @@ def load_cookies_from_browsers(domain_name: str,
             break
         except Exception as e:
             debug.error(f"Error reading cookies from {cookie_fn.__name__} for {domain_name}: {e}")
+    if single_browser == "all":
+        return all_cookies
     return cookies