| 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 |
|
|
|
|
| |
| |
| |
| 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] |
|
|
|
|
| |
| |
| |
| 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.") |
|
|
| |
| self.model = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, api_key=openai_key) |
| self.tavily = TavilyClient(api_key=tavily_key) |
|
|
| |
| 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." |
|
|
| |
| 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() |
|
|
| |
| 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} |
|
|
|
|
| |
| |
| |
| 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) |
|
|
|
|
| |
| |
| |
| if __name__ == "__main__": |
| app = WriterGui() |
| app.launch() |
