app.py

import os
import uvicorn
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from model.main import process_and_analyze
from upload.data import download_pdf_from_url, process_uploaded_file
import shutil

app = FastAPI(title="Agentic Document Extraction", description="API for extracting and analyzing PDF documents")

class URLInput(BaseModel):
    url: str

@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    if not file.filename.endswith('.pdf'):
        raise HTTPException(status_code=400, detail="Only PDF files are allowed!")
    
    # Save uploaded file temporarily
    file_path = f"/tmp/{file.filename}"
    try:
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        
        # Process the uploaded file
        process_uploaded_file(file_path)
        
        if os.path.exists(file_path):
            process_and_analyze(file_path)
            return JSONResponse(content={"message": "Processing successful! Check files in data-extractor and file-upload."})
        else:
            raise HTTPException(status_code=500, detail="Error processing file!")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")
    finally:
        # Clean up temporary file
        if os.path.exists(file_path):
            os.remove(file_path)

@app.post("/process-url")
async def process_url(input: URLInput):
    if not input.url:
        raise HTTPException(status_code=400, detail="URL is required!")
    
    try:
        file_path = download_pdf_from_url(input.url)
        if file_path and os.path.exists(file_path):
            process_and_analyze(file_path)
            return JSONResponse(content={"message": "Processing successful! Check files in data-extractor and file-upload."})
        else:
            raise HTTPException(status_code=500, detail="Error processing file from URL!")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)