app.py

import os
import json
import subprocess
import re
import requests
from datetime import datetime

import gradio as gr
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, TextGenerationPipeline, AutoModel, RagRetriever, AutoModelForSeq2SeqLM
import torch
import tree_sitter
from tree_sitter import Language, Parser
import black
from pylint import lint
from io import StringIO
import sys
from huggingface_hub import Repository, hf_hub_url, HfApi, snapshot_download
import tempfile
import logging
from loguru import logger
logger.add("app.log", format="{time} {level} {message}", level="INFO")

# Constants
MODEL_NAME = "bigscience/bloom"
PROJECT_ROOT = "projects"
AGENT_DIRECTORY = "agents"
AVAILABLE_CODE_GENERATIVE_MODELS = [
    "bigcode/starcoder",
    "Salesforce/codegen-350M-mono",
    "microsoft/CodeGPT-small-py",
    "NinedayWang/PolyCoder-2.7B",
    "facebook/incoder-1B",
]

# Load Models and Resources
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16)
pipe = TextGenerationPipeline(model=model, tokenizer=tokenizer)

# Build Tree-sitter parser libraries (if not already built)
Language.build_library("build/my-languages.so", ["tree-sitter-python", "tree-sitter-javascript"])
PYTHON_LANGUAGE = Language("build/my-languages.so", "python")
JAVASCRIPT_LANGUAGE = Language("build/my-languages.so", "javascript")
parser = Parser()

# Session State Initialization
if 'chat_history' not in gr.State.session_state:
    gr.State.chat_history = []
if 'terminal_history' not in gr.State.session_state:
    gr.State.terminal_history = []
if 'workspace_projects' not in gr.State.session_state:
    gr.State.workspace_projects = {}
if 'available_agents' not in gr.State.session_state:
    gr.State.available_agents = []
if 'current_state' not in gr.State.session_state:
    gr.State.current_state = {
        'toolbox': {},
        'workspace_chat': {}
    }

# Define is_code function
def is_code(message):
    return message.lstrip().startswith("```") or message.lstrip().startswith("code:")

# Define agents variable
agents = ["python", "javascript", "java"]

# Define load_agent_from_file function
def load_agent_from_file(agent_name):
    try:
        with open(os.path.join(AGENT_DIRECTORY, agent_name + ".json"), "r") as f:
            return json.load(f)
    except FileNotFoundError:
        return None

# Define load_pipeline function
def load_pipeline(model_category, model_name):
    return available_models[model_category][model_name]

# Define execute_translation function
def execute_translation(code, target_language, pipe):
    try:
        output = pipe(code, max_length=1000)[0]["generated_text"]
        return output
    except Exception as e:
        logger.error(f"Error in execute_translation function: {e}")
        return "Error: Unable to translate code."

# Refactor using CodeT5+
def execute_refactoring_codet5(code: str) -> str:
    """
    Refactors the provided code using the CodeT5+ model.

    Args:
        code (str): The code to refactor.

    Returns:
        str: The refactored code.
    """
    try:
        refactor_pipe = pipeline(
            "text2text-generation",
            model="Salesforce/codet5p-220m-finetune-Refactor"
        )
        prompt = f"Refactor this Python code:\n{code}"
        output = refactor_pipe(prompt, max_length=1000)[0]["generated_text"]
        return output
    except Exception as e:
        logger.error(f"Error in execute_refactoring_codet5 function: {e}")
        return "Error: Unable to refactor code."

# Chat interface with agent
def chat_interface_with_agent(input_text, agent_name, selected_model):
    """
    Handles interaction with the selected AI agent.
    """
    agent = load_agent_from_file(agent_name)
    if not agent:
        return f"Agent {agent_name} not found."

    agent.pipeline = available_models[selected_model]
    agent_prompt = agent.create_agent_prompt()
    full_prompt = f"{agent_prompt}\n\nUser: {input_text}\nAgent:"

    try:
        response = agent.generate_response(full_prompt)
    except Exception as e:
        logger.error(f"Error generating agent response: {e}")
        response = "Error: Unable to process your request."

    return response

