app.py

import warnings
warnings.filterwarnings("ignore", message=".*TqdmWarning.*")

from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated, List
import operator
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_openai import ChatOpenAI
from pydantic import BaseModel
from tavily import TavilyClient
import gradio as gr


# ----------------------
# Agent State Definition
# ----------------------
class AgentState(TypedDict):
    task: str
    lnode: str
    plan: str
    research_queries: List[str]
    draft: str
    critique: str
    content: List[str]
    revision_number: int
    max_revisions: int
    count: Annotated[int, operator.add]


class Queries(BaseModel):
    queries: List[str]


# ----------------------
# Writer Agent
# ----------------------
class Ewriter:
    def __init__(self, openai_key: str, tavily_key: str):
        if not openai_key or not tavily_key:
            raise ValueError("⚠️ Both OpenAI and Tavily API keys must be provided.")

        # Initialize models with user-provided keys
        self.model = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, api_key=openai_key)
        self.tavily = TavilyClient(api_key=tavily_key)

        # Prompts
        self.PLAN_PROMPT = "You are an expert writer tasked with writing a high-level outline of a short 3-paragraph essay."
        self.RESEARCH_PROMPT = "Generate three research queries to help in writing an essay on the given topic."
        self.WRITER_PROMPT = "You are an essay assistant tasked with writing an excellent 3-paragraph essay."
        self.REFLECTION_PROMPT = "You are a teacher grading an essay. Provide critique and suggestions."

        # Build the workflow graph
        builder = StateGraph(AgentState)
        builder.add_node("planner", self.plan_node)
        builder.add_node("research", self.research_node)
        builder.add_node("generate", self.generation_node)
        builder.add_node("reflect", self.reflection_node)
        builder.set_entry_point("planner")
        builder.add_edge("planner", "research")
        builder.add_edge("research", "generate")
        builder.add_edge("generate", "reflect")
        builder.add_edge("reflect", END)

        self.graph = builder.compile()

    # ----------- Nodes -----------
    def plan_node(self, state: AgentState):
        try:
            response = self.model.invoke([SystemMessage(content=self.PLAN_PROMPT), HumanMessage(content=state['task'])])
            return {"plan": response.content, "lnode": "planner", "count": 1}
        except Exception as e:
            return {"plan": f"Error in planning: {str(e)}", "lnode": "planner", "count": 0}

    def research_node(self, state: AgentState):
        try:
            response = self.model.invoke([SystemMessage(content=self.RESEARCH_PROMPT), HumanMessage(content=state['task'])])
            return {"research_queries": response.content.split('\n'), "lnode": "research", "count": 1}
        except Exception as e:
            return {"research_queries": [f"Error in research: {str(e)}"], "lnode": "research", "count": 0}

    def generation_node(self, state: AgentState):
        try:
            response = self.model.invoke([SystemMessage(content=self.WRITER_PROMPT), HumanMessage(content=state['task'])])
            return {"draft": response.content, "lnode": "generate", "count": 1}
        except Exception as e:
            return {"draft": f"Error in generation: {str(e)}", "lnode": "generate", "count": 0}

    def reflection_node(self, state: AgentState):
        try:
            response = self.model.invoke([SystemMessage(content=self.REFLECTION_PROMPT), HumanMessage(content=state.get("draft", ""))])
            return {"critique": response.content, "lnode": "reflect", "count": 1}
        except Exception as e:
            return {"critique": f"Error in reflection: {str(e)}", "lnode": "reflect", "count": 0}


# ----------------------
# Gradio UI
# ----------------------
class WriterGui:
    def __init__(self):
        self.demo = self.create_interface()

    def run_agent(self, openai_key, tavily_key, topic, revision_number, max_revisions):
        try:
            agent = Ewriter(openai_key, tavily_key)
            config = {'task': topic, 'max_revisions': max_revisions, 'revision_number': revision_number, 'lnode': "", 'count': 0}
            response = agent.graph.invoke(config)
            return response.get("draft", ""), response.get("lnode", ""), response.get("count", 0), response.get("critique", ""), response.get("research_queries", [])
        except Exception as e:
            return f"❌ Error: {str(e)}", "", 0, "", []

    def continue_agent(self, openai_key, tavily_key, topic, revision_number, max_revisions, last_node, current_draft):
        try:
            agent = Ewriter(openai_key, tavily_key)
            config = {'task': topic, 'max_revisions': max_revisions, 'revision_number': revision_number, 'lnode': last_node, 'draft': current_draft, 'count': 0}
            response = agent.graph.invoke(config)
            return response.get("draft", ""), response.get("lnode", ""), response.get("count", 0), response.get("critique", ""), response.get("research_queries", [])
        except Exception as e:
            return f"❌ Error: {str(e)}", "", 0, "", []

    def create_interface(self):
        with gr.Blocks() as demo:
            with gr.Tabs():
                with gr.Tab("Agent"):
                    with gr.Row():
                        openai_input = gr.Textbox(label="🔑 OpenAI API Key", type="password", placeholder="Enter your OpenAI key")
                        tavily_input = gr.Textbox(label="🔑 Tavily API Key", type="password", placeholder="Enter your Tavily key")

                    topic_input = gr.Textbox(label="📘 Essay Topic")
                    last_node = gr.Textbox(label="Last Node", interactive=False)
                    next_node = gr.Textbox(label="Next Node", interactive=False)
                    draft_rev = gr.Textbox(label="Draft Revision", interactive=False)
                    count = gr.Textbox(label="Count", interactive=False)

                    generate_button = gr.Button("Generate Essay", variant="primary")
                    continue_button = gr.Button("Continue Essay")

                    with gr.Row():
                        output_text = gr.Textbox(label="Essay Draft", interactive=False)
                    with gr.Row():
                        critique_text = gr.Textbox(label="Critique", interactive=False)
                    with gr.Row():
                        research_text = gr.Textbox(label="Research Queries", interactive=False)

                    generate_button.click(
                        fn=self.run_agent,
                        inputs=[openai_input, tavily_input, topic_input, gr.State(0), gr.State(2)],
                        outputs=[output_text, last_node, next_node, critique_text, research_text]
                    )
                    continue_button.click(
                        fn=self.continue_agent,
                        inputs=[openai_input, tavily_input, topic_input, gr.State(0), gr.State(2), last_node, draft_rev],
                        outputs=[output_text, last_node, next_node, critique_text, research_text]
                    )
        return demo

    def launch(self):
        self.demo.launch(share=True)


# ----------------------
# Run App
# ----------------------
if __name__ == "__main__":
    app = WriterGui()
    app.launch()