Update app.py

app.py

@@ -2,6 +2,7 @@
 import os
 import io
 import json
+import time
 import traceback
 from typing import List
 from io import BytesIO
@@ -19,16 +20,17 @@ 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")
-proxycurl_api_key = os.getenv("PROXYCURL_API_KEY")
+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 proxycurl_api_key:
-    raise RuntimeError("Missing PROXYCURL_API_KEY")
+if not brightdata_api_key:
+    raise RuntimeError("Missing BRIGHTDATA_API_KEY")
 
 # ───── FastAPI setup ──────────────────────────────────────────
 app = FastAPI(title="Aliro Data Extraction API")
@@ -39,7 +41,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ───── A simple root so GET / won’t 404 ───────────────────────
+# ───── A simple root so GET / won't 404 ───────────────────────
 @app.get("/")
 def read_root():
     return {"message": "Aliro Data Extraction API – POST your files to /extract"}
@@ -94,18 +96,79 @@ def process_buffers(buffers: List[BytesIO]) -> List[str]:
 
     return results
 
-# ───── LinkedIn scrape via Proxycurl ─────────────────────────
+# ───── LinkedIn scrape via BrightData ─────────────────────────
 def scrape_linkedin(url: str) -> str:
     if not url:
         return ""
-    resp = requests.get(
-        "https://nubela.co/proxycurl/api/v2/linkedin",
-        params={"linkedin_profile_url": url},
-        headers={"Authorization": f"Bearer {proxycurl_api_key}"},
-        timeout=10
-    )
-    resp.raise_for_status()
-    return json.dumps(resp.json(), indent=2)
+        
+    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:
@@ -155,6 +218,5 @@ async def extract_endpoint(
 
     except requests.HTTPError as e:
         raise HTTPException(status_code=502, detail=f"LinkedIn scrape failed: {e}")
-    except Exception:
-        tb = traceback.format_exc()
-        raise HTTPException(status_code=500, detail=f"Processing error:\n\n{tb}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")