EmbeddingsFilter: fast relevance filtering without LLM calls

# pip install langchain-openai langchain-community scikit-learn numpy

from langchain_openai import OpenAIEmbeddings
from langchain.schema import Document
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

class EmbeddingsFilter:
    """Filter retrieved docs by embedding similarity before reranking."""
    
    def __init__(self, embeddings_model, similarity_threshold=0.68, min_docs=2):
        self.embeddings = embeddings_model
        self.threshold = similarity_threshold
        self.min_docs = min_docs
    
    def filter_documents(self, query: str, documents: list[Document]) -> list[Document]:
        """Filter docs by cosine similarity to query embedding."""
        if not documents:
            return []
        
        # Embed query
        query_embedding = self.embeddings.embed_query(query)
        
        # Embed all documents (or reuse cached embeddings if available)
        doc_embeddings = []
        for doc in documents:
            if 'embedding' in doc.metadata:
                # Reuse cached embedding
                doc_embeddings.append(doc.metadata['embedding'])
            else:
                # Embed on the fly (avoid this in production for retrieved docs)
                doc_embeddings.append(self.embeddings.embed_query(doc.page_content))
        
        # Compute similarities
        similarities = cosine_similarity([query_embedding], doc_embeddings)[0]
        
        # Filter by threshold
        filtered = [(doc, sim) for doc, sim in zip(documents, similarities) if sim >= self.threshold]
        
        # Safety fallback: ensure minimum docs
        if len(filtered) < self.min_docs:
            # Sort by similarity descending and take top min_docs
            sorted_pairs = sorted(zip(documents, similarities), key=lambda x: x[1], reverse=True)
            filtered = sorted_pairs[:self.min_docs]
        
        # Return docs with similarity metadata for debugging
        result = []
        for doc, sim in filtered:
            doc.metadata['filter_similarity'] = float(sim)
            result.append(doc)
        
        return result


# Example usage
if __name__ == '__main__':
    embeddings = OpenAIEmbeddings(model='text-embedding-3-small')
    
    # Mock retrieved documents (in real RAG, these come from vector DB)
    documents = [
        Document(
            page_content='Python is a programming language used for data science.',
            metadata={'embedding': embeddings.embed_query('Python is a programming language used for data science.'), 'source': 'doc1'}
        ),
        Document(
            page_content='Machine learning requires large datasets.',
            metadata={'embedding': embeddings.embed_query('Machine learning requires large datasets.'), 'source': 'doc2'}
        ),
        Document(
            page_content='The capital of France is Paris.',
            metadata={'embedding': embeddings.embed_query('The capital of France is Paris.'), 'source': 'doc3'}
        ),
        Document(
            page_content='Neural networks learn hierarchical features from data.',
            metadata={'embedding': embeddings.embed_query('Neural networks learn hierarchical features from data.'), 'source': 'doc4'}
        ),
    ]
    
    query = 'How do neural networks process data?'
    
    # Create filter and apply
    filter_obj = EmbeddingsFilter(embeddings, similarity_threshold=0.65, min_docs=2)
    filtered_docs = filter_obj.filter_documents(query, documents)
    
    print(f'Query: {query}')
    print(f'Retrieved: {len(documents)} docs | After filter: {len(filtered_docs)} docs\n')
    for doc in filtered_docs:
        print(f'Source: {doc.metadata["source"]} | Similarity: {doc.metadata["filter_similarity"]:.3f}')
        print(f'  Content: {doc.page_content}\n')