diff --git a/agent_gateway/gateway/gateway.py b/agent_gateway/gateway/gateway.py index 5d49d83..8d8fb01 100644 --- a/agent_gateway/gateway/gateway.py +++ b/agent_gateway/gateway/gateway.py @@ -225,7 +225,7 @@ def _parse_fusion_output(self, raw_answer: str) -> str: is_replan = FUSION_REPLAN in answer if is_replan: - answer = "We couldn't find the information you're looking for. You can try rephrasing your request or validate that the provided tools contain sufficient information." + answer = self._extract_replan_message(raw_answer) if answer is None: raise AgentGatewayError( @@ -257,6 +257,17 @@ def _extract_answer(self, raw_answer): return None + def _extract_replan_message(self, raw_answer): + replan_start = "Action: Replan(" + replan_index = raw_answer.find(replan_start) + if replan_index != -1: + replan_index += len(replan_start) + return raw_answer[replan_index : raw_answer.rfind(")")].strip() + return ( + "We couldn't find the information you're looking for. You can try " + "rephrasing your request or validate that the provided tools contain " + "sufficient information." + ) def _generate_context_for_replanner( self, tasks: Mapping[int, Task], fusion_thought: str ) -> str: