Update app.py

app.py

@@ -1,37 +1,36 @@
 import gradio as gr
-from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import pipeline
 from sentence_transformers import SentenceTransformer, util
 import os
 import requests
-import json
 
-# --- Constants ---
-GITHUB_API_BASE_URL = "https://api.github.com/repos"
-DEFAULT_MODEL = "microsoft/CodeBERT-base"  # Default model for issue resolution
-MAX_RELATED_ISSUES = 3  # Maximum number of related issues to display
-SYSTEM_MESSAGE = "You are GitBot, the Github project guardian angel. You resolve issues and propose implementation of feature requests."
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
-# --- Model Setup ---
-model_name = "enricoros/big-agi"  # Choose your preferred model
+model_name = "enricoros/big-agi"  # You can change this to other models if desired
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-similarity_model = SentenceTransformer('all-mpnet-base-v2')  # For issue similarity
 
-# --- Functions ---
+# Constants for enhanced organization
+GITHUB_API_BASE_URL = "https://api.github.com/repos"
+DEFAULT_MODEL = "microsoft/CodeBERT-base"
+MAX_RELATED_ISSUES = 3
+
+# Load a pre-trained model for sentence similarity
+similarity_model = SentenceTransformer('all-mpnet-base-v2')
 
-def analyze_issues(issue_text: str, model_name: str, severity: str = None, programming_language: str = None) -> dict:
-    """Analyzes an issue description and extracts severity and programming language."""
+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."""
     model = pipeline("text-generation", model=model_name)
     response = model(
-        f"{SYSTEM_MESSAGE}\n{issue_text}\nAssistant: ",
-        max_length=2048,  # Adjust as needed
+        f"{system_message}\n{issue_text}\nAssistant: ",
+        max_length=max_tokens,
         do_sample=True,
-        temperature=0.7,  # Adjust as needed
-        top_k=50,  # Adjust as needed
+        temperature=temperature,
+        top_k=top_p,
     )
     assistant_response = response[0]['generated_text'].strip()
 
-    # Extract information from the response
+    # Extract severity and programming language from the response
     if "Severity" in assistant_response:
         severity = assistant_response.split(":")[1].strip()
     if "Programming Language" in assistant_response:
@@ -44,120 +43,167 @@ def analyze_issues(issue_text: str, model_name: str, severity: str = None, progr
     }
 
 def find_related_issues(issue_text: str, issues: list) -> list:
-    """Finds related issues based on text similarity."""
+    """Finds semantically related issues from a list of issues based on the input issue text."""
     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)
-    return [issues[i] for i, similarity in sorted_issues[:MAX_RELATED_ISSUES]]
+    related_issues = [issues[i] for i, similarity in sorted_issues[:MAX_RELATED_ISSUES]]
+    return related_issues
 
 def fetch_github_issues(github_api_token: str, github_username: str, github_repository: str) -> list:
-    """Fetches issues from the GitHub API."""
+    """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"
     response = requests.get(url, headers=headers)
-    if response.status_code == 200:
-        return response.json()
-    else:
-        return []
+    issues = response.json()
+    return issues
 
 def respond(
-    command: str,
-    history: str,
-    github_api_token: str,
-    github_username: str,
-    github_repository: str,
-    selected_model: str,
-    severity: str,
-    programming_language: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
+    command,
+    history,
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+    github_api_token,
+    github_username,
+    github_repository,
+    selected_model,
+    severity,
+    programming_language,
     *args,
     **kwargs,
 ) -> dict:
-    """Generates a response based on the command, history, and other parameters."""
-    model = pipeline("text-generation", model=selected_model)
-    
-    # Fetch issues if the command is /github
-    if command == "/github":
-        issues = fetch_github_issues(github_api_token, github_username, github_repository)
-        if issues:
-            related_issues = find_related_issues(history, issues)
-            related_issues_text = "\n".join(
-                f"## Related Issue {i+1}: {issue['title']}\n{issue['body']}" for i, issue in enumerate(related_issues)
-            )
-            history += f"\n{related_issues_text}"
-
-    # Generate a response from the LLM
+    """Handles user commands and generates responses using the selected language model."""
+    model = pipeline("text-generation", model="enricoros/big-agi")
     response = model(
-        f"{SYSTEM_MESSAGE}\n{command}\n{history}\n{github_username}/{github_repository}\n{severity}\n{programming_language}\nAssistant: ",
+        f"{system_message}\n{command}\n{history}\n{github_username}/{github_repository}\n{severity}\n{programming_language}\nAssistant: ",
         max_length=max_tokens,
         do_sample=True,
         temperature=temperature,
         top_k=top_p,
     )
     assistant_response = response[0]['generated_text'].strip()
