Add FAISS indexing utilities and enhance text processing functions for improved chunking and validation

utils/faiss_utils.py

@@ -0,0 +1,146 @@
+"""
+FAISS indexing utilities for similarity search.
+
+This module provides utilities for building and searching FAISS indexes.
+"""
+
+from typing import List, Tuple, Any, Dict
+import numpy as np
+import faiss
+from sentence_transformers import SentenceTransformer
+
+from configs.config import Config
+from utils.text_processing import validate_chunk_data
+
+
+def build_faiss_index(embeddings: np.ndarray) -> faiss.IndexHNSWFlat:
+    """
+    Build a FAISS HNSW index from embeddings for similarity search.
+    
+    Args:
+        embeddings: Numpy array of embeddings
+        
+    Returns:
+        FAISS HNSW index
+    """
+    dim = embeddings.shape[1]
+    index = faiss.IndexHNSWFlat(dim, Config.FAISS_NEIGHBORS)
+    index.hnsw.efConstruction = Config.FAISS_EF_CONSTRUCTION
+    index.hnsw.efSearch = Config.FAISS_EF_SEARCH
+    index.add(embeddings)
+    return index
+
+
+def retrieve_similar_chunks(
+    query: str,
+    index: faiss.IndexHNSWFlat,
+    chunks_with_metadata: List[Dict[str, Any]],
+    embedding_model: SentenceTransformer,
+    k: int = None,
+    max_chunk_length: int = None
+) -> List[Tuple[str, float, Dict[str, Any]]]:
+    """
+    Retrieve top k similar chunks to the query from the FAISS index.
+    
+    Args:
+        query: Search query
+        index: FAISS index
+        chunks_with_metadata: List of chunk dictionaries
+        embedding_model: SentenceTransformer model
+        k: Number of chunks to retrieve
+        max_chunk_length: Maximum length for returned chunks
+        
+    Returns:
+        List of tuples (chunk_text, distance, metadata)
+    """
+    if k is None:
+        k = Config.DEFAULT_K_CHUNKS
+    if max_chunk_length is None:
+        max_chunk_length = Config.DEFAULT_CHUNK_SIZE
+    
+    query_embedding = embedding_model.encode([query], convert_to_tensor=True).cpu().numpy()
+    distances, indices = index.search(query_embedding, k)
+    
+    # Ensure indices are within bounds and create mapping for correct distances
+    valid_results = []
+    for idx_pos, chunk_idx in enumerate(indices[0]):
+        if 0 <= chunk_idx < len(chunks_with_metadata):
+            chunk_text = chunks_with_metadata[chunk_idx]["text"][:max_chunk_length]
+            # Only include chunks with meaningful content
+            if chunk_text.strip():  # Skip empty chunks
+                result = (
+                    chunk_text,
+                    distances[0][idx_pos],  # Use original position for correct distance
+                    chunks_with_metadata[chunk_idx]["metadata"]
+                )
+                if validate_chunk_data(result):
+                    valid_results.append(result)
+    
+    return valid_results
+
+
+def search_index_with_validation(
+    query: str,
+    index: faiss.IndexHNSWFlat,
+    chunks_with_metadata: List[Dict[str, Any]],
+    embedding_model: SentenceTransformer,
+    k: int = None,
+    similarity_threshold: float = None
+) -> List[Tuple[str, float, Dict[str, Any]]]:
+    """
+    Search index with additional validation and filtering.
+    
+    Args:
+        query: Search query
+        index: FAISS index
+        chunks_with_metadata: List of chunk dictionaries
+        embedding_model: SentenceTransformer model
+        k: Number of chunks to retrieve
+        similarity_threshold: Threshold for filtering results
+        
+    Returns:
+        List of validated and filtered chunk tuples
+    """
+    if not query or len(query.strip()) < 3:
+        return []
+    
+    if similarity_threshold is None:
+        similarity_threshold = Config.SIMILARITY_THRESHOLD
+    
+    try:
+        # Retrieve similar chunks
+        similar_chunks = retrieve_similar_chunks(
+            query, index, chunks_with_metadata, embedding_model, k
+        )
+        
+        # Filter by similarity threshold
+        filtered_chunks = [
+            chunk for chunk in similar_chunks 
+            if chunk[1] < similarity_threshold
+        ]
+        
+        return filtered_chunks
+        
+    except Exception as e:
+        print(f"Error in index search: {e}")
+        return []
+
+
+def get_index_stats(index: faiss.IndexHNSWFlat) -> Dict[str, Any]:
+    """
+    Get statistics about the FAISS index.
+    
+    Args:
+        index: FAISS index
+        
+    Returns:
+        Dictionary with index statistics
+    """
+    return {
+        "total_vectors": index.ntotal,
+        "dimension": index.d,
+        "index_type": type(index).__name__,
+        "ef_search": index.hnsw.efSearch,
+        "ef_construction": index.hnsw.efConstruction,
+        "is_trained": index.is_trained
+    }

utils/text_processing.py

