Update app.py

app.py

@@ -1,56 +1,19 @@
-import os
-import sys
-import asyncio
+from abc import ABC, abstractmethod
+from typing import Generator, Any
 import logging
-import threading
 import queue
-import gradio as gr
 import httpx
-from typing import Generator, Any, Dict, List
+from gradio import ChatInterface, gr
 
-# -------------------- Configuration --------------------
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
 
-# -------------------- External Model Call --------------------
-async def call_model(prompt: str, model: str = "gpt-4o-mini", 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.
-    """
-    # Use the provided API key or fall back to the environment variable
-    if api_key is None:
-        api_key = os.getenv("OPENAI_API_KEY")
-    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",
-        "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"]
+class Agent(ABC):
+    @abstractmethod
+    async def generate_response(self, prompt: str, api_key: str) -> str:
+        pass
 
-# -------------------- 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 ###
-        """
+class PromptOptimizerAgent(Agent):
+    async def generate_response(self, prompt: str, api_key: str) -> str:
         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. "
@@ -60,210 +23,155 @@ class PromptOptimizerAgent:
             "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 ###"
+            ###User initial prompt###
         )
-        full_prompt = f"{system_prompt}\n{user_prompt}\n<<<"
-        optimized = await call_model(full_prompt, api_key=api_key)
-        return optimized
-
-class OrchestratorAgent:
-    def __init__(self, log_queue: queue.Queue) -> None:
-        self.log_queue = log_queue
-
-    async def generate_plan(self, task: str, api_key: str) -> str:
-        """
-        Generates a detailed, step-by-step plan for completing the given task.
-        """
-        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."
-        )
-        plan = await call_model(prompt, api_key=api_key)
-        return plan
-
-class CoderAgent:
-    async def generate_code(self, instructions: str, api_key: str) -> str:
-        """
-        Generates code based on the given 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"
+        return await call_openai(system_prompt, api_key)
+
+class OrchestratorAgent(Agent):
+    async def generate_response(self, task_message: str, api_key: str) -> str:
+        plan = f"You are an orchestrator agent. The user has provided the task: '{task_message}'. 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."
+        return await call_openai(plan, api_key)
+
+class CoderAgent(Agent):
+    async def generate_response(self, instructions: str, api_key: str) -> str:
+        prompt = f"Implement the task as described in the following plan:\n{instructions}"
+        return await call_openai(prompt, api_key)
+
+class CodeReviewerAgent(Agent):
+    async def generate_response(self, code: str, task: str, api_key: str) -> str:
+        feedback = await call_openai(
+            f"You are a code reviewer agent. Review the provided code: '{code}' and check if it meets the task specifications.",
+            api_key=api_key
         )
-        code = await call_model(prompt, 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.
-        """
-        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"
-        )
-        review = await call_model(prompt, api_key=api_key)
-        return review
-
-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.
-        """
-        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."
-        )
-        documentation = await call_model(prompt, 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:
-    """
-    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.
-    """
-    conversation: List[Dict[str, str]] = []  # List to store each agent's message
-
-    # Step 0: Use Prompt Optimizer to enhance the user's initial prompt.
-    log_queue.put("[Prompt Optimizer]: Received initial task. 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)
-    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.
-    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.")
-    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.
-    reviewer = CodeReviewerAgent()
-    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})...")
+        return feedback
 
-        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}")
+class DocumentationAgent(Agent):
+    async def generate_response(self, code: str, api_key: str) -> str:
+        prompt = f"You are a documentation agent. Generate a brief documentation for the code:\nCode:\n{code}"
+        return await call_openai(prompt, api_key)
 
-        # 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
+async def process_conversation_generator(conversation: list, log_q: queue.Queue, api_key: str) -> Generator[str, None, None]:
+    try:
+        while True:
+            if not conversation or not log_q.get(timeout=0.1):
+                continue
 
-        # 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 :(")
-
-        # 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)
-        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
-
-    # Step 4: Documentation Agent generates documentation for the approved code.
-    doc_agent = DocumentationAgent()
-    log_queue.put("[Documentation Agent]: Generating documentation for the approved code.")
-    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(("result", conversation))
-
-# -------------------- Process Generator for Streaming --------------------
-def process_conversation_generator(task_message: str, api_key: str) -> Generator[str, None, None]:
-    """
-    Wraps the asynchronous multi-agent conversation and yields log messages as they are generated.
-    """
-    log_q: queue.Queue = queue.Queue()
-
-    def run_conversation() -> None:
-        asyncio.run(multi_agent_conversation(task_message, log_q, api_key))
-
-    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."
-            else:
-                yield msg
-        except queue.Empty:
-            continue
-
-    thread.join()
-    if final_result:
-        # Format the final conversation log.
-        conv_text = "\n========== Multi-Agent Conversation ==========\n"
-        for entry in final_result:
-            conv_text += f"[{entry['agent']}]: {entry['message']}\n\n"
-        yield conv_text
+                break
+
+            yield msg
+    except asyncio.CancelledError:
+        pass
+    finally:
+        if log_q.empty():
+            log_q.put("Final conversation complete.")
+
+async def multi_agent_conversation(
+    task_message: str, 
+    log_q: queue.Queue, 
+    api_key: str,
+    additional_inputs=None
+) -> None:
+    if additional_inputs is None:
+        additional_inputs = [gr.Textbox(label="OpenAI API Key (optional)")]
+
+    agents = [
+        PromptOptimizerAgent(),
+        OrchestratorAgent(),
+        CoderAgent(),
+        CodeReviewerAgent(),
+        DocumentationAgent()
+    ]
+
+    log_queue = queue.Queue()
+    run_conversation = None
+
+    async def run_conversation_thread() -> None:
+        nonlocal run_conversation
+        try:
+            if run_conversation is not None:
+                await run_conversation
+        except asyncio.CancelledError:
+            pass
 
-# -------------------- 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.
-    """
-    if not openai_api_key:
-        openai_api_key = os.getenv("OPENAI_API_KEY")
-    yield from process_conversation_generator(message, openai_api_key)
+    thread = asyncio.to_thread(run_conversation_thread)
+    thread.start()
 
