Agent-to-Agent Handover

from langgraph.types import Command
from langchain_core.messages import HumanMessage, AIMessage
from typing import Literal

# Command allows a node to explicitly route to another node
# AND update state in one atomic operation
def billing_classifier(state: dict) -> Command[Literal["billing_agent", "support_agent"]]:
    """Classify intent and immediately route to the right agent."""
    last_msg = state["messages"][-1].content.lower()

    if any(word in last_msg for word in ["invoice", "payment", "charge", "refund"]):
        return Command(
            goto="billing_agent",
            update={
                "messages": [AIMessage(content="Routing to billing specialist...")],
                "routed_to": "billing",
                "reason": "billing keyword detected"
            }
        )
    else:
        return Command(
            goto="support_agent",
            update={
                "messages": [AIMessage(content="Routing to technical support...")],
                "routed_to": "support",
                "reason": "no billing keywords"
            }
        )

from langgraph.prebuilt import create_react_agent
from langgraph.graph import StateGraph, START, END, MessagesState
from langgraph.prebuilt.chat_agent_executor import AgentState
from langchain_openai import ChatOpenAI

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

# create_handoff_tool generates a tool that routes to a named agent
# The LLM can call this tool when it decides to hand off
from langgraph.prebuilt import create_handoff_tool

transfer_to_billing  = create_handoff_tool(agent_name="billing_agent",
    description="Transfer to billing agent for invoice, payment, or refund queries.")
transfer_to_support  = create_handoff_tool(agent_name="support_agent",
    description="Transfer to technical support for product issues or bugs.")
transfer_to_triage   = create_handoff_tool(agent_name="triage_agent",
    description="Return to triage when unsure which specialist is needed.")

# Each agent has its own tools + handoff tools
triage_agent = create_react_agent(
    llm,
    tools=[transfer_to_billing, transfer_to_support],
    state_schema=AgentState,
    name="triage_agent",
)

billing_agent = create_react_agent(
    llm,
    tools=[transfer_to_triage],   # can hand back if wrong agent
    system_prompt="You are a billing specialist. Handle invoices, payments, and refunds.",
    state_schema=AgentState,
    name="billing_agent",
)

support_agent = create_react_agent(
    llm,
    tools=[transfer_to_triage],
    system_prompt="You are a technical support specialist. Debug product issues.",
    state_schema=AgentState,
    name="support_agent",
)

# Wire into a graph
from langgraph.graph import StateGraph
builder = StateGraph(AgentState)
builder.add_node("triage_agent",  triage_agent)
builder.add_node("billing_agent", billing_agent)
builder.add_node("support_agent", support_agent)
builder.add_edge(START, "triage_agent")

from typing import Annotated, Set
from typing_extensions import TypedDict
from langchain_core.messages import BaseMessage, AIMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
import operator

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

class SwarmState(TypedDict):
    messages:      Annotated[list[BaseMessage], add_messages]
    current_agent: str
    visited:       Annotated[Set[str], operator.or_]  # union reducer
    hop_count:     int

def make_swarm_agent(name: str, specialty: str, can_handoff_to: list[str]):
    def agent_node(state: SwarmState) -> dict:
        # Prevent revisiting the same agent (loop detection)
        if name in state["visited"] and state["hop_count"] > 1:
            return {
                "messages": [AIMessage(content=f"[{name}] Already visited; ending here.")],
                "current_agent": "END"
            }

        # Prevent runaway chains
        if state["hop_count"] >= MAX_HOPS:
            return {
                "messages": [AIMessage(content="[Swarm] Max hops reached. Terminating.")],
                "current_agent": "END"
            }

        # Agent decides whether to answer or hand off
        decision_prompt = (
            f"You are {name}, specialised in {specialty}.\n"
            f"Can you answer the user's request? If not, hand off to one of: {can_handoff_to}.\n"
            f"Respond with either your answer OR 'HANDOFF: '"
        )
        response = llm.invoke(state["messages"] + [
            AIMessage(content=decision_prompt)
        ])

        if response.content.startswith("HANDOFF:"):
            next_agent = response.content.split(":")[1].strip()
            return {
                "messages": [AIMessage(content=f"[{name}] Handing off to {next_agent}")],
                "current_agent": next_agent,
                "visited": {name},
                "hop_count": state["hop_count"] + 1,
            }
        else:
            return {
                "messages": [response],
                "current_agent": "END",
                "visited": {name},
            }
    return agent_node

