Update app.py

app.py

@@ -3,47 +3,65 @@ import os
 import matplotlib.pyplot as plt
 import pandas as pd
 from langgraph.graph import StateGraph
-from langgraph.types import Command
-from langchain_core.messages import HumanMessage, AIMessage
+from langgraph.prebuilt import create_agent_executor
+from langchain_core.messages import HumanMessage
 from langchain_anthropic import ChatAnthropic
 
-# Set API key from environment variable (Hugging Face secret)
+# Fallback Command class (since langgraph.types is not available in v0.0.41)
+class Command:
+    def __init__(self, update=None, next=None, goto=None):
+        self.update = update or {}
+        self.next = next
+        self.goto = goto
+
+# Set the API key from Hugging Face secrets
 os.environ["ANTHROPIC_API_KEY"] = os.getenv("ANTHROPIC_API_KEY")
-llm = ChatAnthropic(model="claude-3-5-sonnet-latest")
 
-# System message function
-def make_system_prompt(task: str) -> str:
+# Claude 3.5 Sonnet model
+llm = ChatAnthropic(model="claude-3-5-sonnet-20240229")
+
+# Utility to build system prompts
+def make_system_prompt(suffix: str) -> str:
     return (
-        f"You are an AI assistant. Your task is: {task}. "
-        "If you have the final answer, prefix it with 'FINAL ANSWER'."
+        "You are a helpful AI assistant collaborating with others."
+        " Use your tools to assist. If you can't complete a task, leave it to another agent."
+        " Prefix final output with 'FINAL ANSWER' to signal completion.\n" + suffix
     )
 
-# Node: research
-def research_node(state: dict) -> Command:
-    user_msg = state["messages"][-1]
-    task = user_msg.content
-    system_prompt = make_system_prompt("Research only")
-
-    response = llm.invoke([HumanMessage(content=system_prompt), user_msg])
-    content = response.content
-
-    state["messages"].append(AIMessage(content=content, name="researcher"))
-    goto = "chart_generator" if "FINAL ANSWER" not in content else "__end__"
-    return Command(update={"messages": state["messages"]}, goto=goto)
+# Research agent node
+def research_node(state):
+    agent = create_agent_executor(
+        llm,
+        tools=[],
+        system_message=make_system_prompt("You can only perform research.")
+    )
+    result = agent.invoke(state)
+    last_msg = result["messages"][-1]
 
-# Node: chart generation
-def chart_node(state: dict) -> Command:
-    user_msg = state["messages"][-1]
-    task = user_msg.content
-    system_prompt = make_system_prompt("Generate chart only")
+    # Determine next step
+    goto = "chart_generator" if "FINAL ANSWER" not in last_msg.content else "__end__"
 
-    response = llm.invoke([HumanMessage(content=system_prompt), user_msg])
-    content = response.content
+    result["messages"][-1] = HumanMessage(
+        content=last_msg.content,
+        name="researcher"
+    )
+    return Command(update={"messages": result["messages"]}, goto=goto)
 
-    state["messages"].append(AIMessage(content=content, name="chart_generator"))
-    return Command(update={"messages": state["messages"]}, goto="__end__")
+# Chart generation agent node
+def chart_node(state):
+    agent = create_agent_executor(
+        llm,
+        tools=[],
+        system_message=make_system_prompt("You can only generate charts.")
+    )
+    result = agent.invoke(state)
+    result["messages"][-1] = HumanMessage(
+        content=result["messages"][-1].content,
+        name="chart_generator"
+    )
+    return Command(update={"messages": result["messages"]}, goto="__end__")
 
-# Build LangGraph
+# LangGraph state setup
 workflow = StateGraph(dict)
 workflow.add_node("researcher", research_node)
 workflow.add_node("chart_generator", chart_node)
@@ -52,18 +70,18 @@ workflow.set_finish_point("__end__")
 workflow.add_edge("researcher", "chart_generator")
 graph = workflow.compile()
 
-# Stream function
+# LangGraph runner
 def run_langgraph(input_text):
     try:
-        events = graph.stream({"messages": [HumanMessage(content=input_text)]})
-        final_state = list(events)[-1]
-        messages = final_state["messages"]
-        final_output = messages[-1].content
+        events = graph.stream({"messages": [("user", input_text)]})
+        output = list(events)
+        final_response = output[-1]["messages"][-1].content
 
-        if "FINAL ANSWER" in final_output:
-            # Dummy chart
+        if "FINAL ANSWER" in final_response:
+            # Dummy chart generation
             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")
@@ -71,15 +89,15 @@ def run_langgraph(input_text):
             plt.ylabel("GDP in Trillions USD")
             plt.grid(True)
             plt.tight_layout()
-            chart_path = "gdp_chart.png"
-            plt.savefig(chart_path)
-            return final_output, chart_path
+            plt.savefig("gdp_chart.png")
+
+            return "Chart generated based on FINAL ANSWER.", "gdp_chart.png"
         else:
-            return final_output, None
+            return final_response, 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)