From ac3bd2a5572b941ab2493e1836b3747407139e73 Mon Sep 17 00:00:00 2001
From: hibobmaster <bobmaster@hibobmaster.com>
Date: Tue, 30 May 2023 10:24:29 +0800
Subject: [PATCH] Update bard.py

---
 bard.py | 52 +++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 41 insertions(+), 11 deletions(-)

diff --git a/bard.py b/bard.py
index 44e6deb..0cd1bf9 100644
--- a/bard.py
+++ b/bard.py
@@ -14,6 +14,10 @@ class Bardbot:
     Parameters
         session_id: str
             The __Secure-1PSID cookie.
+        timeout: int
+            Request timeout in seconds.
+        session: requests.Session
+            Requests session object.
     """
 
     __slots__ = [
@@ -23,10 +27,17 @@ class Bardbot:
         "conversation_id",
         "response_id",
         "choice_id",
+        "session_id",
         "session",
+        "timeout",
     ]
 
-    def __init__(self, session_id):
+    def __init__(
+            self, 
+            session_id: str,
+            timeout: int = 20,
+            session: requests.Session = None,
+        ):
         headers = {
             "Host": "bard.google.com",
             "X-Same-Domain": "1",
@@ -39,18 +50,33 @@ class Bardbot:
         self.conversation_id = ""
         self.response_id = ""
         self.choice_id = ""
-        self.session = requests.Session()
+        self.session_id = session_id
+        self.session = session or requests.Session()
         self.session.headers = headers
         self.session.cookies.set("__Secure-1PSID", session_id)
         self.SNlM0e = self.__get_snlm0e()
+        self.timeout = timeout
 
     def __get_snlm0e(self):
-        resp = self.session.get(url="https://bard.google.com/", timeout=10)
         # Find "SNlM0e":"<ID>"
+        if not self.session_id or self.session_id[-1] != ".":
+            raise Exception(
+                "__Secure-1PSID value must end with a single dot. Enter correct __Secure-1PSID value.",
+            )
+        resp = self.session.get(
+            "https://bard.google.com/",
+            timeout=10,
+        )
         if resp.status_code != 200:
-            raise Exception("Could not get Google Bard")
-        SNlM0e = re.search(r"SNlM0e\":\"(.*?)\"", resp.text).group(1)
-        return SNlM0e
+            raise Exception(
+                f"Response code not 200. Response Status is {resp.status_code}",
+            )
+        SNlM0e = re.search(r"SNlM0e\":\"(.*?)\"", resp.text)
+        if not SNlM0e:
+            raise Exception(
+                "SNlM0e value not found in response. Check __Secure-1PSID value.",
+            )
+        return SNlM0e.group(1)
 
     def ask(self, message: str) -> dict:
         """
@@ -60,7 +86,7 @@ class Bardbot:
         """
         # url params
         params = {
-            "bl": "boq_assistant-bard-web-server_20230326.21_p0",
+            "bl": "boq_assistant-bard-web-server_20230523.13_p0",
             "_reqid": str(self._reqid),
             "rt": "c",
         }
@@ -75,19 +101,22 @@ class Bardbot:
             "f.req": json.dumps([None, json.dumps(message_struct)]),
             "at": self.SNlM0e,
         }
-
-        # do the request!
         resp = self.session.post(
             "https://bard.google.com/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate",
             params=params,
             data=data,
-            timeout=120,
+            timeout=self.timeout,
         )
-
         chat_data = json.loads(resp.content.splitlines()[3])[0][2]
         if not chat_data:
             return {"content": f"Google Bard encountered an error: {resp.content}."}
         json_chat_data = json.loads(chat_data)
+        images = set()
+        if len(json_chat_data) >= 3:
+            if len(json_chat_data[4][0]) >= 4:
+                if json_chat_data[4][0][4]:
+                    for img in json_chat_data[4][0][4]:
+                        images.add(img[0][0][0])
         results = {
             "content": json_chat_data[0][0],
             "conversation_id": json_chat_data[1][0],
@@ -95,6 +124,7 @@ class Bardbot:
             "factualityQueries": json_chat_data[3],
             "textQuery": json_chat_data[2][0] if json_chat_data[2] is not None else "",
             "choices": [{"id": i[0], "content": i[1]} for i in json_chat_data[4]],
+            "images": images,
         }
         self.conversation_id = results["conversation_id"]
         self.response_id = results["response_id"]