Create app.py

app.py

@@ -0,0 +1,135 @@
+import warnings
+warnings.filterwarnings("ignore", message=".*TqdmWarning.*")
+from dotenv import load_dotenv
+
+_ = load_dotenv()
+
+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 os
+import gradio as gr
+
+# Define agent state class
+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]
+
+# Define queries class
+class Queries(BaseModel):
+    queries: List[str]
+
+# Writer Agent Class
+class Ewriter():
+    def __init__(self):
+        self.model = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
+        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."
+        self.tavily = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])
+        
+        # Initialize 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)  # Ensure reflect is not a dead-end
+        
+        self.graph = builder.compile()
+
+    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 occurred 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 occurred 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 occurred 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['draft'])])
+            return {"critique": response.content, "lnode": "reflect", "count": 1}
+        except Exception as e:
+            return {"critique": f"Error occurred in reflection: {str(e)}", "lnode": "reflect", "count": 0}
+
+# Gradio UI
+class WriterGui():
+    def __init__(self, graph):
+        self.graph = graph
+        self.demo = self.create_interface()
+    
+    def run_agent(self, topic, revision_number, max_revisions):
+        config = {'task': topic, 'max_revisions': max_revisions, 'revision_number': revision_number, 'lnode': "", 'count': 0}
+        response = self.graph.invoke(config)
+        return response["draft"], response["lnode"], response["count"], response.get("critique", ""), response.get("research_queries", [])
+
+    def continue_agent(self, topic, revision_number, max_revisions, last_node, current_draft):
+        config = {'task': topic, 'max_revisions': max_revisions, 'revision_number': revision_number, 'lnode': last_node, 'draft': current_draft, 'count': 0}
+        response = self.graph.invoke(config)
+        return response["draft"], response["lnode"], response["count"], response.get("critique", ""), response.get("research_queries", [])
+    
+    def create_interface(self):
+        with gr.Blocks() as demo:
+            with gr.Tabs():
+                with gr.Tab("Agent"):
+                    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)
+                    thread = gr.Textbox(label="Thread", 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():
+                        gr.Markdown("**Manage Agent**")
+                    with gr.Row():
+                        output_text = gr.Textbox(label="Live Agent Output", 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=[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=[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 the App
+MultiAgent = Ewriter()
+app = WriterGui(MultiAgent.graph)
+app.launch()