feat: add load_cookies

Author: dsdanielpark

gemini/__init__.py

@@ -13,7 +13,13 @@
 from .src.misc.constants import URLs, Headers
 from .src.misc.decorator import retry, log_method, time_execution, handle_errors
 from .src.misc.exceptions import PackageError, GeminiAPIError, TimeoutError
-from .src.misc.utils import extract_code, upload_image, max_token, max_sentence
+from .src.misc.utils import (
+    extract_code,
+    upload_image,
+    max_token,
+    max_sentence,
+    load_cookies,
+)
 
 from .src.extensions.replit import prepare_replit_data
 

gemini/client.py

@@ -9,7 +9,7 @@
 from requests.exceptions import ConnectionError
 from typing import Optional, Tuple, Dict, Union, List
 
-from .src.misc.utils import upload_image
+from .src.misc.utils import upload_image, load_cookies
 from .src.model.parser.response_parser import ResponseParser
 from .src.model.output import GeminiCandidate, GeminiModelOutput
 from .src.model.parser.custom_parser import ParseMethod1, ParseMethod2
@@ -121,19 +121,7 @@ def _initialize_session(
     def _set_cookies_from_file(self, file_path: str) -> None:
         """Loads cookies from a file and updates the session."""
         try:
-            with open(file_path, "r") as file:
-                if file_path.endswith(".json"):
-                    # Load JSON formatted cookies directly
-                    cookies = json.load(file)
-                else:
-                    content = file.read()
-                    # Attempt to load Python dictionary formatted cookies
-                    try:
-                        content = content.replace("cookies = ", "").replace("'", '"')
-                        cookies = json.loads(content)
-                    except NameError:
-                        # Fallback to converting single quotes to double quotes and loading as JSON
-                        cookies = json.loads(content.replace("'", '"'))
+            cookies = load_cookies(file_path)
         except Exception as e:
             raise Exception(f"Failed to load cookies from {file_path}: {e}")
 

gemini/src/misc/__init__.py

@@ -1,4 +1,4 @@
 from .constants import URLs, Headers
 from .decorator import retry, log_method, time_execution, handle_errors
 from .exceptions import PackageError, GeminiAPIError, TimeoutError
-from .utils import extract_code, upload_image, max_token, max_sentence
+from .utils import extract_code, upload_image, max_token, max_sentence, load_cookies

gemini/src/misc/utils.py

@@ -1,8 +1,41 @@
+import json
 import requests
 from typing import Union
 from gemini.src.misc.constants import Headers
 
 
+def load_cookies(file_path):
+    """
+    Reads a text file and tries to parse it into a Python dictionary.
+    Assumes the file might contain a cookies string or is in JSON format.
+
+    :param file_path: Path to the file containing cookie data
+    :return: A dictionary containing cookie data
+    """
+    with open(file_path, "r") as file:
+        content = file.read().strip()
+
+    # If content includes '=', assume it's a cookie string
+    if "=" in content:
+        try:
+            # Split at the first '=' and prepare for JSON parsing
+            json_str = content.split("=", 1)[1].strip()
+            json_str = json_str.replace(
+                "'", '"'
+            )  # Replace single quotes with double quotes
+            cookies_dict = json.loads(json_str)
+        except json.JSONDecodeError:
+            raise ValueError("The cookie string does not have a valid JSON format.")
+    else:
+        try:
+            # Parse the entire content directly as JSON
+            cookies_dict = json.loads(content)
+        except json.JSONDecodeError:
+            raise ValueError("The file content does not have a valid JSON format.")
+
+    return cookies_dict
+
+
 def extract_code(text: str) -> str:
     """
     Extracts code snippets from the given text.