def swarm_router(state: SwarmState) -> str:
    next = state.get("current_agent", "END")
    return "end" if next == "END" else next

builder = StateGraph(SwarmState)
builder.add_node("billing",  make_swarm_agent("billing",  "invoices & payments", ["support", "compliance"]))
builder.add_node("support",  make_swarm_agent("support",  "technical issues",    ["billing", "compliance"]))
builder.add_node("compliance", make_swarm_agent("compliance", "regulations & policy", ["billing", "support"]))
builder.add_edge(START, "billing")  # entry point
builder.add_conditional_edges("billing",    swarm_router, {"billing": "billing", "support": "support", "compliance": "compliance", "end": END})
builder.add_conditional_edges("support",    swarm_router, {"billing": "billing", "support": "support", "compliance": "compliance", "end": END})
builder.add_conditional_edges("compliance", swarm_router, {"billing": "billing", "support": "support", "compliance": "compliance", "end": END})
swarm_graph = builder.compile()

from pydantic import BaseModel, Field
from typing import List, Optional
from langchain_openai import ChatOpenAI

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

class HandoverPayload(BaseModel):
    """Structured context passed between agents."""
    original_task:    str = Field(description="The user's original goal verbatim")
    summary_so_far:   str = Field(description="2-3 sentence summary of work done")
    key_findings:     List[str] = Field(description="Bullet points of important discoveries")
    remaining_task:   str = Field(description="What the receiving agent needs to do")
    constraints:      List[str] = Field(description="Constraints and requirements to maintain")
    attempted_approaches: List[str] = Field(description="Approaches already tried (to avoid repetition)")

def build_handover_payload(
    original_task: str,
    conversation_history: list,
    remaining_task: str,
) -> HandoverPayload:
    """Use an LLM to extract and structure the handover context."""
    history_text = "\n".join([
        f"{m.type.upper()}: {m.content[:200]}"
        for m in conversation_history[-10:]  # last 10 messages to keep concise
    ])

    structured_llm = llm.with_structured_output(HandoverPayload)
    payload = structured_llm.invoke(
        f"Original task: {original_task}\n\n"
        f"Conversation so far:\n{history_text}\n\n"
        f"Remaining task for next agent: {remaining_task}\n\n"
        "Extract and structure the handover context."
    )
    return payload

# Usage in a handover node
def research_to_writer_handover(state: dict) -> dict:
    payload = build_handover_payload(
        original_task=state["original_task"],
        conversation_history=state["messages"],
        remaining_task="Write a 500-word article based on the research findings",
    )
    # Inject the structured payload as the first message to the writing agent
    handover_message = (
        f"HANDOVER CONTEXT:\n"
        f"Original task: {payload.original_task}\n"
        f"Research summary: {payload.summary_so_far}\n"
        f"Key findings: {chr(10).join(f'- {f}' for f in payload.key_findings)}\n"
        f"Your task: {payload.remaining_task}\n"
        f"Constraints: {', '.join(payload.constraints)}\n"
        f"Do NOT repeat: {', '.join(payload.attempted_approaches)}"
    )
    return {"handover_context": handover_message}

from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver

# Subgraph — isolated worker
class SubState(TypedDict):
    task:   str
    result: str

def sub_worker(state: SubState) -> dict:
    # Perform isolated work
    result = llm.invoke(f"Complete this task: {state['task']}")
    return {"result": result.content}

