Hybrid retrieval: combining BM25 and vector

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.core.retrievers import BM25Retriever
from llama_index.core.postprocessors import SimilarityPostprocessor
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb
import os

# Set up API key
os.environ["OPENAI_API_KEY"] = "your-openai-key"

# Initialize settings
Settings.llm = OpenAI(model="gpt-4o-mini")
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")

# Create sample documents
from llama_index.core.schema import Document

docs = [
    Document(text="Python DataFrame operations are fast with vectorized operations in NumPy."),
    Document(text="Pandas DataFrame provides high-performance data manipulation tools."),
    Document(text="The Apache Spark DataFrame API enables distributed SQL queries across clusters."),
    Document(text="Data structures in Python include lists, tuples, and dictionaries for storage."),
    Document(text="Semantic similarity measures how related two texts are in meaning."),
]

# Initialize Chroma vector store
client = chromadb.Client()
vector_store = ChromaVectorStore(chroma_collection=client.get_or_create_collection("hybrid_demo"))

# Create vector index
vector_index = VectorStoreIndex.from_documents(docs, vector_store=vector_store, show_progress=True)

# Create BM25 retriever
bm25_retriever = BM25Retriever.from_documents(docs, similarity_top_k=3)

# Create vector retriever
vector_retriever = vector_index.as_retriever(similarity_top_k=3)

# Implement hybrid retrieval with reciprocal rank fusion
from llama_index.core.retrievers import BaseRetriever
from llama_index.core.schema import NodeWithScore
from typing import List

class HybridRetriever(BaseRetriever):
    def __init__(self, bm25_retriever, vector_retriever):
        self.bm25_retriever = bm25_retriever
        self.vector_retriever = vector_retriever
        super().__init__()

    def _retrieve(self, query_str: str) -> List[NodeWithScore]:
        # Get results from both retrievers
        bm25_nodes = self.bm25_retriever.retrieve(query_str)
        vector_nodes = self.vector_retriever.retrieve(query_str)

        # Create a dictionary to store reciprocal rank fusion scores
        node_scores = {}
        k = 60  # Constant for RRF formula

        # Add BM25 scores
        for rank, node in enumerate(bm25_nodes):
            node_id = node.node_id
            rrf_score = 1.0 / (k + rank + 1)
            node_scores[node_id] = node_scores.get(node_id, 0) + rrf_score

        # Add vector scores
        for rank, node in enumerate(vector_nodes):
            node_id = node.node_id
            rrf_score = 1.0 / (k + rank + 1)
            node_scores[node_id] = node_scores.get(node_id, 0) + rrf_score

        # Create result list sorted by combined score
        result_nodes = []
        all_nodes = {node.node_id: node for node in bm25_nodes + vector_nodes}

        for node_id in sorted(node_scores.keys(), key=lambda x: node_scores[x], reverse=True):
            node = all_nodes[node_id]
            result_nodes.append(NodeWithScore(node=node.node, score=node_scores[node_id]))

        return result_nodes

# Test hybrid retrieval
hybrid_retriever = HybridRetriever(bm25_retriever, vector_retriever)
query = "DataFrame performance optimization"
results = hybrid_retriever.retrieve(query)

print(f"Hybrid retrieval results for '{query}':\n")
for i, result in enumerate(results, 1):
    print(f"{i}. Score: {result.score:.4f}")
    print(f"   Text: {result.node.get_content()[:70]}...\n")