from haystack import Document, Pipeline
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.document_stores.in_memory import InMemoryDocumentStore


document_store = InMemoryDocumentStore()
document_store.write_documents(
    [
        Document(
            content="Retriever tuning limits how many FAQ documents reach the generator.",
            meta={"title": "retriever tuning"},
        ),
        Document(
            content="FAQ documents explain password reset and account recovery.",
            meta={"title": "account recovery"},
        ),
        Document(
            content="Support agents review FAQ documents before escalating tickets.",
            meta={"title": "support escalation"},
        ),
        Document(
            content="Billing records are exported from the invoice dashboard.",
            meta={"title": "billing export"},
        ),
    ]
)

retriever = InMemoryBM25Retriever(document_store=document_store, top_k=1)


def titles(result):
    return [document.meta["title"] for document in result["documents"]]


def show(label, result, noun="documents"):
    count = len(result["documents"])
    print(f"{label}: {count} {noun} -> {titles(result)}")


default_result = retriever.run(query="FAQ documents")
override_result = retriever.run(query="FAQ documents", top_k=3)

pipeline = Pipeline()
pipeline.add_component("retriever", retriever)
pipeline_result = pipeline.run(
    {"retriever": {"query": "FAQ documents", "top_k": 2}}
)

show("default top_k=1", default_result, "document")
show("run top_k=3", override_result)
show("pipeline top_k=2", pipeline_result["retriever"])