Update app.py

app.py

@@ -1,100 +1,110 @@
 import gradio as gr
 import os
 import matplotlib.pyplot as plt
-from io import BytesIO
-from langgraph.graph import StateGraph, START, END
-from langgraph.graph.message import MessagesState
+import pandas as pd
+from langgraph.graph import StateGraph
 from langgraph.prebuilt import create_react_agent
 from langgraph.types import Command
 from langchain_core.messages import HumanMessage
 from langchain_anthropic import ChatAnthropic
 
-# Set API Key
+# Set the API key from Hugging Face secrets
 os.environ["ANTHROPIC_API_KEY"] = os.getenv("ANTHROPIC_API_KEY")
 
-# Claude LLM
+# Claude 3.5 Sonnet
 llm = ChatAnthropic(model="claude-3-5-sonnet-latest")
 
-# System prompt generator
+# Create the system prompt
 def make_system_prompt(suffix: str) -> str:
     return (
         "You are a helpful AI assistant, collaborating with other assistants."
         " Use the provided tools to progress towards answering the question."
-        " If you are unable to fully answer, that's OK, another assistant with different tools "
+        " If you are unable to fully answer, that's OK—another assistant with different tools"
         " will help where you left off. Execute what you can to make progress."
         " If you or any of the other assistants have the final answer or deliverable,"
         " prefix your response with FINAL ANSWER so the team knows to stop."
         f"\n{suffix}"
     )
 
-# Agent 1: Researcher
-def research_node(state: MessagesState) -> Command[str]:
+# Workflow node: research
+def research_node(state):
     agent = create_react_agent(
         llm,
         tools=[],
         state_modifier=make_system_prompt("You can only do research.")
     )
     result = agent.invoke(state)
-    content = result["messages"][-1].content
-    result["messages"][-1] = HumanMessage(content=content, name="researcher")
-    goto = END if "FINAL ANSWER" in content else "chart_generator"
+    goto = "chart_generator" if "FINAL ANSWER" not in result["messages"][-1].content else "__end__"
+    result["messages"][-1] = HumanMessage(
+        content=result["messages"][-1].content,
+        name="researcher"
+    )
     return Command(update={"messages": result["messages"]}, goto=goto)
 
-# Agent 2: Chart generator
-def chart_node(state: MessagesState) -> Command[str]:
+# Workflow node: chart generation
+def chart_node(state):
     agent = create_react_agent(
         llm,
         tools=[],
         state_modifier=make_system_prompt("You can only generate charts.")
     )
     result = agent.invoke(state)
-    content = result["messages"][-1].content
-    result["messages"][-1] = HumanMessage(content=content, name="chart_generator")
-    goto = END if "FINAL ANSWER" in content else "researcher"
-    return Command(update={"messages": result["messages"]}, goto=goto)
+    result["messages"][-1] = HumanMessage(
+        content=result["messages"][-1].content,
+        name="chart_generator"
+    )
+    return Command(update={"messages": result["messages"]}, goto="__end__")
 
-# LangGraph flow setup
-workflow = StateGraph(MessagesState)
+# LangGraph setup
+workflow = StateGraph(dict)
 workflow.add_node("researcher", research_node)
 workflow.add_node("chart_generator", chart_node)
-workflow.add_edge(START, "researcher")
+workflow.set_entry_point("researcher")
+workflow.set_finish_point("__end__")
 workflow.add_edge("researcher", "chart_generator")
-workflow.add_edge("chart_generator", END)
 graph = workflow.compile()
 
-# Helper: simulate chart from output
-def generate_dummy_chart() -> BytesIO:
-    fig, ax = plt.subplots()
-    ax.plot([2019, 2020, 2021, 2022, 2023], [21.4, 20.9, 22.7, 25.5, 27.6])
-    ax.set_title("Simulated GDP Growth")
-    ax.set_xlabel("Year")
-    ax.set_ylabel("Trillions USD")
-    buf = BytesIO()
-    plt.savefig(buf, format='png')
-    buf.seek(0)
-    return buf
-
-# Run LangGraph
-def run_langgraph(user_input):
+# LangGraph runner
+def run_langgraph(input_text):
     try:
-        events = graph.stream({"messages": [("user", user_input)]}, {"recursion_limit": 150})
-        final_output = ""
+        events = graph.stream({"messages": [("user", input_text)]})
+        output = []
         for event in events:
-            final_output = event["messages"][-1].content
-        if "chart" in user_input.lower():
-            return generate_dummy_chart()
-        return final_output or "No output generated"
+            output.append(event)
+
+        final_response = output[-1]["messages"][-1].content
+
+        if "FINAL ANSWER" in final_response:
+            # Simulated chart creation from dummy data
+            years = [2020, 2021, 2022, 2023, 2024]
+            gdp = [21.4, 22.0, 23.1, 24.8, 26.2]
+
+            plt.figure()
+            plt.plot(years, gdp, marker="o")
+            plt.title("USA GDP Over Last 5 Years")
+            plt.xlabel("Year")
+            plt.ylabel("GDP in Trillions USD")
+            plt.grid(True)
+            plt.tight_layout()
+            plt.savefig("gdp_chart.png")
+
+            return "Chart generated based on FINAL ANSWER.", "gdp_chart.png"
+        else:
+            return final_response, None
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"Error: {str(e)}", None
+
+# Gradio interface
+def process_input(user_input):
+    return run_langgraph(user_input)
 
-# Gradio UI
 interface = gr.Interface(
-    fn=run_langgraph,
-    inputs=gr.Textbox(label="user_input", placeholder="e.g. Get GDP data for the USA over the past 5 years and create a chart."),
-    outputs=gr.Image(type="pil", label="output"),
+    fn=process_input,
+    inputs=gr.Textbox(label="Enter your research task"),
+    outputs=[gr.Textbox(label="Output"), gr.Image(type="filepath", label="Chart")],
     title="LangGraph Research Automation",
     description="Enter a research prompt and view chart output when applicable."
 )
 
 if __name__ == "__main__":
-    interface.launch(server_name="0.0.0.0", server_port=7860)
+    interface.launch()