Spaces:
Paused
Paused
| 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()) |