Update app.py

app.py

@@ -1,21 +1,23 @@
 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, Callable
+from typing import Generator, Any, Dict, List, Optional
 from functools import lru_cache
 
 # -------------------- Configuration --------------------
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
 
 # -------------------- External Model Call (with Caching and Retry) --------------------
-# Removed @lru_cache here, as it caused issues with async and Gradio
 async def call_model(prompt: str, model: str = "gpt-4o", api_key: str = None, max_retries: int = 3) -> str:
-    """Sends a prompt to the OpenAI API endpoint, with caching and retries."""
+    """
+    Sends a prompt to the OpenAI API endpoint with retries and exponential backoff.
+    """
     if api_key is None:
         api_key = os.getenv("OPENAI_API_KEY")
         if api_key is None:
@@ -23,37 +25,33 @@ async def call_model(prompt: str, model: str = "gpt-4o", api_key: str = None, ma
     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}],
+        "Content-Type": "application/json",
     }
+    payload = {"model": model, "messages": [{"role": "user", "content": prompt}]}
 
     for attempt in range(max_retries):
         try:
             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()
+                response_json = response.json()  # Synchronous parsing is acceptable here
                 return response_json["choices"][0]["message"]["content"]
         except httpx.HTTPStatusError as e:
             logging.error(f"HTTP error (attempt {attempt + 1}/{max_retries}): {e}")
-            if e.response.status_code in (502, 503, 504):  # Retry on 502, 503, 504
-                await asyncio.sleep(2 ** attempt)  # Exponential backoff
+            if e.response.status_code in (502, 503, 504):
+                await asyncio.sleep(2 ** attempt)
                 continue
             else:
-                raise  # Re-raise for other HTTP errors
+                raise
         except httpx.RequestError as e:
             logging.error(f"Request error (attempt {attempt + 1}/{max_retries}): {e}")
             await asyncio.sleep(2 ** attempt)
             continue
         except Exception as e:
-            logging.error(f"An unexpected error occurred (attempt {attempt+1}/{max_retries}): {e}")
+            logging.error(f"Unexpected error (attempt {attempt+1}/{max_retries}): {e}")
             raise
     raise Exception(f"Failed to get response from OpenAI API after {max_retries} attempts.")
 
