app.py

from fastapi import FastAPI, File, UploadFile, HTTPException, Query
from fastapi.responses import JSONResponse
import gradio as gr
from docling.document_converter import DocumentConverter
from docling.datamodel.base_models import InputFormat
import tempfile
import os
import json
import re

app = FastAPI(
    title="Docling Document Conversion API",
    description="Convert documents to Markdown and extract basic information",
    version="1.0.0"
)

# Initialize Docling
converter = DocumentConverter()

def extract_fields_from_markdown(markdown_text: str, template_type: str) -> dict:
    """Basic field extraction from markdown using regex patterns"""
    
    if template_type == "payment_advice":
        return {
            "document_type": "payment_advice",
            "payment_document_no": re.search(r'(?:Payment Document|Doc(?:ument)? No\.?)[:\s]+(\w+)', markdown_text, re.I),
            "payment_date": re.search(r'(?:Payment Date)[:\s]+([\d\-\/\.]+)', markdown_text, re.I),
            "total_amount": re.search(r'(?:Total|Amount)[:\s]+([\d,\.]+)', markdown_text, re.I),
            "vendor_code": re.search(r'(?:Vendor Code)[:\s]+(\w+)', markdown_text, re.I),
        }
    elif template_type == "invoice":
        return {
            "document_type": "invoice",
            "bill_no": re.search(r'(?:Bill|Invoice) No\.?[:\s]+(\w+)', markdown_text, re.I),
            "bill_date": re.search(r'(?:Bill|Invoice) Date[:\s]+([\d\-\/\.]+)', markdown_text, re.I),
            "total_amount": re.search(r'(?:Total|Grand Total)[:\s]+([\d,\.]+)', markdown_text, re.I),
            "gst_no": re.search(r'GST No\.?[:\s]+(\w+)', markdown_text, re.I),
        }
    
    return {"document_type": "unknown"}

def clean_extracted_fields(fields: dict) -> dict:
    """Convert regex match objects to strings"""
    cleaned = {}
    for key, value in fields.items():
        if hasattr(value, 'group'):
            cleaned[key] = value.group(1).strip()
        elif value:
            cleaned[key] = value
        else:
            cleaned[key] = None
    return cleaned

@app.post("/api/extract")
async def extract_document(
    file: UploadFile = File(...),
    template_type: str = Query(default="auto", description="Template type: auto, payment_advice, invoice")
):
    """
    Convert document to markdown and extract basic structured data
    """
    try:
        # Save uploaded file temporarily
        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as tmp_file:
            content = await file.read()
            tmp_file.write(content)
            tmp_file_path = tmp_file.name
        
        # Convert document to markdown
        result = converter.convert(tmp_file_path)
        markdown_content = result.document.export_to_markdown()
        
        # Auto-detect template if needed
        if template_type == "auto":
            text_lower = markdown_content.lower()
            if any(keyword in text_lower for keyword in ["payment advice", "payment document", "vendor code"]):
                template_type = "payment_advice"
            elif any(keyword in text_lower for keyword in ["invoice", "bill no", "gst no"]):
                template_type = "invoice"
            else:
                template_type = "general"
        
        # Extract basic fields
        if template_type in ["payment_advice", "invoice"]:
            fields = extract_fields_from_markdown(markdown_content, template_type)
            extracted_data = clean_extracted_fields(fields)
        else:
            extracted_data = {"document_type": template_type}
        
        # Clean up
        os.unlink(tmp_file_path)
        
        return JSONResponse(content={
            "success": True,
            "template_used": template_type,
            "data": extracted_data,
            "markdown": markdown_content,
            "filename": file.filename
        })
    
    except HTTPException:
        raise
    except Exception as e:
        if 'tmp_file_path' in locals():
            try:
                os.unlink(tmp_file_path)
            except:
                pass
        raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")

@app.post("/api/convert")
async def convert_to_markdown(file: UploadFile = File(...)):
    """Simple document to markdown conversion"""
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as tmp_file:
            content = await file.read()
            tmp_file.write(content)
            tmp_file_path = tmp_file.name
        
        result = converter.convert(tmp_file_path)
        markdown_content = result.document.export_to_markdown()
        os.unlink(tmp_file_path)
        
        return JSONResponse(content={
            "success": True,
            "markdown": markdown_content,
            "filename": file.filename
        })
    except Exception as e:
        if 'tmp_file_path' in locals():
            try:
                os.unlink(tmp_file_path)
            except:
                pass
        raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")

@app.get("/")
def read_root():
    """API information"""
    return {
        "message": "Docling Document Conversion API",
        "version": "1.0.0",
        "note": "Advanced extraction API is in beta - currently using markdown conversion with basic field extraction",
        "endpoints": {
            "extract": "/api/extract (POST) - Convert and extract basic fields",
            "convert": "/api/convert (POST) - Convert to markdown only",
            "docs": "/docs - Interactive API documentation",
            "web_ui": "/gradio - Web interface"
        }
    }

@app.get("/health")
def health_check():
    """Health check endpoint"""
    return {"status": "healthy", "service": "docling-api"}

# Gradio interface
def process_document(file):
    """Gradio function to process documents"""
    try:
        if file is None:
            return "Please upload a file"
        
        result = converter.convert(file.name)
        markdown_content = result.document.export_to_markdown()
        
        # Try basic extraction
        text_lower = markdown_content.lower()
        if "invoice" in text_lower or "bill" in text_lower:
            fields = extract_fields_from_markdown(markdown_content, "invoice")
            extracted = clean_extracted_fields(fields)
            return f"✅ Detected: Invoice\n\nExtracted Fields:\n{json.dumps(extracted, indent=2)}\n\n--- Markdown ---\n{markdown_content[:2000]}"
        elif "payment" in text_lower:
            fields = extract_fields_from_markdown(markdown_content, "payment_advice")
            extracted = clean_extracted_fields(fields)
            return f"✅ Detected: Payment Advice\n\nExtracted Fields:\n{json.dumps(extracted, indent=2)}\n\n--- Markdown ---\n{markdown_content[:2000]}"
        else:
            return f"✅ Converted to Markdown\n\n{markdown_content[:3000]}"
            
    except Exception as e:
        return f"❌ Error: {str(e)}"

demo = gr.Interface(
    fn=process_document,
    inputs=gr.File(label="📄 Upload Document (PDF, Image, DOCX)"),
    outputs=gr.Textbox(label="📊 Output", lines=30),
    title="Docling Document Processor",
    description="Upload documents to convert to Markdown and extract basic information. Advanced extraction API coming soon!",
    article="Built with [Docling](https://github.com/docling-project/docling) | API available at `/api/extract` and `/api/convert`",
)

app = gr.mount_gradio_app(app, demo, path="/gradio")

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