Update main.py

main.py

@@ -1,31 +1,44 @@
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
-from pydantic import BaseModel
-from typing import Optional, Tuple, List, Dict
+from transformers import pipeline
+from typing import Tuple
 import io
+import fitz  # PyMuPDF
+from PIL import Image
 import pandas as pd
-import matplotlib.pyplot as plt
-import seaborn as sns
-import base64
-import tempfile
-import logging
-import traceback
 import uvicorn
+from docx import Document
+from pptx import Presentation
+import pytesseract
+import logging
+import re
 from slowapi import Limiter
 from slowapi.util import get_remote_address
 from slowapi.errors import RateLimitExceeded
 from slowapi.middleware import SlowAPIMiddleware
+import matplotlib.pyplot as plt
+import seaborn as sns
+import tempfile
+import base64
+from io import BytesIO
+from typing import Optional
+from pydantic import BaseModel
 
-# Initialize FastAPI app
-app = FastAPI(title="Data Visualization API", version="1.0")
-
-# Rate limiting setup
+# Initialize rate limiter
 limiter = Limiter(key_func=get_remote_address)
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI()
+
+# Apply rate limiting middleware
 app.state.limiter = limiter
 app.add_middleware(SlowAPIMiddleware)
 
-# CORS configuration
+# CORS Configuration
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -35,123 +48,306 @@ app.add_middleware(
 
 # Constants
 MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB
-SUPPORTED_EXCEL_TYPES = {"xlsx", "xls"}
+SUPPORTED_FILE_TYPES = {
+    "docx", "xlsx", "pptx", "pdf", "jpg", "jpeg", "png"
+}
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# Model caching
+summarizer = None
+qa_model = None
+image_captioner = None
 
-class VisualizationRequest(BaseModel):
-    chart_type: str
-    x_column: str
-    y_column: Optional[str] = None
-    hue_column: Optional[str] = None
-    title: Optional[str] = None
-    x_label: Optional[str] = None
-    y_label: Optional[str] = None
-    style: str = "seaborn"
-    width: int = 10
-    height: int = 6
+def get_summarizer():
+    global summarizer
+    if summarizer is None:
+        summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
+    return summarizer
 
-async def validate_excel_file(file: UploadFile) -> Tuple[str, bytes]:
-    """Validate and process uploaded Excel file"""
+def get_qa_model():
+    global qa_model
+    if qa_model is None:
+        qa_model = pipeline("question-answering", model="deepset/roberta-base-squad2")
+    return qa_model
+
+def get_image_captioner():
+    global image_captioner
+    if image_captioner is None:
+        image_captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large")
+    return image_captioner
+
+async def process_uploaded_file(file: UploadFile) -> Tuple[str, bytes]:
+    """Validate and process uploaded file with special handling for each type"""
     if not file.filename:
         raise HTTPException(400, "No filename provided")
 
     file_ext = file.filename.split('.')[-1].lower()
-    if file_ext not in SUPPORTED_EXCEL_TYPES:
-        raise HTTPException(400, f"Unsupported file type. Supported: {', '.join(SUPPORTED_EXCEL_TYPES)}")
+    if file_ext not in SUPPORTED_FILE_TYPES:
+        raise HTTPException(400, f"Unsupported file type. Supported: {', '.join(SUPPORTED_FILE_TYPES)}")
 
     content = await file.read()
     if len(content) > MAX_FILE_SIZE:
         raise HTTPException(413, f"File too large. Max size: {MAX_FILE_SIZE//1024//1024}MB")
 
-    await file.seek(0)
+    # Special validation for PDFs
+    if file_ext == "pdf":
+        try:
+            with fitz.open(stream=content, filetype="pdf") as doc:
+                if doc.is_encrypted:
+                    if not doc.authenticate(""):
+                        raise ValueError("Encrypted PDF - cannot extract text")
+                if len(doc) > 50:
+                    raise ValueError("PDF too large (max 50 pages)")
+        except Exception as e:
+            logger.error(f"PDF validation failed: {str(e)}")
+            raise HTTPException(422, detail=f"Invalid PDF file: {str(e)}")
+
+    await file.seek(0)  # Reset file pointer for processing
     return file_ext, content
 
-def read_excel_with_fallback(content: bytes) -> pd.DataFrame:
-    """Read Excel file with engine fallback"""
+def extract_text(content: bytes, file_ext: str) -> str:
+    """Extract text from various file formats with enhanced support"""
     try:
-        return pd.read_excel(io.BytesIO(content), engine='openpyxl')
+        if file_ext == "docx":
+            doc = Document(io.BytesIO(content))
+            return "\n".join(para.text for para in doc.paragraphs if para.text.strip())
+        
+        elif file_ext in {"xlsx", "xls"}:
+            df = pd.read_excel(io.BytesIO(content), sheet_name=None)
+            all_text = []
+            for sheet_name, sheet_data in df.items():
+                sheet_text = []
+                for column in sheet_data.columns:
+                    sheet_text.extend(sheet_data[column].dropna().astype(str).tolist())
+                all_text.append(f"Sheet: {sheet_name}\n" + "\n".join(sheet_text))
+            return "\n\n".join(all_text)
+        
+        elif file_ext == "pptx":
+            ppt = Presentation(io.BytesIO(content))
+            text = []
+            for slide in ppt.slides:
+                for shape in slide.shapes:
+                    if hasattr(shape, "text") and shape.text.strip():
+                        text.append(shape.text)
+            return "\n".join(text)
+        
+        elif file_ext == "pdf":
+            pdf = fitz.open(stream=content, filetype="pdf")
+            return "\n".join(page.get_text("text") for page in pdf)
+        
+        elif file_ext in {"jpg", "jpeg", "png"}:
+            # First try OCR
+            try:
+                image = Image.open(io.BytesIO(content))
+                text = pytesseract.image_to_string(image, config='--psm 6')
+                if text.strip():
+                    return text
+                
+                # If OCR fails, try image captioning
+                captioner = get_image_captioner()
+                result = captioner(image)
+                return result[0]['generated_text']
+            except Exception as img_e:
+                logger.error(f"Image processing failed: {str(img_e)}")
+                raise ValueError("Could not extract text or caption from image")
+        
     except Exception as e:
-        logger.warning(f"Openpyxl failed, trying xlrd: {str(e)}")
-        try:
-            return pd.read_excel(io.BytesIO(content), engine='xlrd')
-        except Exception as e:
-            raise ValueError(f"Failed to read Excel file: {str(e)}")
+        logger.error(f"Text extraction failed for {file_ext}: {str(e)}")
+        raise HTTPException(422, f"Failed to extract text from {file_ext} file")
 
-def generate_chart(df: pd.DataFrame, request: VisualizationRequest) -> str:
-    """Generate matplotlib/seaborn chart based on request"""
-    plt.style.use(request.style)
-    plt.figure(figsize=(request.width, request.height))
-    
-    chart_type = request.chart_type.lower()
+@app.post("/summarize")
+@limiter.limit("5/minute")
+async def summarize_document(request: Request, file: UploadFile = File(...)):
     try:
-        if chart_type == "line":
-            if request.hue_column:
-                sns.lineplot(data=df, x=request.x_column, y=request.y_column, hue=request.hue_column)
-            else:
-                plt.plot(df[request.x_column], df[request.y_column])
-        elif chart_type == "bar":
-            if request.hue_column:
-                sns.barplot(data=df, x=request.x_column, y=request.y_column, hue=request.hue_column)
-            else:
-                plt.bar(df[request.x_column], df[request.y_column])
-        elif chart_type == "scatter":
-            if request.hue_column:
-                sns.scatterplot(data=df, x=request.x_column, y=request.y_column, hue=request.hue_column)
-            else:
-                plt.scatter(df[request.x_column], df[request.y_column])
-        elif chart_type == "histogram":
-            plt.hist(df[request.x_column], bins=20)
-        else:
-            raise ValueError(f"Unsupported chart type: {chart_type}")
-
-        if request.title:
-            plt.title(request.title)
-        if request.x_label:
-            plt.xlabel(request.x_label)
-        if request.y_label:
-            plt.ylabel(request.y_label)
+        file_ext, content = await process_uploaded_file(file)
+        text = extract_text(content, file_ext)
+        
+        if not text.strip():
+            raise HTTPException(400, "No extractable text found")
+        
+        # Clean and chunk text
+        text = re.sub(r'\s+', ' ', text).strip()
+        chunks = [text[i:i+1000] for i in range(0, len(text), 1000)]
+        
+        # Summarize each chunk
+        summarizer = get_summarizer()
+        summaries = []
+        for chunk in chunks:
+            summary = summarizer(chunk, max_length=150, min_length=50, do_sample=False)[0]["summary_text"]
+            summaries.append(summary)
+        
+        return {"summary": " ".join(summaries)}
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Summarization failed: {str(e)}")
+        raise HTTPException(500, "Document summarization failed")
 
-        plt.tight_layout()
+@app.post("/qa")
+@limiter.limit("5/minute")
+async def question_answering(
+    request: Request,
+    file: UploadFile = File(...),
+    question: str = Form(...),
+    language: str = Form("fr")
+):
+    try:
+        file_ext, content = await process_uploaded_file(file)
+        text = extract_text(content, file_ext)
         
-        # Save to temporary file
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
-            plt.savefig(tmpfile.name, format='png', dpi=300)
-            plt.close()
-            with open(tmpfile.name, "rb") as f:
-                return base64.b64encode(f.read()).decode('utf-8')
+        if not text.strip():
+            raise HTTPException(400, "No extractable text found")
+
+        # Clean and truncate text
+        text = re.sub(r'\s+', ' ', text).strip()[:5000]
+
+        # Theme detection
+        theme_keywords = ["thème", "sujet principal", "quoi le sujet", "theme", "main topic"]
+        if any(kw in question.lower() for kw in theme_keywords):
+            try:
+                summarizer = get_summarizer()
+                summary_output = summarizer(
+                    text,
+                    max_length=min(100, len(text)//4),
+                    min_length=30,
+                    do_sample=False,
+                    truncation=True
+                )
                 
+                theme = summary_output[0].get("summary_text", text[:200] + "...")
+                return {
+                    "question": question,
+                    "answer": f"Le document traite principalement de : {theme}",
+                    "confidence": 0.95,
+                    "language": language
+                }
+            except Exception:
+                theme = text[:200] + ("..." if len(text) > 200 else "")
+                return {
+                    "question": question,
+                    "answer": f"D'après le document : {theme}",
+                    "confidence": 0.7,
+                    "language": language,
+                    "warning": "theme_summary_fallback"
+                }
+
+        # Standard QA
+        qa = get_qa_model()
+        result = qa(question=question, context=text[:3000])
+        
+        return {
+            "question": question,
+            "answer": result["answer"],
+            "confidence": result["score"],
+            "language": language
+        }
+
+    except HTTPException:
+        raise
     except Exception as e:
-        plt.close()
-        raise ValueError(f"Chart generation failed: {str(e)}")
+        logger.error(f"QA processing failed: {str(e)}")
+        raise HTTPException(500, detail=f"Analysis failed: {str(e)}")
+
+@app.exception_handler(RateLimitExceeded)
+async def rate_limit_exceeded_handler(request: Request, exc: RateLimitExceeded):
+    return JSONResponse(
+        status_code=429,
+        content={"detail": "Too many requests. Please try again later."}
+    )
+
+
+
+
+
+    # Add this new Pydantic model for visualization requests
+class VisualizationRequest(BaseModel):
+    chart_type: str
+    x_column: Optional[str] = None
+    y_column: Optional[str] = None
+    hue_column: Optional[str] = None
+    title: Optional[str] = None
+    x_label: Optional[str] = None
+    y_label: Optional[str] = None
+    style: str = "seaborn"  # seaborn or matplotlib
 
-@app.post("/visualize", response_model=Dict[str, str])
+# Add this new function for visualization code generation
+def generate_visualization(df: pd.DataFrame, request: VisualizationRequest) -> str:
+    """Generate and execute visualization code based on request"""
+    plt.style.use(request.style)
+    
+    code_lines = [
+        "import matplotlib.pyplot as plt",
+        "import seaborn as sns",
+        "import pandas as pd",
+        "",
+        "# Data preparation",
+        f"df = pd.DataFrame({df.head().to_dict()})",  # Simplified for demo
+        "",
+        "# Visualization code"
+    ]
+    
+    if request.chart_type == "line":
+        code_lines.append(f"plt.figure(figsize=(10, 6))")
+        if request.hue_column:
+            code_lines.append(f"sns.lineplot(data=df, x='{request.x_column}', y='{request.y_column}', hue='{request.hue_column}')")
+        else:
+            code_lines.append(f"plt.plot(df['{request.x_column}'], df['{request.y_column}'])")
+    elif request.chart_type == "bar":
+        code_lines.append(f"plt.figure(figsize=(10, 6))")
+        if request.hue_column:
+            code_lines.append(f"sns.barplot(data=df, x='{request.x_column}', y='{request.y_column}', hue='{request.hue_column}')")
+        else:
+            code_lines.append(f"plt.bar(df['{request.x_column}'], df['{request.y_column}'])")
+    elif request.chart_type == "scatter":
+        code_lines.append(f"plt.figure(figsize=(10, 6))")
+        if request.hue_column:
+            code_lines.append(f"sns.scatterplot(data=df, x='{request.x_column}', y='{request.y_column}', hue='{request.hue_column}')")
+        else:
+            code_lines.append(f"plt.scatter(df['{request.x_column}'], df['{request.y_column}'])")
+    elif request.chart_type == "histogram":
+        code_lines.append(f"plt.figure(figsize=(10, 6))")
+        code_lines.append(f"plt.hist(df['{request.x_column}'], bins=20)")
+    else:
+        raise ValueError("Unsupported chart type")
+
+    # Add labels and title
+    if request.title:
+        code_lines.append(f"plt.title('{request.title}')")
+    if request.x_label:
+        code_lines.append(f"plt.xlabel('{request.x_label}')")
+    if request.y_label:
+        code_lines.append(f"plt.ylabel('{request.y_label}')")
+    
+    code_lines.append("plt.tight_layout()")
+    code_lines.append("plt.show()")
+    
+    return "\n".join(code_lines)
+
+# Add this new endpoint for visualization
+@app.post("/visualize")
 @limiter.limit("5/minute")
-async def create_visualization(
+async def generate_visualization_from_excel(
     request: Request,
     file: UploadFile = File(...),
     chart_type: str = Form(...),
-    x_column: str = Form(...),
+    x_column: Optional[str] = Form(None),
     y_column: Optional[str] = Form(None),
     hue_column: Optional[str] = Form(None),
     title: Optional[str] = Form(None),
     x_label: Optional[str] = Form(None),
     y_label: Optional[str] = Form(None),
-    style: str = Form("seaborn"),
-    width: int = Form(10),
-    height: int = Form(6)
+    style: str = Form("seaborn")
 ):
     try:
-        # Validate and read file
-        file_ext, content = await validate_excel_file(file)
-        df = read_excel_with_fallback(content)
+        # Validate file
+        file_ext, content = await validate_file(file)
+        if file_ext not in {"xlsx", "xls"}:
+            raise HTTPException(400, "Only Excel files are supported for visualization")
         
-        if df.empty:
-            raise ValueError("Excel file contains no data")
-
-        # Create visualization
+        # Read Excel file
+        df = pd.read_excel(io.BytesIO(content))
+        
+        # Generate visualization request
         vis_request = VisualizationRequest(
             chart_type=chart_type,
             x_column=x_column,
@@ -160,60 +356,60 @@ async def create_visualization(
             title=title,
             x_label=x_label,
             y_label=y_label,
-            style=style,
-            width=width,
-            height=height
+            style=style
         )
         
-        image_base64 = generate_chart(df, vis_request)
+        # Generate and execute the visualization code
+        plt.figure()
+        exec(generate_visualization(df, vis_request), globals(), locals())
+        
+        # Save the plot to a temporary file
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
+            plt.savefig(tmpfile.name, format='png', dpi=300)
+            plt.close()
+            
+            # Read the image back as bytes
+            with open(tmpfile.name, "rb") as f:
+                image_bytes = f.read()
+        
+        # Encode image as base64
+        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
         
         return {
             "status": "success",
-            "image": f"data:image/png;base64,{image_base64}"
+            "image": f"data:image/png;base64,{image_base64}",
+            "code": generate_visualization(df, vis_request)
         }
-
-    except HTTPException as he:
-        raise he
-    except ValueError as ve:
-        logger.error(f"Validation error: {str(ve)}")
-        raise HTTPException(422, detail=str(ve))
+        
+    except HTTPException:
+        raise
     except Exception as e:
-        logger.error(f"Unexpected error: {str(e)}\n{traceback.format_exc()}")
-        raise HTTPException(500, detail="Failed to generate visualization")
+        logger.error(f"Visualization failed: {str(e)}\n{traceback.format_exc()}")
+        raise HTTPException(500, detail=f"Visualization failed: {str(e)}")
 
-@app.post("/get_columns", response_model=Dict[str, List])
+# Add this new endpoint for getting column names
+@app.post("/get_columns")
 @limiter.limit("10/minute")
-async def get_columns(
+async def get_excel_columns(
     request: Request,
     file: UploadFile = File(...)
 ):
     try:
-        file_ext, content = await validate_excel_file(file)
-        df = read_excel_with_fallback(content)
+        file_ext, content = await validate_file(file)
+        if file_ext not in {"xlsx", "xls"}:
+            raise HTTPException(400, "Only Excel files are supported")
         
-        if df.empty:
-            raise ValueError("Excel file contains no data")
-
+        df = pd.read_excel(io.BytesIO(content))
         return {
             "columns": list(df.columns),
-            "sample_data": df.head().replace({float('nan'): None}).to_dict(orient='records')
+            "sample_data": df.head().to_dict(orient='records')
         }
-
-    except HTTPException as he:
-        raise he
-    except ValueError as ve:
-        logger.error(f"Validation error: {str(ve)}")
-        raise HTTPException(422, detail=str(ve))
     except Exception as e:
-        logger.error(f"Unexpected error: {str(e)}\n{traceback.format_exc()}")
-        raise HTTPException(500, detail="Failed to process Excel file")
+        logger.error(f"Column extraction failed: {str(e)}")
+        raise HTTPException(500, detail="Failed to extract columns from Excel file")
 
-@app.exception_handler(RateLimitExceeded)
-async def rate_limit_handler(request: Request, exc: RateLimitExceeded):
-    return JSONResponse(
-        status_code=429,
-        content={"detail": "Too many requests. Please try again later."}
-    )
+
+        
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)