""" Wellfound AI - Main FastAPI Application Automated Excel data completion tool with: - AI web content understanding - Smart address processing - Multi-source contact finding - Breakpoint resume support - Smart timeout detection with auto-skip (5-min threshold per company) """ import asyncio import hashlib import json import logging import os import time import uuid from datetime import datetime from pathlib import Path from typing import Optional import pandas as pd from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request from fastapi.responses import FileResponse, HTMLResponse, JSONResponse from fastapi.staticfiles import StaticFiles from fastapi.middleware.cors import CORSMiddleware from core.excel_handler import ExcelHandler from core.scraper import WebScraper from core.ai_extractor import AIExtractor from core.address_processor import AddressProcessor from core.contact_finder import ContactFinder # ─── Logging Setup ──────────────────────────────────────── logger = logging.getLogger("wellfound_ai") logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s - %(message)s", ) # ─── Timeout Configuration ──────────────────────────────── # Per-company processing timeout in seconds (5 minutes) COMPANY_TIMEOUT_SECONDS = 300 # ─── App Setup ─────────────────────────────────────────── app = FastAPI( title="Wellfound AI - Excel Data Completion", description="AI-powered automated Excel data completion for Wellfound exports", version="1.1.0", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Mount static files (with Docker-safe path resolution) static_dir = Path(__file__).resolve().parent / "static" static_dir.mkdir(parents=True, exist_ok=True) app.mount("/static", StaticFiles(directory=str(static_dir)), name="static") # Data directories (Docker-safe paths) DATA_DIR = Path(__file__).resolve().parent / "data" UPLOAD_DIR = DATA_DIR / "uploads" RESULT_DIR = DATA_DIR / "results" CHECKPOINT_DIR = DATA_DIR / "checkpoints" for d in [UPLOAD_DIR, RESULT_DIR, CHECKPOINT_DIR]: d.mkdir(parents=True, exist_ok=True) # Global processing state processing_tasks: dict = {} excel_handler = ExcelHandler(str(CHECKPOINT_DIR)) address_processor = AddressProcessor() contact_finder = ContactFinder() # ─── Routes ────────────────────────────────────────────── @app.get("/", response_class=HTMLResponse) async def index(): """Serve the main application page.""" template_path = Path(__file__).parent / "templates" / "index.html" if template_path.exists(): return template_path.read_text(encoding="utf-8") return HTMLResponse("

Wellfound AI

Frontend not found.

