import osimport sysfrom dotenv import load_dotenvfrom langchain.docstore.document import Document
from typing import Listfrom rank_bm25 import BM25Okapiimport numpy as np

def create_bm25_index(documents: List[Document]) -> BM25Okapi:"""Create a BM25 index from the given documents.
BM25 (Best Matching 25) is a ranking function used in information retrieval.It's based on the probabilistic retrieval framework and is an improvement over TF-IDF.
Args:documents (List[Document]): List of documents to index.
Returns:BM25Okapi: An index that can be used for BM25 scoring."""# Tokenize each document by splitting on whitespace# This is a simple approach and could be improved with more sophisticated tokenizationtokenized_docs = [doc.page_content.split() for doc in documents]return BM25Okapi(tokenized_docs)

def fusion_retrieval(vectorstore, bm25, query: str, k: int = 5, alpha: float = 0.5) -> List[Document]:"""Perform fusion retrieval combining keyword-based (BM25) and vector-based search.
Args:vectorstore (VectorStore): The vectorstore containing the documents.bm25 (BM25Okapi): Pre-computed BM25 index.query (str): The query string.k (int): The number of documents to retrieve.alpha (float): The weight for vector search scores (1-alpha will be the weight for BM25 scores).
Returns:List[Document]: The top k documents based on the combined scores."""# Step 1: Get all documents from the vectorstoreall_docs = vectorstore.similarity_search("", k=vectorstore.index.ntotal)
# Step 2: Perform BM25 searchbm25_scores = bm25.get_scores(query.split())
# Step 3: Perform vector searchvector_results = vectorstore.similarity_search_with_score(query, k=len(all_docs))
# Step 4: Normalize scoresvector_scores = np.array([score for _, score in vector_results])vector_scores = 1 - (vector_scores - np.min(vector_scores)) / (np.max(vector_scores) - np.min(vector_scores))
bm25_scores = (bm25_scores - np.min(bm25_scores)) / (np.max(bm25_scores) - np.min(bm25_scores))
# Step 5: Combine scorescombined_scores = alpha * vector_scores + (1 - alpha) * bm25_scores
# Step 6: Rank documentssorted_indices = np.argsort(combined_scores)[::-1]
# Step 7: Return top k documentsreturn [all_docs[i] for i in sorted_indices[:k]]