add-assistant #2
|
@ -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()
|
||||
|
|
Loading…
Reference in New Issue