From 1376d76437b4d0a97d4b16b2084f51a8201cccef Mon Sep 17 00:00:00 2001
From: Braelyn Boynton <bboynton97@gmail.com>
Date: Thu, 14 Nov 2024 10:43:51 -0800
Subject: [PATCH] working on code insert

---
 agentstack/generation/gen_utils.py       | 29 ++++++++++++++++++++++++
 agentstack/generation/task_generation.py |  5 ++--
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/agentstack/generation/gen_utils.py b/agentstack/generation/gen_utils.py
index 79f983f..25b2889 100644
--- a/agentstack/generation/gen_utils.py
+++ b/agentstack/generation/gen_utils.py
@@ -1,3 +1,6 @@
+import ast
+
+
 def insert_code_after_tag(file_path, tag, code_to_insert, next_line=False):
     if next_line:
         code_to_insert = ['\n'] + code_to_insert
@@ -16,3 +19,29 @@ def insert_code_after_tag(file_path, tag, code_to_insert, next_line=False):
 
     with open(file_path, 'w') as file:
         file.writelines(lines)
+
+
+def insert_after_tasks(file_path, code_to_insert):
+    with open(file_path, 'r') as file:
+        content = file.read()
+
+    module = ast.parse(content)
+
+    # Track the last task function's end line
+    last_task_end = None
+    for node in ast.walk(module):
+        if isinstance(node, ast.FunctionDef) and \
+                any(isinstance(deco, ast.Name) and deco.id == 'task' for deco in node.decorator_list):
+            last_task_end = node.end_lineno
+
+    if last_task_end is not None:
+        lines = content.split('\n')
+        for i, line in enumerate(code_to_insert):
+            lines.insert(last_task_end + i, line)
+        content = '\n'.join(lines)
+
+        with open(file_path, 'w') as file:
+            file.write(content)
+        return True
+    return False
+
diff --git a/agentstack/generation/task_generation.py b/agentstack/generation/task_generation.py
index ad6c761..7976fe8 100644
--- a/agentstack/generation/task_generation.py
+++ b/agentstack/generation/task_generation.py
@@ -1,6 +1,6 @@
 from typing import Optional
 
-from .gen_utils import insert_code_after_tag
+from .gen_utils import insert_code_after_tag, insert_after_tasks
 from ..utils import verify_agentstack_project, get_framework
 import os
 from ruamel.yaml import YAML
@@ -77,7 +77,6 @@ def generate_crew_task(
 
     # Add task to crew.py
     file_path = 'src/crew.py'
-    tag = '# Task definitions'
     code_to_insert = [
         "@task",
         f"def {name}(self) -> Task:",
@@ -87,4 +86,4 @@ def generate_crew_task(
         ""
     ]
 
-    insert_code_after_tag(file_path, tag, code_to_insert)
+    insert_after_tasks(file_path, code_to_insert)