app.py

from typing import Annotated, TypedDict, List, Dict, Any, Optional
from typing_extensions import TypedDict
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langchain_google_genai import ChatGoogleGenerativeAI
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.graph.message import add_messages
from pydantic import BaseModel, Field
import gradio as gr
import uuid
import asyncio
from datetime import datetime
import os

class EvaluatorOutput(BaseModel):
    feedback: str = Field(description="Feedback on the assistant's response")
    success_criteria_met: bool = Field(description="Whether the success criteria have been met")
    user_input_needed: bool = Field(description="True if more input is needed from the user, or clarifications, or the assistant is stuck")


class State(TypedDict):
    messages: Annotated[List[Any], add_messages]
    success_criteria: str
    feedback_on_work: Optional[str]
    success_criteria_met: bool
    user_input_needed: bool


tools = []

def create_llm(api_key: str, provider: str = "openai"):
    """Create LLM instance based on provider and API key"""
    if provider == "openai":
        return ChatOpenAI(
            model="gpt-4o-mini",
            openai_api_key=api_key,
            temperature=0.7
        )
    elif provider == "gemini":
        return ChatGoogleGenerativeAI(
            model="gemini-pro",
            google_api_key=api_key,
            temperature=0.7
        )
    else:
        raise ValueError(f"Unsupported provider: {provider}")


def worker(state: State, api_key: str, provider: str) -> Dict[str, Any]:
    try:
        worker_llm = create_llm(api_key, provider)
        worker_llm_with_tools = worker_llm.bind_tools(tools)
    except Exception as e:
        return {
            "messages": [AIMessage(content=f"Error setting up LLM: {str(e)}. Please check your API key.")],
        }
    
    system_message = f"""You are TaskMaster AI, a powerful assistant that can use tools to complete tasks.
You keep working on a task until either you have a question or clarification for the user, or the success criteria is met.
You have access to various tools including file operations and more.
The current date and time is {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}

This is the success criteria:
{state['success_criteria']}
You should reply either with a question for the user about this assignment, or with your final response.
If you have a question for the user, you need to reply by clearly stating your question. An example might be:

Question: please clarify whether you want a summary or a detailed answer

If you've finished, reply with the final answer, and don't ask a question; simply reply with the answer.
"""
    
    if state.get("feedback_on_work"):
        system_message += f"""
Previously you thought you completed the assignment, but your reply was rejected because the success criteria was not met.
Here is the feedback on why this was rejected:
{state['feedback_on_work']}
With this feedback, please continue the assignment, ensuring that you meet the success criteria or have a question for the user."""

    found_system_message = False
    messages = state["messages"]
    for message in messages:
        if isinstance(message, SystemMessage):
            message.content = system_message
            found_system_message = True
    
    if not found_system_message:
        messages = [SystemMessage(content=system_message)] + messages
    try:
        response = worker_llm_with_tools.invoke(messages)
    except Exception as e:
        response = AIMessage(content=f"Error during processing: {str(e)}")
    
    return {
        "messages": [response],
    }

def worker_router(state: State) -> str:
    last_message = state["messages"][-1]
    
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        return "tools"
    else:
        return "evaluator"

def format_conversation(messages: List[Any]) -> str:
    conversation = "Conversation history:\n\n"
    for message in messages:
        if isinstance(message, HumanMessage):
            conversation += f"User: {message.content}\n"
        elif isinstance(message, AIMessage):
            text = message.content or "[Tools use]"
            conversation += f"Assistant: {text}\n"
    return conversation

def evaluator(state: State, api_key: str, provider: str) -> State:
    try:
        evaluator_llm = create_llm(api_key, provider)
        evaluator_llm_with_output = evaluator_llm.with_structured_output(EvaluatorOutput)
    except Exception as e:
        return {
            "messages": [AIMessage(content=f"Error setting up evaluator: {str(e)}")],
            "feedback_on_work": f"Error: {str(e)}",
            "success_criteria_met": False,
            "user_input_needed": True
        }
    
    last_response = state["messages"][-1].content

    system_message = """You are an evaluator that determines if a task has been completed successfully by an Assistant.
Assess the Assistant's last response based on the given criteria. Respond with your feedback, and with your decision on whether the success criteria has been met,
and whether more input is needed from the user."""
    
    user_message = f"""You are evaluating a conversation between the User and Assistant. You decide what action to take based on the last response from the Assistant.

The entire conversation with the assistant, with the user's original request and all replies, is:
{format_conversation(state['messages'])}

The success criteria for this assignment is:
{state['success_criteria']}

And the final response from the Assistant that you are evaluating is:
{last_response}

Respond with your feedback, and decide if the success criteria is met by this response.
Also, decide if more user input is required, either because the assistant has a question, needs clarification, or seems to be stuck and unable to answer without help.

The Assistant has access to various tools. If the Assistant says they have performed an action (like writing a file, browsing the web, etc.), then you can assume they have done so.
Overall you should give the Assistant the benefit of the doubt if they say they've done something. But you should reject if you feel that more work should go into this.
"""
    if state["feedback_on_work"]:
        user_message += f"Also, note that in a prior attempt from the Assistant, you provided this feedback: {state['feedback_on_work']}\n"
        user_message += "If you're seeing the Assistant repeating the same mistakes, then consider responding that user input is required."
    
    evaluator_messages = [SystemMessage(content=system_message), HumanMessage(content=user_message)]

    try:
        eval_result = evaluator_llm_with_output.invoke(evaluator_messages)
        new_state = {
            "messages": [AIMessage(content=f"Evaluator Feedback on this answer: {eval_result.feedback}")],
            "feedback_on_work": eval_result.feedback,
            "success_criteria_met": eval_result.success_criteria_met,
            "user_input_needed": eval_result.user_input_needed
        }
    except Exception as e:
        new_state = {
            "messages": [AIMessage(content=f"Error during evaluation: {str(e)}")],
            "feedback_on_work": f"Error: {str(e)}",
            "success_criteria_met": False,
            "user_input_needed": True
        }
    
    return new_state

