retriever.py

from smolagents import Tool
import datasets

class GuestInfoRetrieverTool(Tool):
    name = "guest_info_retriever"
    description = "Retrieves detailed information about gala guests based on their name or relation."
    inputs = {
        "query": {
            "type": "string",
            "description": "The name or relation of the guest you want information about."
        }
    }
    output_type = "string"

    def __init__(self, guest_list, **kwargs):
        super().__init__(**kwargs)
        self.guest_list = guest_list

    def forward(self, query: str):
        query = query.lower()
        # Busca simples por correspondência de texto em qualquer campo
        results = []
        for guest in self.guest_list:
            # Verifica se a query está no nome ou na descrição
            if query in guest['name'].lower() or query in guest['description'].lower() or query in guest['relation'].lower():
                info = (f"Name: {guest['name']}\n"
                        f"Relation: {guest['relation']}\n"
                        f"Description: {guest['description']}\n"
                        f"Email: {guest['email']}")
                results.append(info)
        
        if results:
            # Retorna os 3 primeiros matches
            return "\n\n---\n\n".join(results[:3])
        else:
            return "No matching guest information found."

def load_guest_dataset():
    # Carrega o dataset da Hugging Face
    dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")
    
    # Passamos a lista crua de dicionários para a Tool
    return GuestInfoRetrieverTool(guest_list=list(dataset))