# Implements retrieval functions to support knowledge access.

import datasets
from llama_index.core.schema import Document
from llama_index.retrievers.bm25 import BM25Retriever

# Load dataset
guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")

# Convert dataset entries into Document objects
docs = [
    Document(
        text="\n".join(
            [
                f"Name: {guest['name']}",
                f"Relation: {guest['relation']}",
                f"Description: {guest['description']}",
                f"Email: {guest['email']}",
            ]
        ),
        metadata={"name": guest["name"]},
    )
    for guest in guest_dataset
]

bm25_retriever = BM25Retriever.from_defaults(nodes=docs)


def guest_info_retriever(query: str) -> str:
    """Retrieves detailed info about gala guests based on their name or relation"""
    results = bm25_retriever.retrieve(query)
    if results:
        return "\n\n".join([doc.text for doc in results[:3]])
    else:
        return "No matching guest information found."