-
 # -------------------- Shared Context --------------------
 class Context:
     def __init__(self, original_task: str, optimized_task: Optional[str] = None,
@@ -78,8 +76,13 @@ class Context:
 
 class PromptOptimizerAgent:
     async def optimize_prompt(self, context: Context, api_key: str) -> Context:
-        """Optimizes the user's initial prompt."""
-        system_prompt = "Improve the prompt. Be clear, specific, and complete. Keep original intent. Return ONLY the revised prompt."
+        """
+        Optimizes the user’s original prompt.
+        """
+        system_prompt = (
+            "Improve the prompt. Be clear, specific, and complete. "
+            "Keep original intent. Return ONLY the revised prompt."
+        )
         full_prompt = f"{system_prompt}\n\nUser's prompt:\n{context.original_task}"
         optimized = await call_model(full_prompt, model="gpt-4o", api_key=api_key)
         context.optimized_task = optimized
@@ -87,65 +90,69 @@ class PromptOptimizerAgent:
         return context
 
 class OrchestratorAgent:
-    def __init__(self, log_queue: queue.Queue, human_in_the_loop_event: threading.Event, human_input_queue: queue.Queue) -> None:
+    def __init__(self, log_queue: queue.Queue, human_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_event = human_event
         self.human_input_queue = human_input_queue
 
-    async def generate_plan(self, context: Context, api_key: str, human_feedback: Optional[str] = None) -> Context:
-      """Generates a plan, potentially requesting human feedback."""
-
-      if human_feedback:
-        prompt = (
-            f"You are a planner. Revise/complete the plan for '{context.original_task}' using feedback:\n"
-            f"{human_feedback}\n\nCurrent Plan:\n{context.plan if context.plan else 'No plan yet.'}\n\n"
-            "Output the plan as a numbered list. If unsure, output 'REQUEST_HUMAN_FEEDBACK\\n[Question]'"
-        )
-        plan = await call_model(prompt, model="gpt-4o", api_key=api_key)
-
-      else:
-        prompt = (
-            f"You are a planner. Create a plan for: '{context.optimized_task}'. "
-            "Break down the task. Assign sub-tasks to: Coder, Code Reviewer, Quality Assurance Tester, and Documentation Agent. "
-            "Include review/revision steps. Consider error handling. Include documentation instructions.\n\n"
-             "If unsure, output 'REQUEST_HUMAN_FEEDBACK\\n[Question]'\n\nOutput the plan as a numbered list."
-        )
-        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}")
-
-            #Prepare detailed context for human
-            feedback_request_context = (f"The orchestrator agent is requesting feedback on the following task:\n **{context.optimized_task}**\n\n"
-                                       f"The current plan (if any):\n**{context.plan}**\n\n" if context.plan else "") + f"The specific question is:\n**{question}**"
-
-            self.human_in_the_loop_event.set()  # Signal the human input thread
-
-            human_response = self.get_human_response(feedback_request_context) # Pass context to input function
-            self.human_in_the_loop_event.clear()  # Reset the event
-            self.log_queue.put(f"[Orchestrator]: Received human feedback: {human_response}")
-            context.add_conversation_entry("Orchestrator", f"Plan:\n{plan}\n\nHuman Feedback Requested. Question: {question}")
-            return await self.generate_plan(context, api_key, human_response)  # Recursive call
-
-      context.plan = plan
-      context.add_conversation_entry("Orchestrator", f"Plan:\n{plan}")
-      return context
-
-    def get_human_response(self, feedback_request_context):
-        """Gets human input, using the Gradio queue and event."""
-        self.human_input_queue.put(feedback_request_context)  # Put the question into Gradio
-        human_response = self.human_input_queue.get()       # Get the response
-        return human_response
+    async def generate_plan(self, context: Context, api_key: str) -> Context:
+        """
+        Generates (or revises) a plan using human feedback if necessary.
+        Uses an iterative approach instead of recursion.
+        """
+        while True:
+            if context.plan:
+                prompt = (
+                    f"You are a planner. Revise/complete the plan for '{context.original_task}' using feedback:\n"
+                    f"{context.plan}\n\n"
+                    "If unsure, output 'REQUEST_HUMAN_FEEDBACK\\n[Question]'"
+                )
+            else:
+                prompt = (
+                    f"You are a planner. Create a plan for: '{context.optimized_task}'. "
+                    "Break down the task and assign sub-tasks to: Coder, Code Reviewer, Quality Assurance Tester, and Documentation Agent. "
+                    "Include review/revision steps, error handling, and documentation instructions.\n\n"
+                    "If unsure, output 'REQUEST_HUMAN_FEEDBACK\\n[Question]'"
+                )
+
+            plan = await call_model(prompt, model="gpt-4o", api_key=api_key)
+            context.add_conversation_entry("Orchestrator", f"Plan:\n{plan}")
+
+            # Check if human feedback is requested.
+            if "REQUEST_HUMAN_FEEDBACK" in plan:
+                question = plan.split("REQUEST_HUMAN_FEEDBACK\n", 1)[1].strip()
+                self.log_queue.put("[Orchestrator]: Requesting human feedback...")
+                self.log_queue.put(f"[Orchestrator]: Question for human: {question}")
+
+                # Prepare feedback context and trigger the human feedback event.
+                feedback_request_context = (
+                    f"The orchestrator agent is requesting feedback on the following task:\n"
+                    f"**{context.optimized_task}**\n\n"
+                    f"Current plan:\n**{context.plan or 'None'}**\n\n"
+                    f"Question:\n**{question}**"
+                )
+                self.human_event.set()
+                # Pass the context to the human input handler.
+                self.human_input_queue.put(feedback_request_context)
+                human_response = self.human_input_queue.get()  # Blocking call for human response.
+                self.human_event.clear()
+
+                self.log_queue.put(f"[Orchestrator]: Received human feedback: {human_response}")
+                # Incorporate human feedback into the plan and loop again.
+                context.plan = context.plan + "\n" + human_response if context.plan else human_response
+            else:
+                context.plan = plan
+                break  # Exit loop when no feedback is requested.
+        return context
 
 class CoderAgent:
     async def generate_code(self, context: Context, api_key: str, model: str = "gpt-4o") -> Context:
-        """Generates code based on instructions."""
+        """
+        Generates code based on the provided plan.
+        """
         prompt = (
             "You are a coding agent. Output ONLY the code. "
-            "Adhere to best practices. Include error handling.\n\n"
+            "Adhere to best practices and include error handling.\n\n"
             f"Instructions:\n{context.plan}"
         )
         code = await call_model(prompt, model=model, api_key=api_key)
@@ -155,34 +162,38 @@ class CoderAgent:
 
 class CodeReviewerAgent:
     async def review_code(self, context: Context, api_key: str) -> Context:
-        """Reviews code. Provides concise, actionable feedback or 'APPROVE'."""
+        """
+        Reviews the generated code and returns either actionable feedback or 'APPROVE'.
+        """
         prompt = (
-            "You are a 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"
+            "You are a code reviewer. Provide CONCISE feedback focusing on correctness, efficiency, readability, error handling, and security. "
+            "If the code is acceptable, respond with ONLY 'APPROVE'. Do NOT generate code.\n\n"
             f"Task: {context.optimized_task}\n\nCode:\n{context.code}"
         )
         review = await call_model(prompt, model="gpt-4o", api_key=api_key)
         context.add_conversation_entry("Code Reviewer", f"Review:\n{review}")
 
-        # Structured Feedback (Example)
+        # Check for approval; if not approved, parse feedback.
         if "APPROVE" not in review.upper():
-          structured_review = {"comments": []}
-          #In a real implementation you might use a more advanced parsing technique here
-          for line in review.splitlines():
-            if line.strip(): #Simple example
-              structured_review["comments"].append({"issue": line.strip(), "line_number": "N/A", "severity": "Medium"}) #Dummy data
-          context.review_comments.append(structured_review)
-
+            structured_review = {"comments": []}
+            for line in review.splitlines():
+                if line.strip():
+                    structured_review["comments"].append({
+                        "issue": line.strip(),
+                        "line_number": "N/A",
+                        "severity": "Medium"
+                    })
+            context.review_comments.append(structured_review)
         return context
 
 class QualityAssuranceTesterAgent:
     async def generate_test_cases(self, context: Context, api_key: str) -> Context:
-        """Generates test cases."""
+        """
+        Generates test cases considering edge and error cases.
+        """
         prompt = (
-            "You are a testing agent. Generate test cases. "
-            "Consider edge cases and error scenarios. Output in a clear format.\n\n"
+            "You are a testing agent. Generate comprehensive test cases considering edge cases and error scenarios. "
+            "Output in a clear format.\n\n"
             f"Task: {context.optimized_task}\n\nCode:\n{context.code}"
         )
         test_cases = await call_model(prompt, model="gpt-4o", api_key=api_key)
@@ -191,10 +202,12 @@ class QualityAssuranceTesterAgent:
         return context
 
     async def run_tests(self, context: Context, api_key: str) -> Context:
-        """Runs tests and reports results."""
+        """
+        Runs the generated test cases and compares expected vs. actual outcomes.
+        """
         prompt = (
-            "Run the test cases. Compare actual vs expected output. "
-            "State discrepancies. If all pass, output 'TESTS PASSED'.\n\n"
+            "Run the test cases. Compare actual vs expected outputs and state any discrepancies. "
+            "If all tests pass, output 'TESTS PASSED'.\n\n"
             f"Code:\n{context.code}\n\nTest Cases:\n{context.test_cases}"
         )
         test_results = await call_model(prompt, model="gpt-4o", api_key=api_key)
@@ -204,10 +217,11 @@ class QualityAssuranceTesterAgent:
 
 class DocumentationAgent:
     async def generate_documentation(self, context: Context, api_key: str) -> Context:
-        """Generates documentation, including a --help message."""
+        """
+        Generates concise documentation including a --help message.
+        """
         prompt = (
-            "Generate clear and concise documentation. "
-            "Include a brief description, explanation, and a --help message.\n\n"
+            "Generate clear documentation including a brief description, explanation, and a --help message.\n\n"
             f"Code:\n{context.code}"
         )
         documentation = await call_model(prompt, model="gpt-4o", api_key=api_key)
@@ -215,16 +229,16 @@ class DocumentationAgent:
         context.add_conversation_entry("Documentation Agent", f"Documentation:\n{documentation}")
         return context
 
-# -------------------- Agent Dispatcher (New) --------------------
+# -------------------- Agent Dispatcher --------------------
 
 class AgentDispatcher:
-    def __init__(self, log_queue: queue.Queue, human_in_the_loop_event: threading.Event, human_input_queue: queue.Queue):
+    def __init__(self, log_queue: queue.Queue, human_event: threading.Event, human_input_queue: queue.Queue):
         self.log_queue = log_queue
-        self.human_in_the_loop_event = human_in_the_loop_event
+        self.human_event = human_event
         self.human_input_queue = human_input_queue
         self.agents = {
             "prompt_optimizer": PromptOptimizerAgent(),
-            "orchestrator": OrchestratorAgent(log_queue, human_in_the_loop_event, human_input_queue),
+            "orchestrator": OrchestratorAgent(log_queue, human_event, human_input_queue),
             "coder": CoderAgent(),
             "code_reviewer": CodeReviewerAgent(),
             "qa_tester": QualityAssuranceTesterAgent(),
@@ -232,16 +246,18 @@ class AgentDispatcher:
         }
 
     async def dispatch(self, agent_name: str, context: Context, api_key: str, **kwargs) -> Context:
-        """Dispatches the task to the specified agent."""
+        """
+        Dispatches the task to the specified agent.
+        """
         agent = self.agents.get(agent_name)
         if not agent:
             raise ValueError(f"Unknown agent: {agent_name}")
 
-        self.log_queue.put(f"[{agent_name.replace('_', ' ').title()}]: Starting task...")  # Log here
+        self.log_queue.put(f"[{agent_name.replace('_', ' ').title()}]: Starting task...")
         if agent_name == "prompt_optimizer":
             context = await agent.optimize_prompt(context, api_key)
         elif agent_name == "orchestrator":
-            context = await agent.generate_plan(context, api_key)  #Removed human_feedback
+            context = await agent.generate_plan(context, api_key)
         elif agent_name == "coder":
             context = await agent.generate_code(context, api_key, **kwargs)
         elif agent_name == "code_reviewer":
@@ -254,167 +270,171 @@ class AgentDispatcher:
         elif agent_name == "documentation_agent":
             context = await agent.generate_documentation(context, api_key)
         else:
-          raise ValueError(f"Unknown Agent Name: {agent_name}")
-
+            raise ValueError(f"Unknown Agent Name: {agent_name}")
         return context
-    async def determine_next_agent(self, context:Context, api_key:str) -> str:
-        """Determines the next agent to run based on the current context."""
+
+    async def determine_next_agent(self, context: Context, api_key: str) -> str:
+        """
+        Determines the next agent to run based on the current context.
+        """
         if not context.optimized_task:
             return "prompt_optimizer"
         if not context.plan:
             return "orchestrator"
         if not context.code:
             return "coder"
-        if not context.review_comments or "APPROVE" not in [comment.get('issue',"").upper() for comment_list in context.review_comments for comment in comment_list.get("comments",[])  ]:
+        # Check if any review comment lacks an APPROVE.
+        if not any(
+            "APPROVE" in comment.get("issue", "").upper()
+            for review in context.review_comments
+            for comment in review.get("comments", [])
+        ):
             return "code_reviewer"
         if not context.test_cases:
             return "qa_tester"
-        if not context.test_results or "TESTS PASSED" not in context.test_results.upper() :
+        if not context.test_results or "TESTS PASSED" not in context.test_results.upper():
             return "qa_tester"
         if not context.documentation:
             return "documentation_agent"
 
         return "done"  # All tasks are complete
 
-# -------------------- Multi-Agent Conversation (Refactored) --------------------
-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:
+# -------------------- Multi-Agent Conversation --------------------
+
+async def multi_agent_conversation(task_message: str, log_queue: queue.Queue, api_key: str,
+                                   human_event: threading.Event, human_input_queue: queue.Queue) -> None:
     """
-    Conducts the multi-agent conversation using the AgentDispatcher.
+    Orchestrates the multi-agent conversation.
     """
     context = Context(original_task=task_message)
-    dispatcher = AgentDispatcher(log_queue, human_in_the_loop_event, human_input_queue)
+    dispatcher = AgentDispatcher(log_queue, human_event, human_input_queue)
 
     next_agent = await dispatcher.determine_next_agent(context, api_key)
+    # Prevent endless revisions by tracking coder iterations.
+    coder_iterations = 0
+
     while next_agent != "done":
         if next_agent == "qa_tester":
             if not context.test_cases:
-              context = await dispatcher.dispatch(next_agent, context, api_key, generate_tests=True)
+                context = await dispatcher.dispatch(next_agent, context, api_key, generate_tests=True)
             else:
-              context = await dispatcher.dispatch(next_agent, context, api_key, run_tests=True)
+                context = await dispatcher.dispatch(next_agent, context, api_key, run_tests=True)
         elif next_agent == "coder" and (context.review_comments or context.test_results):
-          #Coder needs a different model after the first coding
-          context = await dispatcher.dispatch(next_agent,context, api_key, model="gpt-3.5-turbo-16k")
+            coder_iterations += 1
+            # Switch to a different model after the first iteration.
+            context = await dispatcher.dispatch(next_agent, context, api_key, model="gpt-3.5-turbo-16k")
+        else:
+            context = await dispatcher.dispatch(next_agent, context, api_key)
+        
+        # Check for approval in code review if applicable.
+        if next_agent == "code_reviewer":
+            approved = any(
+                "APPROVE" in comment.get("issue", "").upper()
+                for review in context.review_comments
+                for comment in review.get("comments", [])
+            )
+            if not approved:
+                # If not approved, we continue with coder to improve the code.
+                next_agent = "coder"
+            else:
+                next_agent = await dispatcher.determine_next_agent(context, api_key)
         else:
-          context = await dispatcher.dispatch(next_agent, context, api_key) # Call the agent
+            next_agent = await dispatcher.determine_next_agent(context, api_key)
 
-        next_agent = await dispatcher.determine_next_agent(context, api_key)
-        if next_agent == "code_reviewer" and context.review_comments and "APPROVE" in [comment.get('issue',"").upper() for comment_list in context.review_comments for comment in comment_list.get("comments",[])  ]:
-          next_agent = await dispatcher.determine_next_agent(context, api_key)
-        # Check for maximum revisions
-        if next_agent == "coder" and len([entry for entry in context.conversation_history if entry["agent"] == "Coder"]) > 5:
-          log_queue.put("Maximum revision iterations reached. Exiting.")
-          break;
+        if next_agent == "coder" and coder_iterations > 5:
+            log_queue.put("Maximum revision iterations reached. Exiting.")
+            break
 
     log_queue.put("Conversation complete.")
     log_queue.put(("result", context.conversation_history))
 
 # -------------------- 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, log_queue: queue.Queue) -> Generator[str, None, None]:
+
+def process_conversation_generator(task_message: str, api_key: str,
+                                   human_event: threading.Event, human_input_queue: queue.Queue,
+                                   log_queue: queue.Queue) -> Generator[str, None, None]:
     """
-    Wraps the conversation, yields log messages, and handles human input within a single thread.
-    Crucially, takes the log_queue as an argument.  Yields Gradio updates.
+    Runs the conversation and yields log messages.
     """
+    # Run the conversation asynchronously.
+    asyncio.run(multi_agent_conversation(task_message, log_queue, api_key, human_event, human_input_queue))
 
-    # Run the multi-agent conversation *synchronously* within this function.
-    asyncio.run(multi_agent_conversation(task_message, log_queue, api_key, human_in_the_loop_event, human_input_queue))
-
-    # Process the log queue and handle human-in-the-loop
     final_result = None
-    while True:  # Loop indefinitely to handle multiple potential human feedback requests.
+    while True:
         try:
-            msg = log_queue.get_nowait()  # Non-blocking get from the log queue.
+            msg = log_queue.get_nowait()
             if isinstance(msg, tuple) and msg[0] == "result":
                 final_result = msg[1]
-                yield gr.Chatbot.update(final_result)  # Update the chatbot with the final result
-                yield "Conversation complete."  # Indicate completion.
-                break  # Exit the loop after processing the final result.
+                yield gr.Chatbot.update(value=final_result, visible=True)
+                yield "Conversation complete."
+                break
             else:
-                yield msg  # Yield the log message.
+                yield msg
         except queue.Empty:
-            pass  # No log message available, continue checking for human input.
-
+            pass
 
-        if human_in_the_loop_event.is_set():
-            yield "Waiting for human feedback..."  # Indicate waiting state.
-            try:
-                feedback_request = human_input_queue.get(
-                    timeout=0.1)  # Get the context/question for feedback.
-                human_interface = get_human_feedback(feedback_request)
-                yield gr.Textbox.update(visible=False), gr.update(visible=True)  # Show feedback UI
-                human_feedback = human_input_queue.get(
-                    timeout=300)  # Wait (block) for human feedback, with a timeout.
-                human_in_the_loop_event.clear()  # Reset the event after getting feedback.
-                yield gr.Textbox.update(visible=True), human_interface.close()  # Hide feedback UI.
+        # If human feedback is requested, yield an appropriate message.
+        if human_event.is_set():
+            yield "Waiting for human feedback..."
+        # Use a short asynchronous sleep to avoid busy-waiting.
+        asyncio.run(asyncio.sleep(0.1))
 
-
-            except queue.Empty:
-                pass
-        # Add a small sleep to avoid busy-waiting and reduce CPU usage.
-        time.sleep(0.1)
-
-
-
-
-def get_human_feedback(placeholder_text):
-    """Gets human input using a Gradio Textbox."""
+def get_human_feedback(placeholder_text: str, human_input_queue: queue.Queue) -> gr.Blocks:
+    """
+    Constructs the Gradio interface to collect human feedback.
+    """
     with gr.Blocks() as human_feedback_interface:
         with gr.Row():
             human_input = gr.Textbox(lines=4, label="Human Feedback", placeholder=placeholder_text)
         with gr.Row():
             submit_button = gr.Button("Submit Feedback")
 
-        def submit_feedback(input_text):
-            # Put the feedback into the shared queue
+        def submit_feedback(input_text: str):
             human_input_queue.put(input_text)
-            return ""  # Clear the input box after submission
-
+            return ""
 
         submit_button.click(fn=submit_feedback, inputs=human_input, outputs=human_input)
-        human_feedback_interface.load(None, [], [])  # Keep interface alive
-
     return human_feedback_interface
 
 # -------------------- 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."""
+def multi_agent_chat(message: str, history: List[Any], openai_api_key: str = None) -> Generator[Any, None, None]:
+    """
+    Gradio chat function that runs the multi-agent conversation.
+    """
     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()  # Use a single queue for both requests and responses
-    log_queue = queue.Queue() #Create log queue here
-
-    yield from process_conversation_generator(message, openai_api_key, human_in_the_loop_event, human_input_queue, log_queue)
-
-
+    human_event = threading.Event()
+    human_input_queue = queue.Queue()
+    log_queue = queue.Queue()
 
+    yield from process_conversation_generator(message, openai_api_key, human_event, human_input_queue, log_queue)
 
 # -------------------- Launch the Chatbot --------------------
 
-# Create the main chat interface
 iface = gr.ChatInterface(
     fn=multi_agent_chat,
-    chatbot=gr.Chatbot(type="feed"), # Use the 'feed' type for a better display of messages
-    additional_inputs=[gr.Textbox(label="OpenAI API Key (optional)", type="password", placeholder="Leave blank to use env variable")],
+    chatbot=gr.Chatbot(type="messages"),
+    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.
-        - Orchestrator can ask for human feedback.
-        - Enter a task; agents will work on it. You may be prompted for input.
-        - Max 5 revisions.
-        - Provide API Key.
-        """
+    description=(
+        "- Collaborative workflow with Human-in-the-Loop.\n"
+        "- Orchestrator can ask for human feedback.\n"
+        "- Enter a task; agents will work on it. You may be prompted for input.\n"
+        "- Max 5 revisions.\n"
+        "- Provide API Key."
+    )
 )
 
-#Need a dummy interface to prevent Gradio errors
-dummy_iface = gr.Interface(lambda x:x, "textbox", "textbox")
+# Dummy interface to prevent Gradio errors.
+dummy_iface = gr.Interface(lambda x: x, "textbox", "textbox")
 
 if __name__ == "__main__":
     demo = gr.TabbedInterface([iface, dummy_iface], ["Chatbot", "Dummy"])
     demo.launch(share=True)
-
-import time #Import the time module