Update app.py

app.py

@@ -1,194 +1,165 @@
 import gradio as gr
 import os
 import matplotlib.pyplot as plt
-import pandas as pd
-from langgraph.graph import StateGraph
-from langchain_core.messages import HumanMessage, AIMessage
+from typing import TypedDict, Annotated, List
+from operator import add
+from langgraph.graph import StateGraph, END
+from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
 import warnings
 warnings.filterwarnings("ignore")
 
-# 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 API key (ensure you add this as a secret in HF Spaces)
+# API key is optional because we use a MockLLM below.
+# If/when you swap in a real Anthropic client, you can reinstate the check.
 api_key = os.getenv("ANTHROPIC_API_KEY")
-if not api_key:
-    raise ValueError("ANTHROPIC_API_KEY environment variable not set")
 
-# Mock LLM implementation that doesn't rely on any external HTTP clients
+
+# ---------------------------------------------------------------------------
+# Mock LLM — does not require network access or an API key.
+# ---------------------------------------------------------------------------
 class MockLLM:
     def __init__(self):
         self.model = "mock-model"
-    
-    def invoke(self, inputs):
-        if isinstance(inputs, dict) and "messages" in inputs:
-            # Process the messages
-            content = ""
-            for msg in inputs["messages"]:
-                if hasattr(msg, "content"):
-                    content += msg.content + "\n"
-                elif isinstance(msg, dict) and "content" in msg:
-                    content += msg["content"] + "\n"
-            
-            # For research queries
-            if "research" in content.lower():
-                return "Based on my research, here's what I found about your topic. FINAL ANSWER: This topic shows trends that would be interesting to visualize in a chart."
-            # For chart generation
-            else:
-                return "I've analyzed the data and created a chart visualization. FINAL ANSWER: The chart shows an upward trend from 2020 to 2024."
-        
-        elif isinstance(inputs, str):
-            if "research" in inputs.lower():
-                return "Based on my research, here's what I found about your topic. FINAL ANSWER: This topic shows trends that would be interesting to visualize in a chart."
-            else:
-                return "I've analyzed the data and created a chart visualization. FINAL ANSWER: The chart shows an upward trend from 2020 to 2024."
-        
-        else:
-            return "I've processed your request. FINAL ANSWER: Here's a summary of what I found."
-
-# Create our mock LLM
+
+    def invoke(self, messages):
+        last_content = ""
+        if isinstance(messages, list) and messages:
+            last = messages[-1]
+            last_content = last.content if hasattr(last, "content") else str(last)
+        elif isinstance(messages, str):
+            last_content = messages
+
+        # Second-pass (chart) request — emit the FINAL ANSWER sentinel.
+        if "chart" in last_content.lower() or "visuali" in last_content.lower():
+            return AIMessage(
+                content=(
+                    "I've analyzed the data and created a chart visualization. "
+                    "FINAL ANSWER: The chart shows an upward trend from 2020 to 2024."
+                )
+            )
+        # First-pass (research) request — stay open-ended.
+        return AIMessage(
+            content=(
+                "Based on my research, here's what I found about your topic. "
+                "This topic shows trends that would be interesting to visualize in a chart."
+            )
+        )
+
+
 llm = MockLLM()
 
