Altium-archimajor-3d-printer-driver-demo/assistant/OpenAIResourceManager.py

from openai import OpenAI
import argparse
import logging
import os
import pandas as pd
import textwrap

class OpenAIResourceManager:
    """
    A class to manage OpenAI resources such as assistants, vector stores, and files.
    Provides methods to delete all resources of each type and display them in tables.
    """

    def __init__(self, api_key: str):
        """
        Initialize the OpenAIResourceManager with the given API key and configure logging.

        :param api_key: The API key for OpenAI.
        """
        self.client = OpenAI(api_key=api_key)
        logging.basicConfig(level=logging.INFO)

    def get_all_assistants(self):
        """
        Retrieve all assistants associated with the OpenAI account.

        :return: A list of assistants.
        """
        try:
            assistants = self.client.beta.assistants.list()
            return assistants.data
        except Exception as e:
            logging.error(f"Failed to retrieve assistants: {e}")
            raise

    def get_all_vector_stores(self):
        """
        Retrieve all vector stores associated with the OpenAI account.

        :return: A list of vector stores.
        """
        try:
            vector_stores = self.client.beta.vector_stores.list()
            return vector_stores.data
        except Exception as e:
            logging.error(f"Failed to retrieve vector stores: {e}")
            raise

    def get_all_files(self):
        """
        Retrieve all files associated with the OpenAI account.

        :return: A list of files.
        """
        try:
            files = self.client.files.list()
            return files.data
        except Exception as e:
            logging.error(f"Failed to retrieve files: {e}")
            raise

    def delete_all_assistants(self):
        """
        Delete all assistants associated with the OpenAI account.
        """
        try:
            assistants = self.get_all_assistants()
            for assistant in assistants:
                self.client.beta.assistants.delete(assistant.id)
                logging.info(f"Deleted assistant: {assistant.id}")
        except Exception as e:
            logging.error(f"Failed to delete assistants: {e}")
            raise

    def delete_all_vector_stores(self):
        """
        Delete all vector stores associated with the OpenAI account.
        """
        try:
            vector_stores = self.get_all_vector_stores()
            for vector_store in vector_stores:
                self.client.beta.vector_stores.delete(vector_store.id)
                logging.info(f"Deleted vector store: {vector_store.id}")
        except Exception as e:
            logging.error(f"Failed to delete vector stores: {e}")
            raise

    def delete_all_files(self):
        """
        Delete all files associated with the OpenAI account.
        """
        try:
            files = self.get_all_files()
            for file in files:
                self.client.files.delete(file.id)
                logging.info(f"Deleted file: {file.id}")
        except Exception as e:
            logging.error(f"Failed to delete files: {e}")
            raise

    def delete_all_resources(self):
        """
        Delete all assistants, vector stores, and files associated with the OpenAI account.
        """
        self.delete_all_assistants()
        self.delete_all_vector_stores()
        self.delete_all_files()

    def wrap_text(self, s, width=30):
        """
        Wrap text to a specific width.

        :param s: The string to wrap.
        :param width: The maximum width of each line.
        :return: The wrapped string.
        """
        return "\n".join(textwrap.wrap(s, width))

    def show_all_assistants(self):
        """
        Display all assistants in a table.
        """
        assistants = self.get_all_assistants()
        assistant_data = [
            {
                k: self.wrap_text(str(v))
                for k, v in assistant.dict().items()
            }
            for assistant in assistants
        ]
        df = pd.DataFrame(assistant_data)
        print("Assistants:")
        print(df.to_markdown(index=False))

    def show_all_vector_stores(self):
        """
        Display all vector stores in a table.
        """
        vector_stores = self.get_all_vector_stores()
        vector_store_data = [
            {k: self.wrap_text(str(v)) for k, v in vector_store.dict().items()}
            for vector_store in vector_stores
        ]
        df = pd.DataFrame(vector_store_data)
        print("Vector Stores:")
        print(df.to_markdown(index=False))

    def show_all_files(self):
        """
        Display all files in a table.
        """
        files = self.get_all_files()
        file_data = [
            {k: self.wrap_text(str(v)) for k, v in file.dict().items()}
            for file in files
        ]
        df = pd.DataFrame(file_data)
        print("Files:")
        print(df.to_markdown(index=False))

    def show_all_resources(self):
        """
        Display all resources in a table.
        """
        self.show_all_assistants()
        print("")
        self.show_all_vector_stores()
        print("")
        self.show_all_files()



def main():
    """
    Main function that either shows or deletes resources based on the command line arguments.
    """
    parser = argparse.ArgumentParser(
        description="Displays or deletes all resources associated with the OpenAI account."
    )
    parser.add_argument(
        "--delete", action="store_true", help="Flag to delete resources instead of showing them."
    )
    args = parser.parse_args()

    # Get the OpenAI API key from the environment variables
    api_key = os.getenv("OPENAI_API_KEY")
    # Create an instance of the OpenAIResourceManager
    manager = OpenAIResourceManager(api_key=api_key)

    # Delete resources conditionally
    if args.delete:
        manager.delete_all_resources()

    manager.show_all_resources()

if __name__ == "__main__":
    main()