"""Find the document chunks most relevant to a question."""
import config
from embed import embed_text
from ingest import collection


def retrieve(question):
    """Return the top-K most relevant documents for the question."""
    q_vec = embed_text(question)
    results = collection.query(query_embeddings=[q_vec], n_results=config.TOP_K)
    return results["documents"][0]