build: 2.0.0 for functions

Author: dsdanielpark

gemini/__init__.py

@@ -1,4 +1,3 @@
-# Copyright 2024 Daniel Park, Antonio Cheang, MIT License
 from os import environ
 from .async_core import GeminiClient
 from .core import Gemini
@@ -11,7 +10,7 @@
 
 gemini_api_key = environ.get("GEMINI_COOKIES")
 
-__version__ = "1.1.2"
+__version__ = "2.0.0"
 __author__ = (
     "daniel park <[email protected]>, antonio cheang <[email protected]>"
 )

gemini/src/async_core.py renamed to gemini/async_core.py

@@ -1,4 +1,3 @@
-# Copyright 2024 Daniel Park, Antonio Cheang, MIT License
 import os
 import re
 import json

gemini/core.py

@@ -1,4 +1,3 @@
-# Copyright 2024 Daniel Park, Antonio Cheang, MIT License
 import os
 import re
 import json
@@ -7,12 +6,13 @@
 import inspect
 import requests
 import urllib.parse
-from typing import Optional, Tuple, Dict
+from typing import Optional, Tuple, Dict, Union
 from requests.exceptions import ConnectionError, RequestException
 
 from .src.parser.custom_parser import ParseMethod1, ParseMethod2
 from .src.parser.response_parser import ResponseParser
 from .src.model.output import GeminiCandidate, GeminiModelOutput
