ingestion/knowledge_extractor.py

"""
GovBridge India — Two-Pass Knowledge Triplet Extraction Engine
Sprint 30: PROJECT INDRA Phase 1.5 — Cascade Intelligence Ingestion

ARCHITECTURE:
  Pass 1 (Global Context): Send first 2000 tokens to Groq to extract
    document title, primary body, and key definitions. This context is
    prepended to every chunk in Pass 2.

  Pass 2 (Reduce): Chunk the document into 800-1200 char windows with
    150-char overlap. Fire all extraction tasks concurrently via
    asyncio.gather() with a Semaphore(GROQ_MAX_CONCURRENCY) gate.
    
  Aggregation: Deduplicate triplets by normalized (subject, predicate, object)
    tuple. Validate all predicates against the 10-predicate ontology.

CONCURRENCY MODEL:
  - asyncio.Semaphore(settings.GROQ_MAX_CONCURRENCY) prevents Groq 429s
  - Exponential backoff (base 2s, max 32s, 4 retries) on rate limit errors
  - Individual chunk failures are logged and skipped (no pipeline crash)

DEPENDENCIES (all already in requirements.txt):
  - groq (async client)
  - pydantic (response validation)
  - asyncio (concurrency)
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import asyncio
import json
import logging
import re
from typing import Optional

from groq import AsyncGroq, RateLimitError, APIStatusError
from pydantic import BaseModel, field_validator, ValidationError

from prompts import (
    GLOBAL_CONTEXT_SYSTEM_PROMPT,
    TRIPLET_EXTRACTION_SYSTEM_PROMPT,
    VALID_PREDICATES,
    build_chunk_user_prompt,
)

logger = logging.getLogger("govbridge.knowledge_extractor")


# ═══════════════════════════════════════════════════════════════
# SECTION 1: Pydantic Response Models (Type Safety Layer)
# ═══════════════════════════════════════════════════════════════

class KnowledgeTriplet(BaseModel):
    """A single validated knowledge triplet."""
    subject: str
    predicate: str
    object: str

    @field_validator("predicate")
    @classmethod
    def validate_predicate(cls, v: str) -> str:
        """Enforce the 10-predicate ontology. Reject hallucinated predicates."""
        normalized = v.strip().upper()
        if normalized not in VALID_PREDICATES:
            raise ValueError(
                f"Invalid predicate '{v}'. Must be one of: {', '.join(sorted(VALID_PREDICATES))}"
            )
        return normalized

    @field_validator("subject", "object")
    @classmethod
    def validate_entity(cls, v: str) -> str:
        """Strip whitespace and reject empty entities."""
        cleaned = v.strip()
        if not cleaned or len(cleaned) < 2:
            raise ValueError(f"Entity too short or empty: '{v}'")
        return cleaned


class ExtractionResponse(BaseModel):
    """Validated LLM extraction response with cognitive exhaust."""
    _reasoning: str = ""
    triplets: list[KnowledgeTriplet] = []


class GlobalContext(BaseModel):
    """Validated global context from Pass 1."""
    document_title: Optional[str] = None
    primary_body: Optional[str] = None
    key_definitions: list[str] = []
    document_date: Optional[str] = None
    document_type: Optional[str] = None


# ═══════════════════════════════════════════════════════════════
# SECTION 2: Text Chunking (Token-Aware Sliding Window)
# ═══════════════════════════════════════════════════════════════

def chunk_text_for_extraction(
    full_text: str,
    chunk_size: int = 1000,
    overlap: int = 150,
) -> list[str]:
    """
    Split document text into overlapping chunks for LLM extraction.
    
    Uses character-level windowing (not token-level) because:
    1. Groq's tokenizer is not locally available
    2. 1000 chars ≈ 250-300 tokens for English/Hindi legal text
    3. Character windowing is deterministic and allocation-free
    
    Args:
        full_text: Complete document text
        chunk_size: Target characters per chunk (800-1200 range)
        overlap: Character overlap to prevent severing cross-boundary relationships
    
    Returns:
        List of text chunks with overlap applied
    """
    if len(full_text) <= chunk_size:
        return [full_text]
    
    chunks: list[str] = []
    start = 0
    
    while start < len(full_text):
        end = start + chunk_size
        
        # Try to break at a paragraph or sentence boundary
        if end < len(full_text):
            # Look for paragraph break within last 200 chars
            para_break = full_text.rfind("\n\n", start + chunk_size - 200, end)
            if para_break > start:
                end = para_break
            else:
                # Fall back to sentence boundary (period + space)
                sentence_break = full_text.rfind(". ", start + chunk_size - 200, end)
                if sentence_break > start:
                    end = sentence_break + 1  # Include the period
        
        chunk = full_text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        
        # Advance with overlap
        start = end - overlap if end < len(full_text) else len(full_text)
    
    return chunks


# ═══════════════════════════════════════════════════════════════
# SECTION 3: Groq API Communication Layer
# ═══════════════════════════════════════════════════════════════

async def _call_groq_with_retry(
    client: AsyncGroq,
    model: str,
    system_prompt: str,
    user_prompt: str,
    semaphore: asyncio.Semaphore,
    max_retries: int = 4,
    base_delay: float = 2.0,
    max_delay: float = 32.0,
) -> Optional[dict]:
    """
    Call Groq API with semaphore-gated concurrency and exponential backoff.
    
    Returns parsed JSON dict on success, None on unrecoverable failure.
    
    Rate Limit Strategy:
      - Semaphore prevents more than N concurrent requests
      - On 429 (rate limit): exponential backoff with jitter
      - On 400/500: log and skip (bad request or server error)
      - On JSON parse failure: log and skip (hallucinated output)
    """
    delay = base_delay
    
    for attempt in range(max_retries + 1):
        async with semaphore:
            try:
                response = await client.chat.completions.create(
                    model=model,
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt},
                    ],
                    response_format={"type": "json_object"},
                    temperature=0.1,
                    max_tokens=2048,
                    stream=False,
                )
                
                raw_content = response.choices[0].message.content or ""
                
                # Parse JSON — the response_format should guarantee valid JSON,
                # but we defend against edge cases.
                try:
                    return json.loads(raw_content)
                except json.JSONDecodeError as e:
                    logger.warning(
                        f"JSON parse failed (attempt {attempt + 1}): {e}\n"
                        f"Raw output: {raw_content[:200]}"
                    )
                    return None
                    
            except RateLimitError:
                if attempt < max_retries:
                    # Add jitter: delay * (0.5 to 1.5)
                    import random
                    jittered = delay * (0.5 + random.random())
                    logger.warning(
                        f"Groq 429 rate limit (attempt {attempt + 1}/{max_retries}). "
                        f"Backing off {jittered:.1f}s"
                    )
                    await asyncio.sleep(jittered)
                    delay = min(delay * 2, max_delay)
                else:
                    logger.error("Groq rate limit exceeded after all retries")
                    return None
                    
            except APIStatusError as e:
                logger.error(f"Groq API error {e.status_code}: {e.message}")
                return None
                
            except Exception as e:
                logger.error(f"Unexpected error calling Groq: {e}")
                return None
    
    return None


# ═══════════════════════════════════════════════════════════════
# SECTION 4: Two-Pass Extraction Pipeline
# ═══════════════════════════════════════════════════════════════

async def extract_global_context(
    client: AsyncGroq,
    model: str,
    full_text: str,
    semaphore: asyncio.Semaphore,
) -> GlobalContext:
    """
    Pass 1: Extract document-level metadata from the first 2000 characters.
    
    This context is prepended to every chunk in Pass 2 to help the LLM
    resolve ambiguous references like "this Act" or "the said Ministry".
    """
    # Take first 2000 chars (roughly first 500 tokens)
    preamble = full_text[:2000]
    
    result = await _call_groq_with_retry(
        client=client,
        model=model,
        system_prompt=GLOBAL_CONTEXT_SYSTEM_PROMPT,
        user_prompt=preamble,
        semaphore=semaphore,
    )
    
    if result is None:
        logger.warning("Global context extraction failed — using empty context")
        return GlobalContext()
    
    try:
        return GlobalContext(**result)
    except ValidationError as e:
        logger.warning(f"Global context validation failed: {e}")
        # Partial extraction: take what we can
        return GlobalContext(
            document_title=result.get("document_title"),
            primary_body=result.get("primary_body"),
        )


async def extract_triplets_from_chunk(
    client: AsyncGroq,
    model: str,
    global_context: dict,
    chunk_text: str,
    chunk_index: int,
    semaphore: asyncio.Semaphore,
) -> list[KnowledgeTriplet]:
    """
    Extract knowledge triplets from a single chunk using Groq.
    
    Returns validated triplets only. Invalid predicates or empty
    entities are silently dropped via Pydantic validation.
    """
    user_prompt = build_chunk_user_prompt(global_context, chunk_text)
    
    result = await _call_groq_with_retry(
        client=client,
        model=model,
        system_prompt=TRIPLET_EXTRACTION_SYSTEM_PROMPT,
        user_prompt=user_prompt,
        semaphore=semaphore,
    )
    
    if result is None:
        logger.warning(f"Chunk {chunk_index}: extraction returned None — skipping")
        return []
    
    reasoning = result.get("_reasoning", "")
    raw_triplets = result.get("triplets", [])
    
    if reasoning:
        logger.debug(f"Chunk {chunk_index} reasoning: {reasoning}")
    
    validated: list[KnowledgeTriplet] = []
    for i, raw in enumerate(raw_triplets):
        try:
            triplet = KnowledgeTriplet(**raw)
            validated.append(triplet)
        except (ValidationError, TypeError) as e:
            logger.warning(
                f"Chunk {chunk_index}, triplet {i}: validation failed — {e}"
            )
    
    return validated


def _normalize_entity(entity: str) -> str:
    """
    Normalize an entity string for deduplication.
    
    Strategy:
      1. Lowercase
      2. Strip leading/trailing whitespace
      3. Collapse multiple spaces
      4. Remove common noise phrases ("the", "said", "hereby")
    """
    s = entity.lower().strip()
    s = re.sub(r'\s+', ' ', s)
    # Remove articles and legal filler that cause false negatives
    for noise in ["the ", "said ", "hereby ", "aforesaid "]:
        if s.startswith(noise):
            s = s[len(noise):]
    return s


def deduplicate_triplets(
    triplets: list[KnowledgeTriplet],
) -> list[KnowledgeTriplet]:
    """
    Deduplicate triplets by normalized (subject, predicate, object) key.
    First occurrence wins (preserves original casing from earliest chunk).
    """
    seen: set[tuple[str, str, str]] = set()
    unique: list[KnowledgeTriplet] = []
    
    for t in triplets:
        key = (
            _normalize_entity(t.subject),
            t.predicate,  # Already normalized by validator
            _normalize_entity(t.object),
        )
        if key not in seen:
            seen.add(key)
            unique.append(t)
    
    return unique


async def extract_knowledge_graph(
    full_text: str,
    groq_api_key: str,
    model: str = "llama-3.3-70b-versatile",
    max_concurrency: int = 5,
) -> tuple[GlobalContext, list[KnowledgeTriplet]]:
    """
    Full Two-Pass Knowledge Extraction Pipeline.
    
    Args:
        full_text: Complete document text (all pages concatenated)
        groq_api_key: Groq API key
        model: Groq model identifier
        max_concurrency: Maximum concurrent API calls
    
    Returns:
        Tuple of (GlobalContext, deduplicated list of KnowledgeTriplet)
    """
    client = AsyncGroq(api_key=groq_api_key)
    semaphore = asyncio.Semaphore(max_concurrency)
    
    # ── Pass 1: Global Context ────────────────────────────────
    logger.info("Pass 1: Extracting global document context...")
    global_ctx = await extract_global_context(
        client, model, full_text, semaphore
    )
    logger.info(
        f"  Title: {global_ctx.document_title}\n"
        f"  Body: {global_ctx.primary_body}\n"
        f"  Definitions: {global_ctx.key_definitions}"
    )
    
    global_ctx_dict = global_ctx.model_dump()
    
    # ── Pass 2: Chunk and extract ─────────────────────────────
    chunks = chunk_text_for_extraction(full_text, chunk_size=1000, overlap=150)
    logger.info(f"Pass 2: Extracting triplets from {len(chunks)} chunks (concurrency={max_concurrency})...")
    
    tasks = [
        extract_triplets_from_chunk(
            client, model, global_ctx_dict, chunk, idx, semaphore
        )
        for idx, chunk in enumerate(chunks)
    ]
    
    # Fire all tasks concurrently (semaphore gates actual API calls)
    results = await asyncio.gather(*tasks, return_exceptions=True)
    
    # Flatten results, skip exceptions
    all_triplets: list[KnowledgeTriplet] = []
    for i, result in enumerate(results):
        if isinstance(result, Exception):
            logger.error(f"Chunk {i} raised exception: {result}")
            continue
        all_triplets.extend(result)
    
    logger.info(f"  Raw triplets extracted: {len(all_triplets)}")
    
    # ── Aggregation: Deduplicate ──────────────────────────────
    unique = deduplicate_triplets(all_triplets)
    logger.info(f"  After deduplication: {len(unique)} unique triplets")
    
    return global_ctx, unique