| | 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. ") |
| | print(agent.system_prompt) |
| | |
| | return agent |
