tools/sub_agents.py

import os
import asyncio
import logging
from abc import ABC

from llama_index.core.agent.workflow import FunctionAgent
from llama_index.llms.openai import OpenAI
from llama_index.core.tools import FunctionTool


from tools.base_tools import get_base_tools
from tools.math_tools import get_math_tools
from tools.web_tools import get_web_tools
from tools.file_tools import get_file_tools

MATH_AGENT_PROMPT = """
You are an AI agent highly skilled in solving math problems using basic operations: addition, subtraction, multiplication, and division. You have access to internal functions that perform these operations precisely.

Your task is to always find the correct final answer by following a clear and reliable solving process:

Carefully read and understand the user's problem.

Break the problem down into smaller, logical steps.

Plan only one step at a time.

For each step:

Clearly explain what operation needs to be performed and why.
Use the appropriate internal tools to carry out the operation.
State the result of each step explicitly before continuing.
Proceed methodically until the problem is fully solved.
Summarize your intermediate results and plan your next step based on the results.

Before answering the question, verify if your answer is correct, and if it is not correct, try to correct it.

At the end, present the final, correct answer clearly.

Accuracy is your highest priority, but your explanations should also be easy for the user to follow. Always solve step-by-step, without skipping any parts, even if the problem looks simple.
"""

WEB_AGENT_PROMPT = """
You are an AI agent specialized in gathering accurate and up-to-date information by performing web searches. When given a user's prompt, you must follow this process:

Carefully read and fully understand the user's query.
Determine what specific information or context is needed to answer it accurately.
Think about the information you need, break it down into as many queries as you see necessary, and formulate your plan.
Plan and perform targeted web searches to gather relevant, reliable information.
Consider doing more web searches to get more information about specific topics within the gathered information as you see fit.
Summarize key findings clearly, identifying trusted sources where appropriate.
Use the gathered information to construct a complete, precise, and helpful final answer.
If multiple sources provide different information, note the differences and recommend the most trustworthy or consistent one.
Always prioritize accuracy, clarity, and usefulness in your final response.
Verify the current time before you check the up-to-date information.
Break the problem down into smaller steps.

Plan only one step at a time.

For each step:

Clearly explain what searches needs to be performed and why.
Use the appropriate internal tools to carry out the tasks.
State the result of each step explicitly before continuing.
Proceed methodically until the problem is fully solved.
Only rely on verified information retrieved from your searches. Be methodical: search first, then answer — never assume facts without confirming them.
"""

FILE_AGENT_PROMPT = """
You are a file agent. You are given a question and you need to answer it. Always use the get_task_id tool to get the task id of the question. Then pass the task id to the transcribe_mp3_file tool to get the transcript of the audio file. Then use the transcribed text to answer the question. It doesnt matter if you dont have the file or the capability to access it. Just pass the exact question along to the transcribe_mp3_file tool to get the transcript which you can then use to answer the question.
"""


class BaseAgent(ABC):
    def __init__(
        self, name, description, prompt, tools, llm_model="gpt-4.1"
    ):
        logging.info(f"{name} Agent initialized.")

        self.llm = OpenAI(
            model=llm_model,
            temperature=0.1,
            api_key=os.getenv("OPENAI_API_KEY"),
        )

        agent_tools = get_base_tools()
        agent_tools.extend(tools)

        self.agent = FunctionAgent(
            llm=self.llm,
            tools=agent_tools,
            name=name,
            description=description,
            system_prompt=prompt,
        )

    async def chat(self, question: str, history: dict) -> str:
        answer = await self(question)
        return answer

    async def __call__(self, question: str) -> str:
        answer = await self.agent.run(user_msg=question, timeout=60)
        print(answer)
        return answer.response.content
    

class MathAgent(BaseAgent):
    def __init__(self):
        super().__init__(
            name="math agent",
            description="A agent specializing in solving math problems.",
            prompt=MATH_AGENT_PROMPT,
            tools=get_math_tools(),
            llm_model="o4-mini",
        )

class WebAgent(BaseAgent):
    def __init__(self):
        super().__init__(
            name="web agent",
            description="A agent specializing in searching the web for information.",
            prompt=WEB_AGENT_PROMPT,
            tools=get_web_tools(),
            llm_model="o4-mini",
        )

class FileAgent(BaseAgent):
    def __init__(self):
        super().__init__(
            name="file agent",
            description="A agent specializing in searching through attached files for information. Always use the get_task_id tool to get the task id of the question. Then pass the task id to the transcribe_mp3_file tool to get the transcript of the audio file. Then use the transcribed text to answer the question.",
            prompt=FILE_AGENT_PROMPT,
            tools=get_file_tools(),
            llm_model="o4-mini",
        )