") @app.get("/api/health") async def health(): """Health check endpoint.""" return { "status": "ok", "timestamp": datetime.now().isoformat(), "version": "1.1.0", } @app.post("/api/upload") async def upload_file(file: UploadFile = File(...)): """Upload Excel file and return metadata.""" if not file.filename.endswith((".xlsx", ".xls")): raise HTTPException(400, "Only Excel files (.xlsx, .xls) are supported") file_id = uuid.uuid4().hex[:12] file_path = UPLOAD_DIR / f"{file_id}_{file.filename}" content = await file.read() file_path.write_bytes(content) try: df = excel_handler.load(str(file_path)) except Exception as e: file_path.unlink(missing_ok=True) raise HTTPException(400, f"Failed to read Excel file: {str(e)}") return { "file_id": file_id, "filename": file.filename, "rows": len(df), "columns": list(df.columns), "missing_data": { col: int(df[col].isnull().sum()) for col in df.columns if df[col].isnull().sum() > 0 }, } @app.post("/api/process") async def process_file(request: Request): """Start processing an uploaded Excel file.""" data = await request.json() file_id = data.get("file_id") api_key = data.get("api_key", "") provider = data.get("provider", "openai") model = data.get("model", "auto") start_row = data.get("start_row", 0) max_rows = data.get("max_rows", 0) # 0 = all use_ai = data.get("use_ai", True) use_web_scraping = data.get("use_web_scraping", True) concurrency = data.get("concurrency", 2) company_timeout = data.get("company_timeout", COMPANY_TIMEOUT_SECONDS) if not file_id: raise HTTPException(400, "file_id is required") # Find uploaded file files = list(UPLOAD_DIR.glob(f"{file_id}_*")) if not files: raise HTTPException(404, "Uploaded file not found") file_path = str(files[0]) # Validate API key if using AI if use_ai and not api_key: raise HTTPException(400, "API key is required when AI extraction is enabled") task_id = uuid.uuid4().hex[:8] processing_tasks[task_id] = { "status": "starting", "progress": 0, "total": 0, "current": "", "errors": [], "skipped": [], # NEW: track skipped companies due to timeout "file_id": file_id, "result_path": None, } # Launch async processing asyncio.create_task( _process_file( task_id=task_id, file_path=file_path, api_key=api_key, provider=provider, model=model, start_row=start_row, max_rows=max_rows, use_ai=use_ai, use_web_scraping=use_web_scraping, concurrency=concurrency, company_timeout=company_timeout, ) ) return {"task_id": task_id, "status": "started"} @app.get("/api/status/{task_id}") async def task_status(task_id: str): """Get processing task status.""" if task_id not in processing_tasks: raise HTTPException(404, "Task not found") task = processing_tasks[task_id] return { "task_id": task_id, "status": task["status"], "progress": task["progress"], "total": task["total"], "current": task["current"], "errors": task["errors"][-10:], # Last 10 errors "error_count": len(task["errors"]), "skipped": task.get("skipped", [])[-10:], # Last 10 skipped entries "skipped_count": len(task.get("skipped", [])), "has_result": task["result_path"] is not None, } @app.get("/api/download/{task_id}") async def download_result(task_id: str): """Download the processed Excel file.""" if task_id not in processing_tasks: raise HTTPException(404, "Task not found") task = processing_tasks[task_id] if not task["result_path"] or not os.path.exists(task["result_path"]): raise HTTPException(404, "Result file not ready yet") return FileResponse( task["result_path"], filename=f"wellfound_processed_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx", media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", ) @app.get("/api/providers") async def get_providers(): """Get available AI providers and models.""" return AIExtractor.PROVIDERS # ─── Processing Logic ──────────────────────────────────── async def _process_file( task_id: str, file_path: str, api_key: str, provider: str, model: str, start_row: int, max_rows: int, use_ai: bool, use_web_scraping: bool, concurrency: int, company_timeout: int = COMPANY_TIMEOUT_SECONDS, ): """Main processing pipeline with smart timeout detection and auto-skip.""" task = processing_tasks[task_id] scraper = None ai_extractor = None try: # Load data task["status"] = "loading" df = excel_handler.load(file_path) # Create working copy result_path = str(RESULT_DIR / f"processed_{task_id}.xlsx") df.to_excel(result_path, index=False, engine="openpyxl") task["result_path"] = result_path total_rows = len(df) if max_rows > 0: total_rows = min(max_rows, total_rows - start_row) else: total_rows = total_rows - start_row task["total"] = total_rows task["status"] = "processing" # Initialize services if use_web_scraping: scraper = WebScraper(headless=True) await scraper.start() if use_ai and api_key: ai_extractor = AIExtractor(provider=provider, api_key=api_key, model=model) # Process rows in batches with concurrency semaphore = asyncio.Semaphore(concurrency) processed = 0 async def process_row(idx: int): nonlocal processed async with semaphore: row = df.iloc[idx].to_dict() company = str(row.get("Company Name", "")) internal_link = str(row.get("Internal Link", "")) external_link = str(row.get("External Link", "")) location = str(row.get("Location", "")) task["current"] = f"Row {idx + 1}/{start_row + total_rows}: {company}" # ── Smart Timeout Wrapper ────────────────────── # Each company gets a bounded time window. If it exceeds the # threshold the row is skipped and the failure is logged so # the batch pipeline never stalls on a single company. try: updates = await asyncio.wait_for( _process_single_company( row=row, idx=idx, company=company, internal_link=internal_link, external_link=external_link, location=location, use_web_scraping=use_web_scraping, scraper=scraper, ai_extractor=ai_extractor, result_path=result_path, ), timeout=company_timeout, ) except asyncio.TimeoutError: # ── Timeout detected → auto-skip ───────── skip_record = { "company": company, "row": idx + 1, "reason": f"Processing exceeded {company_timeout}s timeout", "timestamp": datetime.now().isoformat(), "links": { "internal": internal_link if internal_link != "nan" else None, "external": external_link if external_link != "nan" else None, }, } task["skipped"].append(skip_record) task["errors"].append( f"Row {idx} ({company}): SKIPPED — timeout after {company_timeout}s" ) logger.warning( "⏱ Timeout skip — company '%s' (row %d) exceeded %ds", company, idx + 1, company_timeout, ) # Still count as processed so progress bar advances processed += 1 task["progress"] = processed return # Save updates if updates: excel_handler.save_row(result_path, idx, updates) processed += 1 task["progress"] = processed # Process rows end_row = start_row + total_rows tasks_list = [ process_row(i) for i in range(start_row, min(end_row, len(df))) ] await asyncio.gather(*tasks_list, return_exceptions=True) task["status"] = "completed" task["progress"] = total_rows # Log summary skipped_count = len(task.get("skipped", [])) if skipped_count > 0: logger.info( "Batch complete: %d/%d processed, %d skipped due to timeout", total_rows - skipped_count, total_rows, skipped_count, ) else: logger.info("Batch complete: %d/%d processed, no timeouts", total_rows, total_rows) except Exception as e: task["status"] = "error" task["errors"].append(f"Fatal: {str(e)}") logger.error("Fatal processing error: %s", str(e)) finally: if scraper: await scraper.stop() async def _process_single_company( row: dict, idx: int, company: str, internal_link: str, external_link: str, location: str, use_web_scraping: bool, scraper: Optional[WebScraper], ai_extractor: Optional[AIExtractor], result_path: str, ) -> dict: """Process a single company row — extract funding, contacts, and address. Returns a dict of column updates. Callers should wrap this in ``asyncio.wait_for()`` to enforce per-company timeouts. """ updates = {} # Step 1: Process wellfound page for funding data if use_web_scraping and internal_link and internal_link != "nan": funding_data = await _extract_funding( scraper, ai_extractor, company, internal_link ) if funding_data: for col_key, excel_col in [ ("valuation", "Valuation"), ("rounds", "Rounds"), ("series", "Series"), ("total_raised", "Total Raised"), ]: if funding_data.get(col_key) and pd.isna(row.get(excel_col)): updates[excel_col] = funding_data[col_key] # Step 2: Process company website for contacts and address if use_web_scraping and external_link and external_link != "nan": contact_data = await _extract_contacts( scraper, ai_extractor, company, external_link ) if contact_data: # Contact info if pd.isna(row.get("contact")): contact_email = contact_data.get("contact_email") contact_form = contact_data.get("contact_form_url") contact_str = contact_finder.format_contact_output( contact_email, contact_form ) if contact_str: updates["contact"] = contact_str # Location info if pd.isna(row.get("location.apply")) or pd.isna(row.get("state.apply")): loc_apply, state_apply = address_processor.determine_location_apply( wellfound_location=location, ai_contact_data=contact_data, scraped_addresses=contact_data.get("scraped_addresses", []), ) if loc_apply and pd.isna(row.get("location.apply")): updates["location.apply"] = loc_apply if state_apply and pd.isna(row.get("state.apply")): updates["state.apply"] = state_apply return updates async def _extract_funding( scraper: Optional[WebScraper], ai_extractor: Optional[AIExtractor], company: str, url: str, ) -> dict: """Extract funding information from wellfound page.""" result = {} if not scraper: return result try: page_data = await scraper.fetch_wellfound_page(url) if page_data.get("error"): return result # Get regex-based extraction funding_data = page_data.get("funding_data", {}) # AI-enhanced extraction if ai_extractor and page_data.get("text"): try: ai_data = await ai_extractor.analyze_funding_page( company, page_data["text"], page_data.get("meta"), ) # Merge AI results with regex results, preferring high-confidence AI for key in ["valuation", "total_raised", "rounds", "series"]: ai_val = ai_data.get(key) ai_conf = ai_data.get("confidence", {}).get(key, "low") regex_val = funding_data.get(key) if ai_val and ai_conf == "high": result[key] = ai_val elif regex_val: result[key] = regex_val elif ai_val: result[key] = ai_val except Exception: # Fall back to regex-only for key in ["valuation", "total_raised", "rounds", "series"]: if funding_data.get(key): result[key] = funding_data[key] else: # No AI, use regex results for key in ["valuation", "total_raised", "rounds", "series"]: if funding_data.get(key): result[key] = funding_data[key] except Exception as e: pass # Silently fail on individual row errors return result async def _extract_contacts( scraper: Optional[WebScraper], ai_extractor: Optional[AIExtractor], company: str, url: str, ) -> dict: """Extract contact information from company website.""" result = { "contact_email": None, "contact_form_url": None, "careers_page_url": None, "headquarters_city": None, "headquarters_state": None, "scraped_addresses": [], "us_office_city": None, "us_office_state": None, "hiring_focus_location": None, } if not scraper: return result try: page_data = await scraper.fetch_company_website(url) if page_data.get("error"): return result emails = page_data.get("emails", []) links = page_data.get("links", []) phones = page_data.get("phones", []) addresses = page_data.get("addresses", []) text = page_data.get("text", "") result["scraped_addresses"] = addresses # Contact email best_email = contact_finder.find_best_contact_email(emails, text) result["contact_email"] = best_email # Contact form result["contact_form_url"] = contact_finder.find_contact_form_url( url, links, page_data.get("html", "") ) # Careers page result["careers_page_url"] = contact_finder.find_careers_page_url( url, links, page_data.get("html", "") ) # AI-enhanced extraction for location and additional contacts if ai_extractor and text: try: ai_data = await ai_extractor.analyze_company_website( company, text, emails, links, phones, addresses ) # Use AI data if available if not result["contact_email"] and ai_data.get("contact_email"): result["contact_email"] = ai_data["contact_email"] result["headquarters_city"] = ai_data.get("headquarters_city") result["headquarters_state"] = ai_data.get("headquarters_state") result["us_office_city"] = ai_data.get("us_office_city") result["us_office_state"] = ai_data.get("us_office_state") result["hiring_focus_location"] = ai_data.get("hiring_focus_location") if not result["contact_form_url"]: result["contact_form_url"] = ai_data.get("contact_form_url") if not result["careers_page_url"]: result["careers_page_url"] = ai_data.get("careers_page_url") except Exception: pass except Exception: pass return result # ─── Main ───────────────────────────────────────────────── if __name__ == "__main__": import uvicorn # HF Spaces sets PORT=7860; respect it, otherwise auto-detect hf_port = os.environ.get("PORT") or os.environ.get("HF_SPACE_PORT") if hf_port: port = int(hf_port) print(f"\n HF Spaces detected: using PORT={port}") else: port = 7860 # Try to find an available port locally import socket try: s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) s.bind(("0.0.0.0", port)) s.close() except OSError: for alt in range(7861, 7870): try: s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) s.bind(("0.0.0.0", alt)) s.close() port = alt break except OSError: continue print(f"\n{'='*50}") print(f" Wellfound AI - Server Starting") print(f" URL: http://0.0.0.0:{port}") print(f" Data Dir: {DATA_DIR}") print(f" Company Timeout: {COMPANY_TIMEOUT_SECONDS}s") print(f"{'='*50}\n") uvicorn.run( app, host="0.0.0.0", port=port, log_level="info", # Graceful shutdown for Docker timeout_keep_alive=30, )