Update app.py

app.py

@@ -5,36 +5,33 @@ import threading
 import queue
 import gradio as gr
 import httpx
+import time
 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) --------------------
 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 retries and exponential backoff.
-    """
     if api_key is None:
         api_key = os.getenv("OPENAI_API_KEY")
         if api_key is None:
-            raise ValueError("OpenAI API key not found.")
+            raise ValueError("OpenAI API key not provided.")
     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}]}
-
+    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()  # Synchronous parsing is acceptable here
+                response_json = response.json()
                 return response_json["choices"][0]["message"]["content"]
         except httpx.HTTPStatusError as e:
             logging.error(f"HTTP error (attempt {attempt + 1}/{max_retries}): {e}")
@@ -50,7 +47,7 @@ async def call_model(prompt: str, model: str = "gpt-4o", api_key: str = None, ma
         except Exception as 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.")
+    raise Exception(f"Failed to get response after {max_retries} attempts.")
 
 # -------------------- Shared Context --------------------
 class Context:
@@ -73,12 +70,8 @@ class Context:
         self.conversation_history.append({"agent": agent_name, "message": message})
 
 # -------------------- Agent Classes --------------------
-
 class PromptOptimizerAgent:
     async def optimize_prompt(self, context: Context, api_key: str) -> Context:
-        """
-        Optimizes the user’s original prompt.
-        """
         system_prompt = (
             "Improve the prompt. Be clear, specific, and complete. "
             "Keep original intent. Return ONLY the revised prompt."
@@ -90,21 +83,16 @@ class PromptOptimizerAgent:
         return context
 
 class OrchestratorAgent:
-    def __init__(self, log_queue: queue.Queue, human_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):
         self.log_queue = log_queue
         self.human_event = human_event
         self.human_input_queue = human_input_queue
 
     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"
+                    f"You are a planner. Revise/complete the plan for '{context.original_task}'. "
                     "If unsure, output 'REQUEST_HUMAN_FEEDBACK\\n[Question]'"
                 )
             else:
@@ -114,42 +102,28 @@ class OrchestratorAgent:
                     "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.log_queue.put(f"[Orchestrator]: Question: {question}")
+                feedback_context = (
+                    f"Task: {context.optimized_task}\nCurrent Plan: {context.plan or 'None'}\nQuestion: {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_input_queue.put(feedback_context)
+                human_response = self.human_input_queue.get()  # blocking call waiting 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
+                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.
+                break
         return context
 
 class CoderAgent:
     async def generate_code(self, context: Context, api_key: str, model: str = "gpt-4o") -> Context:
-        """
-        Generates code based on the provided plan.
-        """
         prompt = (
             "You are a coding agent. Output ONLY the code. "
             "Adhere to best practices and include error handling.\n\n"
@@ -162,9 +136,6 @@ class CoderAgent:
 
 class CodeReviewerAgent:
     async def review_code(self, context: Context, api_key: str) -> Context:
-        """
-        Reviews the generated code and returns either actionable feedback or 'APPROVE'.
-        """
         prompt = (
             "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"
@@ -172,8 +143,6 @@ class CodeReviewerAgent:
         )
         review = await call_model(prompt, model="gpt-4o", api_key=api_key)
         context.add_conversation_entry("Code Reviewer", f"Review:\n{review}")
-
-        # Check for approval; if not approved, parse feedback.
         if "APPROVE" not in review.upper():
             structured_review = {"comments": []}
             for line in review.splitlines():
@@ -188,9 +157,6 @@ class CodeReviewerAgent:
 
 class QualityAssuranceTesterAgent:
     async def generate_test_cases(self, context: Context, api_key: str) -> Context:
-        """
-        Generates test cases considering edge and error cases.
-        """
         prompt = (
             "You are a testing agent. Generate comprehensive test cases considering edge cases and error scenarios. "
             "Output in a clear format.\n\n"
@@ -202,9 +168,6 @@ class QualityAssuranceTesterAgent:
         return context
 
     async def run_tests(self, context: Context, api_key: str) -> Context:
-        """
-        Runs the generated test cases and compares expected vs. actual outcomes.
-        """
         prompt = (
             "Run the test cases. Compare actual vs expected outputs and state any discrepancies. "
             "If all tests pass, output 'TESTS PASSED'.\n\n"
@@ -217,9 +180,6 @@ class QualityAssuranceTesterAgent:
 
 class DocumentationAgent:
     async def generate_documentation(self, context: Context, api_key: str) -> Context:
-        """
-        Generates concise documentation including a --help message.
-        """
         prompt = (
             "Generate clear documentation including a brief description, explanation, and a --help message.\n\n"
             f"Code:\n{context.code}"
@@ -230,7 +190,6 @@ class DocumentationAgent:
         return context
 
 # -------------------- Agent Dispatcher --------------------
-
 class AgentDispatcher:
     def __init__(self, log_queue: queue.Queue, human_event: threading.Event, human_input_queue: queue.Queue):
         self.log_queue = log_queue
@@ -246,49 +205,34 @@ class AgentDispatcher:
         }
 
     async def dispatch(self, agent_name: str, context: Context, api_key: str, **kwargs) -> Context:
-        """
-        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...")
         if agent_name == "prompt_optimizer":
