File size: 7,504 Bytes
4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 4878ce8 a9ce647 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 |
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() |