scripts/ai_enrichment.py

import os
import json
import time
from typing import List, Union, Any
from pydantic import BaseModel, Field, validator
from supabase import create_client, Client
from groq import Groq
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from dotenv import load_dotenv
load_dotenv('/workspaces/govbridge/.env')

from config import settings

# --- SUPER-ADVANCED ENGINEERING: TYPE-SAFE NORMALIZATION ---

# Initialize clients
sb: Client = create_client(settings.SUPABASE_URL, settings.SUPABASE_KEY)

api_key = os.environ.get("GROQ_API_KEY")
if not api_key or api_key == "your_groq_api_key_here":
    print("❌ ERROR: GROQ_API_KEY not found. Ensure .env is updated.")
    exit(1)

client = Groq(api_key=api_key)

class AIResponseModel(BaseModel):
    """Strict schema for AI Output with automatic type normalization."""
    summary: str = Field(default="Detailed government document.")
    benefits: Union[str, List[str]] = Field(default="")
    eligibility: Union[str, List[str]] = Field(default="")

    @validator('benefits', 'eligibility', pre=True, always=True)
    def normalize_to_string(cls, v):
        """Converts lists or bullet points into professional strings."""
        if not v or v == "N/A" or v == "None":
            return ""
        if isinstance(v, list):
            # filter out empty strings or N/A
            clean_list = [str(item).strip() for item in v if item and str(item).strip().lower() not in ["n/a", "none", "null"]]
            if not clean_list: return ""
            return "\n".join([f"• {item}" for item in clean_list])
        return str(v).strip()

def get_ai_extraction(text: str, category: str = "") -> AIResponseModel:
    """Robust AI extraction with isolated error handling."""
    prompt = f"""
    SYSTEM: You are a GovBridge AI Data Architect. Extract structured data from this document.
    DOCUMENT CATEGORY: {category}
    DOCUMENT TEXT: {text[:6000]}

    INSTRUCTIONS:
    1. Extract the core summary of the document (2-4 sentences).
    2. List the benefits (as a JSON array of strings). IF THIS IS A PRESS RELEASE OR NO EXPLICIT BENEFITS EXIST, return an empty array []. DO NOT hallucinate.
    3. List eligibility criteria (as a JSON array of strings). IF THIS IS A PRESS RELEASE OR NO EXPLICIT CRITERIA EXIST, return an empty array []. DO NOT hallucinate.

    FORMAT: JSON ONLY.
    {{
      "summary": "...",
      "benefits": ["...", "..."],
      "eligibility": ["...", "..."]
    }}
    """
    try:
        completion = client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model="llama-3.3-70b-versatile",
            response_format={"type": "json_object"}
        )
        content = completion.choices[0].message.content
        if not content:
            return AIResponseModel()
        raw_json = json.loads(content)
        # Pass through Pydantic for automatic cleanup/normalization
        return AIResponseModel(**raw_json)
    except Exception as e:
        print(f"   ⚠️ AI Parsing Error: {e}")
        return AIResponseModel() # Return safe defaults

def run_enrichment():
    print("\n🚀 LAUNCHING WORLD NO. 1 AI ENRICHMENT ENGINE (v2.0 - BULLETPROOF)")
    print("------------------------------------------------------------------")
    
    # 1. Target unverified schemes or schemes with missing/empty benefits
    schemes = sb.table('schemes').select('id, title, source_url, category').eq('is_verified', False).execute()
    total = len(schemes.data)
    print(f"📊 Pipeline detected {total} schemes requiring deep enrichment.")

    import httpx
    from bs4 import BeautifulSoup
    import re

    for idx, s in enumerate(schemes.data, 1):
        print(f"[{idx}/{total}] Processing: {s['title']}...")
        
        try:
            # 2. Extract context from raw chunks
            chunks = sb.table('document_chunks').select('chunk_text').eq('scheme_title', s['title']).execute()
            raw_text = " ".join([c['chunk_text'] for c in chunks.data if c.get('chunk_text')])
            
            # 3. If no chunks, try fetching directly from the source_url
            if not raw_text.strip() and s.get('source_url'):
                try:
                    print(f"   🌐 Fetching raw content from {s['source_url'][:50]}...")
                    with httpx.Client(timeout=15.0, follow_redirects=True) as client:
                        resp = client.get(s['source_url'])
                        if resp.status_code == 200:
                            soup = BeautifulSoup(resp.text, 'html.parser')
                            for script in soup(["script", "style", "nav", "footer"]):
                                script.decompose()
                            raw_text = soup.get_text(separator=' ', strip=True)
                except Exception as e:
                    print(f"   ⚠️ Could not fetch source URL: {e}")

            if not raw_text.strip():
                raw_text = f"General information about the government scheme: {s['title']}"

            # 4. AI Extraction & Safe Normalization
            data = get_ai_extraction(raw_text, category=s.get('category', ''))
            
            # 5. Atomic Database Update
            update_data = {
                'summary': data.summary,
                'benefits': data.benefits,
                'eligibility_text': data.eligibility,
                'full_description': f"{data.summary}\n\nKey Benefits:\n{data.benefits}",
                'is_verified': True
            }
            
            sb.table('schemes').update(update_data).eq('id', s['id']).execute()
            print(f"   ✅ SUCCESS: Data normalized and persisted.")
            
            # Polite delay to respect API limits
            time.sleep(0.5)

        except Exception as loop_err:
            print(f"   ❌ FATAL FOR THIS RECORD: {loop_err}")
            continue

    print("\n------------------------------------------------------------------")
    print("🏁 PIPELINE COMPLETE. ALL SCHEMES ARE NOW DATA-RICH.")

if __name__ == "__main__":
    run_enrichment()