-
-    # Analyze the response for severity and programming language
-    analyzed_data = analyze_issues(assistant_response, selected_model, severity, programming_language)
-
     return {
         'assistant_response': assistant_response,
-        'severity': analyzed_data['severity'],
-        'programming_language': analyzed_data['programming_language'],
+        'severity': severity,
+        'programming_language': programming_language,
     }
 
-# --- Gradio Interface ---
+class MyChatbot(gr.Chatbot):
+    """Custom Chatbot class for enhanced functionality."""
+    def __init__(self, fn, **kwargs):
+        super().__init__(fn, **kwargs)
+        self.issues = []  # Store fetched issues
+
+    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
+                issue = self.issues[issue_number]
+                issue_text = issue['title'] + "\n\n" + 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 '{issue['title']}':\n{resolution['assistant_response']}\n\nRelated Issues:\n{related_issue_text}"
+            except Exception as e:
+                return f"Error analyzing issue: {e}"
+        else:
+            # For other commands or free-form text, simply display the assistant's response
+            return assistant_response
 
 with gr.Blocks() as demo:
-    gr.Markdown("## GitBot: Your GitHub Assistant")
-
     with gr.Row():
         github_api_token = gr.Textbox(label="GitHub API Token", type="password")
         github_username = gr.Textbox(label="GitHub Username")
         github_repository = gr.Textbox(label="GitHub Repository")
 
-    with gr.Row():
-        model_dropdown = gr.Dropdown(
-            choices=["microsoft/CodeBERT-base", "Salesforce/codegen-350M-mono", "enricoros/big-agi"],  # Add more models
-            label="Select Model for Issue Resolution",
-            value=DEFAULT_MODEL,
-        )
+    system_message = gr.Textbox(
+        value="You are GitBot, the Github project guardian angel. You resolve issues and propose implementation of feature requests",
+        label="System message",
+    )
 
-    with gr.Row():
-        severity_dropdown = gr.Dropdown(
-            choices=["Critical", "Major", "Minor", "Trivial"],
-            label="Severity",
-            value=None,
-        )
-        programming_language_textbox = gr.Textbox(label="Programming Language")
+    model_dropdown = gr.Dropdown(
+        choices=[
+            "microsoft/CodeBERT-base",
+            "Salesforce/codegen-45M-mono",
+            
+        ],
+        label="Select Model for Issue Resolution",
+        value=DEFAULT_MODEL,
+    )
 
-    with gr.Row():
-        command_dropdown = gr.Dropdown(
-            choices=[
-                "/github",
-                "/help",
-                "/generate_code",
-                "/explain_concept",
-                "/write_documentation",
-                "/translate_code",
-            ],
-            label="Select Command",
-        )
-
-    chatbot = gr.Chatbot(
+    severity_dropdown = gr.Dropdown(
+        choices=["Critical", "Major", "Minor", "Trivial"],
+        label="Severity",
+        value=None,
+    )
+
+    programming_language_textbox = gr.Textbox(label="Programming Language")
+
+    command_dropdown = gr.Dropdown(
+        choices=[
+            "/github",
+            "/help",
+            "/generate_code",
+            "/explain_concept",
+            "/write_documentation",
+            "/translate_code",
+        ],
+        label="Select Command",
+    )
+
+    chatbot = MyChatbot(
         respond,
         additional_inputs=[
+            command_dropdown,
+            system_message,
+            gr.Slider(minimum=1, maximum=8192, value=2048, step=1, label="Max new tokens"),
+            gr.Slider(minimum=0.1, maximum=4.0, value=0.71, step=0.1, label="Temperature"),
+            gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.95,
+                step=0.01,
+                label="Top-p (nucleus sampling)",
+            ),
             github_api_token,
             github_username,
             github_repository,
             model_dropdown,
             severity_dropdown,
             programming_language_textbox,
-            gr.Slider(minimum=1, maximum=8192, value=2048, step=1, label="Max new tokens"),
-            gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-            gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.01, label="Top-p (nucleus sampling)"),
         ],
     )
 
-    demo.launch(share=True)
+if __name__ == "__main__":
+    demo.queue().launch(
+        share=True,
+        server_name="0.0.0.0",
+        server_port=7860
+    )