setup_vector_db.py

"""
Script to process health.pdf and upload it to Pinecone vector database.
Run this script ONCE to set up the vector database.
"""

import os
import time
import PyPDF2
from dotenv import load_dotenv
from pinecone import Pinecone, ServerlessSpec
from sentence_transformers import SentenceTransformer

load_dotenv()

PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
INDEX_NAME = os.getenv("PINECONE_INDEX_NAME", "doctor-assist")
AWS_REGION = os.getenv("PINECONE_ENVIRONMENT", "us-east-1")

if not PINECONE_API_KEY:
    raise ValueError("❌ PINECONE_API_KEY not found in .env file")

pc = Pinecone(api_key=PINECONE_API_KEY)

print("🔄 Loading embedding model...")
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
EMBEDDING_DIMENSION = 384  # all-MiniLM-L6-v2 output size

def extract_text_from_pdf(pdf_path: str) -> str:
    """Extract text from a PDF file."""
    text = ""

    with open(pdf_path, "rb") as file:
        reader = PyPDF2.PdfReader(file)
        total_pages = len(reader.pages)

        print(f"📄 Processing {total_pages} pages...")

        for i, page in enumerate(reader.pages, start=1):
            page_text = page.extract_text()
            if page_text:
                text += page_text + "\n"

            if i % 10 == 0 or i == total_pages:
                print(f"  ✓ Processed {i}/{total_pages} pages")

    return text


def chunk_text(text: str, chunk_size: int = 500, overlap: int = 100) -> list[str]:
    """Split text into overlapping chunks."""
    words = text.split()
    chunks = []

    step = chunk_size - overlap
    for i in range(0, len(words), step):
        chunk = " ".join(words[i : i + chunk_size])
        if chunk.strip():
            chunks.append(chunk)

    return chunks

def create_index_if_not_exists():
    """Create Pinecone index if it doesn't already exist."""
    try:
        existing_indexes = pc.list_indexes().names()
    except Exception as e:
        raise RuntimeError(f"❌ Failed to list Pinecone indexes: {e}")

    if INDEX_NAME in existing_indexes:
        print(f"✅ Index '{INDEX_NAME}' already exists")
        return

    print(f"🚀 Creating Pinecone index '{INDEX_NAME}' (AWS | {AWS_REGION})")

    pc.create_index(
        name=INDEX_NAME,
        dimension=EMBEDDING_DIMENSION,
        metric="cosine",
        spec=ServerlessSpec(
            cloud="aws",
            region=AWS_REGION
        )
    )

    print("⏳ Waiting for index to be ready...")
    time.sleep(10)


def upload_to_pinecone(chunks: list[str]):
    """Upload text chunks and embeddings to Pinecone."""
    index = pc.Index(INDEX_NAME)

    print(f"⬆️ Uploading {len(chunks)} chunks to Pinecone...")

    batch_size = 100
    for start in range(0, len(chunks), batch_size):
        batch = chunks[start : start + batch_size]

        embeddings = embedding_model.encode(
            batch,
            show_progress_bar=True
        )

        vectors = [
            {
                "id": f"chunk-{start + i}",
                "values": embedding.tolist(),
                "metadata": {"text": chunk},
            }
            for i, (chunk, embedding) in enumerate(zip(batch, embeddings))
        ]

        index.upsert(vectors=vectors)

        print(
            f"  ✓ Uploaded batch "
            f"{start // batch_size + 1}/"
            f"{(len(chunks) + batch_size - 1) // batch_size}"
        )

    print("✅ All vectors uploaded successfully")

def main():
    pdf_path = "health.pdf"

    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"❌ '{pdf_path}' not found")

    create_index_if_not_exists()

    print("🔍 Extracting text from PDF...")
    text = extract_text_from_pdf(pdf_path)

    print("✂️ Chunking text...")
    chunks = chunk_text(text)
    print(f"📦 Created {len(chunks)} chunks")

    upload_to_pinecone(chunks)

    print("\n🎉 Vector database setup complete!")


if __name__ == "__main__":
    main()