import io import datetime from typing import Optional, List from contextlib import asynccontextmanager from fastapi import FastAPI, File, UploadFile, HTTPException, Depends from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from pypdf import PdfReader import pytesseract from PIL import Image from pdf2image import convert_from_bytes import pdfplumber from sqlalchemy.orm import Session # 2. OCR Module def extract_text_from_pdf(file_bytes: bytes) -> str: # Try pdfplumber first for layout preservation try: with pdfplumber.open(io.BytesIO(file_bytes)) as pdf: text = "" for page in pdf.pages: page_text = page.extract_text(layout=True) if page_text: text += page_text + "\n" if text.strip(): return text.strip() except Exception as e: print(f"pdfplumber error: {e}") # Fallback to pypdf try: reader = PdfReader(io.BytesIO(file_bytes)) text = "" for page in reader.pages: page_text = page.extract_text() if page_text: text += page_text + "\n" return text.strip() except Exception as e: print(f"Error reading PDF: {e}") return "" # Internal modules from database import get_db, Email, Vendor, Document from scheduler import start_scheduler, stop_scheduler from mail_service import fetch_and_store_emails # Lifespan for Scheduler @asynccontextmanager async def lifespan(app: FastAPI): # Startup start_scheduler() yield # Shutdown stop_scheduler() app = FastAPI(lifespan=lifespan) # CORS configuration origins = [ "http://localhost", "http://localhost:3000", "http://localhost:5173", # Vite default "http://localhost:4200", # Angular default "http://127.0.0.1:4200", ] app.add_middleware( CORSMiddleware, allow_origins=origins, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Models class LoginRequest(BaseModel): username: str password: str class LoginResponse(BaseModel): token: str message: str class NERResponse(BaseModel): text: str file_path: str def extract_text_from_image(file_bytes: bytes) -> str: try: image = Image.open(io.BytesIO(file_bytes)) text = pytesseract.image_to_string(image) return text.strip() except Exception as e: print(f"Error reading Image: {e}") return "" @app.post("/api/ocr/extract", response_model=NERResponse) async def extract_text(file: UploadFile = File(...)): content = await file.read() filename = file.filename.lower() # Save file for Vision mode file_path = f"uploads/{file.filename}" with open(file_path, "wb") as f: f.write(content) extracted_text = "" if filename.endswith(".pdf"): # Try text extraction first with pdfplumber.open(io.BytesIO(content)) as pdf: try: text = "" for page in pdf.pages: page_text = page.extract_text(layout=True) if page_text: text += page_text + "\n" if text.strip(): extracted_text = text.strip() except Exception: pass if not extracted_text: try: # Fallback to pypdf reader = PdfReader(io.BytesIO(content)) text = "" for page in reader.pages: page_text = page.extract_text() if page_text: text += page_text + "\n" extracted_text = text.strip() except: pass # If text is still empty, it might be a scanned PDF. if not extracted_text.strip(): try: images = convert_from_bytes(content) for i, image in enumerate(images): page_text = pytesseract.image_to_string(image) extracted_text += f"\n--- Page {i+1} ---\n{page_text}" except Exception as e: extracted_text = f"Error processing Scanned PDF: {str(e)}\n\n(Hint: Ensure 'poppler' is installed on your system. Run 'brew install poppler')" elif filename.endswith((".png", ".jpg", ".jpeg", ".tiff", ".bmp")): extracted_text = extract_text_from_image(content) else: raise HTTPException(status_code=400, detail="Unsupported file type") return NERResponse(text=extracted_text, file_path=file_path) # 3. AI Extraction Module from llm_service import extract_data from pdf2image import convert_from_path class AITextRequest(BaseModel): text: Optional[str] = None file_path: Optional[str] = None model_type: str = "text" @app.post("/api/extract/ai") def extract_ai_data(request: AITextRequest): final_image_path = request.file_path if request.model_type == "vision" and request.file_path and request.file_path.endswith(".pdf"): # Convert PDF first page to image try: images = convert_from_path(request.file_path) if images: # Save temp image temp_img_path = request.file_path + ".jpg" images[0].save(temp_img_path, "JPEG") final_image_path = temp_img_path except Exception as e: print(f"Error converting PDF for vision: {e}") data = extract_data(text=request.text, image_path=final_image_path, model_type=request.model_type) return data import zipfile import mimetypes from fastapi.responses import StreamingResponse class EmailDTO(BaseModel): id: int subject: Optional[str] sender: Optional[str] received_date: Optional[datetime.datetime] is_read: bool has_attachments: bool = False class AttachmentDTO(BaseModel): id: int filename: str content_type: str class EmailListResponse(BaseModel): items: List[EmailDTO] total: int class EmailDetailDTO(EmailDTO): body: Optional[str] attachments: List[AttachmentDTO] @app.get("/api/emails", response_model=EmailListResponse) def get_emails(skip: int = 0, limit: int = 50, db: Session = Depends(get_db)): total = db.query(Email).count() emails = db.query(Email).order_by(Email.received_date.desc()).offset(skip).limit(limit).all() items = [ EmailDTO( id=e.id, subject=e.subject, sender=e.sender, received_date=e.received_date, is_read=e.is_read, has_attachments=bool(e.attachments) # Check if list is not empty ) for e in emails ] return EmailListResponse(items=items, total=total) # ... @app.get("/api/emails/{email_id}", response_model=EmailDetailDTO) def get_email_details(email_id: int, db: Session = Depends(get_db)): email = db.query(Email).filter(Email.id == email_id).first() if not email: raise HTTPException(status_code=404, detail="Email not found") # Mark as read if opened if not email.is_read: email.is_read = True db.commit() # Safely handle missing content_type and filename attachments = [] for a in email.attachments: fname = a.filename if not fname or fname == "unknown": ext = mimetypes.guess_extension(a.content_type or "") or ".bin" fname = f"attachment_{a.id}{ext}" attachments.append( AttachmentDTO( id=a.id, filename=fname, content_type=a.content_type or "application/octet-stream" ) ) return EmailDetailDTO( id=email.id, subject=email.subject, sender=email.sender, received_date=email.received_date, is_read=email.is_read, has_attachments=bool(email.attachments), body=email.body, attachments=attachments ) @app.get("/api/emails/{email_id}/download-all") def download_all_attachments(email_id: int, db: Session = Depends(get_db)): email = db.query(Email).filter(Email.id == email_id).first() if not email: raise HTTPException(status_code=404, detail="Email not found") if not email.attachments: raise HTTPException(status_code=400, detail="No attachments found") # Create a zip file in memory zip_buffer = io.BytesIO() with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file: for att in email.attachments: # Use filename from DB, ensure unique names if duplicates exist (simple append here) zip_file.writestr(att.filename, att.file_content) zip_buffer.seek(0) return StreamingResponse( zip_buffer, media_type="application/zip", headers={"Content-Disposition": f"attachment; filename=email-{email_id}-attachments.zip"} ) @app.get("/api/attachments/{attachment_id}") def download_attachment(attachment_id: int, view: bool = False, db: Session = Depends(get_db)): attachment = db.query(DBAttachment).filter(DBAttachment.id == attachment_id).first() if not attachment: raise HTTPException(status_code=404, detail="Attachment not found") disposition = "inline" if view else "attachment" return StreamingResponse( io.BytesIO(attachment.file_content), media_type=attachment.content_type, headers={"Content-Disposition": f"{disposition}; filename={attachment.filename}"} ) @app.post("/api/emails/sync") def sync_emails(): result = fetch_and_store_emails() return result @app.get("/") def read_root(): return {"message": "OCR Backend API is running"} class DocumentSaveRequest(BaseModel): vendor_name: str file_path: str model_type: str data: dict @app.post("/api/documents/save") def save_document(request: DocumentSaveRequest, db: Session = Depends(get_db)): # 1. Find or Create Vendor vendor = db.query(Vendor).filter(Vendor.name == request.vendor_name).first() if not vendor: vendor = Vendor(name=request.vendor_name, default_model=request.model_type) db.add(vendor) db.commit() db.refresh(vendor) else: # Update preference vendor.default_model = request.model_type db.commit() # 2. Save Document filename = request.file_path.split('/')[-1] doc = Document( vendor_id=vendor.id, filename=filename, status="verified", processed_data=request.data ) db.add(doc) db.commit() return {"message": "Document saved and Vendor preference updated", "vendor_id": vendor.id}