Update app.py

app.py

@@ -1,13 +1,17 @@
-import os
-from fastapi import FastAPI, UploadFile, File, HTTPException, Form
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTasks, WebSocket, Request
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import RedirectResponse
+from fastapi_cache import FastAPICache
+from fastapi_cache.backends.redis import RedisBackend
+from fastapi_cache.decorator import cache
 from transformers import pipeline
 import logging
 from PIL import Image
 import io
 from docx import Document
 import fitz  # PyMuPDF
+from pydantic import BaseModel
+import asyncio
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -18,104 +22,47 @@ app = FastAPI()
 # Serve static files (HTML, CSS, JS)
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-# Load a multimodal model for image captioning and visual question answering
+# Load AI models
 multimodal_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-
-# Load a text-based model for summarization and text question answering
 text_pipeline = pipeline("text2text-generation", model="t5-small")
 
+# Initialize Redis cache
+@app.on_event("startup")
+async def startup():
+    FastAPICache.init(RedisBackend("redis://localhost:6379"))
+
+# Root endpoint
 @app.get("/")
 def read_root():
-    # Redirect to the static HTML file
     return RedirectResponse(url="/static/index.html")
 
+# Summarization endpoint
 @app.post("/summarize")
-async def summarize_text(
-    file: UploadFile = File(None),  # Optional file upload
-    text: str = Form(None)  # Optional manual text input
-):
+@cache(expire=300)
+async def summarize_text(file: UploadFile = File(None), text: str = Form(None)):
     if file:
-        logger.info(f"Received document for summarization: {file.filename}")
-        try:
-            # Extract text from the document
-            text = await extract_text_from_file(file)
-        except Exception as e:
-            logger.error(f"Error extracting text from file: {e}")
-            raise HTTPException(status_code=400, detail=str(e))
-    elif text:
-        logger.info("Received manual text for summarization")
-    else:
+        if not file.filename.endswith((".pdf", ".docx")):
+            raise HTTPException(status_code=400, detail="Unsupported file format. Please upload a PDF or DOCX file.")
+        text = await extract_text_from_file(file)
+    elif not text:
         raise HTTPException(status_code=400, detail="No file or text provided")
 
     try:
-        # Use the text pipeline to summarize the text
         summary = text_pipeline(f"summarize: {text}", max_length=100)
-        logger.info(f"Generated summary: {summary[0]['generated_text']}")
-        return {"summary": summary[0]['generated_text']}
+        logger.info(f"Generated summary: {summary[0]['summary_text']}")
+        return {"summary": summary[0]['summary_text']}
     except Exception as e:
         logger.error(f"Error during summarization: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/caption")
-async def caption_image(file: UploadFile = File(...)):
-    logger.info(f"Received image for captioning: {file.filename}")
-    try:
-        # Read the image file
-        image_data = await file.read()
-        image = Image.open(io.BytesIO(image_data))
-
-        # Use the multimodal pipeline to generate a caption for the image
-        caption = multimodal_pipeline(image)
-        logger.info(f"Generated caption: {caption[0]['generated_text']}")
-        return {"caption": caption[0]['generated_text']}
-    except Exception as e:
-        logger.error(f"Error during image captioning: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/answer")
-async def answer_question(
-    file: UploadFile = File(None),  # Optional file upload
-    text: str = Form(None),  # Optional manual text input
-    question: str = Form(...)  # Required question
-):
-    if file:
-        logger.info(f"Received document for question answering: {file.filename}")
-        try:
-            # Extract text from the document
-            text = await extract_text_from_file(file)
-        except Exception as e:
-            logger.error(f"Error extracting text from file: {e}")
-            raise HTTPException(status_code=400, detail=str(e))
-    elif text:
-        logger.info("Received manual text for question answering")
-    else:
-        raise HTTPException(status_code=400, detail="No file or text provided")
-
-    try:
-        # Use the text pipeline to answer the question
-        answer = text_pipeline(f"question: {question} context: {text}")
-        logger.info(f"Generated answer: {answer[0]['generated_text']}")
-        return {"answer": answer[0]['generated_text']}
-    except Exception as e:
-        logger.error(f"Error during question answering: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/vqa")
-async def visual_question_answering(file: UploadFile = File(...), question: str = Form(...)):
-    logger.info(f"Received image for visual question answering: {file.filename}")
-    logger.info(f"Received question: {question}")
-    try:
-        # Read the image file
-        image_data = await file.read()
-        image = Image.open(io.BytesIO(image_data))
-
-        # Use the multimodal pipeline to answer the question about the image
-        answer = multimodal_pipeline(image, question=question)
-        logger.info(f"Generated answer: {answer[0]['generated_text']}")
-        return {"answer": answer[0]['generated_text']}
-    except Exception as e:
-        logger.error(f"Error during visual question answering: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail="An error occurred while processing your request. Please try again.")
+
+# WebSocket for real-time updates
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    while True:
+        data = await websocket.receive_text()
+        summary = text_pipeline(f"summarize: {data}", max_length=100)
+        await websocket.send_text(summary[0]['summary_text'])
 
 # Helper function to extract text from files
 async def extract_text_from_file(file: UploadFile):