from langgraph.graph import StateGraph, START, END
from langgraph.types import Command
from langgraph.graph.message import MessagesState
# ToolNode : 마지막 AI 메세지의 tool_calls를 읽어 실제 툴을 실행해주는 prebuilt(미리 만들어진) 노드
from langgraph.prebuilt import ToolNode, tools_condition
# tool 데코레이터 : handoff tool 생성
from langchain_core.tools import tool
from langchain.chat_models import init_chat_model

# MessagesState를 상속해서 messages + 커스텀 필드 2개를 가진 State를 정의.
# 부모 그래프와 각 서브그래프(에이전트)가 같은 State 스키마를 쓰기 때문에 서브 그래프가 반환한 상태 업데이트가 부모로 그대로 전파됨.
class AgentsState(MessagesState):
    current_agent: str
    transfered_by: str

llm = init_chat_model("openai:gpt-4o")

# 에이전트 팩토리
def make_agent(prompt, tools):
    # agent_node
    def agent_node(state: AgentsState):
        # LLM에 tool 바인딩 -> 툴 스키마를 실어 보내서 LLM이 이런 툴을 쓸 수 있다고 인지하게 됨.
        llm_with_tools = llm.bind_tools(tools)

        # invoke가 돌아오면 AIMessage이고, LLM이 툴을 쓰기로 했다면 .tool_calls에 호출 내역이 들어 있음.
        response = llm_with_tools.invoke(
            f"""
        {prompt}

        You have a tool called 'handoff_tool' use it to transfer to other agent, don't use it to transfer to yourself.

        Conversation History:
        {state["messages"]}
        """
        )
        return {
            # reducer 덕분에 기존 대화에 append
            "messages": [response]
        }

    # ReAct 루프 : agent -> (툴 있으면) tools -> agent -> (툴 없으면) END.
    # Agent 생성
    agent_builder = StateGraph(AgentsState)

    # node 생성
    agent_builder.add_node("agent", agent_node)

    # ToolNode 생성
    # 메세지에 Tool을 호출하는 내용이 있는지 확인하고 있으면 해당 툴 호출.
    agent_builder.add_node(
        "tools",
        ToolNode(tools=tools),
    )

    # edge 생성
    agent_builder.add_edge(START, "agent")
    # tools_condition : 미리 만들어진 함수, 호출되어야 하는 tool이 있는지 확인. 있으면 "tools" 반환. 없으면 END 반환
    agent_builder.add_conditional_edges("agent", tools_condition)
    agent_builder.add_edge("tools", "agent")
    agent_builder.add_edge("agent", END)

    return agent_builder.compile()

# 모든 에이전트가 사용할 수 있는 통제권을 넘기는 Tool 생성
@tool
def handoff_tool(transfer_to: str, transfered_by: str):
    """
    Handoff to another agent.

    Use this tool when the customer speaks a language that you don't understand.

    Possible values for `transfer_to`:
    - `korean_agent`
    - `greek_agent`
    - `spanish_agent`

    Possible values for `transfered_by`:
    - `korean_agent`
    - `greek_agent`
    - `spanish_agent`

    Args:
        transfer_to: The agent to transfer the conversation to
        transfered_by: The agent that transferred the conversation
    """
    if transfer_to == transfered_by:
        return {
            "error": "Stop trying to transfer to yourself and answer the question or i will fire you."
        }

    return Command(
        update={
            "current_agent": transfer_to,
            "transfered_by": transfered_by,
        },
        goto=transfer_to,
        # 이 Tool이 Agent 안에서 호출되므로 Command.PARENT 설정 필요.
        # 이 설정은 transfer_to 노드로 보내고 싶은데, 현재 그래프가 아니라 부모 그래프에서 보내고 싶다는 의미.
        graph=Command.PARENT,
    )

# 네트워크 아키텍처를 구성하는 각 에이전트도 그래프이고 이들을 연결하는 것 또한 그래프.
# 이를 위해 큰 그래프를 먼저 생성하고 하위 그래프는 컴파일된 그래프를 리턴하는 함수로 생성.
graph_builder = StateGraph(AgentsState)

# node로 이루어진 그래프가 있고, make_agent 함수로 Agent를 생성하고 그래프를 만들어서 리턴
graph_builder.add_node(
    "korean_agent",
    make_agent(
        prompt="You're a Korean customer support agent. You only speak and understand Korean.",
        tools=[handoff_tool],
    ),
    destinations=("greek_agent", "spanish_agent"),
)
graph_builder.add_node(
    "greek_agent",
    make_agent(
        prompt="You're a Greek customer support agent. You only speak and understand Greek.",
        tools=[handoff_tool],
    ),
    destinations=("korean_agent", "spanish_agent"),
)
graph_builder.add_node(
    "spanish_agent",
    make_agent(
        prompt="You're a Spanish customer support agent. You only speak and understand Spanish.",
        tools=[handoff_tool],
    ),
    destinations=("greek_agent", "korean_agent"),
)

graph_builder.add_edge(START, "korean_agent")

graph = graph_builder.compile()

{
  "dependencies": ["./graph.py"],
  "graphs": {
    "agent": "./graph.py:graph"
  },
  "env": ".env"
}

from typing import Literal
from pydantic import BaseModel
from langgraph.prebuilt.chat_agent_executor import AgentState

class SupervisorOutput(BaseModel):
    next_agent: Literal["korean_agent", "spanish_agent", "greek_agent", "__end__"]
    reasoning: str

class AgentsState(MessagesState):
    current_agent: str
    transfered_by: str
    reasoning: str

# supervisor 에이전트
# 유저가 보낸 메세지를 보고 누구에게 넘길지 알려주어야 함. (Command 활용, 다른 노드로 전달)
def supervisor(state: AgentState):
    structured_llm = llm.with_structured_output(SupervisorOutput)
    response = structured_llm.invoke(
        f"""
        You are a supervisor that routes conversations to the appropriate language agent.

        Analyse the customers request and the conversation history and decide which agent should handle the conversation.

        The options for the next agent are:
        - greek_agent
        - spanish_agent
        - korean_agent      
        
        <CONVERSATION_HISTORY>
        {state.get("messages", [])}
        </CONVERSATION_HISTORY>

        IMPORTANT:        
        Never transfer to the same agent twice in a row.
        If an agent has replied end the conversation by returning __end__
    """
    )

    return Command(
        goto=response.next_agent,
        update={"reasoning": response.reasoning},
    )

graph_builder = StateGraph(AgentsState)

graph_builder.add_node(
    "supervisor",
    supervisor,
    destinations=(
        "korean_agent",
        "spanish_agent",
        "greek_agent",
        END,
    ),
)

graph_builder.add_node(
    "korean_agent",
    make_agent(
        prompt="You're a Korean customer support agent. You only speak and understand Korean.",
        tools=[],
    ),
)
graph_builder.add_node(
    "greek_agent",
    make_agent(
        prompt="You're a Greek customer support agent. You only speak and understand Greek.",
        tools=[],
    ),
)
graph_builder.add_node(
    "spanish_agent",
    make_agent(
        prompt="You're a Spanish customer support agent. You only speak and understand Spanish.",
        tools=[],
    ),
)

graph_builder.add_edge(START, "supervisor")
# 서브에이전트들은 무조건 supervisor 에이전트에게 돌아가야 함.
graph_builder.add_edge("korean_agent", "supervisor")
graph_builder.add_edge("spanish_agent", "supervisor")
graph_builder.add_edge("greek_agent", "supervisor")

graph = graph_builder.compile()

from typing import Annotated
from langgraph.prebuilt import InjectedState, ToolNode, tools_condition

def make_agent_tool(tool_name, tool_description, system_prompt, tools):

    def agent_node(state: AgentsState):
        llm_with_tools = llm.bind_tools(tools)
        response = llm_with_tools.invoke(
            f"""
        {system_prompt}

        Conversation History:
        {state["messages"]}
        """
        )
        return {"messages": [response]}

    agent_builder = StateGraph(AgentsState)

    agent_builder.add_node("agent", agent_node)
    agent_builder.add_node(
        "tools",
        ToolNode(tools=tools),
    )

    agent_builder.add_edge(START, "agent")
    agent_builder.add_conditional_edges("agent", tools_condition)
    agent_builder.add_edge("tools", "agent")
    agent_builder.add_edge("agent", END)

    agent = agent_builder.compile()

    # tool 커스터마이징 -> tool 이름와 설명을 설정해줌.
    @tool(
        name_or_callable=tool_name,
        description=tool_description,
    )
    # Annotated[dict, InjectedState] -> agent_tool이 호출될 때, 여기에 랭그래프가 현재 state를 주입해줌.
    def agent_tool(state: Annotated[dict, InjectedState]):
        result = agent.invoke(state)
        return result["messages"][-1].content

    return agent_tool

