basic_agent.py

from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
from llama_index.core.tools import FunctionTool
from llama_index.llms.ollama import Ollama
from tools import Tools
from retriever import Retriever
import colorama
from dotenv import load_dotenv
import os

class BasicAgent:
    load_dotenv()
    
    def __init__(self):
        print(colorama.Fore.GREEN + "BasicAgent initialized.")
        LLM_MODEL = os.environ.get("LLM_MODEL")
        SYS_PROMPT = os.environ.get("SYS_PROMPT")
        HF_TOKEN = os.environ.get("HF_TOKEN") 
        
        if SYS_PROMPT is None:
            raise Exception("Missing SYS_PROMPT: add SYS_PROMPT inside your .env")
        
        if HF_TOKEN is None:
            raise Exception("Missing HF_TOKEN: add HF_TOKEN inside your .env")
            
            
        self.agent_llm = self._initialize_llm(LLM_MODEL=LLM_MODEL, HF_TOKEN=HF_TOKEN)
        
        self.agent_tools = self._initialize_agent_tools()
        self.agent_work_flow = self._initialize_workflow()
        self.system_prompt = SYS_PROMPT
        
    async def __call__(self, question: str, sys_prompt=True) -> str:
        '''
            the __call__ method is the method that is triggered when BasicAgent()
            is called the `question` field is expected to be applied to the Agent's 
            LlamaIndex Workflow and produce an LLM response 
            
            the sys_prompt is defaulted to True and is found here: 
            
            [GAIA HuggingFace Paper](https://huggingface.co/spaces/gaia-benchmark/leaderboard)
            
            Args:
                question (str): the desired prompt that the user desires information
                sys_prompt (bool): defaults to True as the GAIA specifications prompt is supplied in the study
            Returns:
                response (str): response from the LLM 
        '''
        if sys_prompt:
            prompt = self.system_prompt + f" QUESTION: {question}"
            agent_output = await self.agent_work_flow.run(prompt)
            response = agent_output.response.blocks[0].text
        else:
            agent_output = await self.agent_work_flow.run(question)
            response = agent_output.response.blocks[0].text
            

        return response
    
    def _initialize_llm(self, LLM_MODEL: str, HF_TOKEN: str):
        """
        initialize the llm brain of the Agent using an Ollama installed
        model. Returns the Ollama object

        Args:
            LLM_MODEL (str): provider model name ex: "llama3.2:latest"

        Returns:
            llm: returns the HuggingFaceInferenceAPI Model
        """
        
        print(colorama.Fore.YELLOW + "🧠 initializing an LLM for the Agent")
        print('-'*15)
        
        
        llm = HuggingFaceInferenceAPI(
            model_name=LLM_MODEL,
            token=HF_TOKEN,
            provider='auto'
            )
        
        return llm
    
    def _initialize_agent_tools(self) -> list[FunctionTool]:
        """
            returns the FunctionalTools from the Tool class calling the
            tool_belt and collecting the tools into a list. also loads
            the Retriever tool `similar_question_tool` for question querying
            
            Args:
                None
            Returns:
                found_tools (list[FunctionTool]): the tools initiated from the Tool agent
        """
        found_tools = []
        
        # tools = Tools(status_updates=True).tool_belt
        
        similar_question_tool = Retriever().similar_question_tool
        found_tools.append(similar_question_tool)
    
        # for tool in tools:
        #     found_tools.append(tool)
            
        return found_tools
    
    def _initialize_workflow(self):
        '''
            creates the BasicAgent Workflow for the agent that will leverage the
            Tools `agent_tools` and the `agent_llm` from the HuggingFaceInterfaceAPI
            
            Args:                
                None
            Returns:
                basic_agent_workflow (AgentWorkflow): leverages the agent_tools, and agent_llm to create a workflow
                                                    that will create the workflow
        '''
        basic_agent_workflow = AgentWorkflow.from_tools_or_functions(
                self.agent_tools,
                llm=self.agent_llm
            )
        
        return basic_agent_workflow