Goodguygregory93/ai-agents-certification-code

1from llama_index.core.agent.workflow import AgentWorkflow2from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI3from llama_index.core.tools import FunctionTool4from llama_index.llms.ollama import Ollama5from tools import Tools6from retriever import Retriever7import colorama8from dotenv import load_dotenv9import os10 11class BasicAgent:12    load_dotenv()13    14    def __init__(self):15        print(colorama.Fore.GREEN + "BasicAgent initialized.")16        LLM_MODEL = os.environ.get("LLM_MODEL")17        SYS_PROMPT = os.environ.get("SYS_PROMPT")18        HF_TOKEN = os.environ.get("HF_TOKEN") 19        20        if SYS_PROMPT is None:21            raise Exception("Missing SYS_PROMPT: add SYS_PROMPT inside your .env")22        23        if HF_TOKEN is None:24            raise Exception("Missing HF_TOKEN: add HF_TOKEN inside your .env")25            26            27        self.agent_llm = self._initialize_llm(LLM_MODEL=LLM_MODEL, HF_TOKEN=HF_TOKEN)28        29        self.agent_tools = self._initialize_agent_tools()30        self.agent_work_flow = self._initialize_workflow()31        self.system_prompt = SYS_PROMPT32        33    async def __call__(self, question: str, sys_prompt=True) -> str:34        '''35            the __call__ method is the method that is triggered when BasicAgent()36            is called the `question` field is expected to be applied to the Agent's 37            LlamaIndex Workflow and produce an LLM response 38            39            the sys_prompt is defaulted to True and is found here: 40            41            [GAIA HuggingFace Paper](https://huggingface.co/spaces/gaia-benchmark/leaderboard)42            43            Args:44                question (str): the desired prompt that the user desires information45                sys_prompt (bool): defaults to True as the GAIA specifications prompt is supplied in the study46            Returns:47                response (str): response from the LLM 48        '''49        if sys_prompt:50            prompt = self.system_prompt + f" QUESTION: {question}"51            agent_output = await self.agent_work_flow.run(prompt)52            response = agent_output.response.blocks[0].text53        else:54            agent_output = await self.agent_work_flow.run(question)55            response = agent_output.response.blocks[0].text56            57 58        return response59    60    def _initialize_llm(self, LLM_MODEL: str, HF_TOKEN: str):61        """62        initialize the llm brain of the Agent using an Ollama installed63        model. Returns the Ollama object64 65        Args:66            LLM_MODEL (str): provider model name ex: "llama3.2:latest"67 68        Returns:69            llm: returns the HuggingFaceInferenceAPI Model70        """71        72        print(colorama.Fore.YELLOW + "🧠 initializing an LLM for the Agent")73        print('-'*15)74        75        76        llm = HuggingFaceInferenceAPI(77            model_name=LLM_MODEL,78            token=HF_TOKEN,79            provider='auto'80            )81        82        return llm83    84    def _initialize_agent_tools(self) -> list[FunctionTool]:85        """86            returns the FunctionalTools from the Tool class calling the87            tool_belt and collecting the tools into a list. also loads88            the Retriever tool `similar_question_tool` for question querying89            90            Args:91                None92            Returns:93                found_tools (list[FunctionTool]): the tools initiated from the Tool agent94        """95        found_tools = []96        97        # tools = Tools(status_updates=True).tool_belt98        99        similar_question_tool = Retriever().similar_question_tool100        found_tools.append(similar_question_tool)101    102        # for tool in tools:103        #     found_tools.append(tool)104            105        return found_tools106    107    def _initialize_workflow(self):108        '''109            creates the BasicAgent Workflow for the agent that will leverage the110            Tools `agent_tools` and the `agent_llm` from the HuggingFaceInterfaceAPI111            112            Args:                113                None114            Returns:115                basic_agent_workflow (AgentWorkflow): leverages the agent_tools, and agent_llm to create a workflow116                                                    that will create the workflow117        '''118        basic_agent_workflow = AgentWorkflow.from_tools_or_functions(119                self.agent_tools,120                llm=self.agent_llm121            )122        123        return basic_agent_workflow