diff --git a/.env.example b/.env.example index 864fe92..3a8a89d 100644 --- a/.env.example +++ b/.env.example @@ -7,6 +7,7 @@ ROOM_ID="!FYCmBSkCRUXXXXXXXXX:matrix.XXX.XXX" # Optional, if not set, bot will w OPENAI_API_KEY="xxxxxxxxxxxxxxxxx" # Optional, for !chat and !gpt command BING_API_ENDPOINT="xxxxxxxxxxxxxxx" # Optional, for !bing command ACCESS_TOKEN="xxxxxxxxxxxxxxxxxxxxx" # Optional, use user_id and password is recommended +BARD_TOKEN="xxxxxxxxxxxxxxxxxxxx", # Optional, for !bard command JAILBREAKENABLED="true" # Optional BING_AUTH_COOKIE="xxxxxxxxxxxxxxxxxxx" # _U cookie, Optional, for Bing Image Creator MARKDOWN_FORMATTED="true" # Optional \ No newline at end of file diff --git a/BingImageGen.py b/BingImageGen.py index 4226ae0..6071de6 100644 --- a/BingImageGen.py +++ b/BingImageGen.py @@ -3,15 +3,12 @@ Code derived from: https://github.com/acheong08/EdgeGPT/blob/f940cecd24a4818015a8b42a2443dd97c3c2a8f4/src/ImageGen.py """ from log import getlogger - -from typing import Union from uuid import uuid4 import os import contextlib import aiohttp import asyncio import random -import time import requests import regex @@ -33,164 +30,6 @@ HEADERS = { "x-forwarded-for": FORWARDED_IP, } -# Error messages -error_timeout = "Your request has timed out." -error_redirect = "Redirect failed" -error_blocked_prompt = ( - "Your prompt has been blocked by Bing. Try to change any bad words and try again." -) -error_noresults = "Could not get results" -error_unsupported_lang = "\nthis language is currently not supported by bing" -error_bad_images = "Bad images" -error_no_images = "No images" -# -sending_message = "Sending request..." -wait_message = "Waiting for results..." -download_message = "\nDownloading images..." - - -def debug(debug_file, text_var): - """helper function for debug""" - with open(f"{debug_file}", "a") as f: - f.write(str(text_var)) - - -class ImageGen: - """ - Image generation by Microsoft Bing - Parameters:3 - auth_cookie: str - """ - - def __init__( - self, auth_cookie: str, debug_file: Union[str, None] = None, quiet: bool = False - ) -> None: - self.session: requests.Session = requests.Session() - self.session.headers = HEADERS - self.session.cookies.set("_U", auth_cookie) - self.quiet = quiet - self.debug_file = debug_file - if self.debug_file: - self.debug = partial(debug, self.debug_file) - - def get_images(self, prompt: str) -> list: - """ - Fetches image links from Bing - Parameters: - prompt: str - """ - if not self.quiet: - print(sending_message) - if self.debug_file: - self.debug(sending_message) - url_encoded_prompt = requests.utils.quote(prompt) - # https://www.bing.com/images/create?q=&rt=3&FORM=GENCRE - url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=4&FORM=GENCRE" - response = self.session.post(url, allow_redirects=False) - # check for content waring message - if "this prompt has been blocked" in response.text.lower(): - if self.debug_file: - self.debug(f"ERROR: {error_blocked_prompt}") - raise Exception( - error_blocked_prompt, - ) - if ( - "we're working hard to offer image creator in more languages" - in response.text.lower() - ): - if self.debug_file: - self.debug(f"ERROR: {error_unsupported_lang}") - raise Exception(error_unsupported_lang) - if response.status_code != 302: - # if rt4 fails, try rt3 - url = f"{BING_URL}/images/create?q={url_encoded_prompt}&rt=3&FORM=GENCRE" - response3 = self.session.post( - url, allow_redirects=False, timeout=200) - if response3.status_code != 302: - if self.debug_file: - self.debug(f"ERROR: {error_redirect}") - print(f"ERROR: {response3.text}") - raise Exception(error_redirect) - response = response3 - # Get redirect URL - redirect_url = response.headers["Location"].replace("&nfy=1", "") - request_id = redirect_url.split("id=")[-1] - self.session.get(f"{BING_URL}{redirect_url}") - # https://www.bing.com/images/create/async/results/{ID}?q={PROMPT} - polling_url = f"{BING_URL}/images/create/async/results/{request_id}?q={url_encoded_prompt}" - # Poll for results - if self.debug_file: - self.debug("Polling and waiting for result") - if not self.quiet: - print("Waiting for results...") - start_wait = time.time() - while True: - if int(time.time() - start_wait) > 200: - if self.debug_file: - self.debug(f"ERROR: {error_timeout}") - raise Exception(error_timeout) - if not self.quiet: - print(".", end="", flush=True) - response = self.session.get(polling_url) - if response.status_code != 200: - if self.debug_file: - self.debug(f"ERROR: {error_noresults}") - raise Exception(error_noresults) - if not response.text or response.text.find("errorMessage") != -1: - time.sleep(1) - continue - else: - break - # Use regex to search for src="" - image_links = regex.findall(r'src="([^"]+)"', response.text) - # Remove size limit - normal_image_links = [link.split("?w=")[0] for link in image_links] - # Remove duplicates - normal_image_links = list(set(normal_image_links)) - - # bad_images = [ - # "https://r.bing.com/rp/in-2zU3AJUdkgFe7ZKv19yPBHVs.png", - # "https://r.bing.com/rp/TX9QuO3WzcCJz1uaaSwQAz39Kb0.jpg", - # ] - # for img in normal_image_links: - # if img in bad_images: - # raise Exception("Bad images") - # No images - if not normal_image_links: - raise Exception(error_no_images) - return normal_image_links - - def save_images(self, links: list, output_dir: str) -> str: - """ - Saves images to output directory - """ - - # image name - image_name = str(uuid4()) - # since matrix only support one media attachment per message, we just need one link - if links: - link = links.pop() - - image_path = os.path.join(output_dir, f"{image_name}.jpeg") - - with contextlib.suppress(FileExistsError): - os.mkdir(output_dir) - try: - with self.session.get(link, stream=True) as response: - # save response to file - response.raise_for_status() - with open( - os.path.join(output_dir, image_path), "wb" - ) as output_file: - for chunk in response.iter_content(chunk_size=8192): - output_file.write(chunk) - return image_path - except requests.exceptions.MissingSchema as url_exception: - raise Exception( - "Inappropriate contents found in the generated images. Please try again or try another prompt.", - ) from url_exception - - class ImageGenAsync: """ Image generation by Microsoft Bing @@ -211,6 +50,17 @@ class ImageGenAsync: async def __aexit__(self, *excinfo) -> None: await self.session.close() + def __del__(self): + try: + loop = asyncio.get_event_loop() + except RuntimeError as e: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(self._close()) + + async def _close(self): + await self.session.close() + async def get_images(self, prompt: str) -> list: """ Fetches image links from Bing diff --git a/askgpt.py b/askgpt.py index 1b1c083..e7b2d86 100644 --- a/askgpt.py +++ b/askgpt.py @@ -6,8 +6,8 @@ logger = getlogger() class askGPT: - def __init__(self): - self.session = aiohttp.ClientSession() + def __init__(self, session: aiohttp.ClientSession): + self.session = session async def oneTimeAsk(self, prompt: str, api_endpoint: str, headers: dict) -> str: jsons = { @@ -19,11 +19,11 @@ class askGPT: }, ], } - max_try = 3 + max_try = 2 while max_try > 0: try: async with self.session.post(url=api_endpoint, - json=jsons, headers=headers, timeout=60) as response: + json=jsons, headers=headers, timeout=120) as response: status_code = response.status if not status_code == 200: # print failed reason diff --git a/bard.py b/bard.py new file mode 100644 index 0000000..44e6deb --- /dev/null +++ b/bard.py @@ -0,0 +1,103 @@ +""" +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. + """ + + __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": "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 = "" + 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/_/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 \ No newline at end of file diff --git a/bing.py b/bing.py index d14b3a7..a06c837 100644 --- a/bing.py +++ b/bing.py @@ -7,13 +7,13 @@ logger = getlogger() class BingBot: - def __init__(self, bing_api_endpoint: str, jailbreakEnabled: bool = False): + def __init__(self, session: aiohttp.ClientSession, bing_api_endpoint: str, jailbreakEnabled: bool = False): self.data = { 'clientOptions.clientToUse': 'bing', } self.bing_api_endpoint = bing_api_endpoint - self.session = aiohttp.ClientSession() + self.session = session self.jailbreakEnabled = jailbreakEnabled @@ -22,10 +22,10 @@ class BingBot: async def ask_bing(self, prompt) -> str: self.data['message'] = prompt - max_try = 3 + max_try = 2 while max_try > 0: try: - resp = await self.session.post(url=self.bing_api_endpoint, json=self.data, timeout=60) + resp = await self.session.post(url=self.bing_api_endpoint, json=self.data, timeout=120) status_code = resp.status body = await resp.read() if not status_code == 200: diff --git a/bot.py b/bot.py index 9764e68..d910957 100644 --- a/bot.py +++ b/bot.py @@ -1,35 +1,26 @@ -import sys import asyncio -import re import os -from functools import partial +import re +import sys import traceback from typing import Optional, Union -from nio import ( - AsyncClient, - MatrixRoom, - RoomMessageText, - InviteMemberEvent, - MegolmEvent, - LoginResponse, - JoinError, - ToDeviceError, - LocalProtocolError, - KeyVerificationEvent, - KeyVerificationStart, - KeyVerificationCancel, - KeyVerificationKey, - KeyVerificationMac, - AsyncClientConfig -) + +import aiohttp +from nio import (AsyncClient, AsyncClientConfig, InviteMemberEvent, JoinError, + KeyVerificationCancel, KeyVerificationEvent, + KeyVerificationKey, KeyVerificationMac, KeyVerificationStart, + LocalProtocolError, LoginResponse, MatrixRoom, MegolmEvent, + RoomMessageText, ToDeviceError) from nio.store.database import SqliteStore + from askgpt import askGPT -from send_message import send_room_message -from v3 import Chatbot -from log import getlogger from bing import BingBot from BingImageGen import ImageGenAsync +from log import getlogger from send_image import send_room_image +from send_message import send_room_message +from v3 import Chatbot +from bard import Bardbot logger = getlogger() @@ -47,6 +38,7 @@ class Bot: bing_api_endpoint: Union[str, None] = None, password: Union[str, None] = None, access_token: Union[str, None] = None, + bard_token: Union[str, None] = None, jailbreakEnabled: Union[bool, None] = True, bing_auth_cookie: Union[str, None] = '', markdown_formatted: Union[bool, None] = False, @@ -64,11 +56,14 @@ class Bot: self.user_id = user_id self.password = password self.access_token = access_token + self.bard_token = bard_token self.device_id = device_id self.room_id = room_id self.api_key = api_key self.chatgpt_api_endpoint = chatgpt_api_endpoint + self.session = aiohttp.ClientSession() + if bing_api_endpoint is None: self.bing_api_endpoint = '' else: @@ -112,16 +107,17 @@ class Bot: self.client.add_to_device_callback( self.to_device_callback, (KeyVerificationEvent, )) - # regular expression to match keyword [!gpt {prompt}] [!chat {prompt}] + # regular expression to match keyword [!gpt {prompt}] [!chat {prompt}] [!bing {prompt}] [!pic {prompt}] [!bard {prompt}] self.gpt_prog = re.compile(r"^\s*!gpt\s*(.+)$") self.chat_prog = re.compile(r"^\s*!chat\s*(.+)$") self.bing_prog = re.compile(r"^\s*!bing\s*(.+)$") + self.bard_prog = re.compile(r"^\s*!bard\s*(.+)$") self.pic_prog = re.compile(r"^\s*!pic\s*(.+)$") self.help_prog = re.compile(r"^\s*!help\s*.*$") # initialize chatbot and chatgpt_api_endpoint if self.api_key != '': - self.chatbot = Chatbot(api_key=self.api_key, timeout=60) + self.chatbot = Chatbot(api_key=self.api_key, timeout=120) self.chatgpt_api_endpoint = self.chatgpt_api_endpoint # request header for !gpt command @@ -136,19 +132,33 @@ class Bot: } # initialize askGPT class - self.askgpt = askGPT() + self.askgpt = askGPT(self.session) # initialize bingbot if self.bing_api_endpoint != '': self.bingbot = BingBot( - bing_api_endpoint, jailbreakEnabled=self.jailbreakEnabled) + self.session, bing_api_endpoint, jailbreakEnabled=self.jailbreakEnabled) # initialize BingImageGenAsync if self.bing_auth_cookie != '': self.imageGen = ImageGenAsync(self.bing_auth_cookie, quiet=True) - # get current event loop - self.loop = asyncio.get_running_loop() + # initialize Bardbot + if bard_token is not None: + self.bardbot = Bardbot(self.bard_token) + + + def __del__(self): + try: + loop = asyncio.get_event_loop() + except RuntimeError as e: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(self._close()) + + + async def _close(self): + await self.session.close() # message_callback RoomMessageText event async def message_callback(self, room: MatrixRoom, event: RoomMessageText) -> None: @@ -194,7 +204,7 @@ class Bot: ) except Exception as e: - logger.error(e) + logger.error(e, exc_info=True) await send_room_message(self.client, room_id, reply_message=str(e)) else: logger.warning("No API_KEY provided") @@ -211,7 +221,7 @@ class Bot: raw_user_message ) except Exception as e: - logger.error(e) + logger.error(e, exc_info=True) await send_room_message(self.client, room_id, reply_message=str(e)) # bing ai @@ -230,6 +240,7 @@ class Bot: ) except Exception as e: + logger.error(e, exc_info=True) await send_room_message(self.client, room_id, reply_message=str(e)) # Image Generation by Microsoft Bing @@ -240,8 +251,28 @@ class Bot: try: await self.pic(room_id, prompt) except Exception as e: + logger.error(e, exc_info=True) await send_room_message(self.client, room_id, reply_message=str(e)) + # Google's Bard + if self.bard_token is not None: + b = self.bard_prog.match(content_body) + if b: + prompt = b.group(1) + try: + await self.bard( + room_id, + reply_to_event_id, + prompt, + sender_id, + raw_user_message + ) + except Exception as e: + logger.error(e, exc_info=True) + await send_room_message(self.client, room_id, reply_message={e}) + + + # help command h = self.help_prog.match(content_body) if h: @@ -491,12 +522,9 @@ class Bot: # !chat command async def chat(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message): - await self.client.room_typing(room_id, timeout=180000) + await self.client.room_typing(room_id, timeout=120000) try: - text = await asyncio.wait_for(self.chatbot.ask_async(prompt), timeout=180) - except TimeoutError as e: - logger.error(f"TimeoutException: {e}", exc_info=True) - raise Exception("Timeout error") + text = await self.chatbot.ask_async(prompt) except Exception as e: raise Exception(e) @@ -509,12 +537,12 @@ class Bot: logger.error(f"Error: {e}", exc_info=True) # !gpt command - async def gpt(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message): + async def gpt(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message) -> None: try: # sending typing state - await self.client.room_typing(room_id, timeout=180000) - # timeout 120s - text = await asyncio.wait_for(self.askgpt.oneTimeAsk(prompt, self.chatgpt_api_endpoint, self.headers), timeout=180) + await self.client.room_typing(room_id, timeout=240000) + # timeout 240s + text = await asyncio.wait_for(self.askgpt.oneTimeAsk(prompt, self.chatgpt_api_endpoint, self.headers), timeout=240) except TimeoutError: logger.error("TimeoutException", exc_info=True) raise Exception("Timeout error") @@ -530,12 +558,12 @@ class Bot: logger.error(f"Error: {e}", exc_info=True) # !bing command - async def bing(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message): + async def bing(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message) -> None: try: # sending typing state await self.client.room_typing(room_id, timeout=180000) - # timeout 120s - text = await asyncio.wait_for(self.bingbot.ask_bing(prompt), timeout=180) + # timeout 240s + text = await asyncio.wait_for(self.bingbot.ask_bing(prompt), timeout=240) except TimeoutError: logger.error("timeoutException", exc_info=True) raise Exception("Timeout error") @@ -548,7 +576,24 @@ class Bot: await send_room_message(self.client, room_id, reply_message=text, reply_to_event_id="", sender_id=sender_id, user_message=raw_user_message, markdown_formatted=self.markdown_formatted) except Exception as e: - logger.error(f"Error: {e}", exc_info=True) + logger.error(e, exc_info=True) + + # !bard command + async def bard(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message) -> None: + try: + # sending typing state + await self.client.room_typing(room_id) + response = await asyncio.to_thread(self.bardbot.ask, prompt) + except Exception as e: + raise Exception(e) + + try: + content = str(response['content']).strip() + await send_room_message(self.client, room_id, reply_message=content, + reply_to_event_id="", sender_id=sender_id, user_message=raw_user_message, markdown_formatted=self.markdown_formatted) + except Exception as e: + logger.error(e, exc_info=True) + # !pic command async def pic(self, room_id, prompt): @@ -561,6 +606,7 @@ class Bot: image_path = await self.imageGen.save_images(links, "images") except Exception as e: logger.error(f"Image Generation error: {e}", exc_info=True) + raise Exception(e) # send image try: @@ -570,7 +616,7 @@ class Bot: logger.error(e, exc_info=True) except Exception as e: - logger.error(f"Error: {e}", exc_info=True) + logger.error(e, exc_info=True) # !help command async def help(self, room_id): @@ -580,12 +626,13 @@ class Bot: help_info = "!gpt [content], generate response without context conversation\n" + \ "!chat [content], chat with context conversation\n" + \ "!bing [content], chat with context conversation powered by Bing AI\n" + \ + "!bard [content], chat with Google's Bard\n" + \ "!pic [prompt], Image generation by Microsoft Bing\n" + \ "!help, help message" await send_room_message(self.client, room_id, reply_message=help_info) except Exception as e: - logger.error(f"Error: {e}", exc_info=True) + logger.error(e, exc_info=True) # bot login async def login(self) -> None: diff --git a/config.json.sample b/config.json.sample index baf6c95..af351a8 100644 --- a/config.json.sample +++ b/config.json.sample @@ -8,6 +8,7 @@ "bing_api_endpoint": "http://api:3000/conversation", "jailbreakEnabled": true, "access_token": "xxxxxxx", + "bard_token": "xxxxxxx", "bing_auth_cookie": "xxxxxxxxxxx", "markdown_formatted": true } diff --git a/main.py b/main.py index cc7d15b..4b3a73e 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,7 @@ +import asyncio import json import os -import asyncio + from bot import Bot from log import getlogger @@ -21,6 +22,7 @@ async def main(): api_key=config.get('api_key'), bing_api_endpoint=config.get('bing_api_endpoint'), access_token=config.get('access_token'), + bard_token=config.get('bard_token'), jailbreakEnabled=config.get('jailbreakEnabled'), bing_auth_cookie=config.get('bing_auth_cookie'), markdown_formatted=config.get('markdown_formatted'), @@ -35,6 +37,7 @@ async def main(): api_key=os.environ.get("OPENAI_API_KEY"), bing_api_endpoint=os.environ.get("BING_API_ENDPOINT"), access_token=os.environ.get("ACCESS_TOKEN"), + bard_token=os.environ.get("BARD_TOKEN"), jailbreakEnabled=os.environ.get( "JAILBREAKENABLED", "false").lower() in ('true', '1', 't'), bing_auth_cookie=os.environ.get("BING_AUTH_COOKIE"), @@ -52,4 +55,6 @@ async def main(): if __name__ == "__main__": logger.info("matrix chatgpt bot start.....") + print("matrix chatgpt bot start.....") asyncio.run(main()) + diff --git a/v3.py b/v3.py index b91dca3..b708bfe 100644 --- a/v3.py +++ b/v3.py @@ -1,7 +1,10 @@ +""" +Code derived from: https://github.com/acheong08/ChatGPT/blob/main/src/revChatGPT/V3.py +""" + import json import os from typing import AsyncGenerator - import httpx import requests import tiktoken