From 4eef6284d897823ee9a92df33a686df134cd79bb Mon Sep 17 00:00:00 2001 From: hibobmaster Date: Sun, 28 May 2023 23:52:44 +0800 Subject: [PATCH] feat: Support chatgpt web session isolation --- .env.example | 2 +- .gitignore | 3 +++ README.md | 2 +- bot.py | 65 +++++++++++++++++++++++++++++----------------------- compose.yaml | 54 +++++++++++++++++++++++++------------------ main.py | 2 +- pandora.py | 4 ++-- 7 files changed, 76 insertions(+), 56 deletions(-) diff --git a/.env.example b/.env.example index 8412f84..1b175cc 100644 --- a/.env.example +++ b/.env.example @@ -5,5 +5,5 @@ OPENAI_API_KEY="sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" BING_API_ENDPOINT="http://api:3000/conversation" BARD_TOKEN="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx." BING_AUTH_COOKIE="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" -PANDORA_API_ENDPOINT="http://127.0.0.1:8008" +PANDORA_API_ENDPOINT="http://pandora:8008" PANDORA_API_MODEL="text-davinci-002-render-sha-mobile" \ No newline at end of file diff --git a/.gitignore b/.gitignore index 188f13e..8827e10 100644 --- a/.gitignore +++ b/.gitignore @@ -134,3 +134,6 @@ dmypy.json # Pyre type checker .pyre/ + +# custom +compose-local-dev.yaml \ No newline at end of file diff --git a/README.md b/README.md index 11efd9c..cbd01eb 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ This is a simple Mattermost Bot that uses OpenAI's GPT API and Bing AI and Googl 1. Support Openai ChatGPT and Bing AI and Google Bard 2. Support Bing Image Creator -3. [pandora](https://github.com/pengzhile/pandora) +3. [pandora](https://github.com/pengzhile/pandora) with Session isolation support ## Installation and Setup diff --git a/bot.py b/bot.py index 37fabb7..a6b4d78 100644 --- a/bot.py +++ b/bot.py @@ -120,9 +120,9 @@ class Bot: if pandora_api_endpoint is not None: self.pandora_api_endpoint = pandora_api_endpoint self.pandora = Pandora( - api_endpoint=pandora_api_endpoint + api_endpoint=pandora_api_endpoint, + clientSession=self.session ) - self.pandora_init() if pandora_api_model is None: self.pandora_api_model = "text-davinci-002-render-sha-mobile" else: @@ -155,24 +155,22 @@ class Bot: self.goon_prog = re.compile(r"^\s*!goon\s*.*$") self.new_prog = re.compile(r"^\s*!new\s*.*$") + self.pandora_data = {} + # close session def __del__(self) -> None: self.driver.disconnect() - async def __aenter__(self): - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - await self.session.close() - def login(self) -> None: self.driver.login() - def pandora_init(self) -> None: - self.conversation_id = None - self.parent_message_id = str(uuid.uuid4()) - self.first_time = True - + def pandora_init(self, user_id: str) -> None: + self.pandora_data[user_id] = { + "conversation_id": None, + "parent_message_id": str(uuid.uuid4()), + "first_time": True + } + async def run(self) -> None: await self.driver.init_websocket(self.websocket_handler) @@ -190,6 +188,9 @@ class Bot: sender_name = response["data"]["sender_name"] raw_message = raw_data_dict["message"] + if user_id not in self.pandora_data: + self.pandora_init(user_id) + try: asyncio.create_task( self.message_callback( @@ -250,32 +251,32 @@ class Bot: if self.talk_prog.match(message): prompt = self.talk_prog.match(message).group(1) try: - if self.conversation_id is not None: + if self.pandora_data[user_id]["conversation_id"] is not None: data = { "prompt": prompt, "model": self.pandora_api_model, - "parent_message_id": self.parent_message_id, - "conversation_id": self.conversation_id, + "parent_message_id": self.pandora_data[user_id]["parent_message_id"], + "conversation_id": self.pandora_data[user_id]["conversation_id"], "stream": False, } else: data = { "prompt": prompt, "model": self.pandora_api_model, - "parent_message_id": self.parent_message_id, + "parent_message_id": self.pandora_data[user_id]["parent_message_id"], "stream": False, } response = await self.pandora.talk(data) - self.conversation_id = response['conversation_id'] - self.parent_message_id = response['message']['id'] + self.pandora_data[user_id]["conversation_id"] = response['conversation_id'] + self.pandora_data[user_id]["parent_message_id"] = response['message']['id'] content = response['message']['content']['parts'][0] - if self.first_time: - self.first_time = False + if self.pandora_data[user_id]["first_time"]: + self.pandora_data[user_id]["first_time"] = False data = { "model": self.pandora_api_model, - "message_id": self.parent_message_id, + "message_id": self.pandora_data[user_id]["parent_message_id"], } - await self.pandora.gen_title(data, self.conversation_id) + await self.pandora.gen_title(data, self.pandora_data[user_id]["conversation_id"]) await asyncio.to_thread( self.send_message, channel_id, f"{content}" @@ -285,17 +286,17 @@ class Bot: raise Exception(e) # !goon command trigger handler - if self.goon_prog.match(message) and self.conversation_id is not None: + if self.goon_prog.match(message) and self.pandora_data[user_id]["conversation_id"] is not None: try: data = { "model": self.pandora_api_model, - "parent_message_id": self.parent_message_id, - "conversation_id": self.conversation_id, + "parent_message_id": self.pandora_data[user_id]["parent_message_id"], + "conversation_id": self.pandora_data[user_id]["conversation_id"], "stream": False, } response = await self.pandora.goon(data) - self.conversation_id = response['conversation_id'] - self.parent_message_id = response['message']['id'] + self.pandora_data[user_id]["conversation_id"] = response['conversation_id'] + self.pandora_data[user_id]["parent_message_id"] = response['message']['id'] content = response['message']['content']['parts'][0] await asyncio.to_thread( self.send_message, channel_id, f"{content}" @@ -306,7 +307,13 @@ class Bot: # !new command trigger handler if self.new_prog.match(message): - self.pandora_init() + self.pandora_init(user_id) + try: + await asyncio.to_thread( + self.send_message, channel_id, "New conversation created, please use !talk to start chatting!" + ) + except Exception: + pass if self.bard_token is not None: # !bard command trigger handler diff --git a/compose.yaml b/compose.yaml index 5ded1a2..e5f774a 100644 --- a/compose.yaml +++ b/compose.yaml @@ -1,23 +1,33 @@ -services: - app: - image: ghcr.io/hibobmaster/mattermost_bot:latest - container_name: mattermost_bot - restart: always - env_file: - - .env - # volumes: - # use env file or config.json - # - ./config.json:/app/config.json - networks: - - mattermost_network - - # api: - # image: hibobmaster/node-chatgpt-api:latest - # container_name: node-chatgpt-api - # volumes: - # - ./settings.js:/var/chatgpt-api/settings.js - # networks: - # - mattermost_network - -networks: +services: + app: + image: ghcr.io/hibobmaster/mattermost_bot:latest + container_name: mattermost_bot + restart: unless-stopped + env_file: + - .env + # volumes: + # use env file or config.json + # - ./config.json:/app/config.json + networks: + - mattermost_network + + # api: + # image: hibobmaster/node-chatgpt-api:latest + # container_name: node-chatgpt-api + # volumes: + # - ./settings.js:/var/chatgpt-api/settings.js + # networks: + # - mattermost_network + + # pandora: + # image: pengzhile/pandora + # container_name: pandora + # restart: unless-stopped + # environment: + # - PANDORA_ACCESS_TOKEN="xxxxxxxxxxxxxx" + # - PANDORA_SERVER="0.0.0.0:8008" + # networks: + # - mattermost_network + +networks: mattermost_network: \ No newline at end of file diff --git a/main.py b/main.py index e85f0b7..a8f2754 100644 --- a/main.py +++ b/main.py @@ -3,7 +3,6 @@ import json import os import asyncio - async def main(): if os.path.exists("config.json"): fp = open("config.json", "r", encoding="utf-8") @@ -51,3 +50,4 @@ async def main(): if __name__ == "__main__": asyncio.run(main()) + diff --git a/pandora.py b/pandora.py index 59af886..7a8a912 100644 --- a/pandora.py +++ b/pandora.py @@ -3,9 +3,9 @@ import uuid import aiohttp import asyncio class Pandora: - def __init__(self, api_endpoint: str) -> None: + def __init__(self, api_endpoint: str, clientSession: aiohttp.ClientSession) -> None: self.api_endpoint = api_endpoint.rstrip('/') - self.session = aiohttp.ClientSession() + self.session = clientSession async def __aenter__(self): return self