app.py

import gradio as gr
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
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."

# --- 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:
    """Analyzes an issue description and extracts severity and programming language."""
    model = pipeline("text-generation", model=model_name)
    response = model(
        f"{SYSTEM_MESSAGE}\n{issue_text}\nAssistant: ",
        max_length=2048,  # Adjust as needed
        do_sample=True,
        temperature=0.7,  # Adjust as needed
        top_k=50,  # Adjust as needed
    )
    assistant_response = response[0]['generated_text'].strip()

    # 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()

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

def find_related_issues(issue_text: str, issues: list) -> list:
    """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]
    sorted_issues = sorted(enumerate(similarities), key=lambda x: x[1], reverse=True)
    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:
    """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)
    if response.status_code == 200:
        return response.json()
    else:
        return []

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,
    *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
    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()

    # 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'],
    }

# --- 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")

    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,
        )

    with gr.Row():
        severity_dropdown = gr.Dropdown(
            choices=["Critical", "Major", "Minor", "Trivial"],
            label="Severity",
            value=None,
        )
        programming_language_textbox = gr.Textbox(label="Programming Language")

    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=[
            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)