app.py

import gradio as gr
import os
import matplotlib.pyplot as plt
import pandas as pd
from langgraph.graph import StateGraph
from langgraph_core.state import MessagesState
from langchain_core.messages import HumanMessage, AIMessage
from langchain_anthropic import ChatAnthropic
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 = os.getenv("ANTHROPIC_API_KEY")
if not api_key:
    raise ValueError("ANTHROPIC_API_KEY environment variable not set")

# 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")

# 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
    
    # Determine next step
    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 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 Command(update={"messages": result["messages"]}, goto="__end__")

# 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")
workflow.set_finish_point("__end__")
workflow.add_edge("researcher", "chart_generator")
graph = workflow.compile()

# LangGraph runner
def run_langgraph(user_input):
    try:
        # 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.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
    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."
)

if __name__ == "__main__":
    interface.launch()