Update app.py

app.py

@@ -1,14 +1,18 @@
 import gradio as gr
 from langgraph.graph import StateGraph, MessagesState, START, END
 from langgraph.types import Command
-from langchain_core.messages import BaseMessage, HumanMessage
+from langchain_core.messages import HumanMessage
 from langgraph.prebuilt import create_react_agent
 from langchain_anthropic import ChatAnthropic
 import os
-os.environ["ANTHROPIC_API_KEY"] = os.getenv("ANTHROPIC_API_KEY")
+import matplotlib.pyplot as plt
+from io import BytesIO
+import base64
 
+# Load API Key
+os.environ["ANTHROPIC_API_KEY"] = os.getenv("ANTHROPIC_API_KEY")
 
-# Define your LangGraph agents and nodes
+# LangGraph setup
 llm = ChatAnthropic(model="claude-3-5-sonnet-latest")
 
 def make_system_prompt(suffix: str) -> str:
@@ -22,35 +26,28 @@ def make_system_prompt(suffix: str) -> str:
         f"\n{suffix}"
     )
 
-# Research agent and node
 def research_node(state: MessagesState) -> Command[str]:
     agent = create_react_agent(
         llm,
-        tools=[],  # Define your tools if needed
+        tools=[],
         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"
-    )
+    result["messages"][-1] = HumanMessage(content=result["messages"][-1].content, name="researcher")
     return Command(update={"messages": result["messages"]}, goto=goto)
 
-# Chart generator agent and node
 def chart_node(state: MessagesState) -> Command[str]:
     agent = create_react_agent(
         llm,
-        tools=[],  # Define your tools if needed
+        tools=[],
         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"
-    )
+    result["messages"][-1] = HumanMessage(content=result["messages"][-1].content, name="chart_generator")
     return Command(update={"messages": result["messages"]}, goto=goto)
 
-# Initialize the LangGraph workflow
 workflow = StateGraph(MessagesState)
 workflow.add_node("researcher", research_node)
 workflow.add_node("chart_generator", chart_node)
@@ -59,7 +56,32 @@ workflow.add_edge("researcher", "chart_generator")
 workflow.add_edge("chart_generator", END)
 graph = workflow.compile()
 
-# Define the function to run the graph
+def extract_chart_data(text):
+    """
+    Try to extract something like:
+    2018: 20
+    2019: 21.5
+    2020: 18
+    """
+    import re
+    matches = re.findall(r'(\d{4})\s*[:\-]?\s*\$?([\d\.]+)', text)
+    if not matches:
+        return None, None
+    years = [m[0] for m in matches]
+    values = [float(m[1]) for m in matches]
+    return years, values
+
+def generate_plot(years, values):
+    fig, ax = plt.subplots()
+    ax.bar(years, values)
+    ax.set_title("Generated Chart")
+    ax.set_xlabel("Year")
+    ax.set_ylabel("Value")
+    buf = BytesIO()
+    plt.savefig(buf, format="png")
+    buf.seek(0)
+    return buf
+
 def run_langgraph(user_input):
     events = graph.stream(
         {"messages": [("user", user_input)]},
@@ -73,20 +95,27 @@ def run_langgraph(user_input):
 
     return final_message or "No output generated"
 
-
-# Create Gradio interface
 def process_input(user_input):
-    result = run_langgraph(user_input)
-    return result
+    result_text = run_langgraph(user_input)
+    years, values = extract_chart_data(result_text)
+
+    if years and values:
+        chart = generate_plot(years, values)
+        return result_text, chart
+    else:
+        return result_text, None
 
 interface = gr.Interface(
     fn=process_input,
     inputs="text",
-    outputs="text",
+    outputs=[
+        gr.Textbox(label="Generated Response"),
+        gr.Image(type="pil", label="Generated Chart")
+    ],
     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 your research task (e.g., 'Get GDP data for the USA over the past 5 years and create a chart.')"
 )
 
-# Launch the Gradio interface
 if __name__ == "__main__":
     interface.launch()
+