from haystack import Document
from haystack.components.joiners import DocumentJoiner


keyword_documents = [
    Document(
        id="doc-hybrid-search",
        content="Hybrid search combines keyword and vector retrieval.",
        score=8.4,
    ),
    Document(
        id="doc-reranker",
        content="A reranker reorders joined retrieval results.",
        score=6.1,
    ),
    Document(
        id="doc-cleaning",
        content="Document cleaning removes repeated headers.",
        score=3.2,
    ),
]

embedding_documents = [
    Document(
        id="doc-hybrid-search",
        content="Hybrid search combines keyword and vector retrieval.",
        score=0.94,
    ),
    Document(
        id="doc-question-answering",
        content="RAG answers questions from retrieved context.",
        score=0.88,
    ),
    Document(
        id="doc-reranker",
        content="A reranker reorders joined retrieval results.",
        score=0.70,
    ),
]

joiner = DocumentJoiner(
    join_mode="reciprocal_rank_fusion",
    top_k=3,
)

result = joiner.run(
    documents=[keyword_documents, embedding_documents]
)
joined_documents = result["documents"]

joined_ids = [document.id for document in joined_documents]
assert joined_ids == [
    "doc-hybrid-search",
    "doc-reranker",
    "doc-question-answering",
]
assert joined_ids.count("doc-hybrid-search") == 1

for position, document in enumerate(joined_documents, start=1):
    print(f"{position}. {document.id} | {document.score:.4f}")