Create app.py

app.py

@@ -0,0 +1,271 @@
+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
+
+# -------------------- 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:
+    """
+    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. "
+            "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 ###"
+        )
+        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"
+        )
+        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})...")
+
+        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}")
+
+        # 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]
+                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
+
+# -------------------- 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)
+
+# -------------------- 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 Primary, Critic, and Documentation agents:
+        "1. Enter a task and observe as your prompt gets magically improved by the Prompt-Enhancing agent before reaching the
+         2. Orchestrator, who then devises a plan and enlists the assistance of a 
+         3. Coder agent that writes code which is iteratively improved upon thanks to the
+         4. Code Reviewer who finally decides if the code should be approved and get sent off to the 
+         5. Code Documentation Generator that will write documentation for your freshly generated code!"
+        "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.launch()