scripts/build_pinecone_DB.py

# -*- coding: utf-8 -*-
"""
build_pinecone_DB.py — Build the Pinecone knowledge base for Serinity.

Run once before starting the app:
    python scripts/build_pinecone_DB.py

IMPORTANT: Ollama must be running and nomic-embed-text must be pulled:
    ollama pull nomic-embed-text
"""

import os
import json
from collections import Counter
from datasets import load_dataset
from langchain_core.documents import Document
from pinecone import Pinecone, ServerlessSpec
from langchain_pinecone import PineconeVectorStore
from langchain_ollama import OllamaEmbeddings
from tqdm import tqdm
from dotenv import load_dotenv

load_dotenv()

# Configuration from .env
OLLAMA_HOST        = os.getenv("OLLAMA_HOST",           "http://localhost:11434")
EMBEDDING_MODEL    = os.getenv("EMBEDDING_MODEL",       "nomic-embed-text")
PINECONE_API_KEY   = os.getenv("PINECONE_API_KEY",      "")
PINECONE_INDEX_NAME= os.getenv("PINECONE_INDEX_NAME",   "serinity-knowledge")
DATASET_NAME       = "Compumacy/Psych_data"
BATCH_SIZE         = 100

# Generic wellness filter markers
GENERIC_WELLNESS_MARKERS = [
    "general wellness",
    "lifestyle advice",
    "wellness discussion",
    "healthy habits",
    "nutrition",
    "exercise tips",
    "sleep hygiene tips",
    "mindfulness general",
    "self-care tips",
]


def _is_generic_wellness(medical_context: str) -> bool:
    """Return True if the row looks like off-topic generic wellness content."""
    if not medical_context:
        return False
    lower = medical_context.lower()
    return any(marker in lower for marker in GENERIC_WELLNESS_MARKERS)


def download_dataset():
    """Download the specified dataset from HuggingFace."""
    print("Downloading dataset from HuggingFace...")
    ds = load_dataset(DATASET_NAME)
    print(f"Downloaded {len(ds['train']):,} rows.")
    return ds


def filter_and_audit(data):
    """
    Sample medical_context values, drop generic wellness rows,
    and print a summary for sanity checking before committing.
    """
    print("\nRunning corpus quality filter...")

    context_counts = Counter()
    dropped_indices = set()
    dropped_samples = []

    for idx, item in enumerate(tqdm(data, desc="Auditing medical_context")):
        metadata = item.get("metadata", {})
        if isinstance(metadata, str):
            try:
                metadata = json.loads(metadata)
            except Exception:
                metadata = {}
        ctx = metadata.get("medical_context", "") or ""
        context_counts[ctx] += 1
        if _is_generic_wellness(ctx):
            dropped_indices.add(idx)
            if len(dropped_samples) < 10:
                dropped_samples.append(ctx)

    print(f"\n{'='*60}")
    print("CORPUS FILTER SUMMARY")
    print(f"{'='*60}")
    print(f"  Total rows:   {len(data):,}")
    print(f"  Rows dropped: {len(dropped_indices):,} ({100*len(dropped_indices)/len(data):.1f}%)")
    print(f"  Rows kept:    {len(data) - len(dropped_indices):,}")
    print("\n  Top 15 medical_context values (before filtering):")
    for ctx, count in context_counts.most_common(15):
        label = ctx[:60] + "..." if len(ctx) > 60 else ctx
        flag = " [DROPPED]" if _is_generic_wellness(ctx) else ""
        print(f"    [{count:>6}] {label}{flag}")
    if dropped_samples:
        print("\n  Sample dropped contexts:")
        for s in dropped_samples:
            print(f"    - {s[:80]}")
    print(f"{'='*60}\n")

    return dropped_indices


def create_documents(data, dropped_indices: set, mode: str = "assistant_only"):
    """
    Build Document objects from the dataset.
    Modes: assistant_only | qa_pairs | both_separate
    """
    print(f"Creating documents (mode: {mode}, skipping {len(dropped_indices):,} filtered rows)...")
    documents = []

    for idx, item in enumerate(tqdm(data, desc="Building documents")):
        if idx in dropped_indices:
            continue

        user_msg      = item.get("user_message",    "") or ""
        assistant_msg = item.get("assistant_message","") or ""
        metadata      = item.get("metadata",        {}) or {}
        if isinstance(metadata, str):
            try:
                metadata = json.loads(metadata)
            except Exception:
                metadata = {}

        doc_metadata = {
            "chunk_id":        idx,
            "source_pdf":      metadata.get("source_pdf",      "unknown"),
            "page_number":     metadata.get("page_number",     -1),
            "confidence_score":metadata.get("confidence_score", 0.0),
            "medical_context": metadata.get("medical_context", ""),
        }

        if mode == "assistant_only":
            if assistant_msg.strip():
                documents.append(Document(page_content=assistant_msg.strip(), metadata=doc_metadata))

        elif mode == "qa_pairs":
            if user_msg and assistant_msg:
                combined = f"Question: {user_msg.strip()}\n\nAnswer: {assistant_msg.strip()}"
                documents.append(Document(page_content=combined, metadata=doc_metadata))

        elif mode == "both_separate":
            if user_msg.strip():
                q_meta = {**doc_metadata, "type": "question"}
                documents.append(Document(page_content=f"Question: {user_msg.strip()}", metadata=q_meta))
            if assistant_msg.strip():
                a_meta = {**doc_metadata, "type": "answer"}
                documents.append(Document(page_content=assistant_msg.strip(), metadata=a_meta))

    print(f"Created {len(documents):,} document objects.")
    return documents


