Update app.py

app.py

@@ -6,28 +6,27 @@ import threading
 import queue
 import gradio as gr
 import httpx
-from typing import Generator, Any, Dict, List
+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-mini", api_key: str = None) -> str:
+async def call_model(prompt: str, model: str = "gpt-4o", api_key: str = None) -> str:
     """
-    Sends a prompt to the OpenAI API endpoint using the specified model (overridden to gpt-4o-mini)
-    and returns the generated response.
+    Sends a prompt to the OpenAI API endpoint.
     """
-    # Use the provided API key or fall back to the environment variable
     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"
     }
-    # Override the model value to always be "gpt-4o-mini"
     payload = {
-        "model": "gpt-4o-mini",
+        "model": model,
         "messages": [{"role": "user", "content": prompt}],
     }
     async with httpx.AsyncClient(timeout=httpx.Timeout(300.0)) as client:
@@ -39,191 +38,214 @@ async def call_model(prompt: str, model: str = "gpt-4o-mini", api_key: str = Non
 # -------------------- Agent Classes --------------------
 class PromptOptimizerAgent:
     async def optimize_prompt(self, user_prompt: str, api_key: str) -> str:
-        """
-        Optimizes the user's initial prompt according to the following instructions:
-        >>> Given the user's initial prompt below the ### characters please enhance it. 
-        1. Start with clear, precise instructions placed at the beginning of the prompt. 
-        2. Include specific details about the desired context, outcome, length, format, and style. 
-        3. Provide examples of the desired output format, if possible. 
-        4. Use appropriate leading words or phrases to guide the desired output, especially if code generation is involved. 
-        5. Avoid any vague or imprecise language. 
-        6. Rather than only stating what not to do, provide guidance on what should be done instead. 
-        Remember to ensure the revised prompt remains true to the user's original intent. <<<
-        ###User initial prompt below ###
-        """
+        """Optimizes the user's initial prompt."""
         system_prompt = (
-            "Given the user's initial prompt below the ### characters please enhance it. "
-            "1. Start with clear, precise instructions placed at the beginning of the prompt. "
-            "2. Include specific details about the desired context, outcome, length, format, and style. "
-            "3. Provide examples of the desired output format, if possible. "
-            "4. Use appropriate leading words or phrases to guide the desired output, especially if code generation is involved. "
-            "5. Avoid any vague or imprecise language. "
-            "6. Rather than only stating what not to do, provide guidance on what should be done instead. "
-            "Remember to ensure the revised prompt remains true to the user's original intent. "
-            "###User initial 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{user_prompt}\n<<<"
-        optimized = await call_model(full_prompt, api_key=api_key)
+        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) -> None:
+    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) -> str:
+    async def generate_plan(self, task: str, api_key: str, human_feedback: Optional[str] = None) -> str:
         """
-        Generates a detailed, step-by-step plan for completing the given task.
+        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 an orchestrator agent. The user has provided the task: '{task}'.\n"
-            "Generate a detailed, step-by-step plan for completing this task by coordinating a coder agent, "
-            "a code reviewer agent, and a documentation agent. List the steps as bullet points."
+            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, api_key=api_key)
+        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) -> str:
-        """
-        Generates code based on the given instructions.
-        """
+    async def generate_code(self, instructions: str, api_key: str, model: str = "gpt-4o") -> str:
+        """Generates code based on instructions."""
         prompt = (
-            "You are a coder agent. Based on the following instructions, generate the requested code. "
-            "Only output the generated code, never any explanations or any other information besides the actual code!\n"
-            f"{instructions}\n"
+            "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, api_key=api_key)
+        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 the provided code to check if it meets the task specifications.
-        NEVER generate any code yourself! Respond only with feedback or with 'APPROVE' if everything is correct.
-        """
+        """Reviews code. Provides concise, actionable feedback or 'APPROVE'."""
         prompt = (
-            "You are a code reviewing agent highly skilled in evaluating code quality. "
-            "Review the provided code and check if it meets the task specifications and properly addresses any provided feedback. "
-            "NEVER generate any code yourself! Respond only with feedback or with 'APPROVE' if everything is correct. "
-            "Do not mention 'APPROVE' before actually approving! Do not request documentation or user guides:\n"
-            f"Task: {task}\n"
-            f"Code:\n{code}\n\n"
+            "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, api_key=api_key)
+        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 clear and concise documentation for the approved code,
-        including a brief and concise --help message.
-        """
+        """Generates documentation, including a --help message."""
         prompt = (
-            "You are a documentation agent. Generate a brief, clear and concise documentation for the following approved code. "
-            "Keep it short and compact, focusing on the main elements, do not include unnecessary extras that limit readability. "
-            "Additionally, generate a brief and concise --help message for the code:\n"
-            f"{code}\n"
-            "Briefly explain what the code does and how it works. Make sure to be clear and concise, do not include unnecessary extras that limit readability."
+            "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, api_key=api_key)
+        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) -> None:
+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 a multi-agent conversation where each agent's response is generated via the external model API.
-    The conversation is logged to the provided queue.
+    Conducts the multi-agent conversation.
     """
-    conversation: List[Dict[str, str]] = []  # List to store each agent's message
+    conversation: List[Dict[str, str]] = []
 
-    # Step 0: Use Prompt Optimizer to enhance the user's initial prompt.
-    log_queue.put("[Prompt Optimizer]: Received initial task. Optimizing prompt...")
+    # 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: Orchestrator generates a plan based on the optimized task.
-    log_queue.put("[Orchestrator]: Received optimized task. Generating plan...")
-    orchestrator = OrchestratorAgent(log_queue)
+    # 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: Coder generates code based on the plan.
+    # Step 2: Generate Code
     coder = CoderAgent()
-    coder_instructions = f"Implement the task as described in the following plan:\n{plan}"
-    log_queue.put("[Coder]: Received coding task from the Orchestrator.")
+    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 Reviewer reviews the generated code.
+    # Step 3: Code Review and Revision
     reviewer = CodeReviewerAgent()
+    tester = QualityAssuranceTesterAgent()
     approval_keyword = "approve"
     revision_iteration = 0
     while True:
-        if revision_iteration == 0:
-            log_queue.put("[Code Reviewer]: Starting review of the generated code...")
-        else:
-            log_queue.put(f"[Code Reviewer]: Reviewing the revised code (Iteration {revision_iteration})...")
-
+        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 feedback (Iteration {revision_iteration}):\n{review}")
+        log_queue.put(f"[Code Reviewer]: Review (Iteration {revision_iteration}):\n{review}")
 
-        # Check if the code has been approved
         if approval_keyword in review.lower():
             log_queue.put("[Code Reviewer]: Code approved.")
-            break  # Exit the loop if approved
+            break
 
-        # If not approved, increment the revision count.
         revision_iteration += 1
-
-        # Kill-switch: After 5 generations without approval, shut down.
         if revision_iteration >= 5:
-            log_queue.put("Unable to solve your task to full satisfaction :(")
-            sys.exit("Unable to solve your task to full satisfaction :(")
+            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}")
 
