Spaces:

Zoey7Web
/

Wellfound

Runtime error

Upload 21 files
28291d7 verified 5 months ago
27.5 kB
"""
Wellfound AI - Main FastAPI Application
Automated Excel data completion tool with:
- AI-powered deep website understanding
- Smart address processing (US Census regions)
- Multi-source contact finding (email priority decision tree)
- Deep multi-page website scanning (careers, contact, about, team, locations)
- Breakpoint resume support
"""

import asyncio
import hashlib
import json
import os
import time
import uuid
from datetime import datetime
from pathlib import Path
from typing import Optional, Dict, Any, List

import pandas as pd
import openpyxl
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, DecisionEngine
from core.address_processor import AddressProcessor  # legacy — kept for /api/decision endpoints
from core.contact_finder import ContactFinder
from core.degradation_manager import DegradationManager, DegradationLogger

# ─── App Setup ───────────────────────────────────────────
app = FastAPI(
    title="Wellfound AI - Excel Data Completion",
    description="AI-powered automated Excel data completion for Wellfound exports",
    version="3.0.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()
degradation_logger = DegradationLogger(str(DATA_DIR / "logs"))
degradation_manager = DegradationManager(degradation_logger)


# ─── 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("<h1>Wellfound AI</h1><p>Frontend not found.</p>")


@app.get("/api/health")
async def health():
    """Health check endpoint."""
    return {
        "status": "ok",
        "timestamp": datetime.now().isoformat(),
        "version": "3.0.0",
        "modules": {
            "excel": True,
            "scraper": True,
            "ai_extractor": True,
            "address_processor": True,
            "contact_finder": True,
            "decision_engine": True,
        },
    }


@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)

    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": [],
        "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,
        )
    )

    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"]),
        "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


@app.get("/api/degradation/summary")
async def degradation_summary():
    """Get degradation statistics for the current session."""
    return degradation_logger.get_session_summary()


@app.get("/api/degradation/recent")
async def degradation_recent(limit: int = 20):
    """Get recent degradation events."""
    return degradation_logger.get_recent_entries(limit)


# ─── Decision Engine Endpoints ───────────────────────────

@app.post("/api/decision/location")
async def decide_location(request: Request):
    """Determine best US location for a company."""
    data = await request.json()
    ai_data = data.get("ai_data", {})
    wellfound_location = data.get("wellfound_location", "")
    website_addresses = data.get("website_addresses", [])

    is_us = DecisionEngine.is_us_office(ai_data)
    city, state = DecisionEngine.determine_best_us_location(
        ai_data, wellfound_location, website_addresses
    )

    return {
        "has_us_office": is_us,
        "location_apply": city,
        "state_apply": state,
    }


@app.post("/api/decision/contact")
async def decide_contact(request: Request):
    """Rank contact methods by priority and return the best one."""
    data = await request.json()
    email = data.get("email")
    contact_form_url = data.get("contact_form_url")
    careers_page_url = data.get("careers_page_url")

    best = DecisionEngine.rank_contact_email(
        email, contact_form_url, careers_page_url
    )

    # Determine which priority level was selected
    if email:
        email_lower = email.lower()
        hr_patterns = [
            "careers@", "career@", "jobs@", "job@",
            "recruiting@", "recruit@", "talent@",
            "hr@", "hiring@", "people@",
        ]
        if any(p in email_lower for p in hr_patterns):
            priority = "hiring_email"
        else:
            priority = "general_email"
    elif contact_form_url:
        priority = "contact_form"
    elif careers_page_url:
        priority = "careers_page"
    else:
        priority = "none"

    return {
        "best_contact": best,
        "priority": priority,
    }


