ingestion/entity_resolver.py

"""
GovBridge India — Entity Resolution for Knowledge Graph Edges
Sprint 30: PROJECT INDRA Phase 1.5

PURPOSE:
  For each extracted triplet's "object" entity (e.g., "Information Technology
  Rules, 2011"), this module searches gazette_chunks to find a matching
  document UUID. If found, the target_node_id is set. If not found,
  target_node_id is NULL and the raw string is stored in metadata.

RESOLUTION STRATEGY (3-tier):
  1. Exact FTS match: websearch_to_tsquery('simple', entity) against chunk_text
  2. Trigram similarity: pg_trgm similarity() > 0.4 against document_title
  3. Unresolved: Insert with NULL target, store in metadata.unresolved_entity

PERFORMANCE:
  - Uses existing GIN FTS index (idx_gazette_chunks_fts)
  - Uses existing GIN trigram index (idx_gazette_chunks_trgm)
  - Single query per entity, batched where possible
"""

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 logging
from typing import Optional
from uuid import UUID

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


async def resolve_entity_to_chunk_id(
    supabase_client,
    entity_name: str,
) -> Optional[str]:
    """
    Attempt to resolve an entity name to a gazette_chunks UUID.
    
    Resolution chain:
      1. Search document_title for trigram similarity > 0.4
         (This catches "IT Rules, 2011" matching "Information Technology Rules, 2011")
      2. Fall back to FTS on chunk_text
      3. Return None if unresolved
    
    Args:
        supabase_client: Initialized Supabase client
        entity_name: The object entity string from the knowledge triplet
    
    Returns:
        UUID string if resolved, None if unresolved
    """
    if not entity_name or len(entity_name) < 3:
        return None
    
    # Strategy 1: Trigram similarity search on document_title
    # This is the most reliable because gazette titles are canonical
    try:
        result = supabase_client.rpc(
            "resolve_entity_by_similarity",
            {"entity_text": entity_name}
        ).execute()
        
        if result.data and len(result.data) > 0:
            match = result.data[0]
            logger.info(
                f"  ✅ Resolved '{entity_name}' → "
                f"'{match['document_title']}' (id={match['id']}, "
                f"similarity={match.get('sim', 'N/A')})"
            )
            return match["id"]
    except Exception as e:
        logger.warning(f"  Trigram resolution RPC failed: {e}")
    
    # Strategy 2: FTS on chunk_text (broader but noisier)
    try:
        # Use ilike for simple substring matching as fallback
        result = (
            supabase_client.table("gazette_chunks")
            .select("id, document_title")
            .ilike("document_title", f"%{entity_name[:50]}%")
            .limit(1)
            .execute()
        )
        
        if result.data and len(result.data) > 0:
            match = result.data[0]
            logger.info(
                f"  ✅ Resolved '{entity_name}' via ilike → '{match['document_title']}'"
            )
            return match["id"]
    except Exception as e:
        logger.warning(f"  ilike resolution failed: {e}")
    
    logger.info(f"  ❌ Unresolved: '{entity_name}' — will insert as dangling edge")
    return None


def build_edge_row(
    source_chunk_id: str,
    predicate: str,
    target_chunk_id: Optional[str],
    object_entity: str,
    subject_entity: str,
    global_context: dict,
) -> dict:
    """
    Build a tensor_edges row for Supabase upsert.
    
    If target_chunk_id is None (unresolved entity), the raw entity
    string is stored in metadata.unresolved_entity for future backfill.
    """
    metadata = {
        "subject_entity": subject_entity,
        "object_entity": object_entity,
        "extraction_model": global_context.get("_extraction_model", "llama-3.3-70b-versatile"),
        "document_title": global_context.get("document_title", "Unknown"),
    }
    
    if target_chunk_id is None:
        metadata["unresolved_entity"] = object_entity
    
    row = {
        "source_node_id": source_chunk_id,
        "target_node_id": target_chunk_id,  # Can be None/NULL
        "edge_type": predicate,
        "metadata": metadata,
    }
    
    return row