matrix_chatgpt_bot/bard.py

133 lines
4.5 KiB
Python
Raw Normal View History

2023-04-13 15:22:24 +00:00
"""
Code derived from: https://github.com/acheong08/Bard/blob/main/src/Bard.py
"""
import random
import string
import re
import json
import requests
class Bardbot:
"""
A class to interact with Google Bard.
Parameters
session_id: str
The __Secure-1PSID cookie.
2023-05-30 02:24:29 +00:00
timeout: int
Request timeout in seconds.
session: requests.Session
Requests session object.
2023-04-13 15:22:24 +00:00
"""
__slots__ = [
"headers",
"_reqid",
"SNlM0e",
"conversation_id",
"response_id",
"choice_id",
2023-05-30 02:24:29 +00:00
"session_id",
2023-04-13 15:22:24 +00:00
"session",
2023-05-30 02:24:29 +00:00
"timeout",
2023-04-13 15:22:24 +00:00
]
2023-05-30 02:24:29 +00:00
def __init__(
self,
session_id: str,
timeout: int = 20,
session: requests.Session = None,
):
2023-04-13 15:22:24 +00:00
headers = {
"Host": "bard.google.com",
"X-Same-Domain": "1",
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.114 Safari/537.36",
"Content-Type": "application/x-www-form-urlencoded;charset=UTF-8",
"Origin": "https://bard.google.com",
"Referer": "https://bard.google.com/",
}
self._reqid = int("".join(random.choices(string.digits, k=4)))
self.conversation_id = ""
self.response_id = ""
self.choice_id = ""
2023-05-30 02:24:29 +00:00
self.session_id = session_id
self.session = session or requests.Session()
2023-04-13 15:22:24 +00:00
self.session.headers = headers
self.session.cookies.set("__Secure-1PSID", session_id)
self.SNlM0e = self.__get_snlm0e()
2023-05-30 02:24:29 +00:00
self.timeout = timeout
2023-04-13 15:22:24 +00:00
def __get_snlm0e(self):
# Find "SNlM0e":"<ID>"
2023-05-30 02:24:29 +00:00
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,
)
2023-04-13 15:22:24 +00:00
if resp.status_code != 200:
2023-05-30 02:24:29 +00:00
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)
2023-04-13 15:22:24 +00:00
def ask(self, message: str) -> dict:
"""
Send a message to Google Bard and return the response.
:param message: The message to send to Google Bard.
:return: A dict containing the response from Google Bard.
"""
# url params
params = {
2023-05-30 02:24:29 +00:00
"bl": "boq_assistant-bard-web-server_20230523.13_p0",
2023-04-13 15:22:24 +00:00
"_reqid": str(self._reqid),
"rt": "c",
}
# message arr -> data["f.req"]. Message is double json stringified
message_struct = [
[message],
None,
[self.conversation_id, self.response_id, self.choice_id],
]
data = {
"f.req": json.dumps([None, json.dumps(message_struct)]),
"at": self.SNlM0e,
}
resp = self.session.post(
"https://bard.google.com/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate",
params=params,
data=data,
2023-05-30 02:24:29 +00:00
timeout=self.timeout,
2023-04-13 15:22:24 +00:00
)
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)
2023-05-30 02:24:29 +00:00
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])
2023-04-13 15:22:24 +00:00
results = {
"content": json_chat_data[0][0],
"conversation_id": json_chat_data[1][0],
"response_id": json_chat_data[1][1],
"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]],
2023-05-30 02:24:29 +00:00
"images": images,
2023-04-13 15:22:24 +00:00
}
self.conversation_id = results["conversation_id"]
self.response_id = results["response_id"]
self.choice_id = results["choices"][0]["id"]
self._reqid += 100000
return results