def route_based_on_evaluation(state: State) -> str:
    if state["success_criteria_met"] or state["user_input_needed"]:
        return "END"
    else:
        return "worker"

def make_thread_id() -> str:
    return str(uuid.uuid4())

async def process_message(message, success_criteria, api_key, provider, history, thread):
    if not api_key.strip():
        return history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": "Please enter your API key to continue."}
        ]
    
    config = {"configurable": {"thread_id": thread}}

    state = {
        "messages": [HumanMessage(content=message)],
        "success_criteria": success_criteria,
        "feedback_on_work": None,
        "success_criteria_met": False,
        "user_input_needed": False
    }
    
    try:
        graph_builder = StateGraph(State)
        graph_builder.add_node("worker", lambda s: worker(s, api_key, provider))
        graph_builder.add_node("tools", ToolNode(tools=tools))
        graph_builder.add_node("evaluator", lambda s: evaluator(s, api_key, provider))
        
        graph_builder.add_conditional_edges("worker", worker_router, {"tools": "tools", "evaluator": "evaluator"})
        graph_builder.add_edge("tools", "worker")
        graph_builder.add_conditional_edges("evaluator", route_based_on_evaluation, {"worker": "worker", "END": END})
        graph_builder.add_edge(START, "worker")
        
        memory = MemorySaver()
        graph = graph_builder.compile(checkpointer=memory)
        
        result = await graph.ainvoke(state, config=config)
        user = {"role": "user", "content": message}
        reply = {"role": "assistant", "content": result["messages"][-2].content}
        feedback = {"role": "assistant", "content": result["messages"][-1].content}
        return history + [user, reply, feedback]
    except Exception as e:
        return history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": f"Error: {str(e)}. Please check your API key and try again."}
        ]

async def reset():
    return "", "", "", "openai", None, make_thread_id()

with gr.Blocks(
    title="TaskMaster AI - Intelligent Task Completion",
    theme=gr.themes.Default(primary_hue="blue")
) as demo:
    gr.Markdown("""
    # TaskMaster AI
    
    **Intelligent Task Completion with Automatic Evaluation**
    
    A LangGraph-powered AI assistant that helps you complete tasks with automatic success criteria evaluation and feedback loops.
    
    ## How it works:
    1. **Enter your API key** (OpenAI or Google Gemini)
    2. **Choose your AI provider**
    3. **Describe your task and success criteria**
    4. **Watch TaskMaster AI work and evaluate itself!**
    
    ## Features:
    - **Multi-agent workflow** with worker and evaluator agents
    - **Automatic success criteria evaluation**
    - **Intelligent feedback loops**
    - **Your API key stays private** (never stored)
    
    ---
    """)
    
    thread = gr.State(make_thread_id())
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### API Configuration")
            api_key = gr.Textbox(
                label="API Key",
                placeholder="Enter your OpenAI or Google Gemini API key",
                type="password",
                show_label=True
            )
            provider = gr.Radio(
                choices=["openai", "gemini"],
                label="AI Provider",
                value="openai",
                show_label=True
            )
            gr.Markdown("""
            **Get API Keys:**
            - [OpenAI API](https://platform.openai.com/api-keys)
            - [Google Gemini API](https://makersuite.google.com/app/apikey)
            """)
        
        with gr.Column(scale=2):
            gr.Markdown("### TaskMaster Interface")
            chatbot = gr.Chatbot(label="TaskMaster AI", height=400, type="messages")
    
    with gr.Group():
        gr.Markdown("### Task Input")
        with gr.Row():
            message = gr.Textbox(
                show_label=False, 
                placeholder="Describe your task (e.g., 'Research the latest AI developments and create a summary')",
                lines=2
            )
        with gr.Row():
            success_criteria = gr.Textbox(
                show_label=False, 
                placeholder="Define success criteria (e.g., 'Provide 3 recent breakthroughs with sources and a concise summary')",
                lines=2
            )
    
    with gr.Row():
        reset_button = gr.Button("New Task", variant="stop")
        go_button = gr.Button("Start Task", variant="primary")
    
    # Event handlers
    message.submit(
        process_message, 
        [message, success_criteria, api_key, provider, chatbot, thread], 
        [chatbot]
    )
    success_criteria.submit(
        process_message, 
        [message, success_criteria, api_key, provider, chatbot, thread], 
        [chatbot]
    )
    go_button.click(
        process_message, 
        [message, success_criteria, api_key, provider, chatbot, thread], 
        [chatbot]
    )
    reset_button.click(
        reset, 
        [], 
        [message, success_criteria, api_key, provider, chatbot, thread]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)