-        # If under the limit, instruct the coder to revise the code.
-        log_queue.put(f"[Orchestrator]: Code not approved. Instructing coder to revise the code (Iteration {revision_iteration}).")
-        update_instructions = f"Please revise the code according to the following feedback. Feedback: {review}"
-        revised_code = await coder.generate_code(update_instructions, api_key=api_key)
+        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 code submitted (Iteration {revision_iteration}):\n{revised_code}")
-        code = revised_code  # Update the code for the next review iteration
+        log_queue.put(f"[Coder]: Revised (Iteration {revision_iteration}):\n{revised_code}")
+        code = revised_code
 
-    # Step 4: Documentation Agent generates documentation for the approved code.
+    # Step 4: Generate Documentation
     doc_agent = DocumentationAgent()
-    log_queue.put("[Documentation Agent]: Generating documentation for the approved code.")
+    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("Multi-agent conversation complete.")
+    log_queue.put("Conversation complete.")
     log_queue.put(("result", conversation))
 
-# -------------------- Process Generator for Streaming --------------------
-def process_conversation_generator(task_message: str, api_key: str) -> Generator[str, None, None]:
+# -------------------- 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 asynchronous multi-agent conversation and yields log messages as they are generated.
+    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))
+        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
-    # Yield log messages as long as the thread is running or the queue is not empty.
     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 "Final conversation complete."
+                yield "Conversation complete."
             else:
                 yield msg
         except queue.Empty:
@@ -231,39 +253,80 @@ def process_conversation_generator(task_message: str, api_key: str) -> Generator
 
     thread.join()
     if final_result:
-        # Format the final conversation log.
-        conv_text = "\n========== Multi-Agent Conversation ==========\n"
+        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.
-    The user's message is interpreted as the task description.
-    An optional OpenAI API key can be provided via the additional input; if not provided, the environment variable is used.
-    This function streams the multi-agent conversation log messages.
-    """
+    """Chat function for Gradio."""
     if not openai_api_key:
         openai_api_key = os.getenv("OPENAI_API_KEY")
-    yield from process_conversation_generator(message, 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 --------------------
-# Use Gradio's ChatInterface with an additional input field for the OpenAI API key.
+
+# 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")],
-    type="messages",
-    title="Actual Multi-Agent Conversation Chatbot",
+    title="Multi-Agent Task Solver with Human-in-the-Loop",
     description="""
-        - Collaborative workflow between Prompt Enhancer, Orchestrator, Coder, Code-Reviewer and Documentation Generator agents.
-        - Enter a task and observe as your prompt gets magically solved! :)
-        - NOTE: The full conversation log will be displayed at the end, showing all the steps taken!
-        - NOTE2: If the Coder is unable to satisfactorily complete the task after five attempts, the script will terminate to prevent endless iterations.
-        - NOTE3: You will have to input your OPENAI_API_KEY at the bottom of the page for this to work!
+        - 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__":
-    iface.launch()
+    demo = gr.TabbedInterface([iface, dummy_iface], ["Chatbot", "Dummy"])
+    demo.launch()