Update app.py

app.py

@@ -1,90 +1,89 @@
 import gradio as gr
 import os
 import matplotlib.pyplot as plt
-from langgraph.graph import StateGraph
-from langgraph.prebuilt import MessagesState, create_react_agent
-from langgraph.types import Command
+import pandas as pd
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt.tool_executor import ToolExecutor
+from langgraph.prebuilt.react import create_react_plan_and_execute
+from langgraph.checkpoint.sqlite import SqliteSaver
+from langgraph.graph.message import add_messages
+
 from langchain_core.messages import HumanMessage
 from langchain_anthropic import ChatAnthropic
 
-# Set API key (ensure you add this as a secret in HF Spaces)
+# Load API Key securely
 os.environ["ANTHROPIC_API_KEY"] = os.getenv("ANTHROPIC_API_KEY")
 
-# Load Claude 3.5 Sonnet model
+# Define the LLM (Claude 3.5 Sonnet)
 llm = ChatAnthropic(model="claude-3-5-sonnet-latest")
 
-# System prompt constructor
+# System prompt modifier
 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 "
-        "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.\n"
+        "You are a helpful AI assistant, collaborating with other assistants."
+        " Use the provided tools to progress toward answering the question."
+        " If you cannot fully answer, another assistant will continue where you left off."
+        " If you or the team has a complete answer, prefix it with FINAL ANSWER.\n"
         f"{suffix}"
     )
 
-# Research phase
-def research_node(state: MessagesState) -> Command[str]:
-    agent = create_react_agent(
-        llm,
-        tools=[],
-        system_message=make_system_prompt("You can only do research.")
-    )
-    result = agent.invoke(state)
-    goto = "chart_generator" if "FINAL ANSWER" not in result["messages"][-1].content else "__end__"
-    result["messages"][-1].name = "researcher"
-    return Command(update={"messages": result["messages"]}, goto=goto)
+# Node 1: Research assistant logic
+def research_node(state: dict) -> dict:
+    messages = state.get("messages", [])
+    prompt = make_system_prompt("You can only do research.")
+    executor = create_react_plan_and_execute(llm=llm, tools=[], system_prompt=prompt)
+    response = executor.invoke(messages)
+    messages.append(HumanMessage(content=response.content, name="researcher"))
+    next_node = "chart_generator" if "FINAL ANSWER" not in response.content else END
+    return {"messages": messages, "next": next_node}
 
-# Chart generation phase
-def chart_node(state: MessagesState) -> Command[str]:
-    agent = create_react_agent(
-        llm,
-        tools=[],
-        system_message=make_system_prompt("You can only generate charts.")
-    )
-    result = agent.invoke(state)
-    result["messages"][-1].name = "chart_generator"
-    return Command(update={"messages": result["messages"]}, goto="__end__")
+# Node 2: Chart assistant logic
+def chart_node(state: dict) -> dict:
+    messages = state.get("messages", [])
+    prompt = make_system_prompt("You can only generate charts.")
+    executor = create_react_plan_and_execute(llm=llm, tools=[], system_prompt=prompt)
+    response = executor.invoke(messages)
+    messages.append(HumanMessage(content=response.content, name="chart_generator"))
+    return {"messages": messages, "next": END}
 
-# Build LangGraph
-workflow = StateGraph(MessagesState)
+# Define LangGraph flow
+workflow = StateGraph(dict)
 workflow.add_node("researcher", research_node)
 workflow.add_node("chart_generator", chart_node)
 workflow.set_entry_point("researcher")
-workflow.set_finish_point("__end__")
-workflow.add_edge("researcher", "chart_generator")
+workflow.set_finish_point(END)
+workflow.add_conditional_edges("researcher", lambda x: x["next"])
+workflow.add_edge("chart_generator", END)
+
 graph = workflow.compile()
 
-# LangGraph runner
-def run_langgraph(user_input):
+# Function to run the graph and optionally return chart
+def run_langgraph(input_text):
     try:
-        events = graph.stream({"messages": [HumanMessage(content=user_input)]})
-        outputs = list(events)
-        final_message = outputs[-1]["messages"][-1].content
+        events = graph.stream({"messages": [HumanMessage(content=input_text)]})
+        output = list(events)[-1]
+        final_content = output["messages"][-1].content
 
-        if "FINAL ANSWER" in final_message:
-            # Simulated chart (you can later parse dynamic values if needed)
+        if "FINAL ANSWER" in final_content:
+            # Example static chart
             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.plot(years, gdp, marker="o")
             plt.title("USA GDP Over Last 5 Years")
             plt.xlabel("Year")
-            plt.ylabel("GDP in Trillions USD")
+            plt.ylabel("GDP in Trillions")
             plt.grid(True)
             plt.tight_layout()
             plt.savefig("gdp_chart.png")
-
-            return "Chart generated based on FINAL ANSWER.", "gdp_chart.png"
+            return "Chart generated based on FINAL ANSWER", "gdp_chart.png"
         else:
-            return final_message, None
+            return final_content, None
+
     except Exception as e:
         return f"Error: {str(e)}", None
 
-# Gradio UI
+# Gradio Interface
 def process_input(user_input):
     return run_langgraph(user_input)