Update app.py

app.py

@@ -1,5 +1,7 @@
 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
 from langgraph.prebuilt import create_react_agent
@@ -7,13 +9,13 @@ from langgraph.types import Command
 from langchain_core.messages import HumanMessage
 from langchain_anthropic import ChatAnthropic
 
-# Set Anthropic API key from Hugging Face secret environment
+# Set API Key
 os.environ["ANTHROPIC_API_KEY"] = os.getenv("ANTHROPIC_API_KEY")
 
-# Initialize Claude 3.5 Sonnet LLM
+# Claude LLM
 llm = ChatAnthropic(model="claude-3-5-sonnet-latest")
 
-# System prompt template function
+# System prompt generator
 def make_system_prompt(suffix: str) -> str:
     return (
         "You are a helpful AI assistant, collaborating with other assistants."
@@ -25,7 +27,7 @@ def make_system_prompt(suffix: str) -> str:
         f"\n{suffix}"
     )
 
-# Research agent logic
+# Agent 1: Researcher
 def research_node(state: MessagesState) -> Command[str]:
     agent = create_react_agent(
         llm,
@@ -33,13 +35,12 @@ def research_node(state: MessagesState) -> Command[str]:
         state_modifier=make_system_prompt("You can only do research.")
     )
     result = agent.invoke(state)
-    goto = END if "FINAL ANSWER" in result["messages"][-1].content else "chart_generator"
-    result["messages"][-1] = HumanMessage(
-        content=result["messages"][-1].content, name="researcher"
-    )
+    content = result["messages"][-1].content
+    result["messages"][-1] = HumanMessage(content=content, name="researcher")
+    goto = END if "FINAL ANSWER" in content else "chart_generator"
     return Command(update={"messages": result["messages"]}, goto=goto)
 
-# Chart generator logic
+# Agent 2: Chart generator
 def chart_node(state: MessagesState) -> Command[str]:
     agent = create_react_agent(
         llm,
@@ -47,13 +48,12 @@ def chart_node(state: MessagesState) -> Command[str]:
         state_modifier=make_system_prompt("You can only generate charts.")
     )
     result = agent.invoke(state)
-    goto = END if "FINAL ANSWER" in result["messages"][-1].content else "researcher"
-    result["messages"][-1] = HumanMessage(
-        content=result["messages"][-1].content, name="chart_generator"
-    )
+    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)
 
-# Build the LangGraph workflow
+# LangGraph flow setup
 workflow = StateGraph(MessagesState)
 workflow.add_node("researcher", research_node)
 workflow.add_node("chart_generator", chart_node)
@@ -62,28 +62,38 @@ workflow.add_edge("researcher", "chart_generator")
 workflow.add_edge("chart_generator", END)
 graph = workflow.compile()
 
-# Function to execute LangGraph flow
-def run_langgraph(user_input):
-    events = graph.stream(
-        {"messages": [("user", user_input)]},
-        {"recursion_limit": 150}
-    )
-    output = []
-    for event in events:
-        output.append(event)
-    return output[-1]["messages"][-1].content if output else "No output generated"
+# 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
 
-# Gradio interface logic
-def process_input(user_input):
-    return run_langgraph(user_input)
+# Run LangGraph
+def run_langgraph(user_input):
+    try:
+        events = graph.stream({"messages": [("user", user_input)]}, {"recursion_limit": 150})
+        final_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"
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-# Launch Gradio app
+# Gradio UI
 interface = gr.Interface(
-    fn=process_input,
-    inputs="text",
-    outputs="text",
+    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"),
     title="LangGraph Research Automation",
-    description="Enter your research task (e.g., 'Get GDP data for the USA over the past 5 years and create a chart.')"
+    description="Enter a research prompt and view chart output when applicable."
 )
 
 if __name__ == "__main__":