import os import sys import asyncio import logging import threading import queue import gradio as gr import httpx from typing import Generator, Any, Dict, List, Optional # -------------------- Configuration -------------------- logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") # -------------------- External Model Call -------------------- async def call_model(prompt: str, model: str = "gpt-4o", api_key: str = None) -> str: """ Sends a prompt to the OpenAI API endpoint. """ if api_key is None: api_key = os.getenv("OPENAI_API_KEY") if api_key is None: raise ValueError("OpenAI API key not found.") url = "https://api.openai.com/v1/chat/completions" headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json" } payload = { "model": model, "messages": [{"role": "user", "content": prompt}], } async with httpx.AsyncClient(timeout=httpx.Timeout(300.0)) as client: response = await client.post(url, headers=headers, json=payload) response.raise_for_status() response_json = response.json() return response_json["choices"][0]["message"]["content"] # -------------------- Agent Classes -------------------- class PromptOptimizerAgent: async def optimize_prompt(self, user_prompt: str, api_key: str) -> str: """Optimizes the user's initial prompt.""" system_prompt = ( "You are a prompt optimization expert. Improve the given user prompt. " "Be clear, specific, and complete. Maintain the user's original intent." "Return ONLY the revised prompt." ) full_prompt = f"{system_prompt}\n\nUser's initial prompt:\n{user_prompt}" optimized = await call_model(full_prompt, model="gpt-4o", api_key=api_key) return optimized class OrchestratorAgent: def __init__(self, log_queue: queue.Queue, human_in_the_loop_event: threading.Event, human_input_queue: queue.Queue) -> None: self.log_queue = log_queue self.human_in_the_loop_event = human_in_the_loop_event self.human_input_queue = human_input_queue async def generate_plan(self, task: str, api_key: str, human_feedback: Optional[str] = None) -> str: """ Generates a plan, potentially requesting human feedback. """ if human_feedback: # Use human feedback if provided prompt = ( f"You are a master planner. You previously generated a partial plan for the task: '{task}'.\n" "You requested human feedback, and here's the feedback you received:\n" f"{human_feedback}\n\n" "Now, complete or revise the plan, incorporating the human feedback. " "Output the plan as a numbered list." ) plan = await call_model(prompt, model="gpt-4o", api_key=api_key) return plan prompt = ( f"You are a master planner. Given the task: '{task}', create a detailed, step-by-step plan. " "Break down the task into sub-tasks. Assign each sub-task to agents: Coder, Code Reviewer, Quality Assurance Tester, and Documentation Agent. " "Include steps for review and revision. Consider potential issues and error handling. " "Include instructions for documentation.\n\n" "HOWEVER, if at ANY point you are unsure how to proceed, you can request human feedback. " "To do this, output ONLY the following phrase (and nothing else): 'REQUEST_HUMAN_FEEDBACK'\n" "Followed by a newline and a clear and concise question for the human. Example:\n\nREQUEST_HUMAN_FEEDBACK\nShould the output be in JSON or XML format?" "\n\nOutput the plan as a numbered list (or as much as you can before requesting feedback)." ) plan = await call_model(prompt, model="gpt-4o", api_key=api_key) if "REQUEST_HUMAN_FEEDBACK" in plan: self.log_queue.put("[Orchestrator]: Requesting human feedback...") question = plan.split("REQUEST_HUMAN_FEEDBACK\n", 1)[1].strip() self.log_queue.put(f"[Orchestrator]: Question for human: {question}") self.human_in_the_loop_event.set() # Signal the human input thread human_response = self.human_input_queue.get() # Wait for human input self.human_in_the_loop_event.clear() # Reset the event self.log_queue.put(f"[Orchestrator]: Received human feedback: {human_response}") return await self.generate_plan(task, api_key, human_response) # Recursive call with feedback return plan class CoderAgent: async def generate_code(self, instructions: str, api_key: str, model: str = "gpt-4o") -> str: """Generates code based on instructions.""" prompt = ( "You are a highly skilled coding agent. Output ONLY the code. " "Adhere to best practices. Include error handling.\n\n" f"Instructions:\n{instructions}" ) code = await call_model(prompt, model=model, api_key=api_key) return code class CodeReviewerAgent: async def review_code(self, code: str, task: str, api_key: str) -> str: """Reviews code. Provides concise, actionable feedback or 'APPROVE'.""" prompt = ( "You are a meticulous code reviewer. Provide CONCISE feedback. " "Focus on correctness, efficiency, readability, error handling, security, and adherence to the task. " "Suggest improvements. If acceptable, respond with ONLY 'APPROVE'. " "Do NOT generate code.\n\n" f"Task: {task}\n\nCode:\n{code}" ) review = await call_model(prompt, model="gpt-4o", api_key=api_key) return review class QualityAssuranceTesterAgent: async def generate_test_cases(self, code: str, task: str, api_key: str) -> str: """Generates test cases.""" prompt = ( "You are a quality assurance testing agent. Generate test cases. " "Consider edge cases and error scenarios. Output in a clear format.\n\n" f"Task: {task}\n\nCode:\n{code}" ) test_cases = await call_model(prompt, model="gpt-4o", api_key=api_key) return test_cases async def run_tests(self, code:str, test_cases:str, api_key:str) -> str: """Runs tests and reports results.""" prompt = ( "Run the generated test cases. Compare actual vs expected output. " "State discrepancies. If all pass, output 'TESTS PASSED'.\n\n" f"Code:\n{code}\n\nTest Cases:\n{test_cases}" ) test_results = await call_model(prompt, model="gpt-4o", api_key=api_key) return test_results class DocumentationAgent: async def generate_documentation(self, code: str, api_key: str) -> str: """Generates documentation, including a --help message.""" prompt = ( "Generate clear and concise documentation. " "Include a brief description, explanation, and a --help message.\n\n" f"Code:\n{code}" ) documentation = await call_model(prompt, model="gpt-4o", api_key=api_key) return documentation # -------------------- Multi-Agent Conversation -------------------- async def multi_agent_conversation(task_message: str, log_queue: queue.Queue, api_key: str, human_in_the_loop_event: threading.Event, human_input_queue: queue.Queue) -> None: """ Conducts the multi-agent conversation. """ conversation: List[Dict[str, str]] = [] # Step 0: Optimize Prompt log_queue.put("[Prompt Optimizer]: Optimizing prompt...") prompt_optimizer = PromptOptimizerAgent() optimized_task = await prompt_optimizer.optimize_prompt(task_message, api_key=api_key) conversation.append({"agent": "Prompt Optimizer", "message": f"Optimized Task:\n{optimized_task}"}) log_queue.put(f"[Prompt Optimizer]: Optimized task prompt:\n{optimized_task}") # Step 1: Generate Plan log_queue.put("[Orchestrator]: Generating plan...") orchestrator = OrchestratorAgent(log_queue, human_in_the_loop_event, human_input_queue) plan = await orchestrator.generate_plan(optimized_task, api_key=api_key) conversation.append({"agent": "Orchestrator", "message": f"Plan:\n{plan}"}) log_queue.put(f"[Orchestrator]: Plan generated:\n{plan}") # Step 2: Generate Code coder = CoderAgent() coder_instructions = f"Implement the task:\n{plan}" log_queue.put("[Coder]: Generating code...") code = await coder.generate_code(coder_instructions, api_key=api_key) conversation.append({"agent": "Coder", "message": f"Code:\n{code}"}) log_queue.put(f"[Coder]: Code generated:\n{code}") # Step 3: Code Review and Revision reviewer = CodeReviewerAgent() tester = QualityAssuranceTesterAgent() approval_keyword = "approve" revision_iteration = 0 while True: log_queue.put(f"[Code Reviewer]: Reviewing code (Iteration {revision_iteration})...") review = await reviewer.review_code(code, optimized_task, api_key=api_key) conversation.append({"agent": "Code Reviewer", "message": f"Review (Iteration {revision_iteration}):\n{review}"}) log_queue.put(f"[Code Reviewer]: Review (Iteration {revision_iteration}):\n{review}") if approval_keyword in review.lower(): log_queue.put("[Code Reviewer]: Code approved.") break revision_iteration += 1 if revision_iteration >= 5: log_queue.put("Unable to solve task satisfactorily.") sys.exit("Unable to solve task satisfactorily.") log_queue.put("[QA Tester]: Generating test cases...") test_cases = await tester.generate_test_cases(code, optimized_task, api_key=api_key) conversation.append({"agent": "QA Tester", "message": f"Test Cases:\n{test_cases}"}) log_queue.put(f"[QA Tester]: Test Cases:\n{test_cases}") log_queue.put("[QA Tester]: Running tests...") test_results = await tester.run_tests(code, test_cases, api_key) conversation.append({"agent": "QA Tester", "message": f"Test Results:\n{test_results}"}) log_queue.put(f"[QA Tester]: Test Results:\n{test_results}") log_queue.put(f"[Orchestrator]: Revising code (Iteration {revision_iteration})...") update_instructions = f"Revise:\nReview:\n{review}\nTests:\n{test_results}\nPlan:\n{plan}" revised_code = await coder.generate_code(update_instructions, api_key=api_key, model="gpt-3.5-turbo-16k") conversation.append({"agent": "Coder", "message": f"Revised Code (Iteration {revision_iteration}):\n{revised_code}"}) log_queue.put(f"[Coder]: Revised (Iteration {revision_iteration}):\n{revised_code}") code = revised_code # Step 4: Generate Documentation doc_agent = DocumentationAgent() log_queue.put("[Documentation Agent]: Generating documentation...") documentation = await doc_agent.generate_documentation(code, api_key=api_key) conversation.append({"agent": "Documentation Agent", "message": f"Documentation:\n{documentation}"}) log_queue.put(f"[Documentation Agent]: Documentation generated:\n{documentation}") log_queue.put("Conversation complete.") log_queue.put(("result", conversation)) # -------------------- Process Generator and Human Input -------------------- def process_conversation_generator(task_message: str, api_key: str, human_in_the_loop_event: threading.Event, human_input_queue: queue.Queue) -> Generator[str, None, None]: """ Wraps the conversation and yields log messages. Handles human input. """ log_q: queue.Queue = queue.Queue() def run_conversation() -> None: asyncio.run(multi_agent_conversation(task_message, log_q, api_key, human_in_the_loop_event, human_input_queue)) thread = threading.Thread(target=run_conversation) thread.start() final_result = None while thread.is_alive() or not log_q.empty(): try: msg = log_q.get(timeout=0.1) if isinstance(msg, tuple) and msg[0] == "result": final_result = msg[1] yield "Conversation complete." else: yield msg except queue.Empty: continue thread.join() if final_result: conv_text = "\n=== Conversation ===\n" for entry in final_result: conv_text += f"[{entry['agent']}]: {entry['message']}\n\n" yield conv_text def get_human_feedback(placeholder_text): """Gets human input using a Gradio Textbox.""" with gr.Blocks() as human_feedback_interface: with gr.Row(): human_input = gr.Textbox(lines=4, placeholder=placeholder_text, label="Human Feedback") with gr.Row(): submit_button = gr.Button("Submit Feedback") feedback_queue = queue.Queue() def submit_feedback(input_text): feedback_queue.put(input_text) return "" submit_button.click(submit_feedback, inputs=human_input, outputs=human_input) human_feedback_interface.load(None, [], []) # This is needed to keep the interface alive return human_feedback_interface, feedback_queue # -------------------- Chat Function for Gradio -------------------- def multi_agent_chat(message: str, history: List[Any], openai_api_key: str = None) -> Generator[str, None, None]: """Chat function for Gradio.""" if not openai_api_key: openai_api_key = os.getenv("OPENAI_API_KEY") if not openai_api_key: yield "Error: API key not provided." return human_in_the_loop_event = threading.Event() human_input_queue = queue.Queue() yield from process_conversation_generator(message, openai_api_key, human_in_the_loop_event, human_input_queue) while human_in_the_loop_event.is_set(): yield "Waiting for human feedback..." placeholder = "Please provide your feedback." human_interface, feedback_queue = get_human_feedback(placeholder) #This is a hacky but currently only working way to make this work with gradio yield gr.Textbox.update(visible=False), gr.update(visible=True) try: human_feedback = feedback_queue.get(timeout=300) # Wait for up to 5 minutes human_input_queue.put(human_feedback) human_in_the_loop_event.clear() yield gr.Textbox.update(visible=True), human_interface.close() yield from process_conversation_generator(message, openai_api_key, human_in_the_loop_event, human_input_queue) except queue.Empty: human_input_queue.put("No feedback provided.") #Timeout human_in_the_loop_event.clear() yield gr.Textbox.update(visible=True), human_interface.close() yield from process_conversation_generator(message, openai_api_key, human_in_the_loop_event, human_input_queue) # -------------------- Launch the Chatbot -------------------- # Create the main chat interface iface = gr.ChatInterface( fn=multi_agent_chat, additional_inputs=[gr.Textbox(label="OpenAI API Key (optional)", type="password", placeholder="Leave blank to use env variable")], title="Multi-Agent Task Solver with Human-in-the-Loop", description=""" - Collaborative workflow with Human-in-the-Loop capability. - The Orchestrator can ask for human feedback if needed. - Enter a task, and the agents will work on it. You may be prompted for input. - Max 5 revision iterations. - Provide your OpenAI API Key below. """ ) #Need a dummy interface to make the human feedback interface update dummy_iface = gr.Interface(lambda x:x, "textbox", "textbox") if __name__ == "__main__": demo = gr.TabbedInterface([iface, dummy_iface], ["Chatbot", "Dummy"]) demo.launch()