From 9eb7f1ffe59aed7c30d2b10cae5cb8803183d629 Mon Sep 17 00:00:00 2001 From: Braelyn Boynton Date: Tue, 3 Dec 2024 00:18:48 -0800 Subject: [PATCH 1/2] preserve comments with AST --- agentstack/generation/tool_generation.py | 33 ++++++++++++++++++++---- examples/web_researcher/src/crew.py | 2 ++ 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/agentstack/generation/tool_generation.py b/agentstack/generation/tool_generation.py index c4e19fc0..aa7fedcf 100644 --- a/agentstack/generation/tool_generation.py +++ b/agentstack/generation/tool_generation.py @@ -378,17 +378,40 @@ def modify_agent_tools( filename = _framework_filename(framework, path) - with open(filename, 'r') as f: - source = f.read() + # Read the original source and store line information + 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 + + # Parse and modify the AST + 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) + + # Generate the modified source without comments modified_source = astor.to_source(modified_tree) - with open(filename, 'w') as f: - f.write(modified_source) \ No newline at end of file + # Split the modified source into lines + modified_lines = modified_source.splitlines() + + # Reinsert comments at appropriate positions + final_lines = [] + for i, line in enumerate(modified_lines, 1): + # If there was a comment at this line number in the original file, add it + if i in comments: + final_lines.append(comments[i]) + final_lines.append(line + '\n') + + # Write the final source back + with open(filename, 'w', encoding='utf-8') as f: + f.write(''.join(final_lines)) \ No newline at end of file diff --git a/examples/web_researcher/src/crew.py b/examples/web_researcher/src/crew.py index 01cfac22..aaf94506 100644 --- a/examples/web_researcher/src/crew.py +++ b/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']) From 194fd90bb553facd656c80afd5d18104ec2fca1c Mon Sep 17 00:00:00 2001 From: Braelyn Boynton Date: Tue, 3 Dec 2024 00:22:34 -0800 Subject: [PATCH 2/2] preserve comments with AST --- agentstack/generation/tool_generation.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/agentstack/generation/tool_generation.py b/agentstack/generation/tool_generation.py index aa7fedcf..d7dd6114 100644 --- a/agentstack/generation/tool_generation.py +++ b/agentstack/generation/tool_generation.py @@ -378,7 +378,6 @@ def modify_agent_tools( filename = _framework_filename(framework, path) - # Read the original source and store line information with open(filename, 'r', encoding='utf-8') as f: source_lines = f.readlines() @@ -389,7 +388,6 @@ def modify_agent_tools( if stripped.startswith('#'): comments[i + 1] = line - # Parse and modify the AST tree = ast.parse(''.join(source_lines)) class ModifierTransformer(ast.NodeTransformer): @@ -397,21 +395,15 @@ def visit_FunctionDef(self, node): return _modify_agent_tools(node, tool_data, operation, agents, base_name) modified_tree = ModifierTransformer().visit(tree) - - # Generate the modified source without comments modified_source = astor.to_source(modified_tree) - - # Split the modified source into lines modified_lines = modified_source.splitlines() - # Reinsert comments at appropriate positions + # Reinsert comments final_lines = [] for i, line in enumerate(modified_lines, 1): - # If there was a comment at this line number in the original file, add it if i in comments: final_lines.append(comments[i]) final_lines.append(line + '\n') - # Write the final source back with open(filename, 'w', encoding='utf-8') as f: f.write(''.join(final_lines)) \ No newline at end of file