Update final.py

final.py

@@ -11,6 +11,7 @@ import sys
 import os
 import asyncio
 import shutil
+import tempfile
 from readfile import prepare_files, USER_FILES_DIR
 from typing import List, Dict, Tuple
 from dotenv import load_dotenv
@@ -91,12 +92,22 @@ async def start():
             content="Please upload a zip file to begin!", accept={"application/zip": [".zip"]}
         ).send()
 
-    zip_file = files[0]  # Assuming only one file is uploaded
-    file_path = os.path.join(USER_FILES_DIR, zip_file.name)
+    zip_file: AskFileResponse = files[0]  # Assuming only one file is uploaded
     uploaded_file_name = zip_file.name
 
-    # Move the uploaded file to the user files directory
-    shutil.move(zip_file.path, file_path)
+    # Get the CHAINLIT_USER_FILES_DIR from environment variables
+    user_files_dir = os.environ.get('CHAINLIT_USER_FILES_DIR', '/tmp/chainlit_user_files')
+
+    # Create a temporary directory to store the uploaded file
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_file_path = os.path.join(temp_dir, zip_file.name)
+        
+        # Copy the uploaded file to the temporary directory
+        shutil.copy(zip_file.path, temp_file_path)
+        
+        # Move the file to the user files directory
+        os.makedirs(user_files_dir, exist_ok=True)
+        shutil.move(temp_file_path, os.path.join(user_files_dir, zip_file.name))
 
     # Let the user know that the system is ready
     await cl.Message(content=f"`{zip_file.name}` uploaded successfully!").send()
@@ -108,8 +119,10 @@ async def process_grading():
     global uploaded_file_name
     if uploaded_file_name:
         try:
+            user_files_dir = os.environ.get('CHAINLIT_USER_FILES_DIR', '/tmp/chainlit_user_files')
+            file_path = os.path.join(user_files_dir, uploaded_file_name)
             # Process the uploaded ZIP file
-            average_similarity, average_score, ref_total_tokens, student_total_tokens, llm_total_tokens = process_data(uploaded_file_name, llm_score_prompt_template)
+            average_similarity, average_score, ref_total_tokens, student_total_tokens, llm_total_tokens = process_data(file_path, llm_score_prompt_template)
 
             # Send results
             await cl.Message(content=f"Processing complete. Results:\n"
@@ -118,6 +131,9 @@ async def process_grading():
                                      f"Reference Total Tokens: {ref_total_tokens}\n"
                                      f"Student Total Tokens: {student_total_tokens}\n"
                                      f"LLM Total Tokens: {llm_total_tokens}").send()
+            
+            # Remove the file after processing
+            os.remove(file_path)
         except Exception as e:
             await cl.Message(content=f"An error occurred while processing the zip file: {str(e)}").send()
     else: