scripts/upload_knowledge_base.py

"""
US Policy Claimer — Knowledge Base Embedding & Upload Pipeline
Parses all knowledge base markdown files, generates embeddings via
Google text-embedding-004, and uploads to Supabase pgvector.
"""

import os
import re
import sys
import time
from pathlib import Path

from dotenv import load_dotenv
from google import genai
from supabase import create_client, Client

# Load .env file from project root
load_dotenv(Path(__file__).parent.parent / ".env")

# ─── Configuration ───────────────────────────────────────────────
SUPABASE_URL = os.environ.get("SUPABASE_URL")
SUPABASE_KEY = os.environ.get("SUPABASE_SERVICE_KEY")
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")

KNOWLEDGE_BASE_DIR = Path(__file__).parent.parent / "knowledge_base"
EMBEDDING_MODEL = "gemini-embedding-001"
EMBEDDING_TASK_TYPE = "RETRIEVAL_DOCUMENT"  # Optimized for document storage
TABLE_NAME = "knowledge_chunks"


def parse_knowledge_file(filepath: Path) -> list[dict]:
    """Parse a knowledge base markdown file into individual chunks."""
    content = filepath.read_text(encoding="utf-8")

    # Strip leading --- (first frontmatter delimiter at start of file)
    content = re.sub(r'^---\n', '', content)

    # Split on YAML frontmatter delimiters (---)
    # Each chunk starts with --- and ends before the next ---
    raw_chunks = re.split(r'\n---\n', content)

    chunks = []
    i = 0
    while i < len(raw_chunks):
        block = raw_chunks[i].strip()

        # Skip empty blocks
        if not block:
            i += 1
            continue

        # Check if this block is a YAML frontmatter block
        if block.startswith("concept_id:"):
            yaml_block = block
            # The next block should be the content
            if i + 1 < len(raw_chunks):
                content_block = raw_chunks[i + 1].strip()
                i += 2
            else:
                i += 1
                continue

            chunk = parse_single_chunk(yaml_block, content_block)
            if chunk:
                chunks.append(chunk)
        else:
            i += 1

    return chunks


def parse_single_chunk(yaml_text: str, content_text: str) -> dict | None:
    """Parse YAML frontmatter and markdown content into a structured dict."""
    try:
        # Parse YAML fields manually (simple key-value extraction)
        def extract_field(text: str, field: str) -> str:
            match = re.search(rf'^{field}:\s*(.+)$', text, re.MULTILINE)
            return match.group(1).strip() if match else ""

        def extract_list(text: str, field: str) -> list[str]:
            match = re.search(rf'^{field}:\s*\[(.+)\]$', text, re.MULTILINE)
            if match:
                items = match.group(1).split(",")
                return [item.strip().strip("'\"") for item in items]
            return []

        concept_id = extract_field(yaml_text, "concept_id")
        if not concept_id:
            return None

        domain = extract_field(yaml_text, "domain")
        jurisdiction = extract_field(yaml_text, "jurisdiction")
        audience = extract_field(yaml_text, "audience")
        tags = extract_list(yaml_text, "tags")

        # Extract title (### heading)
        title_match = re.search(r'^###\s+(.+)$', content_text, re.MULTILINE)
        title = title_match.group(1).strip() if title_match else concept_id

        # Extract semantic summary
        summary_match = re.search(
            r'\*\*Semantic Summary:\*\*\s*\n(.+?)(?=\n\n|\n\*\*)',
            content_text,
            re.DOTALL
        )
        semantic_summary = summary_match.group(1).strip() if summary_match else ""

        return {
            "concept_id": concept_id,
            "domain": domain,
            "jurisdiction": jurisdiction,
            "audience": audience,
            "tags": tags,
            "title": title,
            "semantic_summary": semantic_summary,
            "full_content": content_text,
        }
    except Exception as e:
        print(f"  ⚠️  Error parsing chunk: {e}")
        return None


def generate_embedding(client: genai.Client, text: str) -> list[float]:
    """Generate a 768-dim embedding via Google text-embedding-004."""
    # Combine semantic summary (primary) with full content for richer embedding
    result = client.models.embed_content(
        model=EMBEDDING_MODEL,
        contents=text,
        config={
            "task_type": EMBEDDING_TASK_TYPE,
            "output_dimensionality": 768,
        },
    )
    return result.embeddings[0].values


