Update app.py

app.py

@@ -3,67 +3,102 @@ import os
 import matplotlib.pyplot as plt
 import pandas as pd
 from langgraph.graph import StateGraph
-from langgraph.prebuilt import create_agent_executor
-from langchain_core.messages import HumanMessage
+from langgraph_core.state import MessagesState
+from langchain_core.messages import HumanMessage, AIMessage
 from langchain_anthropic import ChatAnthropic
+import warnings
+warnings.filterwarnings("ignore")
 
-# Fallback Command class (since langgraph.types is not available in v0.0.41)
+# Define a Command class for langgraph 0.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")
+# Set API key (ensure you add this as a secret in HF Spaces)
+api_key = os.getenv("ANTHROPIC_API_KEY")
+if not api_key:
+    raise ValueError("ANTHROPIC_API_KEY environment variable not set")
 
-# Claude 3.5 Sonnet model - remove any default parameters that might be causing issues
-llm = ChatAnthropic(
-    model="claude-3-5-sonnet-20240229",
-    # No additional parameters that might cause conflicts
-)
+# Load Claude 3.5 Sonnet model
+# Using a direct approach to avoid proxies issue
+try:
+    # Explicitly create with minimal parameters
+    llm = ChatAnthropic(api_key=api_key, model="claude-3-5-sonnet-20240229")
+except Exception as e:
+    print(f"Error initializing ChatAnthropic: {e}")
+    # Fallback initialization if needed
+    import anthropic
+    client = anthropic.Anthropic(api_key=api_key)
+    llm = ChatAnthropic(client=client, model="claude-3-5-sonnet-20240229")
 
-# Utility to build system prompts
+# System prompt constructor
 def make_system_prompt(suffix: str) -> str:
     return (
-        "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
+        "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"
+        f"{suffix}"
     )
 
-# Research agent node
+# Research phase
 def research_node(state):
+    # Create a custom research agent using langgraph 0.0.41 compatible approach
+    from langgraph.prebuilt import create_agent_executor
     agent = create_agent_executor(
         llm,
         tools=[],
-        system_message=make_system_prompt("You can only perform research.")
+        system_message=make_system_prompt("You can only do research.")
     )
+    
+    # Process the current state
     result = agent.invoke(state)
-    last_msg = result["messages"][-1]
+    
+    # Check if we have a final answer
+    last_message = result["messages"][-1]
+    content = last_message.content if hasattr(last_message, "content") else last_message
+    
     # Determine next step
-    goto = "chart_generator" if "FINAL ANSWER" not in last_msg.content else "__end__"
-    result["messages"][-1] = HumanMessage(
-        content=last_msg.content,
-        name="researcher"
-    )
+    goto = "chart_generator" if "FINAL ANSWER" not in content else "__end__"
+    
+    # Create an AIMessage with the researcher name
+    if not isinstance(last_message, dict):
+        result["messages"][-1] = AIMessage(content=content, name="researcher")
+    else:
+        result["messages"][-1]["name"] = "researcher"
+    
     return Command(update={"messages": result["messages"]}, goto=goto)
 
-# Chart generation agent node
+# Chart generation phase
 def chart_node(state):
+    # Create a custom chart generator agent
+    from langgraph.prebuilt import create_agent_executor
     agent = create_agent_executor(
         llm,
         tools=[],
         system_message=make_system_prompt("You can only generate charts.")
     )
+    
+    # Process the current state
     result = agent.invoke(state)
-    result["messages"][-1] = HumanMessage(
-        content=result["messages"][-1].content,
-        name="chart_generator"
-    )
+    
+    # Add the chart_generator name
+    last_message = result["messages"][-1]
+    content = last_message.content if hasattr(last_message, "content") else last_message
+    
+    if not isinstance(last_message, dict):
+        result["messages"][-1] = AIMessage(content=content, name="chart_generator")
+    else:
+        result["messages"][-1]["name"] = "chart_generator"
+    
     return Command(update={"messages": result["messages"]}, goto="__end__")
 
-# LangGraph state setup
-workflow = StateGraph(dict)
+# Build LangGraph
+workflow = StateGraph(dict)  # Using dict for state in langgraph 0.0.41
 workflow.add_node("researcher", research_node)
 workflow.add_node("chart_generator", chart_node)
 workflow.set_entry_point("researcher")
@@ -72,17 +107,25 @@ workflow.add_edge("researcher", "chart_generator")
 graph = workflow.compile()
 
 # LangGraph runner
-def run_langgraph(input_text):
+def run_langgraph(user_input):
     try:
-        events = graph.stream({"messages": [HumanMessage(content=input_text)]})
-        output = list(events)
-        final_response = output[-1]["messages"][-1].content
-        if "FINAL ANSWER" in final_response:
-            # Dummy chart generation
+        # Create a human message
+        human_message = HumanMessage(content=user_input)
+        
+        # Stream the events
+        events = graph.stream({"messages": [human_message]})
+        outputs = list(events)
+        
+        # Get the final message
+        final_message = outputs[-1]["messages"][-1]
+        final_content = final_message.content if hasattr(final_message, "content") else final_message
+        
+        if isinstance(final_content, str) and "FINAL ANSWER" in final_content:
+            # Simulated chart (you can later parse dynamic values if needed)
             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")
@@ -91,11 +134,17 @@ def run_langgraph(input_text):
             plt.savefig("gdp_chart.png")
             return "Chart generated based on FINAL ANSWER.", "gdp_chart.png"
         else:
-            return final_response, None
+            if isinstance(final_content, str):
+                return final_content, None
+            else:
+                return str(final_content), None
     except Exception as e:
+        print(f"Error in run_langgraph: {e}")
+        import traceback
+        traceback.print_exc()
         return f"Error: {str(e)}", None
 
-# Gradio interface
+# Gradio UI
 def process_input(user_input):
     return run_langgraph(user_input)