diff --git a/chat.py b/chat.py index a3f9c20..45d618d 100644 --- a/chat.py +++ b/chat.py @@ -24,6 +24,7 @@ AI_MODEL_OPTIONS = [ "gpt-3.5-turbo", "gpt-4", "gpt-4-32k", + "bard", ] st.set_page_config(page_title=PAGE_TITLE, page_icon=PAGE_ICON) @@ -55,7 +56,7 @@ if "user_text" not in st.session_state: def main() -> None: if st.session_state.user_text: - show_conversation(st.session_state.user_text, st.session_state.model, st.session_state.role) + show_conversation() st.session_state.user_text = "" c1, c2, c3 = st.columns(3) diff --git a/requirements.txt b/requirements.txt index 45aa4bf..22325e1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,4 @@ openai==0.27.2 gtts==2.3.1 pip==23.0.1 watchdog==3.0.0 +prompt_toolkit==3.0.38 diff --git a/src/utils/agi/bard.py b/src/utils/agi/bard.py new file mode 100644 index 0000000..c65bbe4 --- /dev/null +++ b/src/utils/agi/bard.py @@ -0,0 +1,176 @@ +""" +Reverse engineering of Google Bard from https://github.com/discordtehe/Bard +""" +import argparse +import json +import random +import re +import string + +import requests +from prompt_toolkit import PromptSession, prompt +from prompt_toolkit.auto_suggest import AutoSuggestFromHistory +from prompt_toolkit.completion import WordCompleter +from prompt_toolkit.history import InMemoryHistory +from prompt_toolkit.key_binding import KeyBindings +from rich.console import Console +from rich.markdown import Markdown + +US_AG = "Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.114 Safari/537.36" + + +def __create_session() -> PromptSession: + return PromptSession(history=InMemoryHistory()) + + +def __create_completer(commands: list, pattern_str: str = "$") -> WordCompleter: + return WordCompleter(words=commands, pattern=re.compile(pattern_str)) + + +def __get_input( + session: PromptSession = None, + completer: WordCompleter = None, + key_bindings: KeyBindings = None, +) -> str: + """ + Multiline input function. + """ + return ( + session.prompt( + completer=completer, + multiline=True, + auto_suggest=AutoSuggestFromHistory(), + key_bindings=key_bindings, + ) + if session + else prompt(multiline=True) + ) + + +class Chatbot: + """ + A class to interact with Google Bard. + Parameters + session_id: str + The __Secure-1PSID cookie. + """ + + __slots__ = [ + "headers", + "_reqid", + "SNlM0e", + "conversation_id", + "response_id", + "choice_id", + "session", + ] + + def __init__(self, session_id): + headers = { + "Host": "bard.google.com", + "X-Same-Domain": "1", + "User-Agent": US_AG, + "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 = "" + self.session = requests.Session() + self.session.headers = headers + self.session.cookies.set("__Secure-1PSID", session_id) + self.SNlM0e = self.__get_snlm0e() + + def __get_snlm0e(self): + resp = self.session.get(url="https://bard.google.com/", timeout=10) + # Find "SNlM0e":"" + if resp.status_code != 200: + raise Exception("Could not get Google Bard") + SNlM0e = re.search(r"SNlM0e\":\"(.*?)\"", resp.text).group(1) + return SNlM0e + + 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 = { + "bl": "boq_assistant-bard-web-server_20230326.21_p0", + "_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, + } + + # do the request! + resp = self.session.post( + "https://bard.google.com/u/1/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate", + params=params, + data=data, + timeout=120, + ) + + 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) + 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]], + } + self.conversation_id = results["conversation_id"] + self.response_id = results["response_id"] + self.choice_id = results["choices"][0]["id"] + self._reqid += 100000 + return results + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--session", + help="__Secure-1PSID cookie.", + type=str, + required=True, + ) + args = parser.parse_args() + + chatbot = Chatbot(args.session) + prompt_session = __create_session() + completions = __create_completer(["!exit", "!reset"]) + console = Console() + try: + while True: + console.print("You:") + user_prompt = __get_input(session=prompt_session, completer=completions) + console.print() + if user_prompt == "!exit": + break + elif user_prompt == "!reset": + chatbot.conversation_id = "" + chatbot.response_id = "" + chatbot.choice_id = "" + continue + print("Bard:") + response = chatbot.ask(user_prompt) + console.print(Markdown(response["content"])) + print() + except KeyboardInterrupt: + print("Exiting...") diff --git a/src/utils/agi/chat_gpt.py b/src/utils/agi/chat_gpt.py index 85545bd..0fab322 100644 --- a/src/utils/agi/chat_gpt.py +++ b/src/utils/agi/chat_gpt.py @@ -6,7 +6,7 @@ import streamlit as st @st.cache_data() -def send_ai_request(ai_model: str, messages: List[dict]) -> dict: +def chat_gpt_request(ai_model: str, messages: List[dict]) -> dict: openai.api_key = st.secrets.api_credentials.api_key logging.warning("messages:") logging.warning(messages) diff --git a/src/utils/conversation.py b/src/utils/conversation.py index b1c6130..d46e01e 100644 --- a/src/utils/conversation.py +++ b/src/utils/conversation.py @@ -2,7 +2,8 @@ import streamlit as st from openai.error import InvalidRequestError, OpenAIError from streamlit_chat import message -from src.utils.agi.chat_gpt import send_ai_request +from src.utils.agi.bard import Chatbot +from src.utils.agi.chat_gpt import chat_gpt_request from src.utils.tts import show_player @@ -42,20 +43,13 @@ def show_chat(ai_content: str, user_text: str) -> None: st.markdown(st.session_state.generated[i]) -def show_conversation(user_content: str, model: str, role: str) -> None: - if st.session_state.messages: - st.session_state.messages.append({"role": "user", "content": user_content}) - else: - st.session_state.messages = [ - {"role": "system", "content": f"{st.session_state.locale.ai_role_prefix} {role}."}, - {"role": "user", "content": user_content}, - ] +def chat_gpt_conversation() -> None: try: - completion = send_ai_request(model, st.session_state.messages) + completion = chat_gpt_request(st.session_state.model, st.session_state.messages) ai_content = completion.get("choices")[0].get("message").get("content") st.session_state.messages.append({"role": "assistant", "content": ai_content}) if ai_content: - show_chat(ai_content, user_content) + show_chat(ai_content, st.session_state.user_text) st.markdown("---") show_player(ai_content) except InvalidRequestError as err: @@ -63,8 +57,28 @@ def show_conversation(user_content: str, model: str, role: str) -> None: st.session_state.messages.pop(1) if len(st.session_state.messages) == 1: st.session_state.user_text = "" - show_conversation(st.session_state.user_text, st.session_state.model, st.session_state.role) + show_conversation() else: st.error(err) except (OpenAIError, UnboundLocalError) as err: st.error(err) + + +def bard_conversation() -> None: + bard = Chatbot(st.secrets.api_credentials.bard_session) + ai_content = bard.ask(st.session_state.user_text) + st.warning(ai_content) + + +def show_conversation() -> None: + if st.session_state.messages: + st.session_state.messages.append({"role": "user", "content": st.session_state.user_text}) + else: + st.session_state.messages = [ + {"role": "system", "content": f"{st.session_state.locale.ai_role_prefix} {st.session_state.role}."}, + {"role": "user", "content": st.session_state.user_text}, + ] + if st.session_state.model == "bard": + bard_conversation() + else: + chat_gpt_conversation()