Enrich Lead

# --- Parameters (you can change these) ---
# Only name+company/university OR linkedin is needed
# You can also give the agent more context following the structure below
LEAD_NAME = "Saarth Shah"
LEAD_COMPANY = "Sixtyfour AI"
LEAD_EMAIL = None
LEAD_LINKEDIN = None

import os, json, time, csv, io
from datetime import datetime
from typing import Dict, Any, List, Optional

import requests
import pandas as pd

API_KEY = os.getenv("SIXTYFOUR_API_KEY", "")  # fallback if already set locally
try:
    from google.colab import userdata
    secret_val = userdata.get("SIXTYFOUR_API_KEY")  # raises if not granted or not created
    if secret_val:  # avoid overwriting with empty
        os.environ["SIXTYFOUR_API_KEY"] = secret_val
        API_KEY = secret_val
except Exception:
    # Not in Colab or secret not available; rely on existing env var
    pass

BASE_URL = "https://api.sixtyfour.ai"
SYNC_ENDPOINT = f"{BASE_URL}/enrich-lead"
ASYNC_ENDPOINT = f"{BASE_URL}/enrich-lead-async"
STATUS_ENDPOINT = f"{BASE_URL}/job-status"
TIMEOUT_S = 600  # We reuse these endpoints and TIMEOUT_S throughout the notebook.

print("✅ Setup complete (key detected, endpoints ready).")

lead_info = {
    "name": LEAD_NAME,
    "company": LEAD_COMPANY,
    "email": LEAD_EMAIL,
    "linkedin": LEAD_LINKEDIN,
}

struct = {
    "name": "The individual's full name",
    "email": "The individual's email address",
    "phone": "The individual's phone number",
    "company": "The company the individual is associated with",
    "title": "The individual's job title",
    "linkedin": "LinkedIn URL for the person",
    "website": "Company website URL",
    "location": "The individual's location and/or company location",
    "industry": "Industry the person operates in",
    "github_url": "url for their github profile",
    "github_notes": "Take detailed notes on their github profile."
}

research_plan = "Review LinkedIn, company website, and public profiles for contact info and role verification."

payload = {
    "lead_info": lead_info,
    "struct": struct,
    "research_plan": research_plan,
}

response = requests.post(
    SYNC_ENDPOINT,
    headers={"x-api-key": API_KEY, "Content-Type": "application/json"},
    json=payload,
    timeout=TIMEOUT_S,
)

if response.ok:
    result = response.json()
    print("✅ Success! Keys:", list(result.keys()))
    print("\n" + "="*50)
    print("STRUCTURED DATA:")
    print("="*50)
    structured_data = result.get("structured_data", {})
    for key, value in structured_data.items():
        print(f"{key}: {value}")
else:
    print("❌ Request failed:", response.status_code, response.text[:800])

# Start async job
start_resp = requests.post(
    ASYNC_ENDPOINT,
    headers={"x-api-key": API_KEY, "Content-Type": "application/json"},
    json=payload,
    timeout=30,
)
start_resp.raise_for_status()
task_id = start_resp.json()["task_id"]
print(f"🔄 Started async job: {task_id}")

# Poll for completion
while True:
    status_resp = requests.get(
        f"{STATUS_ENDPOINT}/{task_id}",
        headers={"x-api-key": API_KEY}
    )
    status = status_resp.json()
    
    if status["status"] == "completed":
        async_result = status["result"]
        print("✅ Async job completed!")
        break
    elif status["status"] == "failed":
        raise RuntimeError(status.get("error", "Async job failed"))
    else:
        print(f"⏳ Status: {status['status']}")
        time.sleep(5)

# Display results
print("\nAsync Results:")
print("="*50)
async_structured_data = async_result.get("structured_data", {})
for key, value in async_structured_data.items():
    print(f"{key}: {value}")