Goodguygregory93/ai-agents-certification-code

1from llama_index.core.schema import Document2import colorama3from llama_index.core.tools import FunctionTool4from llama_index.retrievers.bm25 import BM25Retriever5import json6 7 8class Retriever():9    def __init__(self, confirm_load=False):10        '''11            Ingests the similar questions for the GAIA agent to leverage and supports 12            the exporting of the FunctionalTool for an AgentWorkFlow to leverage13            14            This class handles the detailed logic. of connecting the 15            similar questions and answers for the agent to retrieve16        17            Details:18                leverages the BM25Retriever due to it's ability to search without pre-embedding Documents19                20        '''21        self.confirm_load = confirm_load22        self.question_data = self.load_similar_questions(confirm_load)23        self.bm25_retriever = BM25Retriever.from_defaults(nodes=self.question_data)24        # returns this as an accessible tool for AgentWorkFlows25        self.similar_question_tool = FunctionTool.from_defaults(self.get_similar_question)26    27    28    def filter_questions(self, questions: list[str]) -> list[dict]:29        '''30            filters the incoming questions determining which ones are "level-1" and are potentially 31            associated with the challenge.32 33            Args:34                questions (list[dict]): question objects to filter35            Returns:36                filtered_questions (list[dict]): level 1 questions for similarity37        '''38        filtered_questions = []39        40        for question in questions:41            question_dict = json.loads(question)42            43            if question_dict['Level'] == 1:44                filtered_questions.append(question_dict)45        46        return filtered_questions47 48    49    def load_similar_questions(self, confirm_load=False) -> list[Document]:50        '''51            loads the GAIA dataset for similar questions and returns them as a list of Documents 52            for the BM25Retriever53            54            Args:55                confirm_load (bool): boolean value that controls the printed output of each appended question56            57            Returns:58                docs (list[Document]): returns the Document object for each question for the Retriever59        '''60        # open file with Python61        questions_file = open('./data/metadata.jsonl', 'r')62 63        incoming_questions = questions_file.readlines()64        filtered_questions = self.filter_questions(incoming_questions)65        66        docs = [67            Document(68                text="\n".join([69                    f"question: {filtered_questions[i]['Question']}",70                    f"final_answer: {filtered_questions[i]['Final answer']}",71                ]),72                metadata={"question_number": filtered_questions[i]['task_id']}73            ) for i in range(len(filtered_questions))74        ]75        76        77        if confirm_load:78            79            for question in docs:80                print(colorama.Fore.GREEN + f"✅ Added Question: {question.metadata}")81                82        83        return docs84                85    def get_similar_question(self, query: str) -> str:86        '''87            queries the `bm25_retriever` to determine if the question given in the query88            exists in the potential_questions89            90            Args:91                query (str): string of the potential question that might be supplied for the GAIA Agent92            Returns:93                response (str): a string response that will indicate if the user(s) exist in the guest_book.94        '''95        found_results = self.bm25_retriever.retrieve(query)96        97        if found_results:98            return "/n/n".join([doc.text for doc in found_results[:3]])99        else:100            return "No matching results in the questions database"101