Upload 7 files

app/google_ai_mode.py

@@ -13,6 +13,7 @@ class PageRequest(NamedTuple):
     page_number: int
     image_path: Path
     prompt: str
+    screenshot_path: Path | None = None
 
 
 class GoogleAiModeClient:
@@ -405,6 +406,9 @@ async def _run_single_page_async(
         await _enter_prompt_async(page, request.prompt)
         await _submit_async(page)
         await page.wait_for_timeout(5_000)
+        if request.screenshot_path:
+            request.screenshot_path.parent.mkdir(parents=True, exist_ok=True)
+            await page.screenshot(path=str(request.screenshot_path), full_page=False)
         return await _wait_for_json_response_async(page, response_timeout_seconds)
     finally:
         await page.close()

app/jobs.py

@@ -22,6 +22,10 @@ class JobState:
     total_batches: int = 0
     completed_batches: int = 0
     current_pages: str | None = None
+    active_tabs: int = 0
+    completed_pages: int = 0
+    page_statuses: dict[str, str] = field(default_factory=dict)
+    latest_screenshot: str | None = None
     output_file: str | None = None
     error: str | None = None
     run_dir: str | None = None
@@ -88,9 +92,11 @@ def _run_job(
         image_dir = run_dir / "pages"
         batch_dir = run_dir / "batch-json"
         raw_dir = run_dir / "raw-output"
+        browser_view_dir = run_dir / "browser-view"
         browser_profile_dir = (run_dir / "browser-profile").resolve()
         batch_dir.mkdir(parents=True, exist_ok=True)
         raw_dir.mkdir(parents=True, exist_ok=True)
+        browser_view_dir.mkdir(parents=True, exist_ok=True)
 
         jobs.update(
             job_id,
@@ -119,12 +125,15 @@ def _run_job(
                 message=f"Wave {index}/{len(batches)}: opening AI Mode tabs for pages {current_pages}",
                 completed_batches=index - 1,
                 current_pages=current_pages,
+                active_tabs=len(batch),
                 progress=15 + int(((index - 1) / max(len(batches), 1)) * 70),
             )
 
             requests = []
+            page_statuses = dict(jobs.get(job_id).page_statuses if jobs.get(job_id) else {})
             for image_path in batch:
                 page_number = images.index(image_path) + 1
+                page_statuses[str(page_number)] = "queued"
                 prompt = (
                     f"{BANK_STATEMENT_PROMPT}\n\n"
                     "The JSON block above is only the required output format. "
@@ -134,7 +143,15 @@ def _run_job(
                     f"The uploaded image is PDF page {page_number}. "
                     "Extract transactions only from this uploaded page image."
                 )
-                requests.append(PageRequest(page_number, image_path, prompt))
+                requests.append(
+                    PageRequest(
+                        page_number,
+                        image_path,
+                        prompt,
+                        browser_view_dir / f"page-{page_number:04d}.png",
+                    )
+                )
+            jobs.update(job_id, page_statuses=page_statuses)
 
             wave_profile_dir = (browser_profile_dir / f"wave-{index:03d}").resolve()
             with GoogleAiModeClient(ai_mode_url, wave_profile_dir) as client:
@@ -146,6 +163,13 @@ def _run_job(
                 progress=15 + int(((index - 0.25) / max(len(batches), 1)) * 70),
             )
             for page_number in sorted(responses):
+                page_statuses[str(page_number)] = "parsing"
+                screenshot_path = browser_view_dir / f"page-{page_number:04d}.png"
+                jobs.update(
+                    job_id,
+                    page_statuses=page_statuses,
+                    latest_screenshot=str(screenshot_path) if screenshot_path.exists() else None,
+                )
                 raw_response = responses[page_number]
                 raw_file = raw_dir / f"page-{page_number:04d}-raw.txt"
                 raw_file.write_text(raw_response, encoding="utf-8")
@@ -156,9 +180,17 @@ def _run_job(
                     encoding="utf-8",
                 )
                 page_json_files[page_number] = page_file
+                page_statuses[str(page_number)] = "done"
+                jobs.update(
+                    job_id,
+                    completed_pages=len(page_json_files),
+                    page_statuses=page_statuses,
+                    latest_screenshot=str(screenshot_path) if screenshot_path.exists() else None,
+                )
             jobs.update(
                 job_id,
                 completed_batches=index,
+                active_tabs=0,
                 message=f"Saved page wave {index} of {len(batches)}",
                 progress=15 + int((index / max(len(batches), 1)) * 70),
             )
@@ -173,6 +205,7 @@ def _run_job(
             progress=100,
             output_file=str(final_json),
             current_pages=None,
+            active_tabs=0,
             data_preview=combined["data"][:20],
         )
     except Exception as exc:

app/main.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from tempfile import NamedTemporaryFile
 
 from fastapi import FastAPI, File, Form, HTTPException, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
 
@@ -21,14 +22,45 @@ DEFAULT_BATCH_SIZE = int(os.getenv("PARALLEL_TABS", "10"))
 RUNS_DIR.mkdir(exist_ok=True)
 
 app = FastAPI(title="Bank Statement AI Mode Extractor")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=[
+        origin.strip()
+        for origin in os.getenv("CORS_ORIGINS", "*").split(",")
+        if origin.strip()
+    ],
+    allow_credentials=False,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
 
 
 @app.get("/")
-def index() -> FileResponse:
+def index():
+    if os.getenv("API_ONLY", "false").strip().lower() in {"1", "true", "yes"}:
+        return FileResponse(STATIC_DIR / "backend.html")
     return FileResponse(STATIC_DIR / "index.html")
 
 
+@app.get("/api/info")
+def api_info() -> dict[str, str]:
+    return {
+        "name": "Bank Statement AI Mode Extractor API",
+        "health": "/api/health",
+        "create_job": "POST /api/jobs with multipart fields pdf and optional pwd",
+        "job_status": "/api/jobs/{job_id}",
+        "latest_browser_view": "/api/jobs/{job_id}/browser-view/latest",
+        "download_json": "/api/jobs/{job_id}/download",
+        "download_csv": "/api/jobs/{job_id}/download.csv",
+    }
+
+
+@app.get("/api/health")
+def health() -> dict[str, str]:
+    return {"status": "ok"}
+
+
 @app.post("/api/jobs")
 async def create_job(
     pdf: UploadFile = File(...),
@@ -59,6 +91,25 @@ def get_job(job_id: str) -> dict:
     return asdict(job)
 
 
+@app.get("/api/jobs/{job_id}/browser-view/latest")
+def latest_browser_view(job_id: str) -> FileResponse:
+    job = jobs.get(job_id)
+    if not job or not job.latest_screenshot:
+        raise HTTPException(status_code=404, detail="Browser screenshot is not ready.")
+    return FileResponse(job.latest_screenshot, media_type="image/png")
+
+
+@app.get("/api/jobs/{job_id}/browser-view/page/{page_number}")
+def page_browser_view(job_id: str, page_number: int) -> FileResponse:
+    job = jobs.get(job_id)
+    if not job or not job.run_dir:
+        raise HTTPException(status_code=404, detail="Job not found.")
+    screenshot = Path(job.run_dir) / "browser-view" / f"page-{page_number:04d}.png"
+    if not screenshot.exists():
+        raise HTTPException(status_code=404, detail="Browser screenshot is not ready.")
+    return FileResponse(screenshot, media_type="image/png")
+
+
 @app.get("/api/jobs/{job_id}/download")
 def download_job(job_id: str) -> FileResponse:
     job = jobs.get(job_id)