assistant/ConfigureAssistant.py

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