format code

This commit is contained in:
hibobmaster 2023-04-11 13:42:43 +08:00
parent f5ccc67258
commit da682cb9fd
No known key found for this signature in database
GPG Key ID: 316B77D7914D713C
5 changed files with 56 additions and 48 deletions

View File

@ -49,7 +49,6 @@ 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:
@ -74,7 +73,6 @@ class ImageGen:
if self.debug_file:
self.debug = partial(debug, self.debug_file)
def get_images(self, prompt: str) -> list:
"""
Fetches image links from Bing
@ -106,7 +104,8 @@ class ImageGen:
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)
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}")
@ -165,7 +164,7 @@ class ImageGen:
"""
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
@ -308,7 +307,7 @@ class ImageGenAsync:
async for chunk in response.content.iter_chunked(8192):
output_file.write(chunk)
return f"{output_dir}/{image_name}.jpeg"
except aiohttp.client_exceptions.InvalidURL as url_exception:
raise Exception(
"Inappropriate contents found in the generated images. Please try again or try another prompt.",

6
log.py
View File

@ -14,8 +14,10 @@ def getlogger():
info_handler.setLevel(logging.INFO)
# create formatters
warn_format = logging.Formatter('%(name)s - %(funcName)s - %(levelname)s - %(message)s')
error_format = logging.Formatter('%(asctime)s - %(name)s - %(funcName)s - %(levelname)s - %(message)s')
warn_format = logging.Formatter(
'%(name)s - %(funcName)s - %(levelname)s - %(message)s')
error_format = logging.Formatter(
'%(asctime)s - %(name)s - %(funcName)s - %(levelname)s - %(message)s')
info_format = logging.Formatter('%(message)s')
# set formatter

53
main.py
View File

@ -6,38 +6,41 @@ from log import getlogger
logger = getlogger()
async def main():
if os.path.exists('config.json'):
fp = open('config.json', 'r', encoding="utf8")
config = json.load(fp)
matrix_bot = Bot(homeserver=config.get('homeserver'),
user_id=config.get('user_id') ,
password=config.get('password'),
device_id=config.get('device_id'),
room_id=config.get('room_id'),
api_key=config.get('api_key'),
bing_api_endpoint=config.get('bing_api_endpoint'),
access_token=config.get('access_token'),
jailbreakEnabled=config.get('jailbreakEnabled'),
bing_auth_cookie=config.get('bing_auth_cookie'),
markdown_formatted=config.get('markdown_formatted'),
)
user_id=config.get('user_id'),
password=config.get('password'),
device_id=config.get('device_id'),
room_id=config.get('room_id'),
api_key=config.get('api_key'),
bing_api_endpoint=config.get('bing_api_endpoint'),
access_token=config.get('access_token'),
jailbreakEnabled=config.get('jailbreakEnabled'),
bing_auth_cookie=config.get('bing_auth_cookie'),
markdown_formatted=config.get('markdown_formatted'),
)
else:
matrix_bot = Bot(homeserver=os.environ.get('HOMESERVER'),
user_id=os.environ.get('USER_ID') ,
password=os.environ.get('PASSWORD'),
device_id=os.environ.get("DEVICE_ID"),
room_id=os.environ.get("ROOM_ID"),
api_key=os.environ.get("OPENAI_API_KEY"),
bing_api_endpoint=os.environ.get("BING_API_ENDPOINT"),
access_token=os.environ.get("ACCESS_TOKEN"),
jailbreakEnabled=os.environ.get("JAILBREAKENABLED", "false").lower() in ('true', '1', 't'),
bing_auth_cookie=os.environ.get("BING_AUTH_COOKIE"),
markdown_formatted=os.environ.get("MARKDOWN_FORMATTED", "false").lower() in ('true', '1', 't'),
)
user_id=os.environ.get('USER_ID'),
password=os.environ.get('PASSWORD'),
device_id=os.environ.get("DEVICE_ID"),
room_id=os.environ.get("ROOM_ID"),
api_key=os.environ.get("OPENAI_API_KEY"),
bing_api_endpoint=os.environ.get("BING_API_ENDPOINT"),
access_token=os.environ.get("ACCESS_TOKEN"),
jailbreakEnabled=os.environ.get(
"JAILBREAKENABLED", "false").lower() in ('true', '1', 't'),
bing_auth_cookie=os.environ.get("BING_AUTH_COOKIE"),
markdown_formatted=os.environ.get(
"MARKDOWN_FORMATTED", "false").lower() in ('true', '1', 't'),
)
await matrix_bot.login()
await matrix_bot.sync_encryption_key()
@ -50,5 +53,3 @@ async def main():
if __name__ == "__main__":
logger.info("matrix chatgpt bot start.....")
asyncio.run(main())

