Files
bot-on-anything/model/google/bard_bot.py
2023-04-08 01:21:52 +08:00

67 lines
2.7 KiB
Python

import json
import random
import requests
import re
class BardBot:
BARD_URL = "https://bard.google.com/"
BARD_CHAT_URL = (
"https://bard.google.com/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate"
)
HEADERS = {
"Host": "bard.google.com",
"X-Same-Domain": "1",
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36",
"Content-Type": "application/x-www-form-urlencoded;charset=UTF-8",
"Origin": "https://bard.google.com",
"Referer": "https://bard.google.com/",
}
def __init__(self, session_id: str):
self._reqid = random.randrange(10000,99999)
self.conversation_id = ""
self.response_id = ""
self.choice_id = ""
self.session = requests.Session()
self.session.headers = self.HEADERS
self.session.cookies.set("__Secure-1PSID", session_id)
self.SNlM0e = self.__get_snlm0e()
def __get_snlm0e(self) -> str:
resp = self.session.get(url=self.BARD_URL, timeout=10)
if resp.status_code != 200:
raise Exception("Failed to connect Google Bard")
try:
SNlM0e = re.search(r"SNlM0e\":\"(.*?)\"", resp.text).group(1)
return SNlM0e
except Exception as e:
raise Exception(f"Cookies may be wrong:{e}")
def ask(self, message: str) -> dict[str, str]:
params = {
"bl": "boq_assistant-bard-web-server_20230326.21_p0",
"_reqid": str(self._reqid),
"rt": "c",
}
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}
try:
resp = self.session.post(self.BARD_CHAT_URL, params=params, data=data)
content = json.loads(resp.content.splitlines()[3])[0][2]
if not (content := json.loads(resp.content.splitlines()[3])[0][2]):
return {"content": f"Bard encountered an error: {resp.content}."}
json_data = json.loads(content)
results = {
"content": json_data[0][0],
"conversation_id": json_data[1][0],
"response_id": json_data[1][1],
"reference": json_data[3],
"choices": [{"id": i[0], "content": i[1]} for i in json_data[4]],
}
self.conversation_id = results['conversation_id']
self.response_id = results['response_id']
self.choice_id = results["choices"][0]["id"]
self._reqid += 100000
return results
except Exception as e:
raise Exception(f"Failed to ask Google Bard:{e}")