import os import json import asyncio import gradio as gr import google.generativeai as genai from groq import Groq from mcp import ClientSession, StdioServerParameters from mcp.client.stdio import stdio_client # ---------- Configuration ---------- GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") # free from https://aistudio.google.com GROQ_API_KEY = os.getenv("GROQ_API_KEY") # free from https://console.groq.com genai.configure(api_key=GEMINI_API_KEY) groq_client = Groq(api_key=GROQ_API_KEY) MAX_RETRIES = 3 # ---------- MCP Client Setup ---------- python_server_params = StdioServerParameters( command="python", args=["mcp_python_server.py"], env=None, ) nmap_server_params = StdioServerParameters( command="python", args=["mcp_nmap_server.py"], env=None, ) async def get_mcp_sessions(): """Create persistent MCP sessions for both servers.""" # Python server python_transport = await stdio_client(python_server_params).__aenter__() python_session = await ClientSession(python_transport[0], python_transport[1]).__aenter__() await python_session.initialize() # Nmap server nmap_transport = await stdio_client(nmap_server_params).__aenter__() nmap_session = await ClientSession(nmap_transport[0], nmap_transport[1]).__aenter__() await nmap_session.initialize() return python_session, nmap_session # We'll start the sessions at module level (but must be inside async main of Gradio) # We'll use a global variable set in the main demo function. PYTHON_SESSION = None NMAP_SESSION = None # ---------- Agent Logic ---------- def generate_code_with_gemini(prompt: str, feedback: str = "") -> str: """Generate code using Gemini. If feedback is provided, it's a correction request.""" model = genai.GenerativeModel('gemini-1.5-pro') if feedback: full_prompt = ( f"User request: {prompt}\n\n" f"Previous code was incorrect. Reviewer feedback: {feedback}\n" "Please fix the code and return only the corrected Python code." ) else: full_prompt = ( f"Write Python code to solve the following task. Return only the code, " f"no explanation.\n\nTask: {prompt}" ) response = model.generate_content(full_prompt) # extract code block text = response.text if "```" in text: code = text.split("```")[1] if code.startswith("python"): code = code[6:] return code.strip() return text.strip() async def review_with_groq(code: str, task: str) -> dict: """ Use Groq to review the code, optionally executing it via MCP Python tool. Returns dict with 'pass' (bool) and 'feedback' (str). """ # We'll use Groq's function calling to let the model request code execution. tools = [ { "type": "function", "function": { "name": "execute_code", "description": "Run Python code and return output and any images.", "parameters": { "type": "object", "properties": { "code": {"type": "string", "description": "Python code to run"} }, "required": ["code"] } } } ] messages = [ { "role": "system", "content": ( "You are a strict code reviewer. You can test the code by running it using " "the `execute_code` tool. After testing, output a JSON object with exactly two keys: " "'pass' (boolean) and 'feedback' (string). If the code runs correctly and solves the " "task, set pass=true. Otherwise, set pass=false and give constructive feedback. " "Always run the code before deciding." ) }, { "role": "user", "content": f"TASK: {task}\n\nCODE:\n```python\n{code}\n```\nPlease review." } ] # First call – may return tool call request response = groq_client.chat.completions.create( model="llama-3.1-70b-versatile", messages=messages, tools=tools, tool_choice="auto", max_tokens=1024, temperature=0.1 ) assistant_msg = response.choices[0].message # If tool call requested while assistant_msg.tool_calls: tool_call = assistant_msg.tool_calls[0] if tool_call.function.name == "execute_code": args = json.loads(tool_call.function.arguments) code_to_run = args["code"] # Use MCP Python session to run code result = await PYTHON_SESSION.call_tool("execute_code", {"code": code_to_run}) # Extract text and images from result output_text = "" images = [] for item in result.content: if item.type == "text": output_text += item.text + "\n" elif item.type == "image": images.append(item.data) # base64 # Append assistant tool call and result to messages messages.append(assistant_msg) messages.append({ "role": "tool", "tool_call_id": tool_call.id, "content": output_text + (f"\n[Image generated]" if images else "") }) # Continue with model response = groq_client.chat.completions.create( model="llama-3.1-70b-versatile", messages=messages, tools=tools, max_tokens=512, temperature=0.1 ) assistant_msg = response.choices[0].message else: break # Parse final answer (expected JSON) final_text = assistant_msg.content try: # Sometimes the JSON is wrapped in markdown if "```" in final_text: final_text = final_text.split("```")[1].split("```")[0] if final_text.startswith("json"): final_text = final_text[4:] result = json.loads(final_text) return result except json.JSONDecodeError: # fallback: try to extract pass/fail if "true" in final_text.lower(): return {"pass": True, "feedback": "Code seems correct."} else: return {"pass": False, "feedback": final_text} async def agent_loop(task: str, progress=gr.Progress()): progress(0, desc="Generating code with Gemini...") code = generate_code_with_gemini(task) feedback = "" for i in range(MAX_RETRIES): progress((i+1)/MAX_RETRIES, desc=f"Reviewing iteration {i+1}...") review = await review_with_groq(code, task) if review.get("pass"): progress(1.0, desc="Approved!") return code, review.get("feedback", ""), i+1 feedback = review.get("feedback", "Unknown error") progress(0.6, desc="Revising code...") code = generate_code_with_gemini(task, feedback) progress(1.0, desc="Max retries reached.") return code, feedback, MAX_RETRIES # ---------- Gradio UI ---------- custom_css = """ .gradio-container { max-width: 900px !important; margin: auto; } .chatbot { height: 600px; } """ def format_chat_history(user_message, response_code, review_feedback, retries): return [ {"role": "user", "content": user_message}, {"role": "assistant", "content": f"**Generated Code (attempts: {retries}):**\n```python\n{response_code}\n```\n\n**Review:** {review_feedback}"} ] async def respond(message, history): history.append({"role": "user", "content": message}) # Simple progress tracking not directly in chat function; we'll yield # We'll call agent_loop and stream steps? # For simplicity, we'll run synchronously and update after. code, feedback, retries = await agent_loop(message) assistant_msg = f"**Generated Code** (after {retries} revisions):\n```python\n{code}\n```\n\n**Review:** {feedback}" history.append({"role": "assistant", "content": assistant_msg}) return history def launch_ui(python_session, nmap_session): global PYTHON_SESSION, NMAP_SESSION PYTHON_SESSION = python_session NMAP_SESSION = nmap_session with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo: gr.Markdown("# 🧠 AI Coding Agent with Gemini + Groq + MCP Tools") gr.Markdown("Ask for Python code, and I'll generate, test (with real execution), and review it. Supports matplotlib and nmap security scans via MCP tools.") chatbot = gr.Chatbot(type="messages", height=600, bubble_full_width=False) msg = gr.Textbox(placeholder="Enter your coding task...", label="Your request") clear = gr.Button("Clear") async def user_message(message, history): history = history or [] history.append({"role": "user", "content": message}) yield history, "" # clear input code, feedback, retries = await agent_loop(message) assistant_msg = f"**Generated Code** (attempts: {retries}):\n```python\n{code}\n```\n\n**Review:** {feedback}" history.append({"role": "assistant", "content": assistant_msg}) yield history, "" msg.submit(user_message, [msg, chatbot], [chatbot, msg]) clear.click(lambda: None, None, chatbot, queue=False) return demo async def main(): python_sess, nmap_sess = await get_mcp_sessions() demo = launch_ui(python_sess, nmap_sess) demo.queue() demo.launch(server_name="0.0.0.0", server_port=7860) if __name__ == "__main__": asyncio.run(main())