Enterprise Data Liquidity: The Engineering Framework for an AI-Ready Knowledge Base

from langchain_community.document_loaders import PyPDFDirectoryLoader
from langchain_experimental.text_splitter import SemanticChunker
from langchain_openai import OpenAIEmbeddings
from pinecone import Pinecone, ServerlessSpec
from datetime import datetime, timezone
from pathlib import Path
from typing import List
import hashlib
import logging
import os

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

# --- Configuration ---
PINECONE_INDEX_NAME = "enterprise-knowledge-base"
EMBEDDING_MODEL = "text-embedding-3-large"
EMBEDDING_DIMENSIONS = 3072
BATCH_SIZE = 100  # Pinecone upsert batch size

def initialize_pinecone_index(pc: Pinecone) -> object:
    """Create index if it does not exist, return index object."""
    existing = [idx.name for idx in pc.list_indexes()]
    if PINECONE_INDEX_NAME not in existing:
        logger.info(f"Creating Pinecone index: {PINECONE_INDEX_NAME}")
        pc.create_index(
            name=PINECONE_INDEX_NAME,
            dimension=EMBEDDING_DIMENSIONS,
            metric="cosine",
            spec=ServerlessSpec(cloud="aws", region="us-east-1")
        )
    return pc.Index(PINECONE_INDEX_NAME)

def compute_chunk_id(source: str, page: int, chunk_index: int) -> str:
    """Generate a deterministic, idempotent chunk ID for upsert deduplication."""
    content = f"{source}::p{page}::c{chunk_index}"
    return hashlib.sha256(content.encode()).hexdigest()[:32]

def load_and_chunk_pdfs(pdf_directory: str, embeddings: OpenAIEmbeddings) -> List[dict]:
    """Load PDFs from directory, apply semantic chunking, return list of Pinecone vectors."""
    logger.info(f"Loading PDFs from: {pdf_directory}")

    # Load all PDFs from directory — preserves page metadata
    loader = PyPDFDirectoryLoader(pdf_directory, glob="**/*.pdf", silent_errors=True)
    raw_documents = loader.load()
    logger.info(f"Loaded {len(raw_documents)} pages from {len(set(d.metadata['source'] for d in raw_documents))} documents")

    # Semantic chunker: splits at sentence boundaries, groups by semantic coherence
    # breakpoint_threshold_type="percentile" splits when cosine distance exceeds the 95th percentile
    chunker = SemanticChunker(
        embeddings=embeddings,
        breakpoint_threshold_type="percentile",
        breakpoint_threshold_amount=95
    )

    chunks = chunker.split_documents(raw_documents)
    logger.info(f"Generated {len(chunks)} semantic chunks")

    # Convert to Pinecone vector format
    vectors = []
    for i, chunk in enumerate(chunks):
        source_path = Path(chunk.metadata.get("source", "unknown"))
        page_number = chunk.metadata.get("page", 0)
        chunk_id = compute_chunk_id(source_path.name, page_number, i)

        vectors.append({
            "id": chunk_id,
            "text": chunk.page_content,  # Store text for reconstruction
            "metadata": {
                "source": source_path.name,
                "source_path": str(source_path),
                "page_number": page_number,
                "chunk_id": chunk_id,
                "chunk_index": i,
                "created_at": datetime.now(timezone.utc).isoformat(),
                "document_type": "pdf",
                "char_count": len(chunk.page_content)
            }
        })

    return vectors

def embed_and_upsert(
    vectors: List[dict],
    index: object,
    embeddings: OpenAIEmbeddings
) -> int:
    """Generate embeddings in batches and upsert to Pinecone. Returns count of upserted vectors."""
    upserted_count = 0
    total_batches = (len(vectors) + BATCH_SIZE - 1) // BATCH_SIZE

    for batch_num in range(total_batches):
        batch_start = batch_num * BATCH_SIZE
        batch_end = min(batch_start + BATCH_SIZE, len(vectors))
        batch = vectors[batch_start:batch_end]

        logger.info(f"Embedding batch {batch_num + 1}/{total_batches} ({len(batch)} chunks)")

        # Generate embeddings for this batch
        texts = [v["text"] for v in batch]
        batch_embeddings = embeddings.embed_documents(texts)

        # Build Pinecone upsert payload
        pinecone_vectors = [
            {
                "id": v["id"],
                "values": emb,
                "metadata": {**v["metadata"], "text": v["text"]}  # Store text in metadata for retrieval
            }
            for v, emb in zip(batch, batch_embeddings)
        ]

        # Upsert to Pinecone (idempotent — safe to re-run)
        index.upsert(vectors=pinecone_vectors)
        upserted_count += len(pinecone_vectors)
        logger.info(f"Upserted {upserted_count}/{len(vectors)} vectors")

    return upserted_count

def main(pdf_directory: str = "./documents"):
    # Initialize clients
    pc = Pinecone(api_key=os.environ["PINECONE_API_KEY"])
    embeddings = OpenAIEmbeddings(
        model=EMBEDDING_MODEL,
        dimensions=EMBEDDING_DIMENSIONS,  # Matryoshka: reduce to 256 for cost savings
        openai_api_key=os.environ["OPENAI_API_KEY"]
    )

    # Initialize or connect to Pinecone index
    index = initialize_pinecone_index(pc)
    logger.info(f"Connected to Pinecone index: {PINECONE_INDEX_NAME}")

    # Load, chunk, embed, and upsert
    vectors = load_and_chunk_pdfs(pdf_directory, embeddings)
    if not vectors:
        logger.warning("No chunks generated — check that PDF files exist in the directory")
        return

    total_upserted = embed_and_upsert(vectors, index, embeddings)

    # Report index stats
    stats = index.describe_index_stats()
    logger.info(f"Pipeline complete: {total_upserted} vectors upserted")
    logger.info(f"Index now contains {stats.total_vector_count:,} total vectors")

if __name__ == "__main__":
    import sys
    directory = sys.argv[1] if len(sys.argv) > 1 else "./documents"
    main(pdf_directory=directory)