|
|
from abc import abstractmethod, ABC |
|
|
from smolagents import CodeAgent, Tool |
|
|
from agent.constants import ADDITIONAL_AUTHORIZED_IMPORT |
|
|
from tools.tools_utils import ToolsUtils |
|
|
|
|
|
|
|
|
class BaseAgent(ABC): |
|
|
def __init__(self, model_name: str, tools: list[Tool] | None = None, planning_interval: int = 3, max_steps: int = 12, use_all_custom_tools: bool = True): |
|
|
self.model_name: str = model_name |
|
|
self.planning_interval = planning_interval |
|
|
self.max_steps = max_steps |
|
|
self.use_all_custom_tools = use_all_custom_tools |
|
|
self.tools: list[Tool] = self.init_tools(tools) |
|
|
self.agent = self.init_agent() |
|
|
|
|
|
def __call__(self, question: str) -> str: |
|
|
print(f"Agent received question (first 50 chars): {question[:50]}...") |
|
|
fixed_answer = self.agent.run(question) |
|
|
print(f"Agent returning fixed answer: {fixed_answer}") |
|
|
return fixed_answer |
|
|
|
|
|
@abstractmethod |
|
|
def get_model(self): |
|
|
pass |
|
|
|
|
|
def init_tools(self, tools: list[Tool] | None = None): |
|
|
if tools is None: |
|
|
tools = [] |
|
|
if self.use_all_custom_tools: |
|
|
tools = ToolsUtils.get_default_tools() |
|
|
return tools |
|
|
|
|
|
def add_tool(self, tool: Tool): |
|
|
self.tools.append(tool) |
|
|
|
|
|
def init_agent(self): |
|
|
agent = CodeAgent( |
|
|
model=self.get_model(), |
|
|
tools=[t for t in self.tools], |
|
|
add_base_tools=True, |
|
|
verbosity_level=2, |
|
|
additional_authorized_imports=ADDITIONAL_AUTHORIZED_IMPORT, |
|
|
planning_interval=self.planning_interval, |
|
|
max_steps=self.max_steps |
|
|
) |
|
|
agent.system_prompt += ("\nWhen answering, provide ONLY the precise answer requested. " |
|
|
"Do not include explanations, steps, reasoning, or additional text. " |
|
|
"Be direct and specific. GAIA benchmark requires exact matching answers.") |
|
|
|
|
|
print("Agent initialized.") |
|
|
return agent |
