Open Source RAG Tools & Frameworks

import numpy as np
from typing import List, Dict, Tuple
 
class SimpleBM25:
    """Lightweight BM25 keyword search retriever."""
    def __init__(self, corpus: List[str]):
        self.corpus = corpus
        self.doc_len = [len(doc.split()) for doc in corpus]
        self.avg_doc_len = sum(self.doc_len) / (len(corpus) + 1e-5)
        self.doc_freqs = []
        self.idf = {}
        
        # Calculate term frequencies and document frequencies
        df = {}
        for doc in self.corpus:
            freqs = {}
            for word in doc.lower().split():
                freqs[word] = freqs.get(word, 0) + 1
            self.doc_freqs.append(freqs)
            for word in set(freqs.keys()):
                df[word] = df.get(word, 0) + 1
 
        # Calculate inverse document frequency (IDF)
        for word, count in df.items():
            self.idf[word] = np.log((len(corpus) - count + 0.5) / (count + 0.5) + 1.0)
 
    def retrieve(self, query: str, k1: float = 1.5, b: float = 0.75) -> List[Tuple[float, int]]:
        query_words = query.lower().split()
        scores = []
        for idx, freqs in enumerate(self.doc_freqs):
            score = 0.0
            doc_len = self.doc_len[idx]
            for word in query_words:
                if word in freqs:
                    tf = freqs[word]
                    num = self.idf.get(word, 0.0) * tf * (k1 + 1.0)
                    den = tf + k1 * (1.0 - b + b * (doc_len / self.avg_doc_len))
                    score += num / den
            scores.append((score, idx))
        return sorted(scores, key=lambda x: x[0], reverse=True)
 
# Mocked vector models to enable self-contained compilation
class MockEmbeddingModel:
    def encode(self, texts: List[str]) -> np.ndarray:
        rng = np.random.default_rng(seed=42)
        vectors = rng.random((len(texts), 384))
        norms = np.linalg.norm(vectors, axis=1, keepdims=True)
        return vectors / norms
 
class MockCrossEncoder:
    def predict(self, pairs: List[Tuple[str, str]]) -> np.ndarray:
        scores = []
        for query, doc in pairs:
            overlap = len(set(query.lower().split()) & set(doc.lower().split()))
            score = min(1.0, overlap / (len(query.split()) + 1e-5) + 0.1)
            scores.append(score)
        return np.array(scores)
 
class HybridRerankPipeline:
    def __init__(self, corpus: List[str]):
        self.corpus = corpus
        self.bm25 = SimpleBM25(corpus)
        self.embedder = MockEmbeddingModel()
        self.reranker = MockCrossEncoder()
        self.doc_embeddings = self.embedder.encode(corpus)
 
    def dense_search(self, query: str) -> List[Tuple[float, int]]:
        query_vector = self.embedder.encode([query])[0]
        similarities = np.dot(self.doc_embeddings, query_vector)
        scores = [(float(sim), idx) for idx, sim in enumerate(similarities)]
        return sorted(scores, key=lambda x: x[0], reverse=True)
 
    def reciprocal_rank_fusion(
        self, 
        sparse_res: List[Tuple[float, int]], 
        dense_res: List[Tuple[float, int]], 
        k: int = 60
    ) -> List[Tuple[float, int]]:
        rrf_scores = {}
        for rank, (_, idx) in enumerate(sparse_res):
            rrf_scores[idx] = rrf_scores.get(idx, 0.0) + (1.0 / (k + rank + 1))
        for rank, (_, idx) in enumerate(dense_res):
            rrf_scores[idx] = rrf_scores.get(idx, 0.0) + (1.0 / (k + rank + 1))
        
        sorted_rrf = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
        return [(score, idx) for idx, score in sorted_rrf]
 
    def query(self, query_text: str, top_k: int = 5, top_n: int = 2) -> List[Tuple[str, float]]:
        # 1. Execute dual searches
        sparse_res = self.bm25.retrieve(query_text)[:top_k]
        dense_res = self.dense_search(query_text)[:top_k]
        
        # 2. Fuse candidate ranks via RRF
        fused_candidates = self.reciprocal_rank_fusion(sparse_res, dense_res, k=60)[:top_k]
        
        # 3. Assemble pairs for Cross-Encoder re-scoring
        candidates = [self.corpus[idx] for _, idx in fused_candidates]
        pairs = [(query_text, doc) for doc in candidates]
        rerank_scores = self.reranker.predict(pairs)
        
        # 4. Sort and select top N items
        final_results = [(candidates[i], float(rerank_scores[i])) for i in range(len(candidates))]
        return sorted(final_results, key=lambda x: x[1], reverse=True)[:top_n]