vector_store.py

"""
Vector Store module for PDF ingestion and RAG functionality.
Uses ChromaDB as a simple vector database.
"""

import os
from typing import List
import chromadb
from chromadb.utils import embedding_functions
from pypdf import PdfReader
from openai import OpenAI


# Shared persistent ChromaDB client
_chroma_client = None


def get_chroma_client():
    """Get or create shared persistent ChromaDB client."""
    global _chroma_client
    if _chroma_client is None:
        persist_dir = os.path.join(os.path.dirname(__file__), "chroma_db")
        _chroma_client = chromadb.PersistentClient(path=persist_dir)
    return _chroma_client


class VectorStore:
    def __init__(self, collection_name: str = "documents"):
        self.client = get_chroma_client()
        self.openai_client = OpenAI()

        # Use OpenAI embeddings
        self.embedding_fn = embedding_functions.OpenAIEmbeddingFunction(
            api_key=os.getenv("OPENAI_API_KEY"),
            model_name="text-embedding-3-small"
        )

        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            embedding_function=self.embedding_fn
        )

        print(f"[{collection_name}] Collection has {self.collection.count()} documents")

    def extract_text_from_pdf(self, pdf_path: str) -> str:
        """Extract text content from a PDF file."""
        reader = PdfReader(pdf_path)
        text = ""
        for page in reader.pages:
            text += page.extract_text() + "\n"
        return text

    def chunk_text(self, text: str, chunk_size: int = 500, overlap: int = 50) -> List[str]:
        """Split text into overlapping chunks."""
        chunks = []
        start = 0
        while start < len(text):
            end = start + chunk_size
            chunk = text[start:end]
            if chunk.strip():
                chunks.append(chunk)
            start = end - overlap
        return chunks

    def ingest_pdf(self, pdf_path: str, doc_id_prefix: str = "doc"):
        """Ingest a PDF file into the vector store."""
        text = self.extract_text_from_pdf(pdf_path)
        chunks = self.chunk_text(text)

        ids = [f"{doc_id_prefix}_{i}" for i in range(len(chunks))]
        metadatas = [{"source": pdf_path, "chunk_index": i} for i in range(len(chunks))]

        self.collection.add(
            documents=chunks,
            ids=ids,
            metadatas=metadatas
        )

        print(f"Ingested {len(chunks)} chunks from {pdf_path}")

    def search(self, query: str, n_results: int = 5) -> List[str]:
        """Search for relevant documents based on query."""
        # Get total count to avoid requesting more than available
        total_docs = self.collection.count()
        n_results = min(n_results, total_docs) if total_docs > 0 else n_results

        if total_docs == 0:
            return []

        results = self.collection.query(
            query_texts=[query],
            n_results=n_results
        )

        if results and results["documents"]:
            return results["documents"][0]
        return []

    def get_context(self, query: str, n_results: int = 5) -> str:
        """Get formatted context for RAG."""
        docs = self.search(query, n_results)
        if not docs:
            return "No relevant context found."
        return "\n\n---\n\n".join(docs)