first-aid-rag-assistant / src / search.py
search.py
Raw
"""Pure search utility functions: text search, vector search, and RRF fusion.

Lives at the project root so that both rag_usage.py (production) and
evaluation/evaluation.py (evaluation) can import from here without
creating a circular dependency.
"""


def text_search(query, index, num_results=10):
    boost_dict = {
        "question": 0.5,
        "answer": 1.0,
    }
    return index.search(query, num_results=num_results, boost_dict=boost_dict)


def vector_search(query, embedder, index, num_results=10):
    q_vector = embedder.encode(query)
    return index.search(q_vector, num_results=num_results)


def rrf(result_lists, k=60, num_results=5):
    """Reciprocal Rank Fusion over multiple ranked result lists."""
    scores = {}
    docs = {}

    for results in result_lists:
        for rank, doc in enumerate(results):
            key = doc["id"]
            scores[key] = scores.get(key, 0) + 1 / (k + rank)
            docs[key] = doc

    ranked = sorted(scores, key=scores.get, reverse=True)
    return [docs[key] for key in ranked[:num_results]]