from smolagents import WebSearchTool, CodeAgent, WikipediaSearchTool, LiteLLMModel, VisitWebpageTool, PythonInterpreterTool, FinalAnswerTool, UserInputTool
from agent_tools import *
import os

# load the system prompt from the file
with open("system_prompt.txt", "r", encoding="utf-8") as f:
    system_prompt = f.read()

# Initialize model
api_key = os.environ.get("GOOGLE_API_KEY")
model = LiteLLMModel("gemini/gemini-2.5-flash",api_key=api_key)

# Initialize tools
search_tool = WebSearchTool()
wiki_search_tool = WikipediaSearchTool()
visit_webpage_tool = VisitWebpageTool()
python_interpreter_tool = PythonInterpreterTool()
final_answer_tool = FinalAnswerTool()
user_input_tool = UserInputTool()
calculator_tools = [addition_tool, subtraction_tool, multiplication_tool, division_tool, exponent_tool, modulus_tool]
main_tools = [search_tool, 
              wiki_search_tool, 
              transcriber_tool, 
              visit_webpage_tool, 
              python_interpreter_tool, 
              excel_loader_tool,
              final_answer_tool,
              user_input_tool,
             ]
authorized_imports = ['requests', 'bs4', 'whisper', 'langchain_community', 'langchain', 'markdownify', 'openpyxl', 'tabulate']

# Agent Definitions
calculator_agent = CodeAgent(
    model = model,
    tools = calculator_tools,
    name = "arithmetic_calculator_agent",
    description = "You perform basic arithmetic calculations between two numbers.",
    max_steps = 4,
    verbosity_level=3,
)

main_agent = CodeAgent(
    tools=main_tools,
    model=model,
    additional_authorized_imports=authorized_imports,
    instructions=system_prompt,
    managed_agents=[calculator_agent],
    max_steps=8,
    verbosity_level=3,
)