Update app.py

app.py

@@ -1,8 +1,5 @@
 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
@@ -11,10 +8,12 @@ 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
+
+# ----------------------
+# Agent State Definition
+# ----------------------
 class AgentState(TypedDict):
     task: str
     lnode: str
@@ -27,21 +26,30 @@ class AgentState(TypedDict):
     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)
+
+# ----------------------
+# 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."
-        self.tavily = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])
-        
-        # Initialize Graph
+
+        # Build the workflow graph
         builder = StateGraph(AgentState)
         builder.add_node("planner", self.plan_node)
         builder.add_node("research", self.research_node)
@@ -51,85 +59,108 @@ class Ewriter():
         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
-        
+        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 occurred in planning: {str(e)}", "lnode": "planner", "count": 0}
+            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 occurred in research: {str(e)}", "lnode": "research", "count": 0}
+            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 occurred in generation: {str(e)}", "lnode": "generate", "count": 0}
-    
+            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['draft'])])
+            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 occurred in reflection: {str(e)}", "lnode": "reflect", "count": 0}
+            return {"critique": f"Error in reflection: {str(e)}", "lnode": "reflect", "count": 0}
 
+
+# ----------------------
 # Gradio UI
-class WriterGui():
-    def __init__(self, graph):
-        self.graph = graph
+# ----------------------
+class WriterGui:
+    def __init__(self):
         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 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"):
-                    topic_input = gr.Textbox(label="Essay Topic")
+                    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)
-                    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)
+                        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=[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])
-            
+
+                    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 the App
-MultiAgent = Ewriter()
-app = WriterGui(MultiAgent.graph)
-app.launch()
+
+# ----------------------
+# Run App
+# ----------------------
+if __name__ == "__main__":
+    app = WriterGui()
+    app.launch()