347 lines
10 KiB
Python
347 lines
10 KiB
Python
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}
|