Update app.py

app.py

@@ -3,41 +3,37 @@ from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer, util
 import os
 import requests
+import json
 
-# Constants for enhanced organization
+# --- Constants ---
 GITHUB_API_BASE_URL = "https://api.github.com/repos"
-DEFAULT_MODEL = "microsoft/CodeBERT-base"
-MAX_RELATED_ISSUES = 3
+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."
 
-# Load a pre-trained model for sentence similarity
-similarity_model = SentenceTransformer('all-mpnet-base-v2')
-
-# Define models for issue analysis
-model_name = "enricoros/big-agi"
+# --- Model Setup ---
+model_name = "enricoros/big-agi"  # Choose your preferred model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+similarity_model = SentenceTransformer('all-mpnet-base-v2')  # For issue similarity
+
+# --- Functions ---
 
 def analyze_issues(issue_text: str, model_name: str, severity: str = None, programming_language: str = None) -> dict:
-    # Initialize the model
+    """Analyzes an issue description and extracts severity and programming language."""
     model = pipeline("text-generation", model=model_name)
-
-    # Generate a response
     response = model(
-        f"{issue_text}\nAssistant: ",
-        max_length=512,
+        f"{SYSTEM_MESSAGE}\n{issue_text}\nAssistant: ",
+        max_length=2048,  # Adjust as needed
         do_sample=True,
-        temperature=0.7,
-        top_k=50,
-        top_p=0.9,
+        temperature=0.7,  # Adjust as needed
+        top_k=50,  # Adjust as needed
     )
-
-    # Extract the assistant's response
     assistant_response = response[0]['generated_text'].strip()
 
-    # Analyze the response
+    # Extract information from the response
     if "Severity" in assistant_response:
         severity = assistant_response.split(":")[1].strip()
-
     if "Programming Language" in assistant_response:
         programming_language = assistant_response.split(":")[1].strip()
 
@@ -48,109 +44,120 @@ def analyze_issues(issue_text: str, model_name: str, severity: str = None, progr
     }
 
 def find_related_issues(issue_text: str, issues: list) -> list:
-    # Calculate the similarity between the issue and other issues
+    """Finds related issues based on text similarity."""
     issue_embedding = similarity_model.encode(issue_text)
     similarities = [util.cos_sim(issue_embedding, similarity_model.encode(issue['title'])) for issue in issues]
-
-    # Sort the issues by similarity
     sorted_issues = sorted(enumerate(similarities), key=lambda x: x[1], reverse=True)
-
-    # Select the top related issues
-    related_issues = [issues[i] for i, similarity in sorted_issues[:MAX_RELATED_ISSUES]]
-
-    return related_issues
+    return [issues[i] for i, similarity in sorted_issues[:MAX_RELATED_ISSUES]]
 
 def fetch_github_issues(github_api_token: str, github_username: str, github_repository: str) -> list:
-    # Fetch the issues from the GitHub API
+    """Fetches issues from 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)
-
-    # Parse the JSON response
-    issues = response.json()
-
-    return issues
+    if response.status_code == 200:
+        return response.json()
+    else:
+        return []
 
 def respond(
-    command, history, system_message, max_tokens, temperature, top_p,
-    github_api_token, github_username, github_repository,
-    selected_model, severity, programming_language, *args, **kwargs
+    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,
+    *args,
+    **kwargs,
 ) -> dict:
-    # Initialize the model
-    model = pipeline("text-generation", model="enricoros/big-agi")
-
-    # Generate a response
+    """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
     response = model(
-        f"{system_message}\n{command}\n{history}\n{github_username}/{github_repository}\nSeverity: {severity}\nProgramming Language: {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=50,
-        top_p=top_p,
+        top_k=top_p,
     )
-
-    # Extract the assistant's response
     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': severity,
-        'programming_language': programming_language,
+        'severity': analyzed_data['severity'],
+        'programming_language': analyzed_data['programming_language'],
     }
 
+# --- Gradio Interface ---
+
 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")
 
-    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",
-    )
-
-    model_dropdown = gr.Dropdown(
-        choices=["microsoft/CodeBERT-base", "Salesforce/codegen-45M-mono"],
-        label="Select Model for Issue Resolution",
-        value=DEFAULT_MODEL,
-    )
-
-    severity_dropdown = gr.Dropdown(
-        choices=["Critical", "Major", "Minor", "Trivial"],
-        label="Severity",
-        value=None,
-    )
+    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,
+        )
 
-    programming_language_textbox = gr.Textbox(label="Programming Language")
+    with gr.Row():
+        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",
-    )
+    with gr.Row():
+        command_dropdown = gr.Dropdown(
+            choices=[
+                "/github",
+                "/help",
+                "/generate_code",
+                "/explain_concept",
+                "/write_documentation",
+                "/translate_code",
+            ],
+            label="Select Command",
+        )
 
     chatbot = gr.Chatbot(
         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.7, step=0.1, label="Temperature"),
-            gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, 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)"),
         ],
     )
 
-if __name__ == "__main__":
-    demo.queue().launch(share=True, server_name="0.0.0.0", server_port=7860)
+    demo.launch(share=True)