Merge pull request #80 from AgentOps-AI/ast-with-comments

preserve comments with AST

Author: bboynton97

agentstack/generation/tool_generation.py

@@ -378,17 +378,32 @@ def modify_agent_tools(
 
     filename = _framework_filename(framework, path)
 
-    with open(filename, 'r') as f:
-        source = f.read()
+    with open(filename, 'r', encoding='utf-8') as f:
+        source_lines = f.readlines()
 
-    tree = ast.parse(source)
+    # Create a map of line numbers to comments
+    comments = {}
+    for i, line in enumerate(source_lines):
+        stripped = line.strip()
+        if stripped.startswith('#'):
+            comments[i + 1] = line
+
+    tree = ast.parse(''.join(source_lines))
 
     class ModifierTransformer(ast.NodeTransformer):
         def visit_FunctionDef(self, node):
             return _modify_agent_tools(node, tool_data, operation, agents, base_name)
 
     modified_tree = ModifierTransformer().visit(tree)
     modified_source = astor.to_source(modified_tree)
+    modified_lines = modified_source.splitlines()
+
+    # Reinsert comments
+    final_lines = []
+    for i, line in enumerate(modified_lines, 1):
+        if i in comments:
+            final_lines.append(comments[i])
+        final_lines.append(line + '\n')
 
-    with open(filename, 'w') as f:
-        f.write(modified_source)
+    with open(filename, 'w', encoding='utf-8') as f:
+        f.write(''.join(final_lines))

examples/web_researcher/src/crew.py

@@ -7,6 +7,7 @@
 class WebresearcherCrew:
     """web_researcher crew"""
 
+    # Agent definitions
     @agent
     def content_summarizer(self) ->Agent:
         return Agent(config=self.agents_config['content_summarizer'], tools
@@ -23,6 +24,7 @@ def content_storer(self) ->Agent:
             tools.create_database, tools.execute_sql_ddl, tools.
             run_sql_query], verbose=True)
 
+    # Task definitions
     @task
     def scrape_site(self) ->Task:
         return Task(config=self.tasks_config['scrape_site'])