Source code for OpenAIManager.ThreadsManager

from openai import OpenAI
from openai.types.beta.thread import Thread
import shelve



[docs] class ThreadsManager: """ Manages the storage and retrieval of threads and messages for an OpenAI application. This class handles the interactions between local storage and the remote OpenAI API, specifically focusing on thread management. It provides methods to create, retrieve, and manage threads both locally and remotely. Due to API limitations and the current scope of the class design, it lacks functionality to fetch all threads directly from the remote server, and its capabilities are primarily centered around individual management of known threads. """ def __init__(self, client: OpenAI) -> None: """ Initializes the ThreadsManager with a client for OpenAI API interactions. This constructor sets up the manager with an OpenAI client and initializes local storage for thread tracking. It loads existing thread data from local storage upon instantiation. :param client: The OpenAI client used for API interactions. :type client: openai.OpenAI :return: None """ self.client = client self.local_threads = self.get_threads_local()
[docs] def get_thread_remote(self, thread_id: str) -> Thread: """ Retrieves a specific thread from the remote OpenAI server using its ID. This method attempts to fetch a thread from the OpenAI API. If the thread does not exist or an error occurs during retrieval, an exception is caught and handled. :param thread_id: The unique identifier of the thread to retrieve. :type thread_id: str :return: The retrieved thread or None if not found or an error occurs. :rtype: openai.types.beta.thread.Thread or None """ thread: Thread = None try: thread = self.client.beta.threads.retrieve(thread_id) except Exception as e: print("Unable to locate thread. The following exception was raised: " + str(e)) return thread
[docs] def get_thread_id_local(self, thread_key: str) -> str: """ Retrieves the ID of a local thread based on a provided \"thread key\". This method looks up a thread's ID in the local shelve database using the provided thread key. If the thread key does not exist in the local storage, None is returned. :param thread_key: The key used to identify the thread in local storage. :type thread_key: str :return: The ID of the thread associated with the given key, or None if not found. :rtype: str or None """ thread_id = None if not thread_key is None: thread_dict = self.get_threads_local() thread_id = thread_dict.get(thread_key, None) return thread_id
[docs] def create_thread_remote(self): """ Creates a new thread on the remote OpenAI server. This method attempts to create a new thread using the OpenAI API. If an error occurs during creation, it is printed and the thread is left as None. :return: The created thread or None if an error occurs. :rtype: Thread or None """ thread = None try: thread = self.client.beta.threads.create() except Exception as e: print("Unable to create thread. The following exception was raised: " + str(e)) return thread
[docs] def create_thread_local(self, thread: Thread=None, thread_key: str=None): """ Creates a new thread in local storage or updates an existing one. This method adds a new thread to local storage using the provided thread object and key. If the thread is not provided, a new one is created remotely. If the thread already exists in local storage, it is updated. :param thread: The thread object to store or update locally. :type thread: Thread, optional :param thread_key: The key to associate with the thread in local storage. :type thread_key: str, optional :return: The thread object after creation or update. :rtype: Thread """ if thread is None: print(f"No thread provided. Creating a new thread \'{thread.id}\' at remote.") thread = self.create_thread_remote() if not thread is None: found_thread_id = None with shelve.open("threads_db", writeback=True) as threads_shelf: print("Looking for thread in local shelve db.") found_thread_id = threads_shelf.get(thread_key, None) if found_thread_id is None: print("Did not find existing thread. Creating a new one...") threads_shelf[thread_key] = thread.id else: print("Returning None.") return thread
[docs] def get_threads_local(self) -> dict: """ Retrieves all threads stored in local storage. This method reads the entire local shelve database and returns its contents as a dictionary. The dictionary maps thread keys to thread IDs. :return: A dictionary of thread keys and their corresponding IDs. :rtype: dict """ shelf_dict = {} with shelve.open("threads_db") as threads_shelf: for key in threads_shelf: shelf_dict[key] = threads_shelf[key] self.local_thread_dict = shelf_dict return shelf_dict
[docs] def get_messages_remote(self, thread): """ Retrieves the message history for a given thread from the OpenAI API. This method fetches the message history associated with the specified thread directly from the OpenAI server. The messages are returned in a list, with the latest message at index 0. :param thread: The thread object for which to retrieve messages. :type thread: Thread :return: A list of messages from the thread, latest first. :rtype: list """ raw_msg_list = self.client.beta.threads.messages.list(thread.id) return [msg.content[0].text.value for msg in raw_msg_list]
def __repr__(self) -> str: """ Returns a string representation of the current state of local thread storage. This method generates a string representation of the threads currently stored in local storage, providing a quick overview of the thread keys and their corresponding IDs. :return: A string representation of the local thread storage. :rtype: str """ shelf_dict = self.get_threads_local() return str(shelf_dict)