Update app.py

app.py

@@ -2,48 +2,28 @@ import gradio as gr
 import google.generativeai as genai
 import paramiko
 import io
-import re
-
-def parse_github_url(github_url):
-    if github_url.startswith("git@github.com:"):
-        # SSH format
-        parts = github_url.split(':')[1].split('/')
-        owner = parts[0]
-        repo = parts[1].split('.git')[0]
-        return owner, repo, None, None
-    elif github_url.startswith("https://github.com/"):
-        # HTTPS format
-        parts = github_url.split('/')
-        owner = parts[3]
-        repo = parts[4]
-        branch = parts[6] if len(parts) > 6 else None
-        file_path = '/'.join(parts[7:]) if len(parts) > 7 else None
-        return owner, repo, branch, file_path
-    else:
-        raise ValueError("Invalid GitHub URL format")
 
 def fetch_github_file(github_url, ssh_private_key):
     try:
-        owner, repo, branch, file_path = parse_github_url(github_url)
+        # Parse the GitHub URL
+        parts = github_url.split('/')
+        owner = parts[3]
+        repo = parts[4]
+        branch = parts[6]
+        file_path = '/'.join(parts[7:])
 
         # Set up SSH client
         ssh = paramiko.SSHClient()
         ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
 
         # Load the private key
-        try:
-            private_key = paramiko.RSAKey.from_private_key(io.StringIO(ssh_private_key.strip()))
-        except paramiko.ssh_exception.SSHException as e:
-            return f"Error: Invalid SSH key - {str(e)}"
+        private_key = paramiko.RSAKey.from_private_key(io.StringIO(ssh_private_key))
 
         # Connect to GitHub using SSH
         ssh.connect('github.com', username='git', pkey=private_key)
 
         # Construct the git command to fetch the file content
-        if branch and file_path:
-            git_command = f"git archive --remote=git@github.com:{owner}/{repo}.git {branch} {file_path} | tar -xO"
-        else:
-            git_command = f"git ls-tree -r --name-only HEAD | grep -E '(requirements\.txt|package\.json|Gemfile)'"
+        git_command = f"git archive --remote=git@github.com:{owner}/{repo}.git {branch} {file_path} | tar -xO"
 
         # Execute the command
         stdin, stdout, stderr = ssh.exec_command(git_command)
@@ -56,8 +36,8 @@ def fetch_github_file(github_url, ssh_private_key):
 
         if not file_content:
             return {"error": "Error: File not found or empty."}
-        
-        return {file_path if file_path else "repository_files": file_content}
+
+        return {file_path: file_content}
     except paramiko.AuthenticationException:
         return {"error": "Error: Authentication failed. Please check your SSH key."}
     except paramiko.SSHException as ssh_err:
@@ -94,16 +74,13 @@ def process_input(file, github_url, ssh_private_key, gemini_api_key):
     if file is not None:
         file_content = {file.name: file.read().decode('utf-8')}
     elif github_url and ssh_private_key:
+        if not github_url.startswith("https://github.com/"):
+            return "Error: Invalid GitHub URL. Please use the format: https://github.com/username/repository/blob/branch/path/to/file"
         if not ssh_private_key.strip():
             return "Error: SSH Private Key is empty. Please provide a valid key."
-        if not ssh_private_key.strip().startswith("-----BEGIN"):
-            return "Error: Invalid SSH key format. Please ensure you've copied the entire key, including the BEGIN and END lines."
-        try:
-            file_content = fetch_github_file(github_url, ssh_private_key)
-            if "error" in file_content:
-                return file_content["error"]
-        except ValueError as e:
-            return str(e)
+        file_content = fetch_github_file(github_url, ssh_private_key)
+        if isinstance(file_content, dict) and "error" in file_content:
+            return file_content["error"]
     else:
         return "Error: Please either upload a file OR provide both GitHub URL and SSH Private Key."
     
@@ -118,7 +95,7 @@ iface = gr.Interface(
     fn=process_input,
     inputs=[
         gr.File(label="Upload dependency file (e.g., requirements.txt, package.json, Gemfile)"),
-        gr.Textbox(label="GitHub Repository URL (optional, SSH or HTTPS format)"),
+        gr.Textbox(label="GitHub File URL (optional)"),
         gr.TextArea(label="SSH Private Key (required if using GitHub URL)", lines=10),
         gr.Textbox(label="Gemini API Key", type="password"),
     ],