sub_builder = StateGraph(SubState)
sub_builder.add_node("worker", sub_worker)
sub_builder.add_edge(START, "worker")
sub_builder.add_edge("worker", END)
subgraph = sub_builder.compile()

# Parent graph that calls subgraph and merges results
class ParentState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    task:     str
    result:   str

def call_subgraph(state: ParentState) -> dict:
    sub_result = subgraph.invoke({"task": state["task"], "result": ""})
    # Merge the subgraph's result back into parent state
    return {
        "result": sub_result["result"],
        "messages": [AIMessage(content=f"Subgraph completed: {sub_result['result'][:100]}")]
    }

parent_builder = StateGraph(ParentState)
parent_builder.add_node("call_subgraph", call_subgraph)
parent_builder.add_edge(START, "call_subgraph")
parent_builder.add_edge("call_subgraph", END)
parent_graph = parent_builder.compile(checkpointer=MemorySaver())

import httpx, uuid
from pydantic import BaseModel
from typing import Optional

class CrossServiceHandover(BaseModel):
    correlation_id:  str     # links related calls across services
    source_agent:    str
    target_agent:    str
    original_task:   str
    context_summary: str
    payload:         dict
    callback_url:    Optional[str] = None  # where to POST the result

async def handoff_to_remote_agent(
    target_service_url: str,
    handover: CrossServiceHandover,
) -> dict:
    """Fire a handover to another service and await the response."""
    async with httpx.AsyncClient(timeout=120.0) as client:
        response = await client.post(
            f"{target_service_url}/handover",
            json=handover.model_dump(),
            headers={
                "X-Correlation-ID": handover.correlation_id,
                "X-Source-Agent":   handover.source_agent,
            }
        )
        response.raise_for_status()
        return response.json()

# Usage inside a LangGraph node
async def billing_to_fulfillment_handover(state: dict) -> dict:
    handover = CrossServiceHandover(
        correlation_id  = state.get("correlation_id") or str(uuid.uuid4()),
        source_agent    = "billing_service",
        target_agent    = "fulfillment_service",
        original_task   = state["original_task"],
        context_summary = state["billing_summary"],
        payload         = {"order_id": state["order_id"], "amount": state["amount"]},
    )
    result = await handoff_to_remote_agent(
        "https://fulfillment.internal/api",
        handover
    )
    return {"fulfillment_result": result, "correlation_id": handover.correlation_id}

import pytest
from langchain_core.messages import HumanMessage

def test_triage_routes_to_billing():
    """Billing keywords must reach the billing agent."""
    result = swarm_graph.invoke({
        "messages": [HumanMessage("I need to dispute a charge on my invoice")],
        "current_agent": "billing",
        "visited": set(),
        "hop_count": 0,
    })
    # Last message should come from billing agent
    last_msg = result["messages"][-1]
    assert "billing" in last_msg.content.lower() or "invoice" in last_msg.content.lower()
    assert result["hop_count"] <= MAX_HOPS

def test_handover_preserves_context():
    """Context summary must be non-empty after handover."""
    from langchain_core.messages import AIMessage
    mock_history = [
        HumanMessage("I want to build a RAG pipeline"),
        AIMessage("I found three relevant papers on RAG architectures."),
    ]
    payload = build_handover_payload(
        original_task="Build a RAG pipeline",
        conversation_history=mock_history,
        remaining_task="Write the implementation code",
    )
    assert payload.original_task != ""
    assert len(payload.key_findings) > 0
    assert payload.summary_so_far != ""

def test_swarm_respects_max_hops():
    """Swarm must terminate when MAX_HOPS is reached."""
    result = swarm_graph.invoke({
        "messages": [HumanMessage("Bouncing question that keeps getting handed off")],
        "current_agent": "billing",
        "visited": set(),
        "hop_count": MAX_HOPS - 1,  # start near the limit
    })
    assert result["hop_count"] <= MAX_HOPS
    assert result["current_agent"] in ("END", "billing", "support", "compliance")