refactor(agent_utils.py): extract fallback handler initialization into a separate function to improve code readability and maintainability

fix(ciayn_agent.py): reset fallback handler after executing tool to ensure proper state management
style(fallback_handler.py): format method signature for better readability
This commit is contained in:
Ariel Frischer 2025-02-13 16:53:22 -08:00
parent 63e48db9de
commit 2420dfbb4f
3 changed files with 35 additions and 13 deletions

View File

@ -414,8 +414,10 @@ def run_research_agent(
if agent is not None: if agent is not None:
logger.debug("Research agent created successfully") logger.debug("Research agent created successfully")
fallback_handler = FallbackHandler(config, tools) none_or_fallback_handler = init_fallback_handler(agent, config, tools)
_result = run_agent_with_retry(agent, prompt, run_config, fallback_handler) _result = run_agent_with_retry(
agent, prompt, run_config, none_or_fallback_handler
)
if _result: if _result:
# Log research completion # Log research completion
log_work_event(f"Completed research phase for: {base_task_or_query}") log_work_event(f"Completed research phase for: {base_task_or_query}")
@ -531,8 +533,10 @@ def run_web_research_agent(
console.print(Panel(Markdown(console_message), title="🔬 Researching...")) console.print(Panel(Markdown(console_message), title="🔬 Researching..."))
logger.debug("Web research agent completed successfully") logger.debug("Web research agent completed successfully")
fallback_handler = FallbackHandler(config, tools) none_or_fallback_handler = init_fallback_handler(agent, config, tools)
_result = run_agent_with_retry(agent, prompt, run_config, fallback_handler) _result = run_agent_with_retry(
agent, prompt, run_config, none_or_fallback_handler
)
if _result: if _result:
# Log web research completion # Log web research completion
log_work_event(f"Completed web research phase for: {query}") log_work_event(f"Completed web research phase for: {query}")
@ -637,9 +641,9 @@ def run_planning_agent(
try: try:
print_stage_header("Planning Stage") print_stage_header("Planning Stage")
logger.debug("Planning agent completed successfully") logger.debug("Planning agent completed successfully")
fallback_handler = FallbackHandler(config, tools) none_or_fallback_handler = init_fallback_handler(agent, config, tools)
_result = run_agent_with_retry( _result = run_agent_with_retry(
agent, planning_prompt, run_config, fallback_handler agent, planning_prompt, run_config, none_or_fallback_handler
) )
if _result: if _result:
# Log planning completion # Log planning completion
@ -745,8 +749,10 @@ def run_task_implementation_agent(
try: try:
logger.debug("Implementation agent completed successfully") logger.debug("Implementation agent completed successfully")
fallback_handler = FallbackHandler(config, tools) none_or_fallback_handler = init_fallback_handler(agent, config, tools)
_result = run_agent_with_retry(agent, prompt, run_config, fallback_handler) _result = run_agent_with_retry(
agent, prompt, run_config, none_or_fallback_handler
)
if _result: if _result:
# Log task implementation completion # Log task implementation completion
log_work_event(f"Completed implementation of task: {task}") log_work_event(f"Completed implementation of task: {task}")
@ -846,18 +852,29 @@ def get_agent_type(agent: RAgents) -> Literal["CiaynAgent", "React"]:
Determines the type of the agent. Determines the type of the agent.
Returns "CiaynAgent" if agent is an instance of CiaynAgent, otherwise "React". Returns "CiaynAgent" if agent is an instance of CiaynAgent, otherwise "React".
""" """
if isinstance(agent, CiaynAgent): if isinstance(agent, CiaynAgent):
return "CiaynAgent" return "CiaynAgent"
else: else:
return "React" return "React"
def init_fallback_handler(agent: RAgents, config: Dict[str, Any], tools: List[Any]):
"""
Initialize fallback handler if agent is of type "React"; otherwise return None.
"""
agent_type = get_agent_type(agent)
if agent_type == "React":
return FallbackHandler(config, tools)
return None
def _handle_fallback_response( def _handle_fallback_response(
error: ToolExecutionError, error: ToolExecutionError,
fallback_handler, fallback_handler: Optional[FallbackHandler],
agent: RAgents, agent: RAgents,
agent_type: str, agent_type: str,
msg_list: list msg_list: list,
) -> None: ) -> None:
""" """
Handle fallback response by invoking fallback_handler and updating msg_list. Handle fallback response by invoking fallback_handler and updating msg_list.
@ -921,7 +938,9 @@ def run_agent_with_retry(
logger.debug("Agent run completed successfully") logger.debug("Agent run completed successfully")
return "Agent run completed successfully" return "Agent run completed successfully"
except ToolExecutionError as e: except ToolExecutionError as e:
_handle_fallback_response(e, fallback_handler, agent, agent_type, msg_list) _handle_fallback_response(
e, fallback_handler, agent, agent_type, msg_list
)
continue continue
except FallbackToolExecutionError as e: except FallbackToolExecutionError as e:
msg_list.append( msg_list.append(

View File

@ -288,6 +288,7 @@ class CiaynAgent:
logger.debug(f"Code generated by agent: {response.content}") logger.debug(f"Code generated by agent: {response.content}")
last_result = self._execute_tool(response) last_result = self._execute_tool(response)
self.chat_history.append(response) self.chat_history.append(response)
self.fallback_handler.reset_fallback_handler()
yield {} yield {}
except ToolExecutionError as e: except ToolExecutionError as e:

View File

@ -384,7 +384,9 @@ class FallbackHandler:
tool_calls = response.get("additional_kwargs").get("tool_calls") tool_calls = response.get("additional_kwargs").get("tool_calls")
return tool_calls return tool_calls
def handle_failure_response(self, error: ToolExecutionError, agent, agent_type: str): def handle_failure_response(
self, error: ToolExecutionError, agent, agent_type: str
):
""" """
Handle a tool failure by calling handle_failure and, if a fallback response is returned and the agent type is "React", Handle a tool failure by calling handle_failure and, if a fallback response is returned and the agent type is "React",
return a list of SystemMessage objects wrapping each message from the fallback response. return a list of SystemMessage objects wrapping each message from the fallback response.