diff --git a/cogs/text_service_cog.py b/cogs/text_service_cog.py index 3ca8070..314ab3f 100644 --- a/cogs/text_service_cog.py +++ b/cogs/text_service_cog.py @@ -120,7 +120,6 @@ class GPT3ComCon(discord.Cog, name="GPT3ComCon"): self.full_conversation_history = defaultdict(list) self.summarize = self.model.summarize_conversations - # Pinecone data self.pinecone_service = pinecone_service @@ -232,15 +231,25 @@ class GPT3ComCon(discord.Cog, name="GPT3ComCon"): print("Attempting to load from pickles") # Try to load self.full_conversation_history, self.conversation_threads, and self.conversation_thread_owners from the `pickles` folder try: - with open(EnvService.save_path() / "pickles" / "full_conversation_history.pickle", "rb") as f: + with open( + EnvService.save_path() / "pickles" / "full_conversation_history.pickle", + "rb", + ) as f: self.full_conversation_history = pickle.load(f) print("Loaded full_conversation_history") - with open(EnvService.save_path() / "pickles" / "conversation_threads.pickle", "rb") as f: + with open( + EnvService.save_path() / "pickles" / "conversation_threads.pickle", "rb" + ) as f: self.conversation_threads = pickle.load(f) print("Loaded conversation_threads") - with open(EnvService.save_path() / "pickles" / "conversation_thread_owners.pickle", "rb") as f: + with open( + EnvService.save_path() + / "pickles" + / "conversation_thread_owners.pickle", + "rb", + ) as f: self.conversation_thread_owners = pickle.load(f) print("Loaded conversation_thread_owners") @@ -269,14 +278,17 @@ class GPT3ComCon(discord.Cog, name="GPT3ComCon"): ) print("Commands synced") - # Start an inline async loop that runs every 10 seconds to save the conversation history to a pickle file print("Starting pickle loop") while True: await asyncio.sleep(15) await self.pickle_queue.put( - Pickler(self.full_conversation_history, self.conversation_threads, self.conversation_thread_owners)) - + Pickler( + self.full_conversation_history, + self.conversation_threads, + self.conversation_thread_owners, + ) + ) def check_conversing(self, channel_id, message_content): '''given channel id and a message, return true if it's a conversation thread, false if not, or if the message starts with "~"''' diff --git a/gpt3discord.py b/gpt3discord.py index 3b94f52..0a23d2c 100644 --- a/gpt3discord.py +++ b/gpt3discord.py @@ -79,10 +79,12 @@ asyncio.ensure_future(Deletion.process_deletion_queue(deletion_queue, 1, 1)) # Pickling service for conversation persistence try: - Path(EnvService.save_path()/"pickles").mkdir(exist_ok=True) + Path(EnvService.save_path() / "pickles").mkdir(exist_ok=True) except Exception: traceback.print_exc() - print("Could not start pickle service. Conversation history will not be persistent across restarts.") + print( + "Could not start pickle service. Conversation history will not be persistent across restarts." + ) pickle_queue = asyncio.Queue() asyncio.ensure_future(Pickler.process_pickle_queue(pickle_queue, 5, 1)) diff --git a/services/pickle_service.py b/services/pickle_service.py index 4749e09..4e2f9fe 100644 --- a/services/pickle_service.py +++ b/services/pickle_service.py @@ -10,16 +10,19 @@ from services.environment_service import EnvService class Pickler: - def __init__(self, full_conversation_history, conversation_threads, conversation_thread_owners): + def __init__( + self, + full_conversation_history, + conversation_threads, + conversation_thread_owners, + ): self.full_conversation_history = full_conversation_history self.conversation_threads = conversation_threads self.conversation_thread_owners = conversation_thread_owners # This function will be called by the bot to process the message queue @staticmethod - async def process_pickle_queue( - pickle_queue, PROCESS_WAIT_TIME, EMPTY_WAIT_TIME - ): + async def process_pickle_queue(pickle_queue, PROCESS_WAIT_TIME, EMPTY_WAIT_TIME): while True: try: # If the queue is empty, sleep for a short time before checking again @@ -31,13 +34,26 @@ class Pickler: to_pickle = await pickle_queue.get() # Pickle all the objects inside to_pickle using aiofiles - async with aiofiles.open(EnvService.save_path() / "pickles" / "full_conversation_history.pickle", "wb") as f: + async with aiofiles.open( + EnvService.save_path() + / "pickles" + / "full_conversation_history.pickle", + "wb", + ) as f: await f.write(pickle.dumps(to_pickle.full_conversation_history)) - async with aiofiles.open(EnvService.save_path() / "pickles" / "conversation_threads.pickle", "wb") as f: + async with aiofiles.open( + EnvService.save_path() / "pickles" / "conversation_threads.pickle", + "wb", + ) as f: await f.write(pickle.dumps(to_pickle.conversation_threads)) - async with aiofiles.open(EnvService.save_path() / "pickles" / "conversation_thread_owners.pickle", "wb") as f: + async with aiofiles.open( + EnvService.save_path() + / "pickles" + / "conversation_thread_owners.pickle", + "wb", + ) as f: await f.write(pickle.dumps(to_pickle.conversation_thread_owners)) await asyncio.sleep(PROCESS_WAIT_TIME)