vector_store.py

# vector_store.py
# ================
# Manages the FAISS vector store for document retrieval
# FAISS = Facebook AI Similarity Search
# Converts document chunks into searchable vectors
# No API key needed — runs entirely locally

import os
import pickle
import numpy as np
from typing import List, Dict, Tuple

# We use a simple TF-IDF approach that works on Python 3.14
# This avoids the sentence-transformers/torch compatibility issues
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

class PharmaVectorStore:
    """
    A simple but effective vector store for pharma documents.
    Uses TF-IDF vectors and cosine similarity for retrieval.
    Works on Python 3.14 without torch dependencies.
    """
    
    def __init__(self):
        self.vectorizer = TfidfVectorizer(
            max_features=10000,
            ngram_range=(1, 2),  # Use both single words and pairs
            stop_words='english'
        )
        self.vectors = None
        self.chunks = []
        self.is_fitted = False
    
    def add_documents(self, chunks: List[Dict]):
        """
        Add document chunks to the vector store.
        
        Args:
            chunks: List of chunk dictionaries from document_loader.py
        """
        
        if not chunks:
            print("No chunks to add")
            return
        
        print(f"Adding {len(chunks)} chunks to vector store...")
        
        # Store chunks
        self.chunks.extend(chunks)
        
        # Extract text for vectorisation
        texts = [chunk["text"] for chunk in self.chunks]
        
        # Fit and transform — creates TF-IDF vectors
        self.vectors = self.vectorizer.fit_transform(texts)
        self.is_fitted = True
        
        print(f"Vector store now contains {len(self.chunks)} chunks")
    
    def search(self, query: str, top_k: int = 4) -> List[Dict]:
        """
        Search for the most relevant chunks for a query.
        
        Args:
            query: The user's question
            top_k: Number of top results to return
        
        Returns:
            List of most relevant chunk dictionaries with scores
        """
        
        if not self.is_fitted or len(self.chunks) == 0:
            print("Vector store is empty — add documents first")
            return []
        
        # Vectorise the query
        query_vector = self.vectorizer.transform([query])
        
        # Calculate similarity scores
        similarities = cosine_similarity(query_vector, self.vectors)[0]
        
        # Get top-k most similar chunks
        top_indices = np.argsort(similarities)[::-1][:top_k]
        
        results = []
        for idx in top_indices:
            if similarities[idx] > 0:  # Only return relevant results
                chunk = self.chunks[idx].copy()
                chunk["similarity_score"] = float(similarities[idx])
                results.append(chunk)
        
        print(f"Found {len(results)} relevant chunks for query")
        return results
    
    def clear(self):
        """
        Clear all documents from the vector store.
        Used when loading a new document.
        """
        self.vectorizer = TfidfVectorizer(
            max_features=10000,
            ngram_range=(1, 2),
            stop_words='english'
        )
        self.vectors = None
        self.chunks = []
        self.is_fitted = False
        print("Vector store cleared")
    
    def get_stats(self) -> Dict:
        """
        Get statistics about the vector store.
        """
        return {
            "total_chunks": len(self.chunks),
            "is_fitted": self.is_fitted,
            "unique_sources": len(set(c["source"] for c in self.chunks))
        }
    
    def save(self, path: str):
        """
        Save the vector store to disk.
        """
        with open(path, 'wb') as f:
            pickle.dump({
                'vectorizer': self.vectorizer,
                'vectors': self.vectors,
                'chunks': self.chunks,
                'is_fitted': self.is_fitted
            }, f)
        print(f"Vector store saved to {path}")
    
    def load(self, path: str):
        """
        Load the vector store from disk.
        """
        if not os.path.exists(path):
            print(f"No saved vector store found at {path}")
            return False
        
        with open(path, 'rb') as f:
            data = pickle.load(f)
        
        self.vectorizer = data['vectorizer']
        self.vectors = data['vectors']
        self.chunks = data['chunks']
        self.is_fitted = data['is_fitted']
        print(f"Vector store loaded — {len(self.chunks)} chunks")
        return True


def format_retrieved_chunks(chunks: List[Dict]) -> str:
    """
    Format retrieved chunks as context for Claude.
    
    Args:
        chunks: List of retrieved chunk dictionaries
    
    Returns:
        Formatted string ready to send to Claude
    """
    
    if not chunks:
        return "No relevant context found in the documents."
    
    formatted = []
    for i, chunk in enumerate(chunks, 1):
        formatted.append(f"""
Source {i}: {chunk['source']} — Page {chunk['page_number']}
Relevance: {chunk.get('similarity_score', 0):.2f}
Content: {chunk['text'][:500]}
""")
    
    return "\n---\n".join(formatted)


# Quick test
if __name__ == "__main__":
    from document_loader import chunk_pages
    
    # Create sample chunks
    sample_pages = [
        {
            "text": """KEYTRUDA (pembrolizumab) injection
            INDICATIONS AND USAGE
            KEYTRUDA is indicated for the treatment of patients with 
            unresectable or metastatic melanoma. KEYTRUDA is indicated 
            for the first-line treatment of patients with metastatic 
            non-small cell lung cancer (NSCLC) whose tumors have high 
            PD-L1 expression (TPS ≥50%) with no EGFR or ALK genomic 
            tumor aberrations.
            CONTRAINDICATIONS
            None.
            WARNINGS AND PRECAUTIONS
            Immune-mediated adverse reactions including pneumonitis, 
            colitis, hepatitis, and endocrinopathies may occur.""",
            "page_number": 1,
            "source": "keytruda_label.pdf",
            "file_path": "keytruda_label.pdf"
        }
    ]
    
    chunks = chunk_pages(sample_pages)
    
    # Test vector store
    store = PharmaVectorStore()
    store.add_documents(chunks)
    
    # Test search
    results = store.search("What are the contraindications?")
    
    print(f"\nSearch test passed!")
    print(f"Store stats: {store.get_stats()}")
    print(f"Results found: {len(results)}")
    if results:
        print(f"Top result page: {results[0]['page_number']}")
        print(f"Similarity score: {results[0]['similarity_score']:.3f}")