add-assistant #2

Merged
Ari merged 8 commits from add-assistant into main 2024-08-06 01:14:25 +00:00
3 changed files with 102 additions and 24 deletions
Showing only changes of commit e1220e7165 - Show all commits

View File

@ -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):
"""

View File

@ -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()

View File

@ -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.")