add-assistant #2

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

View File

@ -1,13 +1,19 @@
import os
import io
from openai import OpenAI
class ConfigureAssistant:
"""
A class to configure an OpenAI assistant for aiding designers using the ArchiMajor project.
This class handles the creation of the assistant, the setup of vector stores, and the upload of files for context.
"""
SUPPORTED_FORMATS = {
"c", "cpp", "css", "csv", "docx", "gif", "html", "java", "jpeg", "jpg", "js",
"json", "md", "pdf", "php", "png", "pptx", "py", "rb", "tar", "tex", "ts", "txt",
"webp", "xlsx", "xml", "zip"
}
def __init__(self, root_path: str):
"""
Initialize the ConfigureAssistant with the given root path and OpenAI client.
@ -23,7 +29,7 @@ class ConfigureAssistant:
"""
Create an OpenAI assistant with specific instructions and tools.
"""
AI_ASSISTANT_INSTRUCTIONS = """You are an assistant to aid designers who are either using the ArchiMajor board or designing a new board based off this one. They may ask questions about specific components or interconnects within this project.
AI_ASSISTANT_INSTRUCTIONS = """You are an assistant to aid designers who are either using the ArchiMajor board or designing a new board based off this one. They may ask questions about specific components or interconnects within this project.
You have the following reference documents
1. Schematic (PDF)
@ -58,6 +64,23 @@ class ConfigureAssistant:
file_paths.append(os.path.join(root, file))
return file_paths
def preprocess_file(self, file_path):
"""
Preprocess files based on their type.
:param file_path: The file path to preprocess.
: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:
return os.path.basename(file_path), True
elif extension == "schdoc":
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):
"""
Create a vector store for storing design documents.
@ -74,15 +97,21 @@ class ConfigureAssistant:
:param file_paths: A list of file paths to upload.
"""
for path in file_paths:
new_filename, should_upload = self.preprocess_file(path)
if not should_upload:
continue
try:
with open(path, "rb") as file_stream:
file_content = file_stream.read()
spoofed_file = io.BytesIO(file_content)
spoofed_file.name = new_filename # Spoof the filename
# Upload the file to the vector store
file_batch = self.client.beta.vector_stores.file_batches.upload_and_poll(
vector_store_id=self.vector_store.id, files=[file_stream]
vector_store_id=self.vector_store.id, files=[spoofed_file]
)
print(f"Successfully uploaded: {path}")
print(f"Successfully uploaded: {new_filename}")
except Exception as e:
print(f"Failed to upload: {path} with error: {e}")
print(f"Failed to upload: {new_filename} with error: {e}")
def update_assistant_with_vector_store(self):
"""
@ -113,5 +142,16 @@ if __name__ == "__main__":
root_path = os.path.dirname(os.path.dirname(__file__))
# Create an instance of ConfigureAssistant with the root path
configurator = ConfigureAssistant(root_path=root_path)
# Retrieve file paths
file_paths = configurator.get_file_paths(excluded_folders=[".git"])
# Preprocess and observe the files that will be uploaded
print("Files to be uploaded:")
for path in file_paths:
new_filename, should_upload = configurator.preprocess_file(path)
if should_upload:
print(f"Original: {path}, New: {new_filename}")
# Configure the assistant
configurator.configure()