# ─── 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,
):
    """Main processing pipeline with deep multi-page website scanning."""
    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")

        # Ensure "备注" column exists in the output
        wb = openpyxl.load_workbook(result_path)
        ws = wb.active
        header_row = [cell.value for cell in ws[1]]
        if "备注" not in header_row:
            next_col = len(header_row) + 1
            ws.cell(row=1, column=next_col, value="备注")
            wb.save(result_path)
        wb.close()

        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:
                try:
                    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}"

                    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: Deep scan company website → AI directly decides
                    # location.apply, state.apply, contact
                    # Includes automatic degradation when AI is unavailable/failed
                    if use_web_scraping and external_link and external_link != "nan":
                        ai_result = await _deep_scan_company(
                            scraper, ai_extractor, company, external_link, location, idx
                        )
                        if ai_result:
                            remark = ai_result.get("remark")

                            # Contact
                            if pd.isna(row.get("contact")):
                                contact = ai_result.get("contact")
                                if contact:
                                    updates["contact"] = contact

                            # Location & State — AI decides directly
                            if pd.isna(row.get("location.apply")):
                                loc = ai_result.get("location_apply")
                                if loc:
                                    updates["location.apply"] = loc
                            if pd.isna(row.get("state.apply")):
                                state = ai_result.get("state_apply")
                                if state:
                                    updates["state.apply"] = state

                            # Inject remark if degradation occurred
                            if remark:
                                updates["备注"] = remark

                    # Save updates
                    if updates:
                        excel_handler.save_row(result_path, idx, updates)

                    processed += 1
                    task["progress"] = processed

                except Exception as e:
                    task["errors"].append(f"Row {idx}: {str(e)[:200]}")

        # 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 degradation summary for this processing run
        summary = degradation_logger.get_session_summary()
        if summary["total_degradations"] > 0:
            task["degradation_summary"] = summary
            degradation_logger._file_logger.info(
                "Processing complete | task=%s | total_degradations=%d | affected_rows=%s",
                task_id, summary["total_degradations"], summary["affected_rows"][:20],
            )

    except Exception as e:
        task["status"] = "error"
        task["errors"].append(f"Fatal: {str(e)}")

    finally:
        if scraper:
            await scraper.stop()


