Refactor file upload handling to support single file processing and improve error handling

main.py

@@ -11,10 +11,14 @@ from fastapi.middleware.cors import CORSMiddleware
 app = FastAPI()
 logger = setup_logging()
 
+allow_origins = [
+    "https://*.hf.space",
+    "http://localhost:7860"
+]
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=allow_origins,
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -94,24 +98,40 @@ def extract_address_data(file_stream):
 
 @app.post("/upload")
 async def upload_files(files: list[UploadFile] = File(...)):
-    """Handle multiple file uploads and return processed ZIP"""
+    """Handle file uploads and return processed file(s)"""
     try:
-        zip_buffer = BytesIO()
-        with zipfile.ZipFile(zip_buffer, 'a', zipfile.ZIP_DEFLATED) as zip_file:
-            for file in files:
-                if not allowed_file(file.filename):
-                    continue
-                
-                file_stream = await file.read()
-                output = process_uploaded_file(BytesIO(file_stream))
-                zip_file.writestr(f"processed_{file.filename}", output.getvalue())
-        
-        zip_buffer.seek(0)
-        return StreamingResponse(
-            zip_buffer,
-            media_type="application/zip",
-            headers={"Content-Disposition": "attachment; filename=processed_files.zip"}
-        )
+        if len(files) == 1:
+            # Handle single file
+            file = files[0]
+            if not allowed_file(file.filename):
+                raise HTTPException(status_code=400, detail="Invalid file type")
+            
+            file_stream = await file.read()
+            output = process_uploaded_file(BytesIO(file_stream))
+            
+            return StreamingResponse(
+                output,
+                media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+                headers={"Content-Disposition": f"attachment; filename=processed_{file.filename}"}
+            )
+        else:
+            # Handle multiple files as ZIP
+            zip_buffer = BytesIO()
+            with zipfile.ZipFile(zip_buffer, 'a', zipfile.ZIP_DEFLATED) as zip_file:
+                for file in files:
+                    if not allowed_file(file.filename):
+                        continue
+                    
+                    file_stream = await file.read()
+                    output = process_uploaded_file(BytesIO(file_stream))
+                    zip_file.writestr(f"processed_{file.filename}", output.getvalue())
+            
+            zip_buffer.seek(0)
+            return StreamingResponse(
+                zip_buffer,
+                media_type="application/zip",
+                headers={"Content-Disposition": "attachment; filename=processed_files.zip"}
+            )
         
     except Exception as e:
         logger.error(f"Processing error: {str(e)}")