-# -------------------- Launch the Chatbot --------------------
-# Use Gradio's ChatInterface with an additional input field for the OpenAI API key.
-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",
-    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!
-        """
-)
+    try:
+        conversation = []
+        log_q.put("[Prompt Optimizer]: Received initial task. Optimizing prompt...")
+
+        # Step 0: Use Prompt Optimizer
+        optimized_task = await agents[0].generate_response(task_message, api_key)
+        conversation.append({"agent": "Prompt Optimizer", "message": f"Optimized Task:\n{optimized_task}"})
+        log_q.put(f"[Prompt Optimizer]: Optimized Task:\n{optimized_task}")
+
+        # Step 1: Generate Plan
+        plan = await agents[1].generate_response(optimized_task, api_key)
+        conversation.append({"agent": "Orchestrator", "message": f"Plan:\n{plan}"})
+        log_q.put(f"[Orchestrator]: Plan generated:\n{plan}")
+
+        # Step 2: Generate Code
+        code = await agents[2].generate_response(plan, api_key)
+        conversation.append({"agent": "Coder", "message": f"Code:\n{code}"})
+        log_q.put(f"[Coder]: Code generated:\n{code}")
+
+        # Step 3: Code Review
+        code_review = None
+        iteration = 0
+        while True:
+            if iteration >= 5:
+                log_q.put("[Code Reviewer]: Code not approved after 5 iterations: terminating.")
+                break
+
+            code_review = await agents[3].generate_response(code, plan, api_key)
+            revised_code = await agents[2].generate_response(
+                f"Please revise the code according to the following feedback: {code_review}",
+                api_key
+            )
+            code = revised_code
+            iteration += 1
+
+            if code == revised_code:
+                break
+
+            log_q.put(f"[Code Reviewer]: Feedback received:\n{code_review}")
+            log_q.put(f"[Code Reviewer]: Revised code:\n{revised_code}")
+
+        # Step 4: Documentation
+        doc = await agents[4].generate_response(code, api_key)
+        conversation.append({"agent": "Documentation Agent", "message": f"Documentation:\n{doc}"})
+        log_q.put(f"[Documentation Agent]: Documentation generated:\n{doc}")
+
+    except Exception as e:
+        log_q.put(f"[All Agents]: An error occurred: {str(e)}")
+        logger.error(f"Error in multi_agent_conversation: {str(e)}")
+
+    finally:
+        thread.join()
+
+async def multi_agent_conversation_wrapper(task_message: str, api_key: str) -> None:
+    await multi_agent_conversation(
+        task_message, 
+        log_q=queue.Queue(),
+        api_key=api_key,
+        additional_inputs=[gr.Textbox(label="OpenAI API Key (optional)") if api_key is None else None]
+    )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface = gr.ChatInterface(
+        fn=multi_agent_conversation_wrapper,
+        additional_inputs=[gr.Textbox(label="OpenAI API Key (optional)")],
+        type="messages",
+        title="Actual Multi-Agent Conversation Chatbot",
+        description="""
+            - Collaborative workflow between Prompt Enhancer, Orchestrator, Coder, Code-Reviewer and Documentation Agent agents.
+            - Enter a task description to observe the iterative workflow between the agents.
+            - NOTE: The kill-switch mechanism will terminate after five code rejection iterations to prevent endless loops.
+            - NOTE3: You can input your OPENAI_API_KEY at the bottom of the page for this to work!
+        """,
+    )
+
+    iface.launch()