app.py

import gradio as gr
from transformers import pipeline
from sentence_transformers import SentenceTransformer, util
import os
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")

# Constants for enhanced organization
system_message = "You are GitBot, the Github project guardian angel. You resolve issues and propose implementation of feature requests"
max_tokens = 2048
temperature = 0.71
top_p = 0.95

class InferenceClient:
    def __init__(self):
        pass

    def create_endpoint(self, repo_id, handler_path, model_id, task, description, hyperparameters):
        pass

    def update_endpoint(self, repo_id, handler_path, model_id, task, description, hyperparameters):
        pass

    def delete_endpoint(self, repo_id, handler_path):
        pass

    def list_endpoints(self):
        pass

    def get_endpoint_status(self, repo_id, handler_path):
        pass

    def get_endpoint_logs(self, repo_id, handler_path, num_lines):
        pass

    def get_endpoint_metrics(self, repo_id, handler_path):
        pass

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

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=max_tokens,
        do_sample=True,
        temperature=temperature,
        top_k=top_p,
    )
    assistant_response = response[0]['generated_text'].strip()

    # 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:
        programming_language = assistant_response.split(":")[1].strip()

    return {
        'assistant_response': assistant_response,
        'severity': severity,
        'programming_language': programming_language,
    }

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[: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"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,
    system_message,
    max_tokens,
    temperature,
    top_p,
    github_api_token,
    github_username,
    github_repository,
    selected_model,
    severity,
    programming_language,
    *args,
    **kwargs,
) -> dict:
    """Handles user commands and generates responses using the selected language model."""
    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,
        do_sample=True,
        temperature=temperature,
        top_k=top_p,
    )
    assistant_response = response[0]['generated_text'].strip()
    return {
        'assistant_response': assistant_response,
        'severity': severity,
        'programming_language': programming_language,
    }

with gr.Blocks() as demo:
    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=[
            "mistralai/Mixtral-8x7B-Instruct-v0.1",
            "Gabriel/Swe-review-setfit-model",
            "OpenBMB/multilingual-codeparrot"
        ],
        label="Select Model for Issue Resolution",
        value="OpenBMB/multilingual-codeparrot",
    )

    severity_dropdown = gr.Dropdown(
        choices=["Critical", "Major", "Minor", "Trivial"],
        label="Severity",
        value=None,
    )

    programming_language_textbox = gr.Textbox(label="Programming Language")

    chatbot = MyChatbot(
        respond,
        inputs=[
            gr.Textbox(label="Command"),
            gr.Textbox(label="History"),
            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,
        ],
        outputs=[gr.Textbox(label="Assistant Response")],
    )

    # Add a button to fetch GitHub issues
    fetch_issues_button = gr.Button(label="Fetch Issues")
    fetch_issues_button.click(fn=lambda github_api_token, github_username, github_repository: chatbot.issues, inputs=[github_api_token, github_username, github_repository], outputs=[chatbot])

    # Add a dropdown to select an issue
    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])

if __name__ == "__main__":
    demo.queue().launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860
    )