-            context = await agent.optimize_prompt(context, api_key)
+            context = await self.agents[agent_name].optimize_prompt(context, api_key)
         elif agent_name == "orchestrator":
-            context = await agent.generate_plan(context, api_key)
+            context = await self.agents[agent_name].generate_plan(context, api_key)
         elif agent_name == "coder":
-            context = await agent.generate_code(context, api_key, **kwargs)
+            context = await self.agents[agent_name].generate_code(context, api_key, **kwargs)
         elif agent_name == "code_reviewer":
-            context = await agent.review_code(context, api_key)
+            context = await self.agents[agent_name].review_code(context, api_key)
         elif agent_name == "qa_tester":
             if kwargs.get("generate_tests", False):
-                context = await agent.generate_test_cases(context, api_key)
+                context = await self.agents[agent_name].generate_test_cases(context, api_key)
             elif kwargs.get("run_tests", False):
-                context = await agent.run_tests(context, api_key)
+                context = await self.agents[agent_name].run_tests(context, api_key)
         elif agent_name == "documentation_agent":
-            context = await agent.generate_documentation(context, api_key)
+            context = await self.agents[agent_name].generate_documentation(context, api_key)
         else:
-            raise ValueError(f"Unknown Agent Name: {agent_name}")
+            raise ValueError(f"Unknown agent: {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.
-        """
         if not context.optimized_task:
             return "prompt_optimizer"
         if not context.plan:
             return "orchestrator"
         if not context.code:
             return "coder"
-        # 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", [])
-        ):
+        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"
@@ -296,23 +240,15 @@ class AgentDispatcher:
             return "qa_tester"
         if not context.documentation:
             return "documentation_agent"
-
-        return "done"  # All tasks are complete
+        return "done"
 
 # -------------------- 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:
-    """
-    Orchestrates the multi-agent conversation.
-    """
     context = Context(original_task=task_message)
     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:
@@ -321,34 +257,24 @@ async def multi_agent_conversation(task_message: str, log_queue: queue.Queue, ap
                 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_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", [])
-            )
+            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:
             next_agent = await dispatcher.determine_next_agent(context, api_key)
-
         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 --------------------
+# -------------------- Process Conversation Generator --------------------
 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]:
@@ -359,86 +285,53 @@ def process_conversation_generator(task_message: str, api_key: str,
     def run_conversation():
         asyncio.run(multi_agent_conversation(task_message, log_queue, api_key, human_event, human_input_queue))
     
-    # Start the asynchronous conversation in a separate thread.
     conversation_thread = threading.Thread(target=run_conversation)
     conversation_thread.start()
     
-    # Continuously yield log messages until the conversation is complete.
     while conversation_thread.is_alive() or not log_queue.empty():
         try:
             msg = log_queue.get(timeout=0.1)
             if isinstance(msg, tuple) and msg[0] == "result":
-                # Update the chat box with the final conversation history.
                 yield gr.Chatbot.update(value=msg[1], visible=True)
             else:
                 yield msg
         except queue.Empty:
             pass
-        
         if human_event.is_set():
             yield "Waiting for human feedback..."
-        
-        # Small sleep to prevent busy-waiting.
         time.sleep(0.1)
     
     yield "Conversation complete."
 
-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: str):
-            human_input_queue.put(input_text)
-            return ""
-
-        submit_button.click(fn=submit_feedback, inputs=human_input, outputs=human_input)
-    return human_feedback_interface
-
-# -------------------- 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.
-    """
+# -------------------- Multi-Agent Chat Function --------------------
+def multi_agent_chat(message: str, openai_api_key: str = None) -> Generator[Any, None, None]:
     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_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 --------------------
-
-iface = gr.ChatInterface(
-    fn=multi_agent_chat,
-    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.\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."
-    )
-)
-
-# Dummy interface to prevent Gradio errors.
-dummy_iface = gr.Interface(lambda x: x, "textbox", "textbox")
+# -------------------- Custom Gradio Blocks Interface --------------------
+with gr.Blocks() as demo:
+    gr.Markdown("## Multi-Agent Task Solver with Human-in-the-Loop")
+    
+    with gr.Row():
+        chat_output = gr.Chatbot(label="Conversation")
+    
+    with gr.Row():
+        with gr.Column(scale=8):
+            message_input = gr.Textbox(label="Enter your task", placeholder="Type your task here...", lines=3)
+        with gr.Column(scale=2):
+            api_key_input = gr.Textbox(label="API Key (optional)", type="password", placeholder="Leave blank to use env variable")
+    
+    send_button = gr.Button("Send")
+    
+    # When Send is clicked, the multi_agent_chat generator is called and its output is streamed to the chat.
+    send_button.click(fn=multi_agent_chat, inputs=[message_input, api_key_input], outputs=chat_output, stream=True)
 
 if __name__ == "__main__":
-    demo = gr.TabbedInterface([iface, dummy_iface], ["Chatbot", "Dummy"])
     demo.launch(share=True)