Maverick9876/document-extraction-workbench-cloud
0
1import csv2import os3import shutil4import threading5import time6import queue7from contextlib import contextmanager8from pathlib import Path9 10import gradio as gr11 12from ai_provider import DEFAULT_GEMINI_MODEL, DEFAULT_OPENAI_MODEL13from credential_extractor import run_certificate_extraction14from main_sharepoint_ocr import run_sharepoint_ocr15from name_matcherv2 import MATCH_THRESHOLD, run_name_matching16 17 18APP_DIR = Path(__file__).resolve().parent19ARTIFACTS_DIR = APP_DIR / "artifacts"20ARTIFACTS_DIR.mkdir(exist_ok=True)21RUN_LOCK = threading.Lock()22ARTIFACT_RETENTION_SEC = int(os.getenv("ARTIFACT_RETENTION_SEC", str(6 * 60 * 60)))23 24CSS_FILE = APP_DIR / "styles.css"25CSS_CONTENT = CSS_FILE.read_text(encoding="utf-8") if CSS_FILE.exists() else ""26 27APP_THEME = gr.themes.Base(28 primary_hue=gr.themes.colors.emerald,29 secondary_hue=gr.themes.colors.cyan,30 neutral_hue=gr.themes.colors.slate,31 font=[gr.themes.GoogleFont("Plus Jakarta Sans"), "ui-sans-serif", "system-ui", "sans-serif"],32 font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "ui-monospace", "Consolas", "monospace"],33).set(34 body_background_fill="#f3f7fb",35 body_background_fill_dark="#09121f",36 body_text_color="#102033",37 body_text_color_dark="#e9f1fb",38 body_text_color_subdued="#607286",39 body_text_color_subdued_dark="#95a9bf",40 background_fill_primary="#f3f7fb",41 background_fill_primary_dark="#09121f",42 background_fill_secondary="#ffffff",43 background_fill_secondary_dark="#101b2d",44 block_background_fill="#ffffff",45 block_background_fill_dark="#101b2d",46 block_border_color="#d8e3ef",47 block_border_color_dark="#23364e",48 block_label_background_fill="transparent",49 block_label_background_fill_dark="transparent",50 block_label_text_color="#102033",51 block_label_text_color_dark="#dbe8f7",52 block_title_text_color="#102033",53 block_title_text_color_dark="#f1f6ff",54 button_primary_background_fill="#11715d",55 button_primary_background_fill_dark="#29c6a2",56 button_primary_background_fill_hover="#0d5a4a",57 button_primary_background_fill_hover_dark="#1fa98a",58 button_primary_text_color="#ffffff",59 button_primary_text_color_dark="#05211b",60 button_secondary_background_fill="#eef5fb",61 button_secondary_background_fill_dark="#162438",62 button_secondary_background_fill_hover="#e2edf8",63 button_secondary_background_fill_hover_dark="#20324b",64 button_secondary_text_color="#17304d",65 button_secondary_text_color_dark="#d7e6f8",66 input_background_fill="#fdfefe",67 input_background_fill_dark="#0d1727",68 input_border_color="#c7d7e8",69 input_border_color_dark="#30445f",70 border_color_primary="#c7d7e8",71 border_color_primary_dark="#30445f",72 link_text_color="#11715d",73 link_text_color_dark="#56dfc0",74 link_text_color_hover="#0d5a4a",75 link_text_color_hover_dark="#83eed5",76 block_radius="22px",77 button_large_radius="999px",78 input_radius="16px",79)80 81 82def _timestamp() -> str:83 return time.strftime("%Y%m%d_%H%M%S")84 85 86def _default_model_for_provider(provider: str) -> str:87 return DEFAULT_GEMINI_MODEL if (provider or "").strip().lower() == "gemini" else DEFAULT_OPENAI_MODEL88 89 90def _provider_ui_defaults(provider: str):91 provider = (provider or "openai").strip().lower()92 if provider == "gemini":93 default_model = DEFAULT_GEMINI_MODEL94 note = (95 "**Gemini selected.** Use a Gemini API key and keep the default model unless you need a specific "96 "vision-capable variant."97 )98 else:99 default_model = DEFAULT_OPENAI_MODEL100 note = (101 "**OpenAI selected.** This is the default path for SharePoint OCR and certificate AI fallback. "102 "Faster models can help on large batches, but smaller CSV runs are still the smoothest on free CPU hosting."103 )104 return gr.update(value=default_model, placeholder=f"Default: {default_model}"), note, gr.update(value="")105 106 107def _run_dir(prefix: str) -> Path:108 run_dir = ARTIFACTS_DIR / f"{prefix}_{_timestamp()}_{os.getpid()}"109 run_dir.mkdir(parents=True, exist_ok=True)110 return run_dir111 112 113def _write_sample_csv(path: Path, rows: list[list[str]]) -> str:114 with path.open("w", encoding="utf-8", newline="") as handle:115 writer = csv.writer(handle)116 writer.writerows(rows)117 return str(path)118 119 120def build_sample_ocr_csv() -> str:121 return _write_sample_csv(122 ARTIFACTS_DIR / "sharepoint_ocr_sample.csv",123 [124 ["link"],125 ["https://contoso.sharepoint.com/sites/hr/Shared%20Documents/passport_sample.pdf"],126 ["https://contoso.sharepoint.com/sites/hr/Shared%20Documents/id_card_sample.png"],127 ],128 )129 130 131def build_sample_cert_csv() -> str:132 return _write_sample_csv(133 ARTIFACTS_DIR / "certificate_extraction_sample.csv",134 [135 ["link"],136 ["https://www.credential.net/example-certificate-link-1"],137 ["https://www.credential.net/example-certificate-link-2"],138 ],139 )140 141 142def build_sample_match_csv() -> str:143 return _write_sample_csv(144 ARTIFACTS_DIR / "name_matching_sample.csv",145 [146 ["Email", "Name Entered", "Extracted Database"],147 ["alex@example.com", "Alex Johnson", "Alexander Johnson"],148 ["priya@example.com", "Priya Sharma", "Priya R Sharma"],149 ["omar@example.com", "Omar Khan", "Omar A Khan"],150 ],151 )152 153 154def cleanup_old_artifacts() -> None:155 cutoff = time.time() - ARTIFACT_RETENTION_SEC156 for path in ARTIFACTS_DIR.iterdir():157 try:158 if path.is_dir() and path.stat().st_mtime < cutoff:159 shutil.rmtree(path, ignore_errors=True)160 elif path.is_file() and path.stat().st_mtime < cutoff:161 path.unlink(missing_ok=True)162 except Exception:163 pass164 165 166@contextmanager167def single_run_guard():168 acquired = RUN_LOCK.acquire(blocking=False)169 if not acquired:170 raise gr.Error("Another job is already running in this Space. Wait for it to finish, then try again.")171 try:172 yield173 finally:174 RUN_LOCK.release()175 176 177@contextmanager178def configured_ai_env(provider: str, api_key: str, model: str):179 provider = (provider or "openai").strip().lower()180 if provider not in {"openai", "gemini"}:181 raise ValueError("Provider must be either 'openai' or 'gemini'.")182 183 previous = {184 "AI_PROVIDER": os.environ.get("AI_PROVIDER"),185 "OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY"),186 "OPENAI_MODEL": os.environ.get("OPENAI_MODEL"),187 "GEMINI_API_KEY": os.environ.get("GEMINI_API_KEY"),188 "GEMINI_MODEL": os.environ.get("GEMINI_MODEL"),189 }190 191 os.environ["AI_PROVIDER"] = provider192 os.environ.pop("OPENAI_API_KEY", None)193 os.environ.pop("GEMINI_API_KEY", None)194 195 if provider == "openai":196 os.environ["OPENAI_MODEL"] = (model or DEFAULT_OPENAI_MODEL).strip() or DEFAULT_OPENAI_MODEL197 if api_key.strip():198 os.environ["OPENAI_API_KEY"] = api_key.strip()199 else:200 os.environ["GEMINI_MODEL"] = (model or DEFAULT_GEMINI_MODEL).strip() or DEFAULT_GEMINI_MODEL201 if api_key.strip():202 os.environ["GEMINI_API_KEY"] = api_key.strip()203 204 try:205 yield206 finally:207 for key, value in previous.items():208 if value is None:209 os.environ.pop(key, None)210 else:211 os.environ[key] = value212 213 214def _copy_uploaded_csv(src_path: str, target_dir: Path, target_name: str) -> Path:215 if not src_path:216 raise ValueError("Please upload a CSV file first.")217 src = Path(src_path)218 if not src.exists():219 raise ValueError("Uploaded CSV file could not be found on the server.")220 dst = target_dir / target_name221 shutil.copy2(src, dst)222 return dst223 224 225def _friendly_error(exc: Exception) -> gr.Error:226 message = str(exc).strip() or exc.__class__.__name__227 lowered = message.lower()228 229 if "429" in lowered or "rate limit" in lowered:230 message = "The provider rate limit was hit. Retry in a minute or use a lighter model."231 elif "api key" in lowered and "not set" in lowered:232 message = "Add a valid API key to run this workflow."233 elif "playwright" in lowered or "browser" in lowered:234 message = "Browser automation failed on the hosted runtime. Retry with a smaller CSV or try again shortly."235 elif "csv must include a header column named 'link'" in lowered:236 message = "The uploaded OCR CSV must contain a header column named 'link'."237 elif "missing required columns" in lowered:238 message = "The name matching CSV must include Email, Name Entered, and Extracted Database columns."239 elif "timed out" in lowered or "timeout" in lowered:240 message = "The hosted job timed out. Try a smaller file or fewer large PDFs per run."241 242 return gr.Error(message)243 244 245def run_sharepoint_workflow(csv_file: str, provider: str, api_key: str, model: str):246 if not api_key.strip():247 raise gr.Error("Add your API key to run SharePoint OCR.")248 249 with single_run_guard():250 cleanup_old_artifacts()251 run_dir = _run_dir("sharepoint")252 input_csv = _copy_uploaded_csv(csv_file, run_dir, "sharepoint_input.csv")253 output_dir = run_dir / "outputs"254 output_dir.mkdir(exist_ok=True)255 256 logs = []257 log_queue = queue.Queue()258 259 def status_callback(message: str):260 log_queue.put(str(message))261 262 def progress_callback(done: int, total: int):263 log_queue.put(f"Progress: {done} of {total}")264 265 result_holder = {}266 267 def thread_target():268 try:269 with configured_ai_env(270 provider=provider,271 api_key=api_key,272 model=(model or _default_model_for_provider(provider)),273 ):274 output_path = run_sharepoint_ocr(275 csv_path=str(input_csv),276 output_dir=str(output_dir),277 api_key=api_key.strip(),278 allow_api_key_prompt=False,279 progress_callback=progress_callback,280 status_callback=status_callback,281 show_progress=False,282 )283 result_holder["output_path"] = output_path284 except Exception as exc:285 result_holder["error"] = exc286 finally:287 log_queue.put(None)288 289 t = threading.Thread(target=thread_target)290 t.start()291 292 while t.is_alive() or not log_queue.empty():293 try:294 msg = log_queue.get(timeout=0.2)295 if msg is None:296 break297 logs.append(msg)298 yield None, "\n".join(logs[-25:])299 except queue.Empty:300 pass301 302 t.join()303 if "error" in result_holder:304 raise _friendly_error(result_holder["error"])305 306 output_path = result_holder.get("output_path")307 yield str(output_path), "\n".join(logs[-25:]) or "SharePoint OCR completed."308 309 310def run_certificate_workflow(csv_file: str, provider: str, api_key: str, model: str, use_ai: bool):311 if use_ai and not api_key.strip():312 raise gr.Error("Add your API key or turn off AI fallback for certificate extraction.")313 314 with single_run_guard():315 cleanup_old_artifacts()316 run_dir = _run_dir("certificate")317 input_csv = _copy_uploaded_csv(csv_file, run_dir, "certificate_input.csv")318 output_csv = run_dir / "output_with_certificate_details.csv"319 320 logs = []321 log_queue = queue.Queue()322 323 def status_callback(message: str):324 log_queue.put(str(message))325 326 def progress_callback(done: int, total: int):327 log_queue.put(f"Progress: {done} of {total}")328 329 result_holder = {}330 331 def thread_target():332 try:333 with configured_ai_env(334 provider=provider,335 api_key=api_key,336 model=(model or _default_model_for_provider(provider)),337 ):338 output_path = run_certificate_extraction(339 input_csv=str(input_csv),340 output_csv=str(output_csv),341 progress_callback=progress_callback,342 status_callback=status_callback,343 use_openai_ocr=bool(use_ai),344 openai_api_key=api_key.strip() or None,345 )346 result_holder["output_path"] = output_path347 except Exception as exc:348 result_holder["error"] = exc349 finally:350 log_queue.put(None)351 352 t = threading.Thread(target=thread_target)353 t.start()354 355 while t.is_alive() or not log_queue.empty():356 try:357 msg = log_queue.get(timeout=0.2)358 if msg is None:359 break360 logs.append(msg)361 yield None, "\n".join(logs[-25:])362 except queue.Empty:363 pass364 365 t.join()366 if "error" in result_holder:367 raise _friendly_error(result_holder["error"])368 369 output_path = result_holder.get("output_path")370 yield str(output_path), "\n".join(logs[-25:]) or "Certificate extraction completed."371 372 373def run_name_matching_workflow(csv_file: str, threshold: float):374 try:375 cleanup_old_artifacts()376 with single_run_guard():377 run_dir = _run_dir("matching")378 input_csv = _copy_uploaded_csv(csv_file, run_dir, "name_matching_input.csv")379 output_csv = run_dir / "matched_output.csv"380 output_path = run_name_matching(381 csv_file=str(input_csv),382 output_file=str(output_csv),383 threshold=float(threshold or MATCH_THRESHOLD),384 )385 except gr.Error:386 raise387 except Exception as exc:388 raise _friendly_error(exc)389 390 return str(output_path), f"Name matching completed with threshold {float(threshold or MATCH_THRESHOLD):.2f}."391 392 393with gr.Blocks(394 title="Document Extraction Workbench Cloud",395 theme=APP_THEME,396 css=CSS_CONTENT,397 fill_width=True,398 delete_cache=(3600, 3600 * 6),399) as demo:400 gr.HTML(401 """402 <header class="app-header">403 <div class="header-main">404 <div class="header-logo">405 <svg width="28" height="28" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">406 <path d="M14 2H6C4.9 2 4 2.9 4 4V20C4 21.1 4.9 22 6 22H18C19.1 22 20 21.1 20 20V8L14 2Z" fill="currentColor" fill-opacity="0.15"/>407 <path d="M14 2H6C4.9 2 4.01 2.9 4.01 4L4 20C4 21.1 4.9 22 6 22H18C19.1 22 20 21.1 20 20V8L14 2ZM18 20H6V4H13V9H18V20ZM9 16H15V18H9V16ZM9 12H15V14H9V12ZM9 8H11V10H9V8Z" fill="currentColor"/>408 </svg>409 <h1>Document Extraction Workbench</h1>410 </div>411 <div class="header-badge">Cloud Edition</div>412 </div>413 </header>414 """415 )416 417 with gr.Column(elem_classes=["config-panel"]):418 with gr.Row():419 provider = gr.Dropdown(420 choices=["openai", "gemini"],421 value="openai",422 label="AI Provider",423 )424 model = gr.Textbox(425 label="AI Model",426 value=DEFAULT_OPENAI_MODEL,427 placeholder=f"Default: {DEFAULT_OPENAI_MODEL}",428 )429 api_key = gr.Textbox(430 label="API Key",431 type="password",432 placeholder="Enter API key for session...",433 )434 provider_note = gr.Markdown(435 "**OpenAI selected.** Default path for SharePoint OCR and certificate AI fallback.",436 elem_classes=["provider-note"],437 )438 439 with gr.Tabs(elem_classes=["tool-tabs"]):440 with gr.Tab("SharePoint OCR"):441 with gr.Row(equal_height=True):442 with gr.Column(elem_classes=["panel-card"]):443 sharepoint_input = gr.File(444 label="Upload Input CSV (must include 'link' header)",445 file_types=[".csv"],446 type="filepath",447 )448 with gr.Row():449 sharepoint_sample = gr.File(450 value=build_sample_ocr_csv,451 label="Template CSV",452 interactive=False,453 )454 sharepoint_run = gr.Button("Run SharePoint OCR", variant="primary", elem_classes=["action-btn"])455 with gr.Column(elem_classes=["panel-card"]):456 sharepoint_output = gr.File(label="Result CSV")457 sharepoint_status = gr.Textbox(label="Activity Feed / Logs", lines=12, elem_classes=["log-box"])458 sharepoint_run.click(459 run_sharepoint_workflow,460 inputs=[sharepoint_input, provider, api_key, model],461 outputs=[sharepoint_output, sharepoint_status],462 concurrency_limit=1,463 )464 465 with gr.Tab("Certificate Extraction"):466 with gr.Row(equal_height=True):467 with gr.Column(elem_classes=["panel-card"]):468 cert_input = gr.File(469 label="Upload Input CSV (must include certificate links)",470 file_types=[".csv"],471 type="filepath",472 )473 cert_use_ai = gr.Checkbox(474 value=True,475 label="Enable AI OCR fallback for certificate images",476 )477 with gr.Row():478 cert_sample = gr.File(479 value=build_sample_cert_csv,480 label="Template CSV",481 interactive=False,482 )483 cert_run = gr.Button("Run Certificate Extraction", variant="primary", elem_classes=["action-btn"])484 with gr.Column(elem_classes=["panel-card"]):485 cert_output = gr.File(label="Result CSV")486 cert_status = gr.Textbox(label="Activity Feed / Logs", lines=12, elem_classes=["log-box"])487 cert_run.click(488 run_certificate_workflow,489 inputs=[cert_input, provider, api_key, model, cert_use_ai],490 outputs=[cert_output, cert_status],491 concurrency_limit=1,492 )493 494 with gr.Tab("Name Matching"):495 with gr.Row(equal_height=True):496 with gr.Column(elem_classes=["panel-card"]):497 match_input = gr.File(498 label="Upload Input CSV (Email, Name Entered, Extracted Database headers)",499 file_types=[".csv"],500 type="filepath",501 )502 match_threshold = gr.Slider(503 value=MATCH_THRESHOLD,504 label="Match Threshold",505 minimum=0.0,506 maximum=1.0,507 step=0.05,508 )509 with gr.Row():510 match_sample = gr.File(511 value=build_sample_match_csv,512 label="Template CSV",513 interactive=False,514 )515 match_run = gr.Button("Run Name Matching", variant="primary", elem_classes=["action-btn"])516 with gr.Column(elem_classes=["panel-card"]):517 match_output = gr.File(label="Result CSV")518 match_status = gr.Textbox(label="Activity Feed / Logs", lines=12, elem_classes=["log-box"])519 match_run.click(520 run_name_matching_workflow,521 inputs=[match_input, match_threshold],522 outputs=[match_output, match_status],523 concurrency_limit=1,524 )525 526 gr.HTML(527 """528 <div class="app-footnote">529 Document Extraction Workbench - Dedicated End-User Interface530 </div>531 """532 )533 534 provider.change(535 fn=_provider_ui_defaults,536 inputs=provider,537 outputs=[model, provider_note, api_key],538 queue=False,539 )540 541demo.queue(default_concurrency_limit=1)542 543if __name__ == "__main__":544 demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))545