pull-request.py•7.37 kB
from tkinter import W
from typing import Any, Dict, Optional
import httpx
from mcp.server.fastmcp import FastMCP
import os
from pathlib import Path
from langchain_openai import ChatOpenAI
from github import Github, GithubException
from github.Repository import Repository
from github.PullRequest import PullRequest
from github.GitRef import GitRef
# Initialize FastMCP server
mcp = FastMCP("pull-request")
from dotenv import load_dotenv
load_dotenv(".env")
# Constants
GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
api_key = os.getenv("API_KEY")
api_base = os.getenv("API_BASE")
task_llm = ChatOpenAI(api_key=api_key, model="gpt-4o", base_url=api_base, temperature=0.0)
@mcp.tool()
async def task_planning(task_file: Path, plan_file: Path) -> Dict[str, Any]:
"""
Generate a detailed step-by-step plan based on task description.
Args:
task_file: Path to file containing task description (relative path)
plan_file: Path to output file where plan will be written (relative path)
Returns:
Dictionary containing status and message about the operation
"""
try:
# Read task description from file
with open(task_file, "r", encoding="utf-8") as f:
task_description = f.read()
# Generate plan using LLM
prompt = f"""
Based on the following task description, create a detailed step-by-step plan.
Include:
1. Required libraries and dependencies
2. Implementation steps with code structure
3. Potential exceptions to handle
4. Testing approach
Task Description:
{task_description}
"""
response = task_llm.invoke(prompt).content
# Write plan to output file
with open(plan_file, "w", encoding="utf-8") as f:
f.write(response)
return {"status": "success", "message": f"Plan generated and saved to {plan_file}"}
except FileNotFoundError:
return {"status": "error", "message": f"Task file not found: {task_file}"}
except Exception as e:
return {"status": "error", "message": f"Error generating plan: {str(e)}"}
@mcp.tool()
async def code_for_plan(plan_file: Path, code_file: Path) -> Dict[str, Any]:
"""
Generate Python code based on a plan file.
Args:
plan_file: Path to file containing the plan (relative path)
code_file: Path to output file where generated code will be written (relative path)
Returns:
Dictionary containing status and message about the operation
"""
try:
# Read plan from file
with open(plan_file, "r", encoding="utf-8") as f:
plan_content = f.read()
# Generate code using LLM
prompt = f"""
Based on the following plan, generate Python code that implements all the requirements.
The code should be well-structured, include error handling, and follow best practices.
Plan:
{plan_content}
Provide only the Python code implementation without any explanations.
"""
# Get raw response from the LLM
code_response = task_llm.invoke(prompt).content
# Strip markdown code block delimiters if present
if code_response.startswith("```python"):
code_response = code_response.split("```python", 1)[1]
if code_response.endswith("```"):
code_response = code_response.rsplit("```", 1)[0]
# Clean up any remaining whitespace
code_response = code_response.strip()
# Write code to output file
with open(code_file, "w", encoding="utf-8") as f:
f.write(code_response)
return {"status": "success", "message": f"Code generated and saved to {code_file}"}
except FileNotFoundError:
return {"status": "error", "message": f"Plan file not found: {plan_file}"}
except Exception as e:
return {"status": "error", "message": f"Error generating code: {str(e)}"}
@mcp.tool()
async def create_github_pull_request(
repository: str,
branch_name: str,
file_path: str,
file_content: str,
commit_message: str,
title: str,
body: str,
base_branch: str
) -> Dict[str, Any]:
"""
Create a new GitHub pull request with a file commit.
Args:
repository: Full repository name in format 'owner/repo'
branch_name: Name of the branch to create and use for changes
file_path: Path where the file should be created in the repository
file_content: Content of the file to be committed
commit_message: Message for the commit
title: Title of the pull request
body: Body/description of the pull request
base_branch: Base branch to merge into (default: 'main')
Returns:
Dictionary containing status and message about the operation
"""
github = Github(GITHUB_TOKEN)
try:
# Get repository
repo = github.get_repo(repository)
# Get base branch reference
base_branch_ref = repo.get_branch(base_branch).commit.sha
# Create new branch from base branch
try:
# Try to create new branch
new_branch = repo.create_git_ref(
ref=f'refs/heads/{branch_name}',
sha=base_branch_ref
)
except GithubException as e:
if "Reference already exists" in str(e):
# If branch already exists, get its reference
new_branch = repo.get_git_ref(f'heads/{branch_name}')
else:
return {
"status": "error",
"message": f"Failed to create branch: {str(e)}"
}
# Create blob with file content
blob = repo.create_git_blob(file_content, "utf-8")
# Create tree with the new file
base_tree = repo.get_git_tree(base_branch_ref)
element = repo.create_git_tree(
[repo.create_git_tree_element(
path=file_path,
mode="100644",
type="blob",
sha=blob.sha
)],
base_tree.sha
)
# Create commit
commit = repo.create_git_commit(
message=commit_message,
tree=element.sha,
parents=[base_branch_ref]
)
# Update branch reference
new_branch.edit(sha=commit.sha)
# Create pull request
pr = repo.create_pull(
title=title,
body=body,
head=branch_name,
base=base_branch
)
return {
"status": "success",
"message": f"Pull request created successfully: {pr.html_url}",
"pull_request_url": pr.html_url
}
except GithubException as e:
return {
"status": "error",
"message": f"GitHub API error: {str(e)}"
}
except Exception as e:
return {
"status": "error",
"message": f"Unexpected error: {str(e)}"
}
if __name__ == "__main__":
# Initialize and run the server
mcp.run(transport='stdio')