View File

@ -13,7 +13,7 @@ logger = getlogger()
async def send_room_image(client: AsyncClient,
room_id: str, image: str):
room_id: str, image: str):
"""
image: image path
"""
@ -36,7 +36,8 @@ async def send_room_image(client: AsyncClient,
await client.room_send(
room_id,
message_type="m.room.message",
content={"msgtype": "m.text", "body": f"Failed to generate image. Failure response: {resp}", },
content={"msgtype": "m.text",
"body": f"Failed to generate image. Failure response: {resp}", },
ignore_unverified_devices=True,
)
return
@ -56,4 +57,5 @@ async def send_room_image(client: AsyncClient,
try:
await client.room_send(room_id, message_type="m.room.message", content=content)
except Exception as e:
logger.error(f"Image send of file {image} failed.\n Error: {e}", exc_info=True)
logger.error(
f"Image send of file {image} failed.\n Error: {e}", exc_info=True)

28
v3.py
View File

@ -53,7 +53,8 @@ class Chatbot:
},
)
proxy = (
proxy or os.environ.get("all_proxy") or os.environ.get("ALL_PROXY") or None
proxy or os.environ.get(
"all_proxy") or os.environ.get("ALL_PROXY") or None
)
if proxy:
@ -79,7 +80,6 @@ class Chatbot:
],
}
def add_to_conversation(
self,
message: str,
@ -105,7 +105,6 @@ class Chatbot:
else:
break
def get_token_count(self, convo_id: str = "default") -> int:
"""
Get token count
@ -158,8 +157,10 @@ class Chatbot:
self.__truncate_conversation(convo_id=convo_id)
# Get response
response = self.session.post(
os.environ.get("API_URL") or "https://api.openai.com/v1/chat/completions",
headers={"Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"},
os.environ.get(
"API_URL") or "https://api.openai.com/v1/chat/completions",
headers={
"Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"},
json={
"model": self.engine,
"messages": self.conversation[convo_id],
@ -182,7 +183,7 @@ class Chatbot:
timeout=kwargs.get("timeout", self.timeout),
stream=True,
)
response_role: str = None
full_response: str = ""
for line in response.iter_lines():
@ -205,7 +206,8 @@ class Chatbot:
content = delta["content"]
full_response += content
yield content
self.add_to_conversation(full_response, response_role, convo_id=convo_id)
self.add_to_conversation(
full_response, response_role, convo_id=convo_id)
async def ask_stream_async(
self,
@ -225,8 +227,10 @@ class Chatbot:
# Get response
async with self.aclient.stream(
"post",
os.environ.get("API_URL") or "https://api.openai.com/v1/chat/completions",
headers={"Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"},
os.environ.get(
"API_URL") or "https://api.openai.com/v1/chat/completions",
headers={
"Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"},
json={
"model": self.engine,
"messages": self.conversation[convo_id],
@ -274,7 +278,8 @@ class Chatbot:
content: str = delta["content"]
full_response += content
yield content
self.add_to_conversation(full_response, response_role, convo_id=convo_id)
self.add_to_conversation(
full_response, response_role, convo_id=convo_id)
async def ask_async(
self,
@ -314,11 +319,10 @@ class Chatbot:
full_response: str = "".join(response)
return full_response
def reset(self, convo_id: str = "default", system_prompt: str = None) -> None:
"""
Reset the conversation
"""
self.conversation[convo_id] = [
{"role": "system", "content": system_prompt or self.system_prompt},
]
]