Update app.py

app.py

@@ -6,11 +6,11 @@ import requests
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 from huggingface_hub import InferenceClient, HfApi
+import git
+import gitdb
 
 # If you have a GPU, move the model to it
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device) 
-
 
 # Constants for enhanced organization
 system_message = "You are GitBot, the Github project guardian angel. You resolve issues and propose implementation of feature requests"
@@ -43,33 +43,61 @@ class InferenceClient:
     def get_endpoint_metrics(self, repo_id, handler_path):
         pass
 
-client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
-
-# Load the model and tokenizer
-model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-# Define input prompt
-input_prompt = "(input value = highest-level-quality code content invocation ; True)"
-
-# Tokenize the input prompt
-input_ids = tokenizer(input_prompt, return_tensors="pt", truncation=True)
-
-# Generate the code
-generated_code = model.generate(input_ids.to(model.device))
-
-# Decode the generated code
-generated_code_str = tokenizer.batch_decode(generated_code, skip_special_tokens=True)[0]
-
-# Print the generated code
-print(generated_code_str)
+class MyChatbot(gr.Chatbot):
+    """Custom Chatbot class for enhanced functionality."""
+    def __init__(self, fn, **kwargs):
+        super().__init__(fn, **kwargs)
+        self.issues = []  # Store fetched issues
+        self.current_issue = None  # Store the currently selected issue
 
-DEFAULT_MODEL = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-MAX_RELATED_ISSUES = 5
+    def postprocess(self, y):
+        """Post-processes the response to handle commands and display results."""
+        # Extract the response from the dictionary
+        assistant_response = y['assistant_response']
 
-# Load a pre-trained model for sentence similarity
-similarity_model = SentenceTransformer('all-mpnet-base-v2')
+        # Handle commands
+        if y['command'] == "/github":
+            if not y['github_api_token']:
+                return "Please enter your GitHub API token first."
+            else:
+                try:
+                    self.issues = fetch_github_issues(y['github_api_token'], y['github_username'], y['github_repository'])
+                    issue_list = "\n".join([f"{i+1}. {issue['title']}" for i, issue in enumerate(self.issues)])
+                    return f"Available GitHub Issues:\n{issue_list}\n\nEnter the issue number to analyze:"
+                except Exception as e:
+                    return f"Error fetching GitHub issues: {e}"
+        elif y['command'] == "/help":
+            return """Available commands:
+                - `/github`: Analyze a GitHub issue
+                - `/help`: Show this help message
+                - `/generate_code [code description]`: Generate code based on the description
+                - `/explain_concept [concept]`: Explain a concept
+                - `/write_documentation [topic]`: Write documentation for a given topic
+                - `/translate_code [code] to [target language]`: Translate code to another language"""
+        elif y['command'].isdigit() and self.issues:
+            try:
+                issue_number = int(y['command']) - 1
+                self.current_issue = self.issues[issue_number]  # Store the selected issue
+                issue_text = self.current_issue['title'] + "\n\n" + self.current_issue['body']
+                resolution = analyze_issues(issue_text, y['selected_model'], y['severity'], y['programming_language'])
+                related_issues = find_related_issues(issue_text, self.issues)
+                related_issue_text = "\n".join(
+                    [f"- {issue['title']} (Similarity: {similarity:.2f})" for issue, similarity in related_issues]
+                )
+                return f"Resolution for Issue '{self.current_issue['title']}':\n{resolution['assistant_response']}\n\nRelated Issues:\n{related_issue_text}"
+            except Exception as e:
+                return f"Error analyzing issue: {e}"
+        elif y['command'].startswith("/"):
+            # Handle commands like `/generate_code`, `/explain_concept`, etc.
+            if self.current_issue:
+                # Use the current issue's context for these commands
+                issue_text = self.current_issue['title'] + "\n\n" + self.current_issue['body']
+                return analyze_issues(issue_text, y['selected_model'], y['severity'], y['programming_language'])['assistant_response']
+            else:
+                return "Please select an issue first using `/github`."
+        else:
+            # For free-form text, simply display the assistant's response
+            return assistant_response
 
 def analyze_issues(issue_text: str, model_name: str, severity: str = None, programming_language: str = None) -> str:
     """Analyzes issues and provides solutions using a specified language model."""
