core/ai_extractor.py

"""
AI Extractor - Intelligent data extraction using OpenAI/DeepSeek API.

Handles:
- Deep page content analysis for funding data
- Contact information extraction and validation
- Address parsing and location inference
- Structured data extraction from unstructured text
"""

import json
import re
from typing import Optional, Dict, Any, List

import httpx


class AIExtractor:
    """AI-powered data extraction using LLM APIs."""

    # API endpoints
    PROVIDERS = {
        "openai": {
            "url": "https://api.openai.com/v1/chat/completions",
            "models": ["gpt-4o", "gpt-4o-mini", "gpt-4-turbo", "gpt-3.5-turbo"],
        },
        "deepseek": {
            "url": "https://api.deepseek.com/v1/chat/completions",
            "models": ["deepseek-chat", "deepseek-reasoner"],
        },
    }

    def __init__(self, provider: str = "openai", api_key: str = "", model: str = "auto"):
        if provider not in self.PROVIDERS:
            raise ValueError(f"Unknown provider: {provider}. Use 'openai' or 'deepseek'")

        self.provider = provider
        self.api_key = api_key
        self.base_url = self.PROVIDERS[provider]["url"]

        if model == "auto":
            self.model = self.PROVIDERS[provider]["models"][0]
        else:
            self.model = model

    async def analyze_funding_page(
        self, company_name: str, page_text: str, meta: Dict = None
    ) -> Dict[str, Any]:
        """Extract funding information from a Wellfound-like page."""
        # Truncate text to avoid token limits
        text_sample = page_text[:8000] if page_text else ""

        system_prompt = """You are a precise data extraction AI. Extract funding information from company profile pages.
Return ONLY valid JSON with these fields. Use null for unknown values. Do NOT guess - only extract when confident.

Output format:
{
  "valuation": "string or null - company valuation like '$50M', '$2.5B'",
  "total_raised": "string or null - total funding raised like '$10M'",
  "rounds": "integer or null - number of funding rounds",
  "series": "string or null - latest series like 'Series A', 'Seed', 'Series B'",
  "investors": ["string array - key investors mentioned"],
  "confidence": {
    "valuation": "high|medium|low",
    "total_raised": "high|medium|low", 
    "rounds": "high|medium|low",
    "series": "high|medium|low"
  }
}"""

        user_prompt = f"""Company: {company_name}

Page content:
{text_sample}

Extract funding information from this company profile. Be precise and only return data you can confidently identify."""

        response = await self._call_api(system_prompt, user_prompt)
        return self._parse_json_response(response, "funding")

    async def analyze_company_website(
        self, company_name: str, page_text: str, emails: List[str],
        links: List[Dict], phones: List[str], addresses: List[str]
    ) -> Dict[str, Any]:
        """Extract contact information and address from company website."""
        text_sample = page_text[:10000] if page_text else ""

        system_prompt = """You are a precise contact information extraction AI.
Extract contact details from company website content. Return ONLY valid JSON.

Output format:
{
  "contact_email": "string or null - BEST email for job applications (HR, careers, recruiting, talent). Prioritize: careers@, hr@, jobs@, recruiting@, talent@, people@ over general info@ or contact@",
  "contact_form_url": "string or null - URL of contact/careers form if found",
  "careers_page_url": "string or null - URL of careers/jobs page",
  "headquarters_address": "string or null - full HQ address",
  "headquarters_city": "string or null",
  "headquarters_state": "string or null - US state abbreviation or full name",
  "headquarters_country": "string or null",
  "us_office_city": "string or null - if company has US offices, city of primary US office",
  "us_office_state": "string or null - state of primary US office",
  "other_locations": ["string array - other office locations mentioned"],
  "hiring_focus_location": "string or null - location mentioned most in context of hiring/jobs",
  "contact_confidence": "high|medium|low"
}"""

        user_prompt = f"""Company: {company_name}

Page text content:
{text_sample}

Found emails on page: {json.dumps(emails)}
Contact-related links: {json.dumps(links[:15])}
Phone numbers: {json.dumps(phones[:5])}
Addresses found: {json.dumps(addresses[:5])}

Extract contact information. For contact_email, prioritize HR/recruiting emails over generic ones.
For US office location, if the company has US presence but HQ is elsewhere, find the US office.
Look for location patterns in job listings or 'join us' sections for hiring_focus_location."""

        response = await self._call_api(system_prompt, user_prompt)
        return self._parse_json_response(response, "contact")

    async def verify_and_enhance(
        self, company_name: str, scraped_data: Dict, page_text: str
    ) -> Dict[str, Any]:
        """Verify and enhance extracted data with AI reasoning."""
        text_sample = page_text[:6000] if page_text else ""

        system_prompt = """You are a data verification AI. Review scraped funding data against page content.
Verify accuracy and fill in missing fields when possible. Return ONLY valid JSON.

Output format:
{
  "verified_data": {
    "valuation": "confirmed value or null",
    "total_raised": "confirmed value or null", 
    "rounds": "confirmed integer or null",
    "series": "confirmed value or null"
  },
  "corrections": ["list of corrections made"],
  "verification_notes": "brief note about data confidence"
}"""

        user_prompt = f"""Company: {company_name}

Scraped data: {json.dumps(scraped_data)}

Page content snippet:
{text_sample}

Verify the scraped data against the page content. Correct any errors and fill in missing values where possible."""

        response = await self._call_api(system_prompt, user_prompt)
        return self._parse_json_response(response, "verify")

    async def _call_api(self, system_prompt: str, user_prompt: str) -> str:
        """Make API call to the LLM provider."""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        payload = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            "temperature": 0.1,
            "max_tokens": 1500,
            "response_format": {"type": "json_object"} if self.provider == "openai" else None,
        }

        async with httpx.AsyncClient(timeout=60.0) as client:
            response = await client.post(self.base_url, json=payload, headers=headers)

            if response.status_code != 200:
                raise Exception(f"API error {response.status_code}: {response.text[:500]}")

            data = response.json()
            return data["choices"][0]["message"]["content"]

    def _parse_json_response(self, response: str, extract_type: str) -> Dict[str, Any]:
        """Parse JSON from API response with fallback extraction."""
        # Try direct JSON parse
        try:
            return json.loads(response)
        except json.JSONDecodeError:
            pass

        # Try extracting JSON block
        json_match = re.search(r'\{[\s\S]*\}', response)
        if json_match:
            try:
                return json.loads(json_match.group(0))
            except json.JSONDecodeError:
                pass

        # Return empty structure on failure
        if extract_type == "funding":
            return {
                "valuation": None, "total_raised": None,
                "rounds": None, "series": None, "investors": [],
                "confidence": {"valuation": "low", "total_raised": "low",
                               "rounds": "low", "series": "low"}
            }
        elif extract_type == "contact":
            return {
                "contact_email": None, "contact_form_url": None,
                "careers_page_url": None, "headquarters_address": None,
                "headquarters_city": None, "headquarters_state": None,
                "headquarters_country": None, "us_office_city": None,
                "us_office_state": None, "other_locations": [],
                "hiring_focus_location": None, "contact_confidence": "low"
            }
        else:
            return {"verified_data": {}, "corrections": [], "verification_notes": ""}