Update app.py

app.py

@@ -7,15 +7,50 @@ import queue
 import gradio as gr
 import httpx
 from typing import Generator, Any, Dict, List
-logger = logging.getLogger(__name__)
 
-class Agent(ABC):
-    @abstractmethod
-    async def generate_response(self, prompt: str, api_key: str) -> str:
-        pass
+# -------------------- Configuration --------------------
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
-class PromptOptimizerAgent(Agent):
-    async def generate_response(self, prompt: str, api_key: str) -> str:
+# -------------------- 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"]
+
+# -------------------- 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 ###
+        """
         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. "
@@ -25,155 +60,210 @@ class PromptOptimizerAgent(Agent):
             "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"
+        )
+        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"
         )
-        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
+        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."
         )
-        return feedback
+        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})...")
 
-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)
+        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}")
 
-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
+        # 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
 
+        # 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]
-                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
+                yield "Final conversation complete."
+            else:
+                yield msg
+        except queue.Empty:
+            continue
 
-    thread = asyncio.to_thread(run_conversation_thread)
-    thread.start()
+    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
+
+# -------------------- 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)
 
-    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]
-    )
+# -------------------- 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!
+        """
+)
 
 if __name__ == "__main__":
-    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()
+    iface.launch()