tools/main_agent.py

from llama_index.core.tools import FunctionTool
import asyncio
from tools.sub_agents import MathAgent, WebAgent, BaseAgent

math_agent = MathAgent()
web_agent = WebAgent()

MAIN_AGENT_PROMPT = """
You are the Main Agent, responsible for solving complex tasks by coordinating a team of specialized agents.
When given a task, you must carefully and systematically follow this procedure:

Carefully read the entire task, paying close attention to any specific instructions provided (such as required output format, style, structure, or other constraints).
Break down the task into clear, manageable subproblems.
For each subproblem:
Clearly define what needs to be achieved.
Determine which specialized agent (e.g., Math Agent, Websearch Agent, etc.) is best suited to handle it.
Create a clear, specific instruction for the agent, ensuring it aligns with the overall task requirements.
Delegate each subproblem to the appropriate specialized agent.
Collect and review the agent outputs, checking that they meet the original task's instructions.
Integrate all partial results logically and cohesively.
Format the final answer precisely according to the tasks specified requirements.
Present the complete, correct, and clearly structured final output.
Important:

Always adhere strictly to any instructions related to formatting, style, or special requirements.
Think step-by-step, delegate intelligently, and assemble the solution carefully.
Never solve detailed subproblems yourself — always delegate to the appropriate agent.
Your answer should only be the solution; don't write "Answer:" and don't write a sentence when possible.
Accuracy, completeness, and strict compliance with the user's original instructions are your top priorities.

If the user says they have a file, you should use the file agent to get the information from the file. You dont have to access file directly,so it doesnt matter if you dont have the file or the capability to access it. Just pass the exact question along to the file agent.
"""

async def call_math_agent(question: str) -> str:
    """Call a math agent by delegating part of the question to it."""
    return await math_agent(question=question)


async def call_web_agent(question: str) -> str:
    """Call a search agent by delegating part of the question to it."""
    return await web_agent(question=question)


def get_tools():
    tools = [
        FunctionTool.from_defaults(fn=call_math_agent),
        FunctionTool.from_defaults(fn=call_web_agent),
    ]
    return tools


class MainAgent(BaseAgent):
    def __init__(self):
        super().__init__(
            name="main_agent",
            description="Main agent which creates subtasks and delegates them to other agents. You can use the other agents to look up information in the web, solve math problems, or transcribe audio files. ",
            prompt=MAIN_AGENT_PROMPT,
            tools=get_tools(),
        )

if __name__ == "__main__":
    main_agent = MainAgent()
    asyncio.run(main_agent("I'm making a grocery list for my mom, but she's a professor of botany and she's a real stickler when it comes to categorizing things. I need to add different foods to different categories on the grocery list, but if I make a mistake, she won't buy anything inserted in the wrong category. Here's the list I have so far:\n\nmilk, eggs, flour, whole bean coffee, Oreos, sweet potatoes, fresh basil, plums, green beans, rice, corn, bell pepper, whole allspice, acorns, broccoli, celery, zucchini, lettuce, peanuts\n\nI need to make headings for the fruits and vegetables. Could you please create a list of just the vegetables from my list? If you could do that, then I can figure out how to categorize the rest of the list into the appropriate categories. But remember that my mom is a real stickler, so make sure that no botanical fruits end up on the vegetable list, or she won't get them when she's at the store. Please alphabetize the list of vegetables, and place each item in a comma separated list."))