From a574a6fdaa74c82089dff940fb4f71c130182fb6 Mon Sep 17 00:00:00 2001
From: Alejandro Herrera <149527975+sfc-gh-alherrera@users.noreply.github.com>
Date: Mon, 16 Dec 2024 16:17:42 -0500
Subject: [PATCH] fix: addressing new plan step override bug (#93)

---
 agent_gateway/gateway/output_parser.py | 28 ++++++++++++++++----------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/agent_gateway/gateway/output_parser.py b/agent_gateway/gateway/output_parser.py
index 01b4ed5..db3dc1b 100644
--- a/agent_gateway/gateway/output_parser.py
+++ b/agent_gateway/gateway/output_parser.py
@@ -71,6 +71,7 @@ def parse(self, text: str) -> list[str]:
 ### Helper functions
 def _initialize_task_list(matches):
     new_matches = []
+    new_match_idx = []
     current_index = 1
     index_mapping = {}
 
@@ -82,12 +83,13 @@ def _initialize_task_list(matches):
         if "cortexsearch" in task[2] and i != len(matches) - 2:
             new_step = _create_summarization_step(task[3], current_index)
             new_matches.append(new_step)
+            new_match_idx.append(new_step[1])
             current_index += 1
             index_mapping[task[1]] = str(current_index)
 
         current_index += 1
 
-    return new_matches, index_mapping
+    return new_matches, index_mapping, new_match_idx
 
 
 def _create_summarization_step(context, index):
@@ -101,20 +103,24 @@ def _create_summarization_step(context, index):
     )
 
 
-def _update_task_references(task, index_mapping):
-    updated_string = task[3]
-    updated_string = re.sub(
-        r"\$(\d+)",
-        lambda m: f"${index_mapping.get(m.group(1), m.group(1))}",
-        updated_string,
-    )
-    return (task[0], task[1], task[2], updated_string, "")
+def _update_task_references(task, index_mapping, new_match_ids):
+    if task[1] in new_match_ids:
+        return task
+    else:
+        updated_string = task[3]
+        updated_string = re.sub(
+            r"\$(\d+)",
+            lambda m: f"${index_mapping.get(m.group(1), m.group(1))}",
+            updated_string,
+        )
+        return (task[0], task[1], task[2], updated_string, "")
 
 
 def _update_task_list_with_summarization(matches):
-    new_matches, index_mapping = _initialize_task_list(matches)
+    new_matches, index_mapping, new_match_ids = _initialize_task_list(matches)
     updated_final_matches = [
-        _update_task_references(task, index_mapping) for task in new_matches
+        _update_task_references(task, index_mapping, new_match_ids)
+        for task in new_matches
     ]
     return updated_final_matches