# .env
OPENAI_API_KEY=your_key_here
TAVILY_API_KEY=your_key_here  # For web search tool

from typing import TypedDict, Annotated, List
from langgraph.graph.message import add_messages

class AgentState(TypedDict):
    messages: Annotated[list, add_messages]
    next_agent: str
    task: str
    research_output: str
    final_output: str

from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.tools import tool

# Web search tool
search_tool = TavilySearchResults(max_results=5)

# Custom code execution tool
@tool
def run_python(code: str) -> str:
    """Execute Python code and return the output."""
    import io, contextlib
    output = io.StringIO()
    with contextlib.redirect_stdout(output):
        exec(code)
    return output.getvalue()

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langgraph.prebuilt import create_react_agent

llm = ChatOpenAI(model="gpt-4o", temperature=0)

# --- Research Agent ---
research_agent = create_react_agent(
    llm,
    tools=[search_tool],
    state_modifier=(
        "You are a Research Agent. Your job is to gather accurate, "
        "up-to-date information on the given topic using web search. "
        "Be thorough and cite your sources."
    )
)

# --- Code Agent ---
code_agent = create_react_agent(
    llm,
    tools=[run_python],
    state_modifier=(
        "You are a Code Agent. You write clean, efficient Python code "
        "to solve analytical or data processing tasks. Always test your code."
    )
)

# --- Writer Agent ---
writer_agent = create_react_agent(
    llm,
    tools=[],
    state_modifier=(
        "You are a Writer Agent. Given research and data, you produce "
        "clear, concise, and well-structured written content for technical audiences."
    )
)

from langchain_core.output_parsers import JsonOutputParser
from pydantic import BaseModel

SUPERVISOR_SYSTEM_PROMPT = """
You are a Supervisor orchestrating a team of AI agents. 
Given the current task and conversation, decide who should act next.

Available agents:
- researcher: Gathers information from the web
- coder: Writes and executes Python code
- writer: Produces final written output
- FINISH: The task is complete

Respond ONLY with a JSON object: next
"""

class RouterOutput(BaseModel):
    next: str

def supervisor_node(state: AgentState) -> AgentState:
    messages = [
        {"role": "system", "content": SUPERVISOR_SYSTEM_PROMPT},
        *state["messages"],
    ]
    response = llm.with_structured_output(RouterOutput).invoke(messages)
    return {"next_agent": response.next}

def research_node(state: AgentState) -> AgentState:
    result = research_agent.invoke(state)
    return {
        "messages": result["messages"],
        "research_output": result["messages"][-1].content
    }

def code_node(state: AgentState) -> AgentState:
    result = code_agent.invoke(state)
    return {"messages": result["messages"]}

def writer_node(state: AgentState) -> AgentState:
    result = writer_agent.invoke(state)
    return {
        "messages": result["messages"],
        "final_output": result["messages"][-1].content
    }

from langgraph.graph import StateGraph, END

# Initialize graph
workflow = StateGraph(AgentState)

# Add nodes
workflow.add_node("supervisor", supervisor_node)
workflow.add_node("researcher", research_node)
workflow.add_node("coder", code_node)
workflow.add_node("writer", writer_node)

# Set entry point
workflow.set_entry_point("supervisor")

# Conditional routing from supervisor
workflow.add_conditional_edges(
    "supervisor",
    lambda state: state["next_agent"],
    {
        "researcher": "researcher",
        "coder": "coder",
        "writer": "writer",
        "FINISH": END,
    }
)

# All agents report back to supervisor
workflow.add_edge("researcher", "supervisor")
workflow.add_edge("coder", "supervisor")
workflow.add_edge("writer", "supervisor")

# Compile
graph = workflow.compile()

initial_state = {
    "messages": [
        {"role": "user", "content": (
            "Research the top 3 vector databases for RAG in 2024, "
            "write Python code to benchmark their query latency, "
            "then write a summary report of your findings."
        )}
    ],
    "task": "Vector DB research and benchmark",
    "next_agent": "",
    "research_output": "",
    "final_output": "",
}

# Stream the execution
for step in graph.stream(initial_state, {"recursion_limit": 20}):
    for node_name, output in step.items():
        print(f"\n{'='*50}")
        print(f"Node: {node_name}")
        if "messages" in output:
            print(output["messages"][-1].content[:500])

from langgraph.checkpoint.memory import MemorySaver

memory = MemorySaver()
graph = workflow.compile(checkpointer=memory)

# Each run is tied to a thread_id
config = {"configurable": {"thread_id": "session_001"}}

result = graph.invoke(initial_state, config=config)

# Resume from checkpoint later
follow_up = {
    "messages": [{"role": "user", "content": "Now compare pricing for those databases."}]
}
result2 = graph.invoke(follow_up, config=config)  # Remembers full prior context

from langgraph.checkpoint.memory import MemorySaver

graph = workflow.compile(
    checkpointer=MemorySaver(),
    interrupt_before=["writer"]  # Pause before writer runs
)

# After review, resume
graph.invoke(None, config=config)  # Pass None to continue from checkpoint

# Integrate LangSmith for full tracing
import os
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "multi-agent-prod"