@@ -0,0 +1,196 @@
+"""
+Utility functions for text processing and embeddings.
+
+This module contains utility functions for text processing, tokenization,
+chunking, and embedding operations.
+"""
+
+from typing import List, Dict, Any, Tuple
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from langchain_community.document_loaders import PyMuPDFLoader
+from langchain.schema import Document
+
+from configs.config import Config
+
+
+def estimate_tokens(text: str) -> int:
+    """
+    Estimate the number of tokens in a text (rough approximation).
+    
+    Args:
+        text: Input text
+        
+    Returns:
+        Estimated number of tokens
+    """
+    return len(text) // 4
+
+
+def process_pdf_file(file_path: str) -> List[Document]:
+    """
+    Load a PDF file and extract its text with metadata.
+    
+    Args:
+        file_path: Path to the PDF file
+        
+    Returns:
+        List of Document objects with metadata
+        
+    Raises:
+        FileNotFoundError: If the file doesn't exist
+    """
+    import os
+    if not os.path.exists(file_path):
+        raise FileNotFoundError(f"The file {file_path} does not exist.")
+    
+    loader = PyMuPDFLoader(file_path)
+    documents = loader.load()
+    return documents
+
+
+def chunk_text(documents: List[Document], max_length: int = None) -> List[Dict[str, Any]]:
+    """
+    Split documents into chunks with metadata.
+    
+    Args:
+        documents: List of Document objects
+        max_length: Maximum chunk length in tokens
+        
+    Returns:
+        List of chunk dictionaries with text and metadata
+    """
+    if max_length is None:
+        max_length = Config.DEFAULT_CHUNK_SIZE
+    
+    chunks = []
+    
+    for doc in documents:
+        text = doc.page_content
+        metadata = doc.metadata
+        paragraphs = text.split("\n\n")
+        current_chunk = ""
+        current_metadata = metadata.copy()
+        
+        for paragraph in paragraphs:
+            # Skip very short paragraphs
+            if len(paragraph.strip()) < Config.MIN_PARAGRAPH_LENGTH:
+                continue
+                
+            if estimate_tokens(current_chunk + paragraph) <= max_length // 4:
+                current_chunk += paragraph + "\n\n"
+            else:
+                # Only add chunks with meaningful content
+                if current_chunk.strip() and len(current_chunk.strip()) > Config.MIN_CHUNK_LENGTH:
+                    chunks.append({
+                        "text": current_chunk.strip(), 
+                        "metadata": current_metadata
+                    })
+                current_chunk = paragraph + "\n\n"
+        
+        # Add the last chunk if it has meaningful content
+        if current_chunk.strip() and len(current_chunk.strip()) > Config.MIN_CHUNK_LENGTH:
+            chunks.append({
+                "text": current_chunk.strip(), 
+                "metadata": current_metadata
+            })
+    
+    return chunks
+
+
+def create_embeddings(chunks: List[Dict[str, Any]], model: SentenceTransformer) -> Tuple[np.ndarray, List[Dict[str, Any]]]:
+    """
+    Create embeddings for a list of chunk texts.
+    
+    Args:
+        chunks: List of chunk dictionaries
+        model: SentenceTransformer model
+        
+    Returns:
+        Tuple of (embeddings array, chunks)
+    """
+    texts = [chunk["text"] for chunk in chunks]
+    embeddings = model.encode(texts, show_progress_bar=True, convert_to_tensor=True)
+    return embeddings.cpu().numpy(), chunks
+
+
+def filter_relevant_chunks(chunks_data: List[Tuple], threshold: float = None) -> List[Tuple]:
+    """
+    Filter chunks based on similarity threshold.
+    
+    Args:
+        chunks_data: List of (text, score, metadata) tuples
+        threshold: Similarity threshold (lower is more similar)
+        
+    Returns:
+        Filtered list of chunks
+    """
+    if threshold is None:
+        threshold = Config.SIMILARITY_THRESHOLD
+    
+    return [chunk for chunk in chunks_data if len(chunk) >= 3 and chunk[1] < threshold]
+
+
+def prepare_context_from_chunks(context_chunks: List[Tuple], max_tokens: int = None) -> str:
+    """
+    Prepare context string from chunk data.
+    
+    Args:
+        context_chunks: List of (text, score, metadata) tuples
+        max_tokens: Maximum tokens for context
+        
+    Returns:
+        Formatted context string
+    """
+    if max_tokens is None:
+        max_tokens = Config.MAX_CONTEXT_TOKENS
+    
+    # Sort chunks by relevance (lower distance = more relevant)
+    sorted_chunks = sorted(context_chunks, key=lambda x: x[1]) if context_chunks else []
+    
+    # Filter out chunks with very high distance scores (low similarity)
+    relevant_chunks = filter_relevant_chunks(sorted_chunks)
+    
+    context = ""
+    total_tokens = 0
+    
+    for chunk, _, _ in relevant_chunks:
+        if chunk and chunk.strip():
+            chunk_tokens = estimate_tokens(chunk)
+            if total_tokens + chunk_tokens <= max_tokens:
+                context += chunk + "\n\n"
+                total_tokens += chunk_tokens
+            else:
+                break
+    
+    return context.strip() if context else "No initial context provided from preliminary search."
+
+
+def validate_chunk_data(chunk_data: Any) -> bool:
+    """
+    Validate chunk data structure.
+    
+    Args:
+        chunk_data: Chunk data to validate
+        
+    Returns:
+        True if valid, False otherwise
+    """
+    if not isinstance(chunk_data, (list, tuple)):
+        return False
+    
+    if len(chunk_data) < 3:
+        return False
+    
+    text, score, metadata = chunk_data[0], chunk_data[1], chunk_data[2]
+    
+    if not isinstance(text, str) or not text.strip():
+        return False
+    
+    if not isinstance(score, (int, float)):
+        return False
+    
+    if not isinstance(metadata, dict):
+        return False
+    
+    return True