agent.py

from tools import visit_webpage, image_generation_tool, image_diplay_tool
from smolagents import (
    CodeAgent,
    ToolCallingAgent,
    HfApiModel,
    ManagedAgent,
    DuckDuckGoSearchTool,
    LiteLLMModel,
    PythonInterpreterTool
)
from groq import Groq
import os

def multi_agent_framework(model_id):
    # model = HfApiModel(model_id)
    api_key = os.getenv("GROQ_API_KEY")
    model = LiteLLMModel(
                model_id,
                api_base="https://api.groq.com/openai/v1",
                api_key=api_key
                )
    model.flatten_messages_as_text = True

    web_agent = ToolCallingAgent(
        tools=[DuckDuckGoSearchTool(), visit_webpage],
        model=model,
        max_steps=3,
    )

    managed_web_agent = ManagedAgent(
        agent=web_agent,
        name="search",
        description="Runs web searches for you. Give it your query as an argument.",
    )

    head_counsellor_agent = CodeAgent(
        tools=[PythonInterpreterTool()],
        model=model,
        managed_agents=[managed_web_agent],
        additional_authorized_imports=["time", "numpy", "pandas", "requests"],
    )

    return head_counsellor_agent