# app.py import os import io import json import time import traceback from typing import List from io import BytesIO from fastapi import FastAPI, UploadFile, File, Form, HTTPException from fastapi.middleware.cors import CORSMiddleware from dotenv import load_dotenv import PyPDF2 import openai import requests from pdf2image import convert_from_bytes from google.cloud import vision_v1 as vision # ───── Load .env (locally) ──────────────────────────────────── # In HF Spaces you'll set these as Secrets under Settings → Variables & secrets load_dotenv() openai.api_key = os.getenv("OPENAI_API_KEY") gcv_api_key = os.getenv("GCV_API_KEY") brightdata_api_key = os.getenv("BRIGHTDATA_API_KEY") brightdata_dataset_id = "gd_l1viktl72bvl7bjuj0" if not openai.api_key: raise RuntimeError("Missing OPENAI_API_KEY") if not gcv_api_key: raise RuntimeError("Missing GCV_API_KEY") if not brightdata_api_key: raise RuntimeError("Missing BRIGHTDATA_API_KEY") # ───── FastAPI setup ────────────────────────────────────────── app = FastAPI(title="Aliro Data Extraction API") app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) # ───── A simple root so GET / won't 404 ─────────────────────── @app.get("/") def read_root(): return {"message": "Aliro Data Extraction API – POST your files to /extract"} # ───── PDF/OCR processing ───────────────────────────────────── def process_buffers(buffers: List[BytesIO]) -> List[str]: client = vision.ImageAnnotatorClient(client_options={"api_key": gcv_api_key}) results: List[str] = [] for buf in buffers: name = getattr(buf, "name", "").lower() data = buf.getvalue() # PDF: try text → fallback to page-by-page OCR if name.endswith(".pdf"): # 1) PyPDF2 text try: reader = PyPDF2.PdfReader(io.BytesIO(data)) text = "".join(page.extract_text() or "" for page in reader.pages) if text.strip(): results.append(text) continue except: pass # 2) Fallback: render pages → OCR try: for img in convert_from_bytes(data): img_buf = io.BytesIO() img.save(img_buf, format="JPEG") resp = client.text_detection(image=vision.Image(content=img_buf.getvalue())) desc = resp.text_annotations[0].description if resp.text_annotations else "" if desc: results.append(desc) continue except: pass # Image → always OCR if name.endswith((".png", ".jpg", ".jpeg")): try: resp = client.text_detection(image=vision.Image(content=data)) desc = resp.text_annotations[0].description if resp.text_annotations else "" if desc: results.append(desc) continue except: pass # Otherwise unsupported results.append(f"[Unsupported file type: {name}]") return results # ───── LinkedIn scrape via BrightData ───────────────────────── def scrape_linkedin(url: str) -> str: if not url: return "" try: # Step 1: Trigger data collection trigger_url = "https://api.brightdata.com/datasets/v3/trigger" headers = { "Authorization": f"Bearer {brightdata_api_key}", "Content-Type": "application/json", } params = { "dataset_id": brightdata_dataset_id, "include_errors": "true", } trigger_response = requests.post( trigger_url, headers=headers, params=params, json=[{"url": url}], timeout=30 ) if not trigger_response.ok: raise Exception(f"BrightData trigger failed: {trigger_response.status_code} - {trigger_response.text}") trigger_data = trigger_response.json() snapshot_id = trigger_data.get('snapshot_id') if not snapshot_id: raise Exception("No snapshot_id received from BrightData") print(f"BrightData collection triggered, snapshot_id: {snapshot_id}") # Step 2: Poll for completion (max 5 minutes) max_attempts = 30 for attempt in range(max_attempts): time.sleep(10) # Wait 10 seconds between checks progress_url = f"https://api.brightdata.com/datasets/v3/progress/{snapshot_id}" progress_response = requests.get(progress_url, headers=headers, timeout=30) if progress_response.ok: progress_data = progress_response.json() status = progress_data.get('status') print(f"Progress check {attempt + 1}: status = {status}") if status == 'ready': # Step 3: Fetch results result_url = f"https://api.brightdata.com/datasets/v3/snapshot/{snapshot_id}" result_params = {"format": "json"} result_response = requests.get(result_url, headers=headers, params=result_params, timeout=30) if result_response.ok: linkedin_data = result_response.json() return json.dumps(linkedin_data, indent=2) else: raise Exception(f"Failed to fetch results: {result_response.status_code} - {result_response.text}") elif status == 'failed': raise Exception("BrightData collection failed") # Continue polling if status is 'running' else: print(f"Progress check failed: {progress_response.status_code}") raise Exception("Timeout waiting for BrightData collection to complete") except Exception as e: raise Exception(f"LinkedIn scraping error: {str(e)}") # ───── Summarize via OpenAI ───────────────────────────────── def make_summary(chunks: List[str]) -> str: if not chunks: return "No data extracted." prompt = "\n\n---\n\n".join(chunks) resp = openai.ChatCompletion.create( model="gpt-4o", messages=[ { "role": "system", "content": ( "Extract and structure personal, educational, and professional details " "into a clear, logical and relational hierarchy. Translate non-English " "content, remove duplicates, and output in English." ) }, {"role": "user", "content": f"Summarize the following data:\n\n{prompt}"}, ], max_tokens=1500, ) return resp.choices[0].message.content.strip() # ───── The POST /extract endpoint ───────────────────────────── @app.post("/extract") async def extract_endpoint( files: List[UploadFile] = File(default=[]), linkedin_url: str = Form(default="") ): try: # Wrap each UploadFile in a BytesIO (so we can peek at .name & .getvalue()) buffers: List[BytesIO] = [] for f in files: data = await f.read() bio = BytesIO(data) bio.name = f.filename buffers.append(bio) # 1) OCR/PDF text texts = process_buffers(buffers) # 2) LinkedIn JSON if linkedin_url: texts.append(scrape_linkedin(linkedin_url)) # 3) Summarize summary = make_summary(texts) return {"summary": summary} except requests.HTTPError as e: raise HTTPException(status_code=502, detail=f"LinkedIn scrape failed: {e}") except Exception as e: raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")