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"), "Mixtral-8x7B-v0.1": pipeline("text-generation", model="mistralai/Mixtral-8x7B-v0.1"), "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"), #... more code generation models }, "Code Translation": { "Python to JavaScript": (lambda code, pipe=pipeline("translation", model="transformersbook/codeparrot-translation-en-java"): execute_translation(code, "javascript", pipe), []), # Pipeline for Python to JavaScript "Python to C++": (lambda code, pipe=pipeline("text-generation", model="konodyuk/codeparrot-small-trans-py-cpp"): execute_translation(code, "cpp", pipe), []), # Pipeline for Python to C++ #... more language pairs }, #... 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()