tools = [
    make_agent_tool(
        tool_name="korean_agent",
        tool_description="Use this when the user is speaking korean",
        system_prompt="You're a korean customer support agent you speak in korean",
        tools=[],
    ),
    make_agent_tool(
        tool_name="spanish_agent",
        tool_description="Use this when the user is speaking spanish",
        system_prompt="You're a spanish customer support agent you speak in spanish",
        tools=[],
    ),
    make_agent_tool(
        tool_name="greek_agent",
        tool_description="Use this when the user is speaking greek",
        system_prompt="You're a greek customer support agent you speak in greek",
        tools=[],
    ),
]

def supervisor(state: AgentState):
    llm_with_tools = llm.bind_tools(tools=tools)
    result = llm_with_tools.invoke(state["messages"])
    return {
        "messages": [result],
    }


graph_builder = StateGraph(AgentsState)

graph_builder.add_node("supervisor", supervisor)
graph_builder.add_node("tools", ToolNode(tools=tools))


graph_builder.add_edge(START, "supervisor")
graph_builder.add_conditional_edges("supervisor", tools_condition)
graph_builder.add_edge("tools", "supervisor")
graph_builder.add_edge("supervisor", END)


graph = graph_builder.compile()

pip install langgraph-supervisor langchain-openai

from langchain.chat_models import init_chat_model
from langgraph.prebuilt import create_react_agent
from langgraph_supervisor import create_supervisor

MODEL = "openai:gpt-4o-mini"

history_agent = create_react_agent(
    model=MODEL,
    tools=[],
    name="history_agent",
    prompt="You are a history expert. You only answer questions about history."
)
geography_agent = create_react_agent(
    model=MODEL,
    tools=[],
    name="geography_agent",
    prompt="You are a geography expert. You only answer questions about geography."
)

supervisor = create_supervisor(
    agents=[
        history_agent,
        geography_agent
    ],
    model=init_chat_model(MODEL),
    prompt="""
    You are a supervisor that routes student questions to the appropriate subject expert.
    You manage a history agent and a geography agent.
    Analyze the student's question and assign it to the correct expert based on the subject matter:
    - history_agent: For historical events, dates, historical figures
    - geography_agent: For locations, rivers, mountains, countries
    """
).compile()

questions = [
    "....",
    "....",
    "...."    
]

for question in questions:
    result = supervisor.invoke({
        "messages": [
            {"role": "user", "content": question}
        ]
    })
    if result["messages"]:
        for message in result["messages"]:
            message.pretty_print()

pip install langgraph-swarm

from langchain_openai import ChatOpenAI

from langgraph.checkpoint.memory import InMemorySaver
from langchain.agents import create_agent
from langgraph_swarm import create_handoff_tool, create_swarm

model = ChatOpenAI(model="gpt-4o")

def add(a: int, b: int) -> int:
    """Add two numbers"""
    return a + b

alice = create_agent(
    model,
    tools=[
        add,
        create_handoff_tool(
            agent_name="Bob",
            description="Transfer to Bob",
        ),
    ],
    system_prompt="You are Alice, an addition expert.",
    name="Alice",
)

bob = create_agent(
    model,
    tools=[
        create_handoff_tool(
            agent_name="Alice",
            description="Transfer to Alice, she can help with math",
        ),
    ],
    system_prompt="You are Bob, you speak like a pirate.",
    name="Bob",
)

checkpointer = InMemorySaver()
workflow = create_swarm(
    [alice, bob],
    default_active_agent="Alice"
)
app = workflow.compile(checkpointer=checkpointer)

config = {"configurable": {"thread_id": "1"}}
turn_1 = app.invoke(
    {"messages": [{"role": "user", "content": "i'd like to speak to Bob"}]},
    config,
)
print(turn_1)
turn_2 = app.invoke(
    {"messages": [{"role": "user", "content": "what's 5 + 7?"}]},
    config,
)
print(turn_2)