Skip to content

Commit

Permalink
fix: addressing new plan step override bug (#93)
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-alherrera authored Dec 16, 2024
1 parent a96bbb9 commit a574a6f
Showing 1 changed file with 17 additions and 11 deletions.
28 changes: 17 additions & 11 deletions agent_gateway/gateway/output_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {}

Expand All @@ -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):
Expand All @@ -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

Expand Down

0 comments on commit a574a6f

Please sign in to comment.