def build_pinecone_db(documents: list):
    """
    Push documents to Pinecone using Ollama embeddings.
    """
    if not documents:
        raise ValueError("No documents to index — check filter settings.")

    print(f"\nInitializing OllamaEmbeddings ({EMBEDDING_MODEL})...")
    embeddings = OllamaEmbeddings(
        model=EMBEDDING_MODEL,
        base_url=OLLAMA_HOST,
        keep_alive=-1,
    )

    print("Verifying Ollama connection...")
    test_vec = embeddings.embed_query("test")
    print(f"Embedding model ready (dimension: {len(test_vec)}).")

    if not PINECONE_API_KEY:
        raise ValueError("PINECONE_API_KEY is missing from .env")

    print(f"\nInitializing Pinecone Client and pushing to index: {PINECONE_INDEX_NAME} ...")
    pc = Pinecone(api_key=PINECONE_API_KEY)
    
    existing_indexes = [index_info["name"] for index_info in pc.list_indexes()]
    if PINECONE_INDEX_NAME not in existing_indexes:
        print(f"Index '{PINECONE_INDEX_NAME}' not found. Creating it with dimension {len(test_vec)}...")
        pc.create_index(
            name=PINECONE_INDEX_NAME,
            dimension=len(test_vec),
            metric="cosine",
            spec=ServerlessSpec(
                cloud="aws",
                region="us-east-1"
            )
        )
        print("Index created successfully.")
        
    index = pc.Index(PINECONE_INDEX_NAME)
    
    vectorstore = PineconeVectorStore(
        index=index,
        embedding=embeddings,
        text_key="text"
    )

    # For a fresh build, we start from 0. 
    # Pinecone stats can be accessed but keeping track of exact batches is harder due to distributed nature.
    # We will just upload all provided documents.
    start_doc = 0

    print(f"Index:          {PINECONE_INDEX_NAME}")
    print(f"Total docs:     {len(documents):,}")
    print(f"Batch size:     {BATCH_SIZE}")
    print("This will take a while (embedding locally and uploading to Pinecone)...\n")

    remaining_docs = documents[start_doc:]
    total_batches  = (len(remaining_docs) + BATCH_SIZE - 1) // BATCH_SIZE

    for i in tqdm(range(0, len(remaining_docs), BATCH_SIZE),
                  total=total_batches, desc="Embedding and Uploading batches"):
        batch = remaining_docs[i : i + BATCH_SIZE]
        vectorstore.add_documents(batch)

    print(f"\nPinecone build complete! Vectors uploaded.")
    return vectorstore


def main():
    import sys
    resuming = "--resume" in sys.argv

    print("=" * 60)
    print("  Serinity — LOCAL KNOWLEDGE BASE BUILDER")
    print("  Target: Pinecone + nomic-embed-text (Ollama)")
    if resuming:
        print("  MODE: RESUME (skipping already-embedded batches)")
    print("=" * 60)

    print("\nChoose indexing mode:")
    print("  1. Assistant messages only (default — cleanest clinical answers)")
    print("  2. Q&A pairs (combined question + answer)")
    print("  3. Both separate (individual question and answer docs)")
    choice = input("\nEnter choice (1/2/3) [1]: ").strip() or "1"
    mode_map = {"1": "assistant_only", "2": "qa_pairs", "3": "both_separate"}
    mode = mode_map.get(choice, "assistant_only")

    try:
        dataset     = download_dataset()
        data        = dataset["train"]
        dropped_idx = filter_and_audit(data)

        if not resuming:
            print("\nProceed with the filter above? (Ctrl+C to abort, Enter to continue)")
            input()
        else:
            print("\n(Resume mode — skipping filter confirmation, using same filter as before.)")

        documents = create_documents(data, dropped_idx, mode=mode)
        build_pinecone_db(documents)
        print("\nKnowledge base is ready. Start the app with: uvicorn main:app --reload")
    except KeyboardInterrupt:
        print("\nAborted by user. Run with --resume to continue from where you stopped.")
    except Exception as e:
        print(f"\nERROR: {e}")
        raise


if __name__ == "__main__":
    main()