Nioi's picture
trying deepseek
a210bde
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)