Introduction
This article outlines the construction of an advanced LLM-based agent designed to assist with code refactoring and design, supporting both Python and JavaScript/TypeScript. The agent's primary function is to analyze code, identify and suggest refactorings, prioritize them, and then, upon user approval, implement these changes. It will incorporate robust versioning mechanisms and leverage external tools for enhanced functionality. We will develop this agent in two progressive stages: first, by focusing on refactorings confined to a single file, and subsequently, by extending its capabilities to handle project-wide transformations across multiple files. The emphasis throughout will be on the agent's internal workings and its interaction with various components.
The foundation of our refactoring agent is an intelligent orchestrator, acting as the central control unit. This orchestrator manages the flow of information, facilitates interaction with the user, invokes the Large Language Model (LLM), and coordinates a suite of specialized tools. Its operational model follows a "plan, act, observe" loop: it first formulates a plan based on the user's request and its understanding of the code, then executes actions (such as calling the LLM or external tools), and finally observes the results to refine its plan or proceed with further steps.
Integrating the Large Language Model is the cornerstone of the agent's intelligence. For maximum flexibility, the agent is designed to support both remote and local LLMs. Remote LLMs, typically accessed via an API (e.g., OpenAI's models), offer high performance and access to the latest models. Local LLMs provide greater control over data privacy and can be deployed on diverse hardware. To ensure compatibility with all major GPU types, including Nvidia, AMD, and Apple Silicon, open-source solutions like `llama.cpp` or `Ollama` are excellent choices. These projects provide highly optimized C++ implementations that can leverage various GPU architectures for accelerated inference, making local LLM deployment accessible across a broad spectrum of employee workstations. The agent will abstract the LLM interaction through a common interface, enabling seamless switching between different LLM providers or local models without altering the agent's core logic.
The agent's ability to interact with the external environment is facilitated by a comprehensive tooling system. This system empowers the LLM to perform actions beyond mere text generation, such as reading and writing files, executing Git commands, and searching the web. Each tool is essentially a function or a class method that the LLM can invoke, accompanied by a clear description of its purpose, expected inputs, and anticipated outputs. For file system operations, the agent can wrap standard operating system commands or Python's built-in `os` and `io` modules. Git integration, crucial for version control, will be managed by a library like `gitpython` for Python, which offers a convenient API for programmatic interaction with Git repositories. Web search functionality can be implemented by integrating with open-source search APIs or by using libraries like `requests` to query search engines programmatically, always adhering to API usage policies.
Understanding the code is paramount for effective refactoring. The agent employs a dedicated Code Analysis module that utilizes Abstract Syntax Trees (ASTs) and static analysis techniques. An AST provides a hierarchical, tree-like representation of the source code's abstract syntactic structure, which is invaluable for understanding the code's organization and relationships without executing it. For Python, the built-in `ast` module enables parsing Python code into an AST, which can then be traversed to identify patterns, extract specific information, or even programmatically manipulate the code structure. For JavaScript and TypeScript, tools like `typescript-eslint` parser or `tree-sitter` offer similar capabilities, allowing the agent to build a detailed understanding of the code's syntax and semantics. Furthermore, static analysis tools such as Pylint for Python or ESLint for JavaScript/TypeScript can be integrated to automatically identify common code quality issues, potential bugs, and predefined refactoring opportunities based on established rules.
Versioning is a critical safety mechanism for any code modification agent. The agent must integrate tightly with Git to meticulously track all changes, providing users with the ability to review, revert, or undo refactorings with ease. Before applying any significant modification, the agent should create a new Git branch or, at a minimum, commit the current state of the repository. This ensures that a reliable rollback point is always available. The `gitpython` library in Python provides all the necessary functionalities to programmatically stage changes, commit them with descriptive messages, create new branches, and revert to previous commits. This robust integration guarantees that even if an LLM-generated refactoring introduces unintended issues, the user can effortlessly undo the changes without any data loss.
One of the significant challenges when working with LLMs on large codebases is the inherent context window limitation. LLMs can only process a finite amount of text at any given time. To overcome this, the agent employs several sophisticated strategies. Summarization involves condensing large code blocks or entire file contents into shorter, yet informative, summaries before feeding them to the LLM. Retrieval-Augmented Generation (RAG) is a powerful technique where the entire codebase (or relevant portions) is indexed and stored in a vector database. When the LLM requires specific information about a particular part of the code or project, the agent can perform a semantic search against this database. The database then retrieves only the most semantically relevant code snippets or documentation, which are subsequently provided to the LLM as part of its context. A sliding window approach can be utilized for very long files, processing them in chunks and maintaining a condensed context of previously processed sections. Finally, the agent can leverage function calling or tool use, allowing the LLM to explicitly "ask" for specific files, function definitions, or class implementations as needed, rather than receiving the entire project upfront, thereby optimizing context usage.
Stage 1: Single-File Refactoring
In the initial stage, the LLM-based agent focuses on refactorings that are entirely confined to a single file. This approach significantly reduces complexity by limiting the scope of changes and dependencies, making it an ideal starting point for development.
The process begins with user input and initial code loading. The user might explicitly suggest a specific refactoring, such as "extract this block into a new function," or they might simply ask the agent to "find possible refactorings in this file." The agent's first action is to load the content of the specified code file.
Here is an example of how the agent's internal `FileHandler` tool might read a Python file:
--------------------------------------------------------------------------------
Agent Code Snippet: FileHandler Tool - Reading a File
--------------------------------------------------------------------------------
import os
class FileHandler:
"""
A tool for the agent to interact with the file system.
Provides methods for reading and writing file contents.
"""
def read_file(self, file_path):
"""
Reads the content of a specified file.
Parameters:
file_path (str): The path to the file.
Returns:
str: The content of the file, or None if an error occurs.
"""
if not os.path.exists(file_path):
print(f"Agent Error: File not found at {file_path}")
return None
try:
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
print(f"Agent Action: Successfully read file: {file_path}")
return content
except Exception as e:
print(f"Agent Error reading file {file_path}: {e}")
return None
def write_file(self, file_path, content):
"""
Writes content to a specified file, overwriting existing content.
Parameters:
file_path (str): The path to the file.
content (str): The content to write.
Returns:
bool: True if successful, False otherwise.
"""
try:
with open(file_path, 'w', encoding='utf-8') as f:
f.write(content)
print(f"Agent Action: Successfully wrote to file: {file_path}")
return True
except Exception as e:
print(f"Agent Error writing to file {file_path}: {e}")
return False
# Example usage within the agent's orchestrator:
# file_handler = FileHandler()
# code_content = file_handler.read_file("my_application_file.py")
# if code_content:
# print("Agent: File content loaded, proceeding to analysis.")
--------------------------------------------------------------------------------
Once the code is loaded, the agent proceeds with analysis and refactoring suggestion. The agent's internal Code Analysis module parses the file's content into an AST. For Python, this involves using the `ast` module. The AST allows the agent to construct a detailed understanding of the code's structure, identifying functions, classes, variables, and their interrelationships within the single file.
A snippet showing the agent's `CodeAnalyzer` tool parsing Python AST:
--------------------------------------------------------------------------------
Agent Code Snippet: CodeAnalyzer Tool - Parsing Python AST
--------------------------------------------------------------------------------
import ast
class CodeAnalyzer:
"""
A tool for the agent to perform static code analysis.
Provides methods for parsing code into ASTs and potentially
extracting information from them.
"""
def parse_python_code_to_ast(self, code_content):
"""
Parses Python code content into an Abstract Syntax Tree (AST).
Parameters:
code_content (str): The Python code as a string.
Returns:
ast.Module: The root node of the AST, or None if parsing fails.
"""
try:
tree = ast.parse(code_content)
print("Agent Action: Successfully parsed code into AST.")
return tree
except SyntaxError as e:
print(f"Agent Error: Syntax error in code during AST parsing: {e}")
return None
except Exception as e:
print(f"Agent Error: An unexpected error occurred during AST parsing: {e}")
return None
def extract_function_names(self, ast_tree):
"""
Extracts names of functions defined in the AST.
(A simple example of AST traversal for demonstration)
"""
if not ast_tree:
return []
function_names = []
for node in ast.walk(ast_tree):
if isinstance(node, ast.FunctionDef):
function_names.append(node.name)
print(f"Agent Action: Extracted function names: {function_names}")
return function_names
# Example usage within the agent's orchestrator:
# code_analyzer = CodeAnalyzer()
# application_code = "def greet(name):\n print(f'Hello, {name}!')\nclass MyClass:\n def method_a(self):\n pass"
# ast_tree = code_analyzer.parse_python_code_to_ast(application_code)
# if ast_tree:
# function_names = code_analyzer.extract_function_names(ast_tree)
# print(f"Agent: Identified functions: {function_names}")
--------------------------------------------------------------------------------
The parsed AST, along with the original code content and the user's prompt, is then formatted into a comprehensive prompt for the LLM. This is where prompt engineering becomes crucial. The agent's orchestrator constructs a prompt that instructs the LLM to analyze the code, identify refactoring opportunities (if not explicitly provided by the user), explain its reasoning, and suggest a prioritized list of refactorings. The LLM might suggest refactorings such as "extract method," "rename variable," "introduce constant," or "simplify conditional expression."
An example of how the agent's orchestrator might construct a prompt for refactoring suggestion:
--------------------------------------------------------------------------------
Agent Code Snippet: Orchestrator - Constructing LLM Prompt
--------------------------------------------------------------------------------
class AgentOrchestrator:
"""
The central brain of the agent, managing workflow and LLM interaction.
"""
def __init__(self, llm_connector, file_handler, code_analyzer, git_manager):
self.llm_connector = llm_connector # Handles actual LLM API calls or local inference
self.file_handler = file_handler
self.code_analyzer = code_analyzer
self.git_manager = git_manager
def generate_refactoring_suggestions(self, code_content, user_request=None):
"""
Constructs a prompt for the LLM to get refactoring suggestions.
"""
prompt_template = """
You are an expert software engineer assistant specializing in Python refactoring.
Your task is to analyze the provided Python code and identify potential refactoring opportunities.
For each suggestion, provide a brief explanation of why it's beneficial and a proposed priority (High, Medium, Low).
If the user has suggested a specific refactoring, evaluate its feasibility and provide a plan for its implementation.
Code to analyze:
```python
{code_content}
```
User's specific request (if any):
{user_request_text}
Please output your suggestions in a clear, structured format, prioritizing the most impactful changes first.
"""
user_request_text = user_request if user_request else "No specific request. Find general improvements."
full_prompt = prompt_template.format(
code_content=code_content,
user_request_text=user_request_text
)
print("Agent Action: Constructed LLM prompt for refactoring suggestions.")
# In a real scenario, self.llm_connector.query(full_prompt) would be called here.
# For this example, we return the prompt for inspection.
return full_prompt
# Example usage within the agent's orchestrator:
# # Assume llm_connector, file_handler, code_analyzer, git_manager are initialized
# orchestrator = AgentOrchestrator(llm_connector, file_handler, code_analyzer, git_manager)
# application_code = "def process_data(data):\n # complex logic here\n pass"
# user_prompt = "This process_data function is too long, can you suggest splitting it?"
# llm_prompt = orchestrator.generate_refactoring_suggestions(application_code, user_prompt)
# print(f"Agent: LLM prompt generated:\n{llm_prompt}")
--------------------------------------------------------------------------------
Upon receiving the LLM's suggestions, the agent's orchestrator presents them to the user for review and approval. The user can then select which refactorings to implement or ask for further details. Once approved, the agent instructs the LLM to generate the actual refactored code. This involves providing the original code, the chosen refactoring, and potentially relevant AST nodes to guide the LLM. The LLM then outputs the modified code.
Applying the changes and committing them to version control is the next crucial step. Before modifying the file, the agent ensures a safe rollback point by interacting with its `GitManager` tool.
Here is a snippet demonstrating how the agent's `GitManager` tool would apply changes and commit them using `gitpython`:
--------------------------------------------------------------------------------
Agent Code Snippet: GitManager Tool - Applying and Committing Changes
--------------------------------------------------------------------------------
import git
import os
import shutil # For temporary repo setup in running example
class GitManager:
"""
A tool for the agent to interact with Git repositories.
Provides methods for committing changes, creating branches, and reverting.
"""
def __init__(self, repo_path):
self.repo_path = repo_path
self.repo = None
self._initialize_repo()
def _initialize_repo(self):
"""Initializes or loads the Git repository."""
if not os.path.exists(self.repo_path):
print(f"Agent Git: Creating new repository at {self.repo_path}")
os.makedirs(self.repo_path)
self.repo = git.Repo.init(self.repo_path)
else:
try:
self.repo = git.Repo(self.repo_path)
print(f"Agent Git: Loaded existing repository at {self.repo_path}")
except git.InvalidGitRepositoryError:
print(f"Agent Git Error: {self.repo_path} is not a valid Git repository. Initializing.")
self.repo = git.Repo.init(self.repo_path)
# Ensure there's at least one commit for branch operations
if not self.repo.heads:
print("Agent Git: Initializing empty repository with a commit.")
dummy_file_path = os.path.join(self.repo_path, ".gitkeep")
with open(dummy_file_path, "w") as f:
f.write("")
self.repo.index.add([".gitkeep"])
self.repo.index.commit("Initial commit by agent.")
os.remove(dummy_file_path) # Clean up dummy file
self.repo.index.add([".gitkeep"]) # Stage deletion
self.repo.index.commit("Cleaned up .gitkeep.")
def apply_and_commit_changes(self, file_path_relative_to_repo, new_content, commit_message):
"""
Applies new content to a file, stages it, and commits it to the Git repository.
Parameters:
file_path_relative_to_repo (str): The path to the file relative to the repo root.
new_content (str): The new content to write to the file.
commit_message (str): The Git commit message.
Returns:
bool: True if successful, False otherwise.
"""
if not self.repo:
print("Agent Git Error: Repository not initialized.")
return False
full_file_path = os.path.join(self.repo_path, file_path_relative_to_repo)
# Ensure directory exists for the file
os.makedirs(os.path.dirname(full_file_path), exist_ok=True)
try:
# Write the new content to the file
with open(full_file_path, 'w', encoding='utf-8') as f:
f.write(new_content)
print(f"Agent Git Action: File {file_path_relative_to_repo} updated with new content.")
# Stage the changes
self.repo.index.add([file_path_relative_to_repo])
print(f"Agent Git Action: File {file_path_relative_to_repo} staged.")
# Commit the changes
commit = self.repo.index.commit(commit_message)
print(f"Agent Git Action: Changes committed with message: '{commit_message}' (SHA: {commit.hexsha})")
return True
except Exception as e:
print(f"Agent Git Error applying and committing changes: {e}")
return False
def create_branch(self, branch_name):
"""Creates and checks out a new Git branch."""
if not self.repo:
print("Agent Git Error: Repository not initialized.")
return False
try:
new_branch = self.repo.create_head(branch_name)
new_branch.checkout()
print(f"Agent Git Action: Created and checked out new branch '{branch_name}'.")
return True
except Exception as e:
print(f"Agent Git Error creating branch: {e}")
return False
def revert_last_commit(self):
"""Reverts the last commit in the current branch."""
if not self.repo:
print("Agent Git Error: Repository not initialized.")
return False
try:
self.repo.head.reset('HEAD~1', index=True, working_tree=True)
print("Agent Git Action: Reverted to the previous commit.")
return True
except Exception as e:
print(f"Agent Git Error reverting last commit: {e}")
return False
# Example usage within the agent's orchestrator:
# # For demonstration, let's set up a temporary repo
# temp_repo_dir = "temp_agent_repo"
# if os.path.exists(temp_repo_dir):
# shutil.rmtree(temp_repo_dir) # Clean up previous run
# git_manager = GitManager(temp_repo_dir)
#
# # Create a dummy file for initial commit
# git_manager.apply_and_commit_changes("initial_file.py", "print('hello')", "Initial commit")
#
# # Now apply a refactoring
# git_manager.create_branch("feature/refactor-greeting")
# refactored_code = "def greet():\n print('hello world')"
# git_manager.apply_and_commit_changes("initial_file.py", refactored_code, "Refactor: extracted greet function")
#
# # If user wants to undo
# # git_manager.revert_last_commit()
#
# shutil.rmtree(temp_repo_dir) # Clean up
--------------------------------------------------------------------------------
The agent's `GitManager` tool also provides an undo mechanism. If the user is dissatisfied with a refactoring, the agent can use Git commands to revert the last commit or switch back to a previous branch. This is typically achieved by calling `repo.head.reset('HEAD~1', index=True, working_tree=True)` to revert the last commit, or by checking out a specific commit hash or branch using `repo.git.checkout('branch_name')`.
Stage 2: Multi-File Refactoring
The second stage extends the agent's capabilities to handle refactorings that span multiple files within a project. This requires a more sophisticated understanding of the project structure, inter-file dependencies, and a robust context management strategy.
Project setup and Git repository integration are crucial. The agent assumes the project is available as a Git repository. It starts by either cloning the repository (if remote) or initializing a `gitpython` `Repo` object if the repository already exists locally. This allows the agent to access all files, track changes, and manage branches effectively for complex multi-file refactorings. The `GitManager` tool, as shown previously, handles this initialization and interaction.
Advanced context management becomes paramount for multi-file projects. Since feeding the entire codebase to the LLM is impractical due to context window limitations, the agent needs intelligent ways to provide only the most relevant information. Retrieval-Augmented Generation (RAG) is a key technique here. The agent can build an embedding database of the entire codebase. Each significant code entity (e.g., function, class, or even code block) can be converted into a vector embedding using models like `sentence-transformers`. These embeddings are then stored in a vector database (e.g., FAISS, ChromaDB, Weaviate). When the LLM needs to understand a specific part of the code or identify cross-file dependencies, the agent queries this vector database with a semantic query. The database then returns the most semantically similar code snippets, which are subsequently passed to the LLM as additional context. This ensures that the LLM receives only the most pertinent information without being overwhelmed by irrelevant code.
An example of how the agent's `KnowledgeBaseManager` might conceptually build a code embedding database:
--------------------------------------------------------------------------------
Agent Code Snippet: KnowledgeBaseManager - Conceptual Code Embedding
--------------------------------------------------------------------------------
# This is a conceptual example. In a real scenario, you would use a library
# like 'sentence-transformers' for embeddings and a vector database like
# 'ChromaDB' or 'FAISS' for storage and retrieval.
class KnowledgeBaseManager:
"""
A tool for the agent to manage and query a knowledge base of code embeddings.
This enables RAG for overcoming context window limitations.
"""
def __init__(self):
self.code_chunks = [] # Stores (chunk_id, text, file_path)
self.embeddings = {} # Stores {chunk_id: embedding_vector}
# self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2') # Actual model
# self.vector_db = ChromaDB.Client() # Actual vector database
def index_project_code(self, project_root_path, file_handler):
"""
Conceptually indexes code files in a project by creating chunks and embeddings.
In a real implementation, this would involve:
1. Iterating through all relevant code files (e.g., .py, .js, .ts) using file_handler.
2. Splitting each file into logical chunks (functions, classes, significant blocks).
3. Generating an embedding vector for each chunk using an LLM or embedding model.
4. Storing the chunk text and its embedding in a vector database.
"""
print(f"Agent KB: Scanning project at {project_root_path} for code files to index...")
# Placeholder for actual file scanning and chunking logic
# In reality, this would use file_handler.read_file and code_analyzer
# to parse and chunk code.
mock_chunks_data = [
{"id": "func_calc_sum", "text": "def calculate_sum(a, b):\n return a + b", "file": "src/utils.py"},
{"id": "class_data_proc", "text": "class DataProcessor:\n def __init__(self):\n pass\n def process(self, data):\n return data.upper()", "file": "src/processor.py"},
{"id": "main_logic", "text": "if __name__ == '__main__':\n processor = DataProcessor()\n result = processor.process('hello')", "file": "main.py"},
]
for i, chunk_data in enumerate(mock_chunks_data):
chunk_id = f"chunk_{i}_{chunk_data['id']}"
self.code_chunks.append({"id": chunk_id, "text": chunk_data["text"], "file": chunk_data["file"]})
# In a real system: embedding = self.embedding_model.encode(chunk_data["text"])
# self.embeddings[chunk_id] = embedding
# self.vector_db.add(documents=[chunk_data["text"]], metadatas=[{"file": chunk_data["file"]}], ids=[chunk_id])
print(f"Agent KB: Indexed chunk '{chunk_id}' from {chunk_data['file']}.")
print("Agent KB: Code embedding database (conceptual) indexing complete.")
def retrieve_relevant_code(self, query_text, top_k=3):
"""
Conceptually retrieves code chunks semantically similar to the query.
"""
print(f"Agent KB: Retrieving relevant code for query: '{query_text}'")
# In a real system: query_embedding = self.embedding_model.encode(query_text)
# results = self.vector_db.query(query_embeddings=[query_embedding], n_results=top_k)
# For demonstration, return some mock relevant chunks
mock_relevant_chunks = [
self.code_chunks[0], # calculate_sum
self.code_chunks[1] # DataProcessor
]
print(f"Agent KB: Retrieved {len(mock_relevant_chunks)} relevant code chunks.")
return mock_relevant_chunks
# Example usage within the agent's orchestrator:
# # Assume file_handler is initialized
# kb_manager = KnowledgeBaseManager()
# kb_manager.index_project_code("path/to/project", file_handler)
#
# # When LLM needs context for 'data processing logic'
# relevant_context = kb_manager.retrieve_relevant_code("how to process data in this project")
# # This context would then be appended to the LLM's prompt.
--------------------------------------------------------------------------------
Cross-file analysis and refactoring involve understanding how changes in one file impact others. For instance, renaming a function might require updating all its call sites across the entire project. The agent uses its `CodeAnalyzer` capabilities to build a conceptual dependency graph, mapping function calls, class instantiations, and variable usages across files. When a multi-file refactoring is proposed, the agent can query its embedding database (`KnowledgeBaseManager`) for relevant code snippets and then use the LLM to identify all affected files and lines of code. The LLM's response would include a detailed plan for modifying each affected file.
Coordinated code modifications are then performed. After the LLM identifies all necessary changes across multiple files, the agent's orchestrator, using its `FileHandler` and `GitManager` tools, applies these changes systematically. This might involve iterating through a list of files, applying specific modifications to each, and then staging and committing them as a single logical change or a series of closely related commits. For example, if a function is moved from `src/utils.py` to `src/new_utils/math_ops.py`, the agent would modify `src/utils.py` to remove the function, create `src/new_utils/math_ops.py` with the function's definition, and then update all files that imported the function to import it from its new location. All these changes would be part of a single, atomic Git commit to maintain consistency and traceability.
An example of how the agent's orchestrator might coordinate multi-file changes:
--------------------------------------------------------------------------------
Agent Code Snippet: Orchestrator - Coordinating Multi-File Refactoring
--------------------------------------------------------------------------------
import os
class AgentOrchestrator:
"""
The central brain of the agent, managing workflow and LLM interaction.
(Expanded from previous snippet)
"""
def __init__(self, llm_connector, file_handler, code_analyzer, git_manager, kb_manager):
self.llm_connector = llm_connector
self.file_handler = file_handler
self.code_analyzer = code_analyzer
self.git_manager = git_manager
self.kb_manager = kb_manager
# ... (previous methods like generate_refactoring_suggestions) ...
def execute_multi_file_refactoring_plan(self, project_root, refactoring_plan, commit_message):
"""
Executes a refactoring plan that spans multiple files.
The plan is assumed to be generated by the LLM.
Parameters:
project_root (str): The root directory of the project.
refactoring_plan (list): A list of dictionaries, each describing a file change.
Example: [{"file": "old_file.py", "action": "delete"},
{"file": "new_file.py", "action": "create", "content": "..."}]
commit_message (str): The Git commit message for these changes.
Returns:
bool: True if all changes were applied and committed, False otherwise.
"""
print(f"Agent Orchestrator: Starting multi-file refactoring in {project_root}...")
# Create a new branch for safety before starting multi-file changes
branch_name = f"refactor_multi_file_{os.urandom(4).hex()}" # Unique branch name
if not self.git_manager.create_branch(branch_name):
print("Agent Orchestrator Error: Failed to create a new branch for refactoring.")
return False
files_to_stage_for_commit = []
for change in refactoring_plan:
file_path_relative = change["file"]
full_file_path = os.path.join(project_root, file_path_relative)
action = change["action"]
content = change.get("content") # Content only for 'create' or 'modify' actions
try:
if action == "create":
if not self.file_handler.write_file(full_file_path, content):
raise Exception(f"Failed to create file {file_path_relative}")
files_to_stage_for_commit.append(file_path_relative)
elif action == "modify":
if not self.file_handler.write_file(full_file_path, content):
raise Exception(f"Failed to modify file {file_path_relative}")
files_to_stage_for_commit.append(file_path_relative)
elif action == "delete":
if os.path.exists(full_file_path):
os.remove(full_file_path)
print(f"Agent Orchestrator: Deleted file: {file_path_relative}")
files_to_stage_for_commit.append(file_path_relative)
else:
print(f"Agent Orchestrator Warning: Attempted to delete non-existent file {file_path_relative}")
else:
print(f"Agent Orchestrator Error: Unknown action: {action} for file {file_path_relative}")
return False
except Exception as e:
print(f"Agent Orchestrator Error applying change to {file_path_relative}: {e}")
# Revert branch if error occurs during plan execution
print("Agent Orchestrator: Attempting to revert branch due to error.")
self.git_manager.revert_last_commit() # Revert to state before branch creation
return False
if files_to_stage_for_commit:
if self.git_manager.apply_and_commit_changes_internal(files_to_stage_for_commit, commit_message):
print("Agent Orchestrator: All multi-file changes committed successfully.")
return True
else:
print("Agent Orchestrator Error: Failed to commit multi-file changes.")
return False
else:
print("Agent Orchestrator: No files were staged for commit.")
return True # No changes to commit, still considered successful if no errors
# Helper method for GitManager to stage multiple files, used by orchestrator
def apply_and_commit_changes_internal(self, file_paths_relative_to_repo, commit_message):
"""
Internal method for GitManager to stage multiple files and commit.
Used by the orchestrator for multi-file operations.
"""
if not self.git_manager.repo:
print("Agent Git Error: Repository not initialized for multi-file commit.")
return False
try:
self.git_manager.repo.index.add(file_paths_relative_to_repo)
commit = self.git_manager.repo.index.commit(commit_message)
print(f"Agent Git Action: Multi-file changes committed with message: '{commit_message}' (SHA: {commit.hexsha})")
return True
except Exception as e:
print(f"Agent Git Error during multi-file commit: {e}")
return False
# Example usage (conceptual within the orchestrator):
# # Assume orchestrator with all tools initialized
# project_path = "/path/to/your/multi_file_project"
# refactoring_plan_from_llm = [
# {"file": "src/utils.py", "action": "modify", "content": "# Updated utils.py content, function moved\n"},
# {"file": "src/new_math.py", "action": "create", "content": "def add_numbers(a, b):\n return a + b\n"},
# {"file": "src/main.py", "action": "modify", "content": "from src.new_math import add_numbers\n# ... rest of main.py ...\n"}
# ]
# orchestrator.execute_multi_file_refactoring_plan(project_path, refactoring_plan_from_llm, "Refactor: Moved math functions to new module")
--------------------------------------------------------------------------------
Challenges and Future Enhancements
Building such an agent presents several challenges. Ensuring the correctness of LLM-generated code is paramount; integrating automated testing frameworks (like `pytest` for Python or `Jest` for JavaScript) to run tests before and after refactoring can help validate changes. The agent would use its `FileHandler` to read test files, and then a `ProcessExecutor` tool to run the test commands, reporting results back to the LLM. Performance optimization, especially for large codebases and local LLM inference, requires careful resource management and efficient RAG implementations. Future enhancements could include supporting more complex refactoring patterns, integrating with Integrated Development Environments (IDEs) for a smoother user experience, and incorporating user feedback loops to continuously improve the agent's refactoring intelligence.
Addendum: Full Running Example of the Agent
This addendum provides a complete running example of the agent's core logic, demonstrating a simple single-file refactoring scenario. The LLM interaction is simulated by providing predefined responses, allowing the agent's workflow to be fully executed and observed.
--------------------------------------------------------------------------------
Addendum: Agent Running Example - Single-File Refactoring
--------------------------------------------------------------------------------
import os
import shutil
import git # Used by GitManager
# --- Agent Tools Definitions (as described in the article) ---
class FileHandler:
"""
A tool for the agent to interact with the file system.
Provides methods for reading and writing file contents.
"""
def read_file(self, file_path):
if not os.path.exists(file_path):
print(f"Agent.FileHandler Error: File not found at {file_path}")
return None
try:
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
print(f"Agent.FileHandler Action: Successfully read file: {file_path}")
return content
except Exception as e:
print(f"Agent.FileHandler Error reading file {file_path}: {e}")
return None
def write_file(self, file_path, content):
try:
# Ensure directory exists for the file
os.makedirs(os.path.dirname(file_path), exist_ok=True)
with open(file_path, 'w', encoding='utf-8') as f:
f.write(content)
print(f"Agent.FileHandler Action: Successfully wrote to file: {file_path}")
return True
except Exception as e:
print(f"Agent.FileHandler Error writing to file {file_path}: {e}")
return False
class CodeAnalyzer:
"""
A tool for the agent to perform static code analysis.
Provides methods for parsing code into ASTs and potentially
extracting information from them.
"""
def parse_python_code_to_ast(self, code_content):
import ast
try:
tree = ast.parse(code_content)
print("Agent.CodeAnalyzer Action: Successfully parsed code into AST.")
return tree
except SyntaxError as e:
print(f"Agent.CodeAnalyzer Error: Syntax error in code during AST parsing: {e}")
return None
except Exception as e:
print(f"Agent.CodeAnalyzer Error: An unexpected error occurred during AST parsing: {e}")
return None
def extract_function_names(self, ast_tree):
import ast
if not ast_tree:
return []
function_names = []
for node in ast.walk(ast_tree):
if isinstance(node, ast.FunctionDef):
function_names.append(node.name)
print(f"Agent.CodeAnalyzer Action: Extracted function names: {function_names}")
return function_names
class GitManager:
"""
A tool for the agent to interact with Git repositories.
Provides methods for committing changes, creating branches, and reverting.
"""
def __init__(self, repo_path):
self.repo_path = repo_path
self.repo = None
self._initialize_repo()
def _initialize_repo(self):
"""Initializes or loads the Git repository."""
if not os.path.exists(self.repo_path):
print(f"Agent.GitManager: Creating new repository at {self.repo_path}")
os.makedirs(self.repo_path)
self.repo = git.Repo.init(self.repo_path)
else:
try:
self.repo = git.Repo(self.repo_path)
print(f"Agent.GitManager: Loaded existing repository at {self.repo_path}")
except git.InvalidGitRepositoryError:
print(f"Agent.GitManager Error: {self.repo_path} is not a valid Git repository. Initializing.")
self.repo = git.Repo.init(self.repo_path)
# Ensure there's at least one commit for branch operations
if not self.repo.heads:
print("Agent.GitManager: Initializing empty repository with a commit.")
dummy_file_path = os.path.join(self.repo_path, ".gitkeep")
with open(dummy_file_path, "w") as f:
f.write("")
self.repo.index.add([".gitkeep"])
self.repo.index.commit("Initial commit by agent.")
os.remove(dummy_file_path) # Clean up dummy file
self.repo.index.add([".gitkeep"]) # Stage deletion
self.repo.index.commit("Cleaned up .gitkeep.")
def apply_and_commit_changes(self, file_path_relative_to_repo, new_content, commit_message):
"""
Applies new content to a file, stages it, and commits it to the Git repository.
This method uses the FileHandler to write content.
"""
if not self.repo:
print("Agent.GitManager Error: Repository not initialized.")
return False
full_file_path = os.path.join(self.repo_path, file_path_relative_to_repo)
# Use FileHandler to write the content
if not FileHandler().write_file(full_file_path, new_content):
return False
try:
# Stage the changes
self.repo.index.add([file_path_relative_to_repo])
print(f"Agent.GitManager Action: File {file_path_relative_to_repo} staged.")
# Commit the changes
commit = self.repo.index.commit(commit_message)
print(f"Agent.GitManager Action: Changes committed with message: '{commit_message}' (SHA: {commit.hexsha})")
return True
except Exception as e:
print(f"Agent.GitManager Error applying and committing changes: {e}")
return False
def create_branch(self, branch_name):
"""Creates and checks out a new Git branch."""
if not self.repo:
print("Agent.GitManager Error: Repository not initialized.")
return False
try:
new_branch = self.repo.create_head(branch_name)
new_branch.checkout()
print(f"Agent.GitManager Action: Created and checked out new branch '{branch_name}'.")
return True
except Exception as e:
print(f"Agent.GitManager Error creating branch: {e}")
return False
def revert_last_commit(self):
"""Reverts the last commit in the current branch."""
if not self.repo:
print("Agent.GitManager Error: Repository not initialized.")
return False
try:
self.repo.head.reset('HEAD~1', index=True, working_tree=True)
print("Agent.GitManager Action: Reverted to the previous commit.")
return True
except Exception as e:
print(f"Agent.GitManager Error reverting last commit: {e}")
return False
def get_current_file_content(self, file_path_relative_to_repo):
"""Retrieves the current content of a file from the repository's working tree."""
full_file_path = os.path.join(self.repo_path, file_path_relative_to_repo)
return FileHandler().read_file(full_file_path)
def apply_and_commit_changes_internal(self, file_paths_relative_to_repo, commit_message):
"""
Internal method for GitManager to stage multiple files and commit.
Used by the orchestrator for multi-file operations.
"""
if not self.repo:
print("Agent.GitManager Error: Repository not initialized for multi-file commit.")
return False
try:
self.repo.index.add(file_paths_relative_to_repo)
commit = self.repo.index.commit(commit_message)
print(f"Agent.GitManager Action: Multi-file changes committed with message: '{commit_message}' (SHA: {commit.hexsha})")
return True
except Exception as e:
print(f"Agent.GitManager Error during multi-file commit: {e}")
return False
class LLMConnector:
"""
A simulated tool for the agent to interact with an LLM.
In a real scenario, this would make API calls or run local inference.
For this example, it provides predefined responses.
"""
def __init__(self):
self.simulated_responses = {
"suggest_refactoring_prompt": {
"response": """
Here are refactoring suggestions for 'process_data':
1. **Refactoring Type:** Extract Method
**Description:** The data validation logic is complex and can be extracted into a separate private method '_validate_item'. This improves readability and testability.
**Priority:** High
**Proposed Code Change (Conceptual):**
- Original: complex validation block inside loop
- Refactored: call `self._validate_item(item)`
- New method: `_validate_item(self, item, config)`
2. **Refactoring Type:** Introduce Constant
**Description:** The string "Input raw_data must be a list." is a magic string. It could be a class constant.
**Priority:** Low
""",
"refactored_code": """
import logging
class DataProcessor:
def __init__(self, config):
self.config = config
def _validate_item(self, item):
# Extracted validation logic
if not isinstance(item, str):
logging.warning(f"Skipping non-string item: {item}")
return False
is_valid_length = len(item) >= self.config.get("min_length", 5)
try:
numeric_value = int(item)
is_valid_value = numeric_value <= self.config.get("max_value", 100)
except ValueError:
is_valid_value = False
if not (is_valid_length and is_valid_value):
logging.warning(f"Item '{item}' failed validation (length valid: {is_valid_length}, value valid: {is_valid_value}).")
return False
return True
def process_data(self, raw_data):
logging.info("Starting data processing.")
if not isinstance(raw_data, list):
logging.error("Input raw_data must be a list.")
return []
processed_results = []
for item in raw_data:
if self._validate_item(item):
transformed_item = item.upper() # Example transformation
processed_results.append(transformed_item)
logging.debug(f"Item '{item}' processed successfully.")
logging.info("Finished data processing.")
return processed_results
"""
},
"multi_file_refactoring_plan": {
"plan": [
{"file": "src/simple_app.py", "action": "modify", "content": """
import logging
from utils.data_validators import validate_numeric_item
class DataProcessor:
def __init__(self, config):
self.config = config
def process_data(self, raw_data):
logging.info("Starting data processing.")
if not isinstance(raw_data, list):
logging.error("Input raw_data must be a list.")
return []
processed_results = []
for item in raw_data:
if validate_numeric_item(item, self.config): # Using external validator
transformed_item = item.upper()
processed_results.append(transformed_item)
logging.debug(f"Item '{item}' processed successfully.")
else:
logging.warning(f"Item '{item}' failed validation.")
logging.info("Finished data processing.")
return processed_results
"""},
{"file": "utils/data_validators.py", "action": "create", "content": """
import logging
def validate_numeric_item(item, config):
if not isinstance(item, str):
logging.warning(f"Validator: Skipping non-string item: {item}")
return False
is_valid_length = len(item) >= config.get("min_length", 5)
try:
numeric_value = int(item)
is_valid_value = numeric_value <= config.get("max_value", 100)
except ValueError:
is_valid_value = False
if not (is_valid_length and is_valid_value):
logging.warning(f"Validator: Item '{item}' failed validation (length valid: {is_valid_length}, value valid: {is_valid_value}).")
return False
return True
"""}
],
"commit_message": "Refactor: Extracted numeric item validation to utils/data_validators.py"
}
}
def query_llm(self, prompt_type, **kwargs):
"""
Simulates querying the LLM based on a prompt type.
"""
print(f"Agent.LLMConnector Action: Simulating LLM query for type: {prompt_type}")
if prompt_type == "suggest_refactoring_prompt":
return self.simulated_responses["suggest_refactoring_prompt"]
elif prompt_type == "get_refactored_code":
# For simplicity, assume the LLM directly provides the refactored code
# after a suggestion is chosen. In reality, this would be another prompt.
return self.simulated_responses["suggest_refactoring_prompt"]["refactored_code"]
elif prompt_type == "multi_file_refactoring_plan":
return self.simulated_responses["multi_file_refactoring_plan"]
else:
print(f"Agent.LLMConnector Error: Unknown prompt type '{prompt_type}'.")
return None
class KnowledgeBaseManager:
"""
A simulated tool for the agent to manage and query a knowledge base of code embeddings.
For this running example, it's simplified as it's not directly used in the single-file flow.
"""
def __init__(self):
print("Agent.KnowledgeBaseManager: Initialized (simulated).")
def index_project_code(self, project_root_path, file_handler):
print(f"Agent.KnowledgeBaseManager: Simulating indexing project at {project_root_path}.")
# In a real scenario, this would populate a vector DB
pass
def retrieve_relevant_code(self, query_text, top_k=3):
print(f"Agent.KnowledgeBaseManager: Simulating retrieval for '{query_text}'.")
return [] # Return empty for this single-file example
# --- Agent Orchestrator Definition ---
class AgentOrchestrator:
"""
The central brain of the agent, managing workflow and LLM interaction.
"""
def __init__(self, repo_path):
self.file_handler = FileHandler()
self.code_analyzer = CodeAnalyzer()
self.git_manager = GitManager(repo_path)
self.llm_connector = LLMConnector() # Simulated LLM
self.kb_manager = KnowledgeBaseManager() # Simulated KB
def run_single_file_refactoring(self, target_file_relative_path, user_request=None):
"""
Orchestrates a single-file refactoring workflow.
"""
print("\n--- Agent Orchestrator: Starting Single-File Refactoring ---")
full_target_file_path = os.path.join(self.git_manager.repo_path, target_file_relative_path)
# 1. Read the code file
initial_code_content = self.file_handler.read_file(full_target_file_path)
if not initial_code_content:
print("Agent Orchestrator Error: Could not read target file.")
return False
# 2. Parse AST for initial understanding
ast_tree = self.code_analyzer.parse_python_code_to_ast(initial_code_content)
if not ast_tree:
print("Agent Orchestrator Error: Could not parse AST.")
return False
# 3. Generate refactoring suggestions using LLM (simulated)
print("\nAgent Orchestrator: Requesting refactoring suggestions from LLM...")
llm_suggestion_response = self.llm_connector.query_llm(
"suggest_refactoring_prompt",
code_content=initial_code_content,
user_request=user_request
)
if not llm_suggestion_response:
print("Agent Orchestrator Error: LLM did not provide suggestions.")
return False
print("\n--- LLM Suggested Refactorings ---")
print(llm_suggestion_response["response"])
print("----------------------------------")
# 4. User approval (simulated as automatic acceptance of first suggestion)
print("\nAgent Orchestrator: User approves the 'Extract Method' refactoring.")
# 5. Get refactored code from LLM (simulated)
print("Agent Orchestrator: Requesting refactored code from LLM...")
refactored_code_content = self.llm_connector.query_llm("get_refactored_code")
if not refactored_code_content:
print("Agent Orchestrator Error: LLM did not provide refactored code.")
return False
# 6. Apply changes and commit
print("\nAgent Orchestrator: Applying refactored code and committing changes.")
self.git_manager.create_branch(f"refactor/{os.path.basename(target_file_relative_path).replace('.', '_')}")
commit_msg = "Refactor: Extracted _validate_item method from process_data."
if self.git_manager.apply_and_commit_changes(target_file_relative_path, refactored_code_content, commit_msg):
print(f"\nAgent Orchestrator: Refactoring of {target_file_relative_path} completed and committed.")
print("Current content of the refactored file:")
print(self.git_manager.get_current_file_content(target_file_relative_path))
return True
else:
print("Agent Orchestrator Error: Failed to apply and commit refactoring.")
return False
def run_multi_file_refactoring(self, project_root_path, user_request=None):
"""
Orchestrates a multi-file refactoring workflow.
"""
print("\n--- Agent Orchestrator: Starting Multi-File Refactoring ---")
# 1. Index project code for RAG (simulated)
self.kb_manager.index_project_code(project_root_path, self.file_handler)
# 2. Simulate LLM generating a multi-file refactoring plan
print("\nAgent Orchestrator: Requesting multi-file refactoring plan from LLM (simulated)...")
llm_plan_response = self.llm_connector.query_llm("multi_file_refactoring_plan")
if not llm_plan_response:
print("Agent Orchestrator Error: LLM did not provide a multi-file plan.")
return False
refactoring_plan = llm_plan_response["plan"]
commit_message = llm_plan_response["commit_message"]
print("\n--- LLM Generated Multi-File Refactoring Plan ---")
for change in refactoring_plan:
print(f"- File: {change['file']}, Action: {change['action']}")
print(f"Commit Message: {commit_message}")
print("--------------------------------------------------")
# 3. User approval (simulated as automatic acceptance)
print("\nAgent Orchestrator: User approves the multi-file refactoring plan.")
# 4. Execute the multi-file refactoring plan
print("\nAgent Orchestrator: Executing multi-file refactoring plan.")
if self.execute_multi_file_refactoring_plan(project_root_path, refactoring_plan, commit_message):
print("\nAgent Orchestrator: Multi-file refactoring completed and committed.")
# Verify changes (optional)
for change in refactoring_plan:
print(f"Content of {change['file']} after refactoring:")
print(self.file_handler.read_file(os.path.join(project_root_path, change['file'])))
return True
else:
print("Agent Orchestrator Error: Failed to execute multi-file refactoring plan.")
return False
# Helper method from GitManager, moved here for orchestrator's direct use
def execute_multi_file_refactoring_plan(self, project_root, refactoring_plan, commit_message):
print(f"Agent Orchestrator: Executing multi-file plan in {project_root}...")
branch_name = f"refactor_multi_file_{os.urandom(4).hex()}"
if not self.git_manager.create_branch(branch_name):
print("Agent Orchestrator Error: Failed to create a new branch for multi-file refactoring.")
return False
files_to_stage_for_commit = []
for change in refactoring_plan:
file_path_relative = change["file"]
full_file_path = os.path.join(project_root, file_path_relative)
action = change["action"]
content = change.get("content")
try:
if action == "create":
if not self.file_handler.write_file(full_file_path, content):
raise Exception(f"Failed to create file {file_path_relative}")
files_to_stage_for_commit.append(file_path_relative)
elif action == "modify":
if not self.file_handler.write_file(full_file_path, content):
raise Exception(f"Failed to modify file {file_path_relative}")
files_to_stage_for_commit.append(file_path_relative)
elif action == "delete":
if os.path.exists(full_file_path):
os.remove(full_file_path)
print(f"Agent Orchestrator: Deleted file: {file_path_relative}")
files_to_stage_for_commit.append(file_path_relative)
else:
print(f"Agent Orchestrator Warning: Attempted to delete non-existent file {file_path_relative}")
else:
print(f"Agent Orchestrator Error: Unknown action: {action} for file {file_path_relative}")
return False
except Exception as e:
print(f"Agent Orchestrator Error applying change to {file_path_relative}: {e}")
print("Agent Orchestrator: Attempting to revert branch due to error.")
self.git_manager.revert_last_commit()
return False
if files_to_stage_for_commit:
if self.git_manager.apply_and_commit_changes_internal(files_to_stage_for_commit, commit_message):
print("Agent Orchestrator: All multi-file changes committed successfully.")
return True
else:
print("Agent Orchestrator Error: Failed to commit multi-file changes.")
return False
else:
print("Agent Orchestrator: No files were staged for commit.")
return True
# --- Main Execution for Running Example ---
if __name__ == "__main__":
# Setup a temporary Git repository for the agent to work in
temp_repo_dir = "agent_workspace"
if os.path.exists(temp_repo_dir):
shutil.rmtree(temp_repo_dir)
os.makedirs(temp_repo_dir)
print(f"Setup: Created temporary agent workspace at {temp_repo_dir}")
# --- Initial Application Code (Target for Refactoring) ---
# This is the 'application code' that the agent will refactor.
# It's kept simple to focus on the agent's actions.
initial_app_code = """
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
class DataProcessor:
def __init__(self, config):
self.config = config
def process_data(self, raw_data):
logging.info("Starting data processing.")
if not isinstance(raw_data, list):
logging.error("Input raw_data must be a list.")
return []
processed_results = []
for item in raw_data:
# Complex validation logic here
if not isinstance(item, str):
logging.warning(f"Skipping non-string item: {item}")
continue
is_valid_length = len(item) >= self.config.get("min_length", 5)
try:
numeric_value = int(item)
is_valid_value = numeric_value <= self.config.get("max_value", 100)
except ValueError:
numeric_value = None
is_valid_value = False
if is_valid_length and is_valid_value:
transformed_item = item.upper() # Example transformation
processed_results.append(transformed_item)
logging.debug(f"Item '{item}' processed successfully.")
else:
logging.warning(f"Item '{item}' failed validation (length valid: {is_valid_length}, value valid: {is_valid_value}).")
logging.info("Finished data processing.")
return processed_results
"""
target_app_file = "src/simple_app.py"
full_target_app_file_path = os.path.join(temp_repo_dir, target_app_file)
# Initialize the agent
agent = AgentOrchestrator(temp_repo_dir)
# Write the initial application code to the repo and commit it
# This simulates the existing codebase the agent works on.
print("\n--- Initializing Application Code in Agent Workspace ---")
agent.file_handler.write_file(full_target_app_file_path, initial_app_code)
agent.git_manager.apply_and_commit_changes(target_app_file, initial_app_code, "Initial commit of simple_app.py")
print("\nInitial application code:")
print(agent.git_manager.get_current_file_content(target_app_file))
print("-------------------------------------------------------")
# --- Stage 1: Single-File Refactoring Demonstration ---
print("\n#######################################################")
print("# DEMONSTRATING STAGE 1: SINGLE-FILE REFACTORING #")
print("#######################################################")
agent.run_single_file_refactoring(target_app_file, "Extract the validation logic into a separate method.")
# --- Demonstrating Undo ---
print("\n--- Agent Orchestrator: User decides to undo the last refactoring ---")
if agent.git_manager.revert_last_commit():
print(f"\nAgent Orchestrator: Successfully reverted last commit. {target_app_file} is now:")
print(agent.git_manager.get_current_file_content(target_app_file))
else:
print("\nAgent Orchestrator Error: Failed to undo last commit.")
print("-------------------------------------------------------")
# Re-apply the single-file refactoring so we have the refactored state for multi-file
print("\n--- Agent Orchestrator: Re-applying single-file refactoring for multi-file stage ---")
agent.run_single_file_refactoring(target_app_file, "Extract the validation logic into a separate method.")
# --- Stage 2: Multi-File Refactoring Demonstration ---
print("\n\n#######################################################")
print("# DEMONSTRATING STAGE 2: MULTI-FILE REFACTORING #")
print("#######################################################")
# This will simulate moving the _validate_item logic to a new file
# and updating simple_app.py to import it.
agent.run_multi_file_refactoring(temp_repo_dir, "Move _validate_item to a new utils file.")
# Clean up the temporary repository
print(f"\nCleanup: Removing temporary agent workspace at {temp_repo_dir}")
shutil.rmtree(temp_repo_dir)
print("Cleanup: Done.")
--------------------------------------------------------------------------------