# Available models
available_models = {
    "Code Generation & Completion": {
        "Salesforce CodeGen-350M (Mono)": pipeline("text-generation", model="Salesforce/codegen-350M-mono"),
        "BigCode StarCoder": pipeline("text-generation", model="bigcode/starcoder"),
        "CodeGPT-small-py": pipeline("text-generation", model="microsoft/CodeGPT-small-py"),
        "PolyCoder-2.7B": pipeline("text-generation", model="NinedayWang/PolyCoder-2.7B"),
        "InCoder-1B": pipeline("text-generation", model="facebook/incoder-1B"),
    },
    "Code Translation": {
        "Python to JavaScript": (lambda code, pipe=pipeline("translation", model="transformersbook/codeparrot-translation-en-java"): execute_translation(code, "javascript", pipe), []),
        "Python to C++": (lambda code, pipe=pipeline("text-generation", model="konodyuk/codeparrot-small-trans-py-cpp"): execute_translation(code, "cpp", pipe), []),
    },
    # ... other categories
}

# Gradio interface with tabs
with gr.Blocks(title="AI Power Tools for Developers") as demo:
    # --- State ---
    code = gr.State("")  # Use gr.State to store code across tabs
    task_dropdown = gr.State(list(available_models.keys())[0])  # Initialize task dropdown
    model_dropdown = gr.State(
        list(available_models[task_dropdown.value].keys())[0]
    )  # Initialize model dropdown

    def update_model_dropdown(selected_task):
        models_for_task = list(available_models[selected_task].keys())
        return gr.Dropdown.update(choices=models_for_task)

    with gr.Tab("Chat & Code"):
        chatbot = gr.Chatbot(elem_id="chatbot")
        msg = gr.Textbox(label="Enter your message", placeholder="Type your message here...")
        clear = gr.ClearButton([msg, chatbot])

        def user(message, history):
            if is_code(message):
                response = ""  # Initialize response
                task = message.split()[0].lower()  # Extract task keyword

                # Use the selected model or a default one
                model_category = task_dropdown.value
                model_name = model_dropdown.value
                pipeline = load_pipeline(model_category, model_name)

                if task in agents:
                    agent = load_agent_from_file(task)
                    try:
                        response = agent.generate_response(message)
                    except Exception as e:
                        logger.error(f"Error executing agent {task}: {e}")
                        response = f"Error executing agent {task}: {e}"
                else:
                    response = "Invalid command or task not found."
            else:
                # Process as natural language request
                response = pipe(message, max_length=1000)[0]["generated_text"]

            return response, history + [(message, response)]

        msg.change(user, inputs=[msg, chatbot], outputs=[chatbot, chatbot])
        clear.click(lambda: None, None, chatbot, queue=False)

    # Model Selection Tab
    with gr.Tab("Model Selection"):
        task_dropdown.render()
        model_dropdown.render()
        task_dropdown.change(update_model_dropdown, task_dropdown, model_dropdown)

    # Workspace Tab
    with gr.Tab("Workspace"):
        with gr.Row():
            with gr.Column():
                code.render()
                file_output = gr.File(label="Save File As...", interactive=False)
            with gr.Column():
                output = gr.Textbox(label="Output")

        run_btn = gr.Button(value="Run Code")
        upload_btn = gr.UploadButton("Upload Python File", file_types=[".py"])
        save_button = gr.Button(value="Save Code")

        def run_code(code_str):
            try:
                # Save code to a temporary file
                with open("temp_code.py", "w") as f:
                    f.write(code_str)

                # Execute the code using subprocess
                process = subprocess.Popen(["python", "temp_code.py"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                output, error = process.communicate()

                # Return the output and error messages
                if error:
                    return "Error: " + error.decode("utf-8")
                else:
                    return output.decode("utf-8")

            except Exception as e:
                logger.error(f"Error running code: {e}")
                return f"Error running code: {e}"

        def upload_file(file):
            with open("uploaded_code.py", "wb") as f:
                f.write(file.file.getvalue())
            return "File uploaded successfully!"

        def save_code(code_str):
            file_output.value = code_str
            return file_output

        run_btn.click(run_code, inputs=[code], outputs=[output])
        upload_btn.click(upload_file, inputs=[upload_btn], outputs=[output])
        save_button.click(save_code, inputs=[code], outputs=[file_output])

    demo.launch()