import os
from smolagents import CodeAgent, DuckDuckGoSearchTool, OpenAIServerModel, VisitWebpageTool, LiteLLMModel
from tools.attached_files import csv_reader, excel_reader, transcribe_audio
from tools.basic_math import *
from tools.browser import arvix_search, wiki_search

def get_agent(provider:str) -> CodeAgent:
  search_tool = DuckDuckGoSearchTool(max_results=5)
  web_page_tool = VisitWebpageTool(max_output_length=1_000_000)

  codestral_key = os.getenv('CODESTRAL_API_KEY')
  groq_key = os.getenv('GROQ_API_KEY')

  model = None

  if provider=='groq':
    model = LiteLLMModel(
      model_id="groq/deepseek-r1-distill-llama-70b",
      api_base="https://api.groq.com/openai/v1/chat/completions",
      api_key=groq_key)

  else:
    model = OpenAIServerModel(
      model_id='codestral-latest',
      api_base="https://codestral.mistral.ai/v1/chat/completions",
      api_key=codestral_key)

  return CodeAgent(tools=[search_tool, web_page_tool, add, subtract, multiply, divide, modulus, rounder, power, square_root, csv_reader, excel_reader, transcribe_audio, wiki_search, arvix_search], model=model, additional_authorized_imports=['random', 'time'], add_base_tools=True)