app.py

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 functools import lru_cache

# -------------------- Configuration --------------------
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

# -------------------- External Model Call (with Caching) --------------------
@lru_cache(maxsize=128)  # Cache up to 128 responses
async def call_model(prompt: str, model: str = "gpt-4o", api_key: str = None) -> str:
    """Sends a prompt to the OpenAI API endpoint, with caching."""
    if api_key is None:
        api_key = os.getenv("OPENAI_API_KEY")
        if api_key is None:
            raise ValueError("OpenAI API key not found.")
    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}],
    }
    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"]

# -------------------- Shared Context --------------------
class Context:
    def __init__(self, original_task: str, optimized_task: Optional[str] = None,
                 plan: Optional[str] = None, code: Optional[str] = None,
                 review_comments: Optional[List[Dict[str, str]]] = None,
                 test_cases: Optional[str] = None, test_results: Optional[str] = None,
                 documentation: Optional[str] = None, conversation_history: Optional[List[Dict[str, str]]] = None):
        self.original_task = original_task
        self.optimized_task = optimized_task
        self.plan = plan
        self.code = code
        self.review_comments = review_comments or []
        self.test_cases = test_cases
        self.test_results = test_results
        self.documentation = documentation
        self.conversation_history = conversation_history or []

    def add_conversation_entry(self, agent_name: str, message: str):
        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 initial 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
        context.add_conversation_entry("Prompt Optimizer", f"Optimized Task:\n{optimized}")
        return context

class OrchestratorAgent:
    def __init__(self, log_queue: queue.Queue, human_in_the_loop_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_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

class CoderAgent:
    async def generate_code(self, context: Context, api_key: str, model: str = "gpt-4o") -> Context:
        """Generates code based on instructions."""
        prompt = (
            "You are a coding agent. Output ONLY the code. "
            "Adhere to best practices. Include error handling.\n\n"
            f"Instructions:\n{context.plan}"
        )
        code = await call_model(prompt, model=model, api_key=api_key)
        context.code = code
        context.add_conversation_entry("Coder", f"Code:\n{code}")
        return context

class CodeReviewerAgent:
    async def review_code(self, context: Context, api_key: str) -> Context:
        """Reviews code. Provides concise, 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"
            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)
        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)

        return context

class QualityAssuranceTesterAgent:
    async def generate_test_cases(self, context: Context, api_key: str) -> Context:
        """Generates test cases."""
        prompt = (
            "You are a testing agent. Generate test cases. "
            "Consider 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)
        context.test_cases = test_cases
        context.add_conversation_entry("QA Tester", f"Test Cases:\n{test_cases}")
        return context

    async def run_tests(self, context: Context, api_key: str) -> Context:
        """Runs tests and reports results."""
        prompt = (
            "Run the test cases. Compare actual vs expected output. "
            "State discrepancies. If all 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)
        context.test_results = test_results
        context.add_conversation_entry("QA Tester", f"Test Results:\n{test_results}")
        return context

class DocumentationAgent:
    async def generate_documentation(self, context: Context, api_key: str) -> Context:
        """Generates documentation, including a --help message."""
        prompt = (
            "Generate clear and concise documentation. "
            "Include 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)
        context.documentation = documentation
        context.add_conversation_entry("Documentation Agent", f"Documentation:\n{documentation}")
        return context

# -------------------- Agent Dispatcher (New) --------------------

class AgentDispatcher:
    def __init__(self, log_queue: queue.Queue, human_in_the_loop_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_input_queue = human_input_queue
        self.agents = {
            "prompt_optimizer": PromptOptimizerAgent(),
            "orchestrator": OrchestratorAgent(log_queue, human_in_the_loop_event, human_input_queue),
            "coder": CoderAgent(),
            "code_reviewer": CodeReviewerAgent(),
            "qa_tester": QualityAssuranceTesterAgent(),
            "documentation_agent": DocumentationAgent(),
        }

    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)
        elif agent_name == "orchestrator":
            context = await agent.generate_plan(context, api_key)  #Removed human_feedback
        elif agent_name == "coder":
            context = await agent.generate_code(context, api_key, **kwargs)
        elif agent_name == "code_reviewer":
            context = await agent.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)
            elif kwargs.get("run_tests", False):
                context = await agent.run_tests(context, api_key)
        elif agent_name == "documentation_agent":
            context = await agent.generate_documentation(context, api_key)
        else:
          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."""
        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",[])  ]:
            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() :
            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:
    """
    Conducts the multi-agent conversation using the AgentDispatcher.
    """
    context = Context(original_task=task_message)
    dispatcher = AgentDispatcher(log_queue, human_in_the_loop_event, human_input_queue)

    next_agent = await dispatcher.determine_next_agent(context, api_key)
    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)
            else:
              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")
        else:
          context = await dispatcher.dispatch(next_agent, context, api_key) # Call the agent

        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;

    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) -> Generator[str, None, None]:
    """
    Wraps the conversation and yields log messages. Handles human input.
    """
    log_q: queue.Queue = queue.Queue()

    def run_conversation() -> None:
        asyncio.run(multi_agent_conversation(task_message, log_q, api_key, human_in_the_loop_event, human_input_queue))

    thread = threading.Thread(target=run_conversation)
    thread.start()

    final_result = None
    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 "Conversation complete."
            else:
                yield msg
        except queue.Empty:
            continue

    thread.join()
    if final_result:
        conv_text = "\n=== Conversation ===\n"
        for entry in final_result:
            conv_text += f"[{entry['agent']}]: {entry['message']}\n\n"
        yield conv_text

def get_human_feedback(placeholder_text):
    """Gets human input using a Gradio Textbox."""
    with gr.Blocks() as human_feedback_interface:
        with gr.Row():
            human_input = gr.Textbox(lines=4, label="Human Feedback", placeholder=placeholder_text) #Removed placeholder
        with gr.Row():
            submit_button = gr.Button("Submit Feedback")

        feedback_queue = queue.Queue()

        def submit_feedback(input_text):
            feedback_queue.put(input_text)
            return ""

        submit_button.click(submit_feedback, inputs=human_input, outputs=human_input)
        human_feedback_interface.load(None, [], [])  # Keep interface alive

    return human_feedback_interface, feedback_queue

# -------------------- 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."""
    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() #For receiving the feedback request

    yield from process_conversation_generator(message, openai_api_key, human_in_the_loop_event, human_input_queue)

    while human_in_the_loop_event.is_set():
        yield "Waiting for human feedback..."
        try:
          feedback_request = human_input_queue.get(timeout=0.1) #Non-blocking, check for feedback request

          human_interface, feedback_queue = get_human_feedback(feedback_request)

          #This is a hacky but currently only working way to make this work with gradio
          yield gr.Textbox.update(visible=False), gr.update(visible=True)
          human_feedback = feedback_queue.get(timeout=300)  # Wait for up to 5 minutes
          human_input_queue.put(human_feedback) #Put feedback where Orchestrator can find it.
          human_in_the_loop_event.clear()
          yield gr.Textbox.update(visible=True), human_interface.close() #Hide human input box
          yield from process_conversation_generator(message, openai_api_key, human_in_the_loop_event, human_input_queue)

        except queue.Empty: #If we get here, there was NO human feedback request, so skip.
          continue #Go back to the top of the while loop

# -------------------- Launch the Chatbot --------------------

# Create the main chat interface
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")],
    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.
        """
)

#Need a 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()