benchmarks/lowq2_reconstruction/PRfunctions.py

0001 #!/usr/bin/env python3
0002 """
0003 Script to create GitHub PR suggestions for ONNX model updates
0004 """
0005 import requests
0006 import base64
0007 import xml.etree.ElementTree as ET
0008 from pathlib import Path
0009 from datetime import datetime, timezone
0010 import os
0011
0012 # =============================================================================
0013 # Utility Functions
0014 # =============================================================================
0015
0016 def parse_repository(repository):
0017     """Parse repository string into owner and name"""
0018     try:
0019         owner, name = repository.split('/')
0020         return owner, name
0021     except ValueError:
0022         print(f"❌ Invalid repository format: {repository}. Expected: 'owner/name'")
0023         return None, None
0024
0025 # =============================================================================
0026 # GitHub API Functions
0027 # =============================================================================
0028
0029 def get_pr_info(repo_owner, repo_name, pr_number, github_token=None):
0030     """Get PR information including head SHA"""
0031     print(f"Getting PR information for #{pr_number}...")
0032
0033     headers = {
0034         'Accept': 'application/vnd.github+json'
0035     }
0036
0037     if github_token:
0038         headers['Authorization'] = f'token {github_token}'
0039
0040     url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/pulls/{pr_number}"
0041     response = requests.get(url, headers=headers)
0042
0043     if response.status_code == 200:
0044         pr_data = response.json()
0045         print(f"✅ Found PR: {pr_data['title']}")
0046         print(f"   Branch: {pr_data['head']['ref']}")
0047         print(f"   SHA: {pr_data['head']['sha']}")
0048         return pr_data
0049     else:
0050         print(f"❌ Failed to get PR info: {response.status_code}")
0051         print(f"   Response: {response.text}")
0052         return None
0053
0054 def get_file_content(repo_owner, repo_name, file_path, sha, github_token=None):
0055     """Get file content from GitHub at a specific commit"""
0056     print(f"Getting content for {file_path} at commit {sha[:8]}...")
0057
0058     headers = {
0059         'Accept': 'application/vnd.github+json'
0060     }
0061
0062     if github_token:
0063         headers['Authorization'] = f'token {github_token}'
0064
0065     url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/contents/{file_path}?ref={sha}"
0066     response = requests.get(url, headers=headers)
0067
0068     if response.status_code == 200:
0069         content_data = response.json()
0070         if 'content' in content_data:
0071             # GitHub returns base64-encoded content
0072             content = base64.b64decode(content_data['content']).decode('utf-8')
0073             print(f"✅ Got file content ({len(content)} characters)")
0074             return content
0075         else:
0076             print(f"❌ No content found in response")
0077             return None
0078     elif response.status_code == 404:
0079         print(f"❌ File not found: {file_path}")
0080         return None
0081     else:
0082         print(f"❌ Failed to get file content: {response.status_code}")
0083         print(f"   Response: {response.text}")
0084         return None
0085
0086 # =============================================================================
0087 # XML Processing Functions
0088 # =============================================================================
0089
0090 def find_and_update_epic_fileloader_url(content, calibration_file, new_url, plugin_name='epic_FileLoader'):
0091     """Find the line with FileLoader URL and return line number and suggested change"""
0092     print(f"Parsing XML to find {plugin_name} URL for {calibration_file}...")
0093
0094     try:
0095         # Parse the XML content
0096         root = ET.fromstring(content)
0097         calibration_filename = Path(calibration_file).name
0098
0099         # Find all plugin elements with the specified name
0100         file_loader_plugins = root.findall(f".//plugin[@name='{plugin_name}']")
0101
0102         if not file_loader_plugins:
0103             print(f"❌ No {plugin_name} plugin found")
0104             return None, None
0105
0106         print(f"✅ Found {len(file_loader_plugins)} {plugin_name} plugin(s)")
0107
0108         # Look through each FileLoader plugin
0109         for plugin in file_loader_plugins:
0110             print(f"   Checking plugin...")
0111
0112             # Find all arg elements with value starting with "url:"
0113             for arg in plugin.findall("arg"):
0114                 value = arg.get('value', '')
0115
0116                 if value.startswith('url:') and calibration_filename in value:
0117                     print(f"✅ Found matching URL arg: {value}")
0118
0119                     # Find the line number of this change in the original content
0120                     line_number = find_line_number_of_change(content, value)
0121
0122                     if line_number:
0123                         # Get the original line and create the suggested replacement
0124                         lines = content.split('\n')
0125                         original_line = lines[line_number - 1]  # Convert to 0-based index
0126                         suggested_line = original_line.replace(value, f'url:{new_url}')
0127
0128                         print(f"✅ Line {line_number}: {original_line.strip()}")
0129                         print(f"   Suggested: {suggested_line.strip()}")
0130
0131                         return line_number, suggested_line
0132                     else:
0133                         print("❌ Could not find line number for the URL")
0134                         return None, None
0135
0136         print(f"❌ No matching URL argument found in {plugin_name} plugins")
0137         return None, None
0138
0139     except ET.ParseError as e:
0140         print(f"❌ XML parsing error: {e}")
0141         return None, None
0142
0143 def find_line_number_of_change(original_content, old_value):
0144     """Find the line number where the change occurred"""
0145     lines = original_content.split('\n')
0146
0147     for i, line in enumerate(lines, 1):
0148         if old_value in line:
0149             return i
0150
0151     return None
0152
0153 # =============================================================================
0154 # GitHub PR Comment Functions
0155 # =============================================================================
0156
0157 def create_pr_suggestion(repo_owner, repo_name, pr_number, calibration_file, xml_file, line_number, suggested_line, head_sha, github_token, artifacts_url=''):
0158     """Create a PR comment with proposed changes"""
0159     print(f"Creating PR comment with calibration update for #{pr_number}...")
0160
0161     headers = {
0162         'Accept': 'application/vnd.github+json',
0163         'Authorization': f'token {github_token}'
0164     }
0165
0166     bot_comment_base = f"🤖 **Automated Calibration `{calibration_file}` Update**"
0167
0168     # Check for existing comments from bot
0169     existing_comment_id = find_existing_bot_comment_general(repo_owner, repo_name, pr_number, bot_comment_base, github_token)
0170
0171     # Generate timestamp for the comment
0172     timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
0173
0174     # Get current line content for context
0175     content = get_file_content(repo_owner, repo_name, xml_file, head_sha, github_token)
0176     lines = content.split('\n') if content else []
0177     current_line = lines[line_number - 1].strip() if line_number <= len(lines) else "Line not found"
0178
0179     comment_body = f"""{bot_comment_base}{' (Updated)' if existing_comment_id else ''}
0180
0181 A new calibration has been generated and is ready for use.
0182
0183 **File:** `{xml_file}`
0184 **Line:** {line_number}
0185 **Last updated:** {timestamp}
0186
0187 **Current line:**
0188 ```xml
0189 {current_line}
0190 ```
0191
0192 **Proposed change:**
0193 ```xml
0194 {suggested_line.strip()}
0195 ```
0196
0197 Please update the calibration URL in `{xml_file}` at line {line_number}."""
0198
0199     # Add artifacts link if provided
0200     if artifacts_url:
0201         comment_body += f"\n\n---\n\n### 📊 Review Results\n\nPlease review the artifacts here: {artifacts_url}"
0202
0203     # Create or update comment via GitHub REST API (no gh CLI)
0204     if existing_comment_id:
0205         print(f"Updating existing comment {existing_comment_id}...")
0206         update_url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/issues/comments/{existing_comment_id}"
0207         update_data = {'body': comment_body}
0208         response = requests.patch(update_url, headers=headers, json=update_data)
0209         if response.status_code == 200:
0210             print("✅ Existing PR comment updated successfully")
0211             return response.json()
0212         else:
0213             print(f"❌ Failed to update existing comment: {response.status_code}\n{response.text}")
0214             return None
0215     else:
0216         print("Creating new PR comment...")
0217         comment_url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/issues/{pr_number}/comments"
0218         comment_data = {'body': comment_body}
0219         response = requests.post(comment_url, headers=headers, json=comment_data)
0220         if response.status_code == 201:
0221             print("✅ New PR comment created successfully")
0222             return response.json()
0223         else:
0224             print(f"❌ Failed to create PR comment: {response.status_code}\n{response.text}")
0225             return None
0226
0227 def find_existing_bot_comment_general(repo_owner, repo_name, pr_number, bot_comment_base, github_token):
0228     """Find existing bot comment (general PR comment, not line-specific)"""
0229     headers = {
0230         'Accept': 'application/vnd.github+json',
0231         'Authorization': f'token {github_token}'
0232     }
0233
0234     # Get all general comments for the PR
0235     comments_url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/issues/{pr_number}/comments"
0236     response = requests.get(comments_url, headers=headers)
0237
0238     if response.status_code != 200:
0239         print(f"❌ Failed to get PR comments: {response.status_code}")
0240         return None
0241
0242     # Look for existing bot comment
0243     for comment in response.json():
0244         if bot_comment_base in comment.get('body', ''):
0245             print(f"✅ Found existing bot comment: {comment['id']}")
0246             return comment['id']
0247
0248     print("No existing bot comment found")
0249     return None