Update Summarization/main.py

Summarization/main.py

@@ -43,21 +43,74 @@ async def caption(file: UploadFile = File(...)):
     except Exception as e:
         return JSONResponse({"error": str(e)}, status_code=500)
 
+from fastapi import FastAPI, UploadFile, File, Form, Request
+from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
+from fastapi.templating import Jinja2Templates
+from fastapi.middleware.cors import CORSMiddleware
+import tempfile
+import os
+from pathlib import Path
+import mimetypes  # Add this import
+
+from .app import summarize_document
+from .appImage import caption_image
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+current_dir = Path(__file__).parent
+templates = Jinja2Templates(directory=current_dir.parent / "templates")
+
+@app.get("/", response_class=HTMLResponse)
+async def serve_home(request: Request):
+    return templates.TemplateResponse("HomeS.html", {"request": request})
+
+@app.post("/summarize/")
+async def summarize_api(file: UploadFile = File(...), length: str = Form("medium")):
+    try:
+        result = await summarize_document(file, length)
+        return JSONResponse(result)
+    except Exception as e:
+        return JSONResponse({"detail": str(e)}, status_code=500)
+
+@app.post("/imagecaption/")
+async def caption(file: UploadFile = File(...)):
+    try:
+        result = await caption_image(file)
+        return JSONResponse(result)
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
+
 @app.get("/files/{filename}")
-async def serve_file(filename: str):
-    filepath = os.path.join(tempfile.gettempdir(), filename)
-    if filename.endswith('.mp3'):
-        return FileResponse(filepath, media_type="audio/mpeg")
 async def serve_file(filename: str):
     filepath = os.path.join(tempfile.gettempdir(), filename)
     
     if not os.path.exists(filepath):
         return JSONResponse({"error": "File not found"}, status_code=404)
     
-    # Determine media type for proper browser handling
+    # Force download for PDF files
+    if filename.endswith('.pdf'):
+        return FileResponse(
+            filepath,
+            media_type="application/pdf",
+            headers={"Content-Disposition": f"attachment; filename={filename}"}
+        )
+    
+    # Handle audio files
+    if filename.endswith('.mp3'):
+        return FileResponse(filepath, media_type="audio/mpeg")
+    
+    # Default handling for other files
     media_type, _ = mimetypes.guess_type(filename)
     return FileResponse(
         filepath,
         media_type=media_type or "application/octet-stream",
-        filename=filename 
+        filename=filename
     )