File size: 1,184 Bytes
d4f6460 74a8ea8 050ce04 726d2d9 38e5771 74a8ea8 69c427f 74a8ea8 c5eab2c 74a8ea8 a210bde 74a8ea8 ab0b2e4 74a8ea8 3ad9743 050ce04 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
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) |