Update main.py

main.py

@@ -1,14 +1,18 @@
-from fastapi import FastAPI, File, UploadFile, Form,  HTTPException
+from fastapi import FastAPI, File, UploadFile, Form, HTTPException
 from fastapi.responses import PlainTextResponse, FileResponse, HTMLResponse
 from fastapi.requests import Request
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
 from typing import Optional
 from starlette.exceptions import HTTPException as StarletteHTTPException
+
+# Import your custom AI service
 from services.client import generate_code_comments_and_docs
 
 # ✅ Initialize FastAPI app
 app = FastAPI()
+
+# Add CORS middleware to allow cross-origin requests if needed
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -17,28 +21,33 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Load static data
+# ✅ Mount the static folder so FastAPI can serve HTML, CSS, and JS files
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
 
+# ✅ Serve the main frontend page
 @app.get("/", response_class=HTMLResponse)
 async def serve_index():
     return FileResponse("static/index.html")
 
-# ✅ API route
+
+# ✅ Main API route to handle code generation
 @app.post("/generate-docs")
 async def generate_docs(
     file: Optional[UploadFile] = File(None),
     code: Optional[str] = Form(None)
-    ):
+):
     try:
+        # Check if neither file nor code was provided
         if not file and not code:
             raise HTTPException(status_code=400, detail="You must provide either a file or raw code.")
 
+        # If raw text code is provided
         elif code:
             response = generate_code_comments_and_docs(code)
             return response
         
+        # If a file is uploaded
         elif file:
             content = await file.read()
             code = content.decode("utf-8")
@@ -51,8 +60,10 @@ async def generate_docs(
     except Exception as ex:
         return PlainTextResponse(f"❌ Server error: {str(ex)}", status_code=500)
 
+
+# ✅ Custom 404 Error Handler to serve your not-found.html
 @app.exception_handler(StarletteHTTPException)
 async def custom_404_handler(request: Request, exc: StarletteHTTPException):
     if exc.status_code == 404:
-        return FileResponse("static/not-found.html", media_type="text/html")
-    return HTMLResponse(content=str(exc.detail), status_code=exc.status_code)
+        return FileResponse("static/not-found.html")
+    return PlainTextResponse(str(exc.detail), status_code=exc.status_code)