@@ -97,20 +125,46 @@ def analyze_issues(issue_text: str, model_name: str, severity: str = None, progr
 
 def find_related_issues(issue_text: str, issues: list) -> list:
     """Finds semantically related issues from a list of issues based on the input issue text."""
+    similarity_model = SentenceTransformer('all-mpnet-base-v2')
     issue_embedding = similarity_model.encode(issue_text)
     similarities = [util.cos_sim(issue_embedding, similarity_model.encode(issue['title'])) for issue in issues]
     sorted_issues = sorted(enumerate(similarities), key=lambda x: x[1], reverse=True)
-    related_issues = [issues[i] for i, similarity in sorted_issues[:MAX_RELATED_ISSUES]]
+    related_issues = [issues[i] for i, similarity in sorted_issues[:5]]
     return related_issues
 
 def fetch_github_issues(github_api_token: str, github_username: str, github_repository: str) -> list:
     """Fetches issues from a specified GitHub repository using the GitHub API."""
     headers = {'Authorization': f'token {github_api_token}'}
-    url = f"{GITHUB_API_BASE_URL}/{github_username}/{github_repository}/issues"
+    url = f"https://api.github.com/repos/{github_username}/{github_repository}/issues"
     response = requests.get(url, headers=headers)
     issues = response.json()
     return issues
 
+def push_to_repo(github_api_token: str, github_username: str, github_repository: str, commit_message: str, commit_file: str):
+    """Pushes changes to a GitHub repository."""
+    try:
+        repo = git.Repo.clone_from(f"https://github.com/{github_username}/{github_repository}.git", f"{github_repository}")
+        repo.git.add(commit_file)
+        repo.git.commit(m=commit_message)
+        repo.git.push()
+        return f"Changes pushed to {github_repository} successfully."
+    except gitdb.exc.InvalidGitRepositoryError:
+        return f"Invalid repository: {github_repository}"
+    except Exception as e:
+        return f"Error pushing to repository: {e}"
+
+def resolve_issue(github_api_token: str, github_username: str, github_repository: str, issue_number: int, resolution: str):
+    """Resolves an issue by pushing a commit with the resolution."""
+    try:
+        issue_text = fetch_github_issues(github_api_token, github_username, github_repository)[issue_number]['body']
+        commit_message = f"Resolved issue {issue_number}: {issue_text}"
+        commit_file = "resolution.txt"
+        with open(commit_file, "w") as f:
+            f.write(resolution)
+        return push_to_repo(github_api_token, github_username, github_repository, commit_message, commit_file)
+    except Exception as e:
+        return f"Error resolving issue: {e}"
+
 def respond(
     command,
     history,
@@ -128,7 +182,7 @@ def respond(
     **kwargs,
 ) -> dict:
     """Handles user commands and generates responses using the selected language model."""
-    model = pipeline("text-generation", model="enricoros/big-agi")
+    model = pipeline("text-generation", model="OpenBMB/multilingual-codeparrot")
     response = model(
         f"{system_message}\n{command}\n{history}\n{github_username}/{github_repository}\n{severity}\n{programming_language}\nAssistant: ",
         max_length=max_tokens,
@@ -143,62 +197,6 @@ def respond(
         'programming_language': programming_language,
     }
 
-class MyChatbot(gr.Chatbot):
-    """Custom Chatbot class for enhanced functionality."""
-    def __init__(self, fn, **kwargs):
-        super().__init__(fn, **kwargs)
-        self.issues = []  # Store fetched issues
-        self.current_issue = None  # Store the currently selected issue
-
-    def postprocess(self, y):
-        """Post-processes the response to handle commands and display results."""
-        # Extract the response from the dictionary
-        assistant_response = y['assistant_response']
-
-        # Handle commands
-        if y['command'] == "/github":
-            if not y['github_api_token']:
-                return "Please enter your GitHub API token first."
-            else:
-                try:
-                    self.issues = fetch_github_issues(y['github_api_token'], y['github_username'], y['github_repository'])
-                    issue_list = "\n".join([f"{i+1}. {issue['title']}" for i, issue in enumerate(self.issues)])
-                    return f"Available GitHub Issues:\n{issue_list}\n\nEnter the issue number to analyze:"
-                except Exception as e:
-                    return f"Error fetching GitHub issues: {e}"
-        elif y['command'] == "/help":
-            return """Available commands:
-                - `/github`: Analyze a GitHub issue
-                - `/help`: Show this help message
-                - `/generate_code [code description]`: Generate code based on the description
-                - `/explain_concept [concept]`: Explain a concept
-                - `/write_documentation [topic]`: Write documentation for a given topic
-                - `/translate_code [code] to [target language]`: Translate code to another language"""
-        elif y['command'].isdigit() and self.issues:
-            try:
-                issue_number = int(y['command']) - 1
-                self.current_issue = self.issues[issue_number]  # Store the selected issue
-                issue_text = self.current_issue['title'] + "\n\n" + self.current_issue['body']
-                resolution = analyze_issues(issue_text, y['selected_model'], y['severity'], y['programming_language'])
-                related_issues = find_related_issues(issue_text, self.issues)
-                related_issue_text = "\n".join(
-                    [f"- {issue['title']} (Similarity: {similarity:.2f})" for issue, similarity in related_issues]
-                )
-                return f"Resolution for Issue '{self.current_issue['title']}':\n{resolution['assistant_response']}\n\nRelated Issues:\n{related_issue_text}"
-            except Exception as e:
-                return f"Error analyzing issue: {e}"
-        elif y['command'].startswith("/"):
-            # Handle commands like `/generate_code`, `/explain_concept`, etc.
-            if self.current_issue:
-                # Use the current issue's context for these commands
-                issue_text = self.current_issue['title'] + "\n\n" + self.current_issue['body']
-                return analyze_issues(issue_text, y['selected_model'], y['severity'], y['programming_language'])['assistant_response']
-            else:
-                return "Please select an issue first using `/github`."
-        else:
-            # For free-form text, simply display the assistant's response
-            return assistant_response
-
 with gr.Blocks() as demo:
     with gr.Row():
         github_api_token = gr.Textbox(label="GitHub API Token", type="password")
@@ -217,7 +215,7 @@ with gr.Blocks() as demo:
             "OpenBMB/multilingual-codeparrot"
         ],
         label="Select Model for Issue Resolution",
-        value=DEFAULT_MODEL,
+        value="OpenBMB/multilingual-codeparrot",
     )
 
     severity_dropdown = gr.Dropdown(
@@ -258,6 +256,10 @@ with gr.Blocks() as demo:
     issue_dropdown = gr.Dropdown(label="Select Issue", choices=[], interactive=True)
     issue_dropdown.change(fn=lambda issue_number, chatbot: chatbot.postprocess(issue_number), inputs=[issue_dropdown, chatbot], outputs=[chatbot])
 
+    # Add a button to resolve an issue
+    resolve_issue_button = gr.Button(label="Resolve Issue")
+    resolve_issue_button.click(fn=lambda github_api_token, github_username, github_repository, issue_number, resolution: resolve_issue(github_api_token, github_username, github_repository, issue_number, resolution), inputs=[github_api_token, github_username, github_repository, gr.Number(label="Issue Number"), chatbot], outputs=[chatbot])
+
     # Connect the chatbot input to the issue dropdown
     chatbot.input.change(fn=lambda chatbot, github_api_token, github_username, github_repository: chatbot.postprocess("/github"), inputs=[chatbot, github_api_token, github_username, github_repository], outputs=[chatbot])
 
@@ -266,4 +268,4 @@ if __name__ == "__main__":
         share=True,
         server_name="0.0.0.0",
         server_port=7860
-    )
+    )