Update backend/app/main.py

backend/app/main.py

@@ -1,9 +1,18 @@
 import os
-from fastapi import FastAPI
+import time
+from typing import List, Dict
+
+from fastapi import FastAPI, UploadFile, File, Depends
 from fastapi.middleware.cors import CORSMiddleware
-from .db import Base, engine
+from fastapi.staticfiles import StaticFiles
+from sqlalchemy.orm import Session
+
+from .db import Base, engine, SessionLocal
+from .models import ExtractionRecord
+from .schemas import ExtractionRecordBase, ExtractionStage
+from .openrouter_client import extract_fields_from_document
 
-# Ensure data directory exists (for SQLite file)
+# Ensure data dir exists for SQLite
 os.makedirs("data", exist_ok=True)
 
 # Create tables
@@ -11,7 +20,7 @@ Base.metadata.create_all(bind=engine)
 
 app = FastAPI(title="Document Capture Demo – Backend")
 
-# Allow CORS from anywhere (fine for a demo Space)
+# CORS (for safety we allow all; you can tighten later)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -21,10 +30,165 @@ app.add_middleware(
 )
 
 
+def get_db():
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()
+
+
 @app.get("/ping")
 def ping():
+    """Healthcheck."""
+    return {"status": "ok", "message": "backend alive"}
+
+
+def make_stages(total_ms: int, status: str) -> Dict[str, ExtractionStage]:
     """
-    Simple health check endpoint.
-    You can hit this from the browser to confirm backend is running.
+    Build synthetic stage timing data for the History UI.
+    For now we just split total_ms into 4 stages.
     """
-    return {"status": "ok", "message": "backend alive"}
+    if total_ms <= 0:
+        total_ms = 1000
+
+    return {
+        "uploading": ExtractionStage(
+            time=int(total_ms * 0.15),
+            status="completed",
+            variation="normal",
+        ),
+        "aiAnalysis": ExtractionStage(
+            time=int(total_ms * 0.55),
+            status="completed" if status == "completed" else "failed",
+            variation="normal",
+        ),
+        "dataExtraction": ExtractionStage(
+            time=int(total_ms * 0.2),
+            status="completed" if status == "completed" else "skipped",
+            variation="fast",
+        ),
+        "outputRendering": ExtractionStage(
+            time=int(total_ms * 0.1),
+            status="completed" if status == "completed" else "skipped",
+            variation="normal",
+        ),
+    }
+
+
+@app.post("/api/extract")
+async def extract_document(
+    file: UploadFile = File(...),
+    db: Session = Depends(get_db),
+):
+    """
+    Main extraction endpoint used by the Dashboard.
+    1) Read the uploaded file
+    2) Call OpenRouter + Qwen3-VL
+    3) Store a record in SQLite
+    4) Return extraction result + metadata
+    """
+    start = time.time()
+    content = await file.read()
+    content_type = file.content_type or "application/octet-stream"
+    size_mb = len(content) / 1024 / 1024
+    size_str = f"{size_mb:.2f} MB"
+
+    try:
+        extracted = await extract_fields_from_document(content, content_type, file.filename)
+        total_ms = int((time.time() - start) * 1000)
+
+        confidence = float(extracted.get("confidence", 90))
+        fields = extracted.get("fields", {})
+        fields_extracted = len(fields) if isinstance(fields, dict) else 0
+
+        status = "completed"
+        error_message = None
+    except Exception as e:
+        total_ms = int((time.time() - start) * 1000)
+        confidence = 0.0
+        fields = {}
+        fields_extracted = 0
+        status = "failed"
+        error_message = str(e)
+
+    # Save record to DB
+    rec = ExtractionRecord(
+        file_name=file.filename,
+        file_type=content_type,
+        file_size=size_str,
+        status=status,
+        confidence=confidence,
+        fields_extracted=fields_extracted,
+        total_time_ms=total_ms,
+        raw_output=str(fields),
+        error_message=error_message,
+    )
+    db.add(rec)
+    db.commit()
+    db.refresh(rec)
+
+    stages = make_stages(total_ms, status)
+
+    # Response shape that frontend will consume
+    return {
+        "id": rec.id,
+        "fileName": rec.file_name,
+        "fileType": rec.file_type,
+        "fileSize": rec.file_size,
+        "status": status,
+        "confidence": confidence,
+        "fieldsExtracted": fields_extracted,
+        "totalTime": total_ms,
+        "fields": fields,
+        "stages": {k: v.dict() for k, v in stages.items()},
+        "errorMessage": error_message,
+    }
+
+
+@app.get("/api/history", response_model=List[ExtractionRecordBase])
+def get_history(db: Session = Depends(get_db)):
+    """
+    Used by the History page.
+    Returns last 100 records, with synthetic stage data.
+    """
+    recs = (
+        db.query(ExtractionRecord)
+        .order_by(ExtractionRecord.created_at.desc())
+        .limit(100)
+        .all()
+    )
+
+    output: List[ExtractionRecordBase] = []
+    for r in recs:
+        stages = make_stages(r.total_time_ms or 1000, r.status or "completed")
+        output.append(
+            ExtractionRecordBase(
+                id=r.id,
+                fileName=r.file_name,
+                fileType=r.file_type or "",
+                fileSize=r.file_size or "",
+                extractedAt=r.created_at,
+                status=r.status or "completed",
+                confidence=r.confidence or 0.0,
+                fieldsExtracted=r.fields_extracted or 0,
+                totalTime=r.total_time_ms or 0,
+                stages=stages,
+                errorMessage=r.error_message,
+            )
+        )
+    return output
+
+
+# Static frontend mounting (used after we build React)
+# Dockerfile copies the Vite build into backend/frontend_dist
+frontend_dir = os.path.join(
+    os.path.dirname(os.path.dirname(__file__)), "frontend_dist"
+)
+
+if os.path.isdir(frontend_dir):
+    app.mount(
+        "/",
+        StaticFiles(directory=frontend_dir, html=True),
+        name="frontend",
+    )