add-assistant #2
|
@ -9,9 +9,10 @@ class ConfigureAssistant:
|
|||
"""
|
||||
|
||||
SUPPORTED_FORMATS = {
|
||||
"c", "cpp", "css", "csv", "docx", "gif", "html", "java", "jpeg", "jpg", "js",
|
||||
"c", "cpp", "css", "docx", "gif", "html", "java", "jpeg", "jpg", "js",
|
||||
"json", "md", "pdf", "php", "png", "pptx", "py", "rb", "tar", "tex", "ts", "txt",
|
||||
"webp", "xlsx", "xml", "zip"
|
||||
"webp", "xlsx", "xml", "zip",
|
||||
# "csv", # CSV is supported but not actually parsed so we're going to treat it as text
|
||||
}
|
||||
|
||||
def __init__(self, root_path: str):
|
||||
|
@ -72,14 +73,15 @@ class ConfigureAssistant:
|
|||
:return: The new filename and a flag indicating if the file should be uploaded.
|
||||
"""
|
||||
extension = file_path.split(".")[-1].lower()
|
||||
if extension in self.SUPPORTED_FORMATS:
|
||||
|
||||
# TO DO: Preprocess Outjob and PcbDoc files into something OpenAI (or future vector DB) can understand
|
||||
excluded_extensions = ["schdoc", "exe", "so", "dll", "outjob", "pcbdoc", "png"]
|
||||
|
||||
if extension in self.SUPPORTED_FORMATS and extension not in excluded_extensions:
|
||||
return os.path.basename(file_path), True
|
||||
elif extension == "schdoc":
|
||||
else:
|
||||
print(f"Skipping unsupported file: {file_path}")
|
||||
return file_path, False
|
||||
else:
|
||||
new_filename = os.path.basename(file_path) + ".txt"
|
||||
return new_filename, True
|
||||
|
||||
def create_vector_store(self):
|
||||
"""
|
||||
|
|
|
@ -1,12 +1,12 @@
|
|||
import os
|
||||
import logging
|
||||
import pandas as pd
|
||||
from openai import OpenAI
|
||||
|
||||
|
||||
class OpenAIResourceManager:
|
||||
"""
|
||||
A class to manage OpenAI resources such as assistants, vector stores, and files.
|
||||
Provides methods to delete all resources of each type.
|
||||
Provides methods to delete all resources of each type and display them in tables.
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: str):
|
||||
|
@ -18,15 +18,52 @@ class OpenAIResourceManager:
|
|||
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:
|
||||
# Retrieve the list of all assistants
|
||||
assistants = self.client.beta.assistants.list()
|
||||
# Loop through each assistant and delete it
|
||||
for assistant in assistants.data:
|
||||
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:
|
||||
|
@ -38,10 +75,8 @@ class OpenAIResourceManager:
|
|||
Delete all vector stores associated with the OpenAI account.
|
||||
"""
|
||||
try:
|
||||
# Retrieve the list of all vector stores
|
||||
vector_stores = self.client.beta.vector_stores.list()
|
||||
# Loop through each vector store and delete it
|
||||
for vector_store in vector_stores.data:
|
||||
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:
|
||||
|
@ -53,10 +88,8 @@ class OpenAIResourceManager:
|
|||
Delete all files associated with the OpenAI account.
|
||||
"""
|
||||
try:
|
||||
# Retrieve the list of all files
|
||||
files = self.client.files.list()
|
||||
# Loop through each file and delete it
|
||||
for file in files.data:
|
||||
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:
|
||||
|
@ -71,11 +104,54 @@ class OpenAIResourceManager:
|
|||
self.delete_all_vector_stores()
|
||||
self.delete_all_files()
|
||||
|
||||
def truncate_string(self, s, max_length=50):
|
||||
"""
|
||||
Truncate a string to a maximum length with ellipsis.
|
||||
|
||||
:param s: The string to truncate.
|
||||
:param max_length: The maximum length of the string.
|
||||
:return: The truncated string.
|
||||
"""
|
||||
return (s[:max_length] + '...') if len(s) > max_length else s
|
||||
|
||||
def show_all_assistants(self):
|
||||
"""
|
||||
Display all assistants in a table.
|
||||
"""
|
||||
assistants = self.get_all_assistants()
|
||||
assistant_data = [{k: self.truncate_string(str(v), max_length=25) 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.truncate_string(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.truncate_string(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))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 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 all resources
|
||||
manager.delete_all_resources()
|
||||
|
||||
# Show all resources in tables
|
||||
manager.show_all_assistants()
|
||||
manager.show_all_vector_stores()
|
||||
manager.show_all_files()
|
||||
|
|
|
@ -172,7 +172,7 @@ def main(query: str, assistant_id: str, context: str, use_streaming: bool):
|
|||
if __name__ == "__main__":
|
||||
# Default query and context
|
||||
DEFAULT_QUERY = "What are you capable of as an assistant?"
|
||||
DEFAULT_CONTEXT = "Use your vector store to answer questions about the Arty A7 Evaluation Board. Take time to understand the context and introspect yourself. If you don't know the answer simply respond with 'I don't know'. It is NEVER okay to return an empty response."
|
||||
DEFAULT_CONTEXT = "Use your vector store to answer questions about the ArchiMajor Board. Take time to understand the context and introspect. If you don't know the answer simply respond with 'I don't know'. It is NEVER okay to return an empty response."
|
||||
|
||||
# Parse command line arguments
|
||||
parser = argparse.ArgumentParser(description="Run an assistant query.")
|
||||
|
|
Loading…
Reference in New Issue