async def _deep_scan_company(
    scraper: WebScraper,
    ai_extractor: Optional[AIExtractor],
    company: str,
    url: str,
    wellfound_location: str = "",
    row_idx: int = 0,
) -> dict:
    """Deep scan a company website and determine location.apply, state.apply,
    and contact — with automatic degradation when AI fails.
    Pipeline:
    1. Fetch homepage → extract emails, links, phones, addresses, text
    2. Discover and follow key sub-pages (careers, contact, about, locations, team)
    3. Aggregate ALL content from all pages
    4. Send aggregated content to AI with the decision prompt → get final fields
    5. If AI fails → DegradationManager handles fallback:
       - contact_finder for email scoring + URL detection
       - address_processor for address parsing
       - "备注" column tagged with source: "non-ai" + failure reason
       - Error logged with full traceability
    Returns:
        Dict with: location_apply, state_apply, contact, remark (nullable).
    """
    result = {
        "location_apply": None,
        "state_apply": None,
        "contact": None,
        "remark": None,
    }

    try:
        # ── Step 1: Fetch homepage ──
        homepage = await scraper.fetch_company_website(url)
        if homepage.get("error"):
            # Scraping itself failed — no data for AI or fallback
            await degradation_manager.handle_degradation(
                row_idx=row_idx,
                company=company,
                ai_extractor=ai_extractor,
                scraped_data={"base_url": url},
                wellfound_location=wellfound_location,
            )
            return result

        homepage_emails = homepage.get("emails", [])
        homepage_links = homepage.get("links", [])
        homepage_phones = homepage.get("phones", [])
        homepage_addresses = homepage.get("addresses", [])
        homepage_text = homepage.get("text", "")

        all_emails = list(homepage_emails)
        all_addresses = list(homepage_addresses)
        all_text_parts = [homepage_text]

        # ── Step 2: Discover and visit sub-pages ──
        sub_pages_to_visit = []
        visited_urls = set()

        for link in homepage_links:
            link_url = link.get("url", "")
            link_text = (link.get("text", "") or "").lower()
            href_lower = link_url.lower()
            resolved = contact_finder._resolve_url(url, link_url)

            if resolved in visited_urls or resolved == url:
                continue
            if not resolved.startswith(("http://", "https://")):
                continue
            if any(domain in resolved for domain in [
                "facebook.com", "twitter.com", "x.com", "linkedin.com",
                "instagram.com", "youtube.com", "github.com",
            ]):
                continue

            visited_urls.add(resolved)

            if any(kw in href_lower or kw in link_text for kw in
                   ["career", "job", "join", "work-with", "opening", "position", "talent"]):
                sub_pages_to_visit.append(("careers", resolved))
            elif any(kw in href_lower or kw in link_text for kw in
                     ["contact", "reach", "get-in-touch", "connect"]):
                sub_pages_to_visit.append(("contact", resolved))
            elif any(kw in href_lower or kw in link_text for kw in
                     ["about", "company", "our-story", "who-we-are"]):
                sub_pages_to_visit.append(("about", resolved))
            elif any(kw in href_lower or kw in link_text for kw in
                     ["team", "people", "leadership"]):
                sub_pages_to_visit.append(("team", resolved))
            elif any(kw in href_lower or kw in link_text for kw in
                     ["location", "office", "headquarter"]):
                sub_pages_to_visit.append(("location", resolved))

        page_priority = {"careers": 0, "contact": 1, "about": 2, "location": 3, "team": 4}
        sub_pages_to_visit.sort(key=lambda x: page_priority.get(x[0], 99))

        for page_type, page_url in sub_pages_to_visit[:4]:
            try:
                sub_page = await scraper.fetch_page(page_url)
                if sub_page.get("text") and not sub_page.get("error"):
                    all_text_parts.append(sub_page["text"])
                    all_emails.extend(sub_page.get("emails", []))
                    sub_addrs = sub_page.get("addresses", [])
                    if sub_addrs:
                        all_addresses.extend(sub_addrs)
            except Exception:
                continue

        # Deduplicate emails
        seen = set()
        unique_emails = []
        for e in all_emails:
            el = e.lower()
            if el not in seen:
                seen.add(el)
                unique_emails.append(e)

        # Deduplicate addresses
        seen_addrs = set()
        unique_addrs = []
        for a in all_addresses:
            al = a.strip().lower()
            if al not in seen_addrs:
                seen_addrs.add(al)
                unique_addrs.append(a)

        # ── Step 3: AI deep analysis ──
        combined_text = "\n\n---\n\n".join(all_text_parts)

        # Build scraped data dict for degradation manager
        scraped_data = {
            "emails": unique_emails,
            "links": homepage_links,
            "addresses": unique_addrs,
            "html": homepage.get("html", ""),
            "base_url": url,
            "combined_text": combined_text,
        }

        if ai_extractor and combined_text:
            try:
                ai_data = await ai_extractor.analyze_company_website(
                    company_name=company,
                    page_text=combined_text,
                    emails=unique_emails,
                    links=homepage_links,
                    phones=homepage_phones,
                    addresses=unique_addrs,
                    wellfound_location=wellfound_location,
                )

                # Check if AI returned useful results
                failure_type = degradation_manager.detect_failure(
                    ai_extractor=ai_extractor,
                    ai_result=ai_data,
                )

                if failure_type is None:
                    # AI succeeded — use its results directly
                    result["location_apply"] = ai_data.get("location_apply")
                    result["state_apply"] = ai_data.get("state_apply")
                    result["contact"] = ai_data.get("contact")
                    return result

                # AI returned empty — trigger degradation
                degraded = await degradation_manager.handle_degradation(
                    row_idx=row_idx,
                    company=company,
                    ai_extractor=ai_extractor,
                    ai_result=ai_data,
                    scraped_data=scraped_data,
                    wellfound_location=wellfound_location,
                )
                return degraded

            except Exception as exc:
                # AI call failed with exception — trigger degradation
                degraded = await degradation_manager.handle_degradation(
                    row_idx=row_idx,
                    company=company,
                    ai_extractor=ai_extractor,
                    exc=exc,
                    scraped_data=scraped_data,
                    wellfound_location=wellfound_location,
                )
                return degraded

        # ── Step 4: AI not enabled — use degradation directly ──
        degraded = await degradation_manager.handle_degradation(
            row_idx=row_idx,
            company=company,
            ai_extractor=ai_extractor,  # None if not enabled
            scraped_data=scraped_data,
            wellfound_location=wellfound_location,
        )
        return degraded

    except Exception as outer_exc:
        # Outer exception (e.g., scraper crash) — still log
        try:
            await degradation_manager.handle_degradation(
                row_idx=row_idx,
                company=company,
                ai_extractor=ai_extractor,
                exc=outer_exc,
                scraped_data={"base_url": url},
                wellfound_location=wellfound_location,
            )
        except Exception:
            pass  # Don't let logging failures crash the pipeline

    return result


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:
        pass  # Silently fail on individual row errors

    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"{'='*50}\n")

    uvicorn.run(
        app,
        host="0.0.0.0",
        port=port,
        log_level="info",
        timeout_keep_alive=30,
    )