-# System prompt constructor
-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"
-        f"{suffix}"
-    )
-
-# 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 do research.")
-    )
-    
-    # Process the current state
-    result = agent.invoke(state)
-    
-    # Check if we have a final answer
-    last_message = result["messages"][-1]
-    content = last_message.content if hasattr(last_message, "content") else last_message
-    
-    # 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"
-    
-    # In langgraph 0.0.41 style, return the updated state
-    return {"messages": result["messages"], "has_final_answer": "FINAL ANSWER" in content}
-
-# 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)
-    
-    # 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 the updated state
-    return {"messages": result["messages"], "has_final_answer": True}
-
-# Define the router function for the graph
-def router(state):
-    # Check if we have a final answer
-    return "chart_generator" if not state.get("has_final_answer", False) else "end"
-
-# Simplified end node that doesn't modify the state
-def end_node(state):
-    return state
-
-# Build a proper LangGraph for v0.0.41
-workflow = StateGraph({"messages": [], "has_final_answer": False})
+
+# ---------------------------------------------------------------------------
+# Graph state — TypedDict is required by StateGraph (not a dict instance).
+# `add` is used as the reducer so that node returns of {"messages": [msg]}
+# append rather than overwrite.
+# ---------------------------------------------------------------------------
+class GraphState(TypedDict):
+    messages: Annotated[List[BaseMessage], add]
+    has_final_answer: bool
+
+
+# ---------------------------------------------------------------------------
+# Nodes
+# ---------------------------------------------------------------------------
+def research_node(state: GraphState):
+    response = llm.invoke(state["messages"])
+    response.name = "researcher"
+    return {
+        "messages": [response],
+        "has_final_answer": "FINAL ANSWER" in response.content,
+    }
+
+
+def chart_node(state: GraphState):
+    # Nudge the LLM toward the chart-generation branch.
+    prompt = state["messages"] + [HumanMessage(content="Now create the chart visualization.")]
+    response = llm.invoke(prompt)
+    response.name = "chart_generator"
+    return {
+        "messages": [response],
+        "has_final_answer": True,
+    }
+
+
+# Router returns the key used in the conditional-edge mapping below.
+def router(state: GraphState) -> str:
+    return "end" if state.get("has_final_answer", False) else "chart_generator"
+
+
+# ---------------------------------------------------------------------------
+# Build the graph — use END (not a custom "end" node) to terminate.
+# ---------------------------------------------------------------------------
+workflow = StateGraph(GraphState)
 workflow.add_node("researcher", research_node)
 workflow.add_node("chart_generator", chart_node)
-workflow.add_node("end", end_node)
 workflow.set_entry_point("researcher")
-workflow.add_conditional_edges("researcher", router)
-workflow.add_edge("chart_generator", "end")
+workflow.add_conditional_edges(
+    "researcher",
+    router,
+    {
+        "chart_generator": "chart_generator",
+        "end": END,
+    },
+)
+workflow.add_edge("chart_generator", END)
 graph = workflow.compile()
 
-# LangGraph runner
-def run_langgraph(user_input):
+
+# ---------------------------------------------------------------------------
+# Runner
+# ---------------------------------------------------------------------------
+def run_langgraph(user_input: str):
     try:
-        # Create a human message
-        human_message = HumanMessage(content=user_input)
-        
-        # Run the workflow with proper initialization
-        events = graph.stream({"messages": [human_message], "has_final_answer": False})
-        outputs = list(events)
-        
-        # Get the final message
-        final_state = outputs[-1]
+        initial_state: GraphState = {
+            "messages": [HumanMessage(content=user_input)],
+            "has_final_answer": False,
+        }
+        final_state = graph.invoke(initial_state)
         final_messages = final_state.get("messages", [])
-        
+
         if not final_messages:
             return "Error: No messages in the final state.", None
-            
+
         final_message = final_messages[-1]
-        final_content = final_message.content if hasattr(final_message, "content") else final_message
-        
+        final_content = (
+            final_message.content if hasattr(final_message, "content") else str(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")
             plt.grid(True)
             plt.tight_layout()
             plt.savefig("gdp_chart.png")
-            return "Chart generated based on FINAL ANSWER.", "gdp_chart.png"
-        else:
-            if isinstance(final_content, str):
-                return final_content, None
-            else:
-                return str(final_content), None
+            plt.close()
+            return f"{final_content}\n\nChart generated based on FINAL ANSWER.", "gdp_chart.png"
+
+        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 UI
+# ---------------------------------------------------------------------------
 def process_input(user_input):
     return run_langgraph(user_input)
 
+
 interface = gr.Interface(
     fn=process_input,
     inputs=gr.Textbox(label="Enter your research task"),
     outputs=[gr.Textbox(label="Output"), gr.Image(type="filepath", label="Chart")],
     title="LangGraph Research Automation",
-    description="Enter a research prompt and view chart output when applicable."
+    description="Enter a research prompt and view chart output when applicable.",
 )
 
 if __name__ == "__main__":