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":"" + 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"]