+from .src.misc.utils import upload_image
 from .src.misc.constants import (
     HEADERS,
     HOST,
@@ -52,12 +52,16 @@ def __init__(
         auto_cookies: bool = False,
         timeout: int = 30,
         proxies: Optional[dict] = None,
+        rcid: str = None,
     ) -> None:
         """
         Initializes the Gemini object with session, cookies, and other configurations.
         """
         self._nonce = None
         self._sid = None
+        self._rcid = rcid or None
+        self._rid = None
+        self._cid = None
         self.auto_cookies = auto_cookies
         self.cookie_fp = cookie_fp
         self.cookies = cookies
@@ -69,6 +73,14 @@ def __init__(
         self._reqid = int("".join(random.choices(string.digits, k=7)))
         self.parser = ResponseParser(cookies=self.cookies)
 
+    @property
+    def rcid(self):
+        return self._rcid
+
+    @rcid.setter
+    def rcid(self, value):
+        self._rcid = value
+
     def _initialize_session(
         self,
     ) -> requests.Session:
@@ -87,15 +99,15 @@ def _initialize_session(
         if self.cookies:
             session.cookies.update(self.cookies)
         elif self.cookie_fp:
-            self._load_cookies_from_file(self.cookie_fp)
+            self._set_cookies_from_file(self.cookie_fp)
         elif self.auto_cookies == True:
             self._set_cookies_automatically()
 
         self._set_sid_and_nonce()
 
         return session
 
-    def _load_cookies_from_file(self, file_path: str) -> None:
+    def _set_cookies_from_file(self, file_path: str) -> None:
         """Loads cookies from a file and updates the session."""
         try:
             if file_path.endswith(".json"):
@@ -120,15 +132,17 @@ def _set_sid_and_nonce(self):
             response = requests.get(f"{HOST}/app", cookies=self.cookies)
             response.raise_for_status()
 
-            sid_match, nonce_match = self.extract_sid_nonce(response.text)
+            sid_match = re.search(r'"FdrFJe":"([\d-]+)"', response.text)
+            nonce_match = re.search(r'"SNlM0e":"(.*?)"', response.text)
 
-            if sid_match and nonce_match:
+            if sid_match:
                 self._sid = sid_match.group(1)
-                self._nonce = nonce_match.group(1)
             else:
                 raise ValueError(
                     "Failed to parse SID or SNlM0e nonce from the response.\nRefresh the Gemini web page or access Gemini in a new incognito browser to resend cookies."
                 )
+            if nonce_match:
+                self._nonce = nonce_match.group(1)
 
         except requests.RequestException as e:
             raise ConnectionError(f"Request failed: {e}")
@@ -137,12 +151,6 @@ def _set_sid_and_nonce(self):
         except Exception as e:
             raise RuntimeError(f"An unexpected error occurred: {e}")
 
-    @staticmethod
-    def extract_sid_nonce(response_text):
-        sid_match = re.search(r'"FdrFJe":"([\d-]+)"', response_text)
-        nonce_match = re.search(r'"SNlM0e":"(.*?)"', response_text)
-        return sid_match, nonce_match
-
     def _construct_params(self, sid: str) -> str:
         """
         Constructs URL-encoded parameters for a request.
@@ -163,7 +171,9 @@ def _construct_params(self, sid: str) -> str:
             }
         )
 
-    def _construct_payload(self, prompt: str, nonce: str) -> str:
+    def _construct_payload(
+        self, prompt: str, image: Union[bytes, str], nonce: str
+    ) -> str:
         """
         Constructs URL-encoded payload for a request.
 
@@ -177,15 +187,35 @@ def _construct_payload(self, prompt: str, nonce: str) -> str:
         return urllib.parse.urlencode(
             {
                 "at": nonce,
-                "f.req": json.dumps([None, json.dumps([[prompt], None, None])]),
-            }
+                "f.req": json.dumps(
+                    [
+                        None,
+                        json.dumps(
+                            [
+                                image
+                                and [
+                                    prompt,
+                                    0,
+                                    None,
+                                    [[[upload_image(image), 1]]],
+                                ]
+                                or [prompt],
+                                None,
+                                [self._cid, self._rid, self._rcid],
+                            ]
+                        ),
+                    ]
+                ),
+            },
         )
 
-    def send_request(self, prompt: str) -> Tuple[str, int]:
+    def send_request(
+        self, prompt: str, image: Union[bytes, str] = None
+    ) -> Tuple[str, int]:
         """Sends a request and returns the response text and status code."""
         try:
             params = self._construct_params(self._sid)
-            data = self._construct_payload(prompt, self._nonce)
+            data = self._construct_payload(prompt, image, self._nonce)
             response = self.session.post(
                 POST_ENDPOINT,
                 params=params,
@@ -200,7 +230,7 @@ def send_request(self, prompt: str) -> Tuple[str, int]:
             # self._update_cookies_from_browser()
             self._set_sid_and_nonce()
             params = self._construct_params(self._sid)
-            data = self._construct_payload(prompt, self._nonce)
+            data = self._construct_payload(prompt, image, self._nonce)
             response = self.session.post(
                 POST_ENDPOINT,
                 params=params,
@@ -217,9 +247,11 @@ def send_request(self, prompt: str) -> Tuple[str, int]:
             response.raise_for_status()
         return response.text, response.status_code
 
-    def generate_content(self, prompt: str) -> GeminiModelOutput:
+    def generate_content(
+        self, prompt: str, image: Union[bytes, str] = None
+    ) -> GeminiModelOutput:
         """Generates content based on the prompt and returns a GeminiModelOutput object."""
-        response_text, response_status_code = self.send_request(prompt)
+        response_text, response_status_code = self.send_request(prompt, image)
         if response_status_code != 200:
             raise ValueError(f"Response status: {response_status_code}")
         else:
@@ -231,7 +263,12 @@ def generate_content(self, prompt: str) -> GeminiModelOutput:
     def _create_model_output(self, parsed_response: dict) -> GeminiModelOutput:
         candidates = self.collect_candidates(parsed_response)
         metadata = parsed_response.get("metadata", [])
-
+        try:
+            self._cid = metadata[0]
+            self._rid = metadata[1]
+            # self._rcid = candidates["candidates"][0]["rcid"]
+        except:
+            pass
         return GeminiModelOutput(
             metadata=metadata,
             candidates=candidates,

gemini/src/misc/utils.py

@@ -1,7 +1,8 @@
 # Copyright 2024 Daniel Park, Antonio Cheang, MIT License
 import json
 import requests
-from .constants import IMG_UPLOAD_HEADERS, REPLIT_SUPPORT_PROGRAM_LANGUAGES
+from typing import Union
+from .constants import REPLIT_SUPPORT_PROGRAM_LANGUAGES
 
 
 def extract_code(text: str, language: str) -> str:
@@ -51,36 +52,31 @@ def extract_code(text: str, language: str) -> str:
         return text
 
 
-def upload_image(image: bytes, filename: str = "Photo.jpg"):
+def upload_image(file: Union[bytes, str]) -> str:
     """
     Upload image into bard bucket on Google API, do not need session.
 
     Returns:
         str: relative URL of image.
     """
-    resp = requests.options("https://content-push.googleapis.com/upload/")
-    resp.raise_for_status()
-    size = len(image)
-
-    headers = IMG_UPLOAD_HEADERS
-    headers["size"] = str(size)
-    headers["x-goog-upload-command"] = "start"
-
-    data = f"File name: {filename}"
-    resp = requests.post(
-        "https://content-push.googleapis.com/upload/", headers=headers, data=data
+    if isinstance(file, str):
+        with open(file, "rb") as f:
+            file_data = f.read()
+    else:
+        file_data = file
+
+    response = requests.post(
+        url="https://content-push.googleapis.com/upload/",
+        headers={
+            "Push-ID": "feeds/mcudyrk2a4khkz",
+            "Content-Type": "application/octet-stream",
+        },
+        data=file_data,
+        allow_redirects=True,
     )
-    resp.raise_for_status()
-    upload_url = resp.headers["X-Goog-Upload-Url"]
-    resp = requests.options(upload_url, headers=headers)
-    resp.raise_for_status()
-    headers["x-goog-upload-command"] = "upload, finalize"
-
-    # It can be that we need to check returned offset
-    headers["X-Goog-Upload-Offset"] = "0"
-    resp = requests.post(upload_url, headers=headers, data=image)
-    resp.raise_for_status()
-    return resp.text
+    response.raise_for_status()
+
+    return response.text
 
 
 def max_token(text: str, n: int) -> str:

gemini/src/model/output.py

@@ -21,6 +21,12 @@ class GeminiModelOutput(BaseModel):
     chosen: int = 0
     response_dict: Optional[dict] = None
 
+    def __setattr__(self, name, value):
+        if name == "chosen":
+            self.chosen = value
+        else:
+            super().__setattr__(name, value)
+
     @property
     def rcid(self) -> str:
         """The RCID (Response Choice ID) of the chosen candidate."""

gemini/src/parser/response_parser.py

@@ -11,41 +11,65 @@ def parse(self, response_text):
 
     def parse_response_text(self, response_text):
         body = self._extract_body(response_text)
+
         if not body or not body[4]:
             raise ValueError(
                 "Failed to parse response body. Data structure is invalid."
             )
 
         candidates = self._parse_candidates(body[4])
+        metadata = body[1]
+
+        if body[0] is None:
+            prompt_class = None
+            prompt_candidates = []
+        else:
+            prompt_class = body[0][0]
+            prompt_candidates = body[0][1:] if len(body[0]) > 1 else []
+
         if not candidates:
             raise ValueError(
                 "Failed to generate contents. No output data found in response."
             )
 
         return {
-            "metadata": body[1],
+            "metadata": metadata,
+            "prompt_class": prompt_class,
+            "prompt_candidates": prompt_candidates,
             "candidates": candidates,
         }
 
     def _extract_body(self, response_text):
         try:
-            body = json.loads(json.loads(max(response_text.split("\n"), key=len))[0][2])
+            body = json.loads(
+                json.loads(response_text.lstrip("')]}'\n\n").split("\n")[1])[0][2]
+            )
             if not body[4]:
                 body = json.loads(
-                    json.loads(max(response_text.split("\n"), key=len))[4][2]
+                    json.loads(response_text.lstrip("')]}'\n\n").split("\n")[1])[4][2]
                 )
             return body
         except:
-            body = json.loads(json.loads(response_text.text.split("\n")[2])[0][2])
-            if not body[4]:
-                body = json.loads(json.loads(response_text.text.split("\n")[2])[4][2])
-            else:
-                raise ValueError(
-                    "Failed to generate contents. Invalid response data received."
-                )
-            return body
+            try:
+                body = json.loads(json.loads(response_text.text.split("\n")[2])[0][2])
+                if not body[4]:
+                    body = json.loads(
+                        json.loads(response_text.text.split("\n")[2])[4][2]
+                    )
+                else:
+                    raise ValueError(
+                        "Failed to generate contents. Invalid response data received."
+                    )
+                return body
+            except:
+                max_response = max(response_text.split("\n"), key=len)
+                body = json.loads(json.loads(max_response)[0][2])
+                if not body[4]:
+                    body = json.loads(json.loads(max_response)[4][2])
+                return body
 
     def _parse_candidates(self, candidates_data):
+        candidates_list = []
         for candidate_data in candidates_data:
             web_images = self._parse_web_images(candidate_data[4])
             generated_images = self._parse_generated_images(candidate_data[12])
@@ -55,7 +79,8 @@ def _parse_candidates(self, candidates_data):
                 "web_images": web_images,
                 "generated_images": generated_images,
             }
-        return candidate_dict
+            candidates_list.append(candidate_dict)
+        return candidates_list
 
     def _parse_web_images(self, images_data):
         if not images_data:
@@ -75,7 +100,7 @@ def _parse_generated_images(self, images_data):
         return [
             {
                 "url": image[0][3][3],
-                "title": f"[Generated Image {image[3][6]}]",
+                "title": f"[GeneratedImage {image[3][6]}]",
                 "alt": image[3][5][i] if len(image[3][5]) > i else image[3][5][0],
                 "cookies": self.cookies,
             }