Top-K recherche de documents

from langchain_core.documents import Document
docs = [
    Document(page_content='Python est un langage de programmation'),
    Document(page_content='Java est un autre langage'),
    Document(page_content='Python et Django pour le web'),
    Document(page_content='Les chats sont mignons'),
]
result = search_top_k('Python programmation', docs, 2)
assert len(result) == 2, f'Attendu 2 résultats, obtenu {len(result)}'
assert all(isinstance(r, Document) for r in result), 'Les résultats doivent etre des Documents'
top1 = result[0].page_content
assert 'Python' in top1, f'Le premier résultat devrait mentionner Python: {top1}'
result_all = search_top_k('Python', docs, 10)
assert len(result_all) == 4, 'Si k > nombre de docs, renvoyer tous les docs'

Solution officielle

from langchain_core.documents import Document
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

def search_top_k(query, documents, k):
    texts = [doc.page_content for doc in documents]
    vectorizer = TfidfVectorizer()
    doc_vectors = vectorizer.fit_transform(texts)
    query_vector = vectorizer.transform([query])
    scores = cosine_similarity(query_vector, doc_vectors)[0]
    top_indices = scores.argsort()[::-1][:k]
    return [documents[i] for i in top_indices]