def main():
    # ─── Validate environment ────────────────────────────────────
    if not SUPABASE_KEY:
        print("❌ SUPABASE_SERVICE_KEY environment variable not set.")
        print("   Export it: export SUPABASE_SERVICE_KEY='your-service-role-key'")
        sys.exit(1)

    if not GEMINI_API_KEY:
        print("❌ GEMINI_API_KEY environment variable not set.")
        print("   Export it: export GEMINI_API_KEY='your-gemini-api-key'")
        sys.exit(1)

    if not KNOWLEDGE_BASE_DIR.exists():
        print(f"❌ Knowledge base directory not found: {KNOWLEDGE_BASE_DIR}")
        sys.exit(1)

    # ─── Initialize clients ──────────────────────────────────────
    print("🔌 Connecting to Supabase...")
    supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)

    print("🤖 Initializing Gemini embedding client...")
    gemini_client = genai.Client(api_key=GEMINI_API_KEY)

    # ─── Parse all knowledge base files ──────────────────────────
    md_files = sorted(KNOWLEDGE_BASE_DIR.glob("*.md"))
    print(f"\n📂 Found {len(md_files)} knowledge base files:")
    for f in md_files:
        print(f"   • {f.name}")

    all_chunks = []
    for filepath in md_files:
        chunks = parse_knowledge_file(filepath)
        print(f"   ✅ {filepath.name}: {len(chunks)} chunks parsed")
        all_chunks.extend(chunks)

    print(f"\n📊 Total chunks parsed: {len(all_chunks)}")

    # ─── Generate embeddings and upload ──────────────────────────
    print(f"\n🚀 Generating embeddings and uploading to Supabase...\n")

    success_count = 0
    error_count = 0

    for i, chunk in enumerate(all_chunks, 1):
        concept_id = chunk["concept_id"]
        try:
            # Create embedding text: combine summary + title for best retrieval
            embedding_text = f"{chunk['title']}. {chunk['semantic_summary']}"

            # Generate embedding
            embedding = generate_embedding(gemini_client, embedding_text)

            # Upsert to Supabase (insert or update if concept_id exists)
            row = {
                "concept_id": chunk["concept_id"],
                "domain": chunk["domain"],
                "jurisdiction": chunk["jurisdiction"],
                "audience": chunk["audience"],
                "tags": chunk["tags"],
                "title": chunk["title"],
                "semantic_summary": chunk["semantic_summary"],
                "full_content": chunk["full_content"],
                "embedding": embedding,
            }

            supabase.table(TABLE_NAME).upsert(
                row, on_conflict="concept_id"
            ).execute()

            print(f"   [{i:02d}/{len(all_chunks)}] ✅ {concept_id}")
            success_count += 1

            # Small delay to respect rate limits
            time.sleep(0.1)

        except Exception as e:
            print(f"   [{i:02d}/{len(all_chunks)}] ❌ {concept_id}: {e}")
            error_count += 1

    # ─── Summary ─────────────────────────────────────────────────
    print(f"\n{'='*50}")
    print(f"📊 Upload Complete!")
    print(f"   ✅ Successful: {success_count}")
    print(f"   ❌ Errors:     {error_count}")
    print(f"   📦 Total:      {len(all_chunks)}")
    print(f"{'='*50}")

    # ─── Verification query ──────────────────────────────────────
    print(f"\n🔍 Verification: Testing similarity search...")
    try:
        test_query = "What happens when my insurance denies a claim?"
        test_embedding = generate_embedding(gemini_client, test_query)

        # Use the search function we created
        result = supabase.rpc("search_knowledge", {
            "query_embedding": test_embedding,
            "match_count": 3,
        }).execute()

        if result.data:
            print(f"   Query: \"{test_query}\"")
            print(f"   Top 3 results:")
            for r in result.data:
                sim = r.get("similarity", 0)
                print(f"     • [{sim:.4f}] {r['title'][:80]}")
        else:
            print("   ⚠️  No results returned. Check the table and search function.")

    except Exception as e:
        print(f"   ⚠️  Verification query failed: {e}")
        print(f"   (This is OK — the data is uploaded. The search function can be tested later.)")


if __name__ == "__main__":
    main()