Escalation

"""Cookbook 04 — Escalation pattern.

A request flows up a tiered support stack. Each tier either
``resolve``\\ s the request (terminating the workflow) or escalates
to the next tier. Demonstrates two routing idioms together:

* Typed ``Handoff`` returns for escalation — each tier's
  ``escalate_to_*`` tool returns a ``Handoff(target=...)`` carrying
  the next tier's name. Mirrors classic AG2's
  ``ReplyResult(target=AgentTarget(...))`` pattern.
* ``set_context`` + ``ToolCalled`` graph rule for termination —
  the ``resolve`` tool writes the answer into
  ``context_vars["resolution"]`` and a ``ToolCalled("resolve")``
  graph rule fires the terminate transition.
"""

import asyncio

from dotenv import load_dotenv

from autogen.beta import Agent
from autogen.beta.config import AnthropicConfig
from autogen.beta.knowledge import MemoryKnowledgeStore
from autogen.beta.network import (
    EV_PACKET,
    EV_SESSION_CLOSED,
    EV_TEXT,
    WORKFLOW_TYPE,
    AgentTarget,
    FromSpeaker,
    Handoff,
    Hub,
    HubClient,
    LocalLink,
    Passport,
    Resume,
    SessionInject,
    TerminateTarget,
    ToolCalled,
    Transition,
    TransitionGraph,
)
from autogen.beta.network.workflow_helpers import set_context
from autogen.beta.testing import TestConfig

load_dotenv()

async def escalate_to_tier2(reason: str) -> Handoff:
    """Escalate to the tier-2 specialist. The returned Handoff
    carries the target's Passport.name; the framework resolves
    it and routes the next turn there."""
    print(f"  [tool] escalate_to_tier2({reason!r})")
    return Handoff(target="tier2", reason=reason)

async def escalate_to_senior(reason: str) -> Handoff:
    """Escalate to senior support."""
    print(f"  [tool] escalate_to_senior({reason!r})")
    return Handoff(target="senior", reason=reason)

async def resolve(answer: str, session: SessionInject) -> str:
    """Resolve the request. Stores the answer in
    context_vars['resolution'] and the graph's
    ToolCalled('resolve') rule terminates the workflow."""
    if session is None:
        return "no session"
    print(f"  [tool] resolve({answer[:60]!r}…)")
    await set_context(session, "resolution", answer)
    return "resolved"

async def main() -> None:
    config = AnthropicConfig(model="claude-sonnet-4-6")

    hub_obj = await Hub.open(MemoryKnowledgeStore(), ttl_sweep_interval=0)
    link = LocalLink(hub_obj)

    user_hc = HubClient(link, hub=hub_obj)
    tier1_hc = HubClient(link, hub=hub_obj)
    tier2_hc = HubClient(link, hub=hub_obj)
    senior_hc = HubClient(link, hub=hub_obj)

    user_agent = Agent("user", config=TestConfig())

    tier1_agent = Agent(
        "tier1",
        prompt=(
            "You are tier-1 support. You only handle simple FAQ-like "
            "questions: business hours, general policies, account "
            "lookups. ANY technical diagnosis, billing dispute, or "
            "specialist topic is OUT OF SCOPE for you — call "
            "`escalate_to_tier2(reason)` immediately. Otherwise call "
            "`resolve(answer)` with a 1-line answer.\n"
            "\n"
            "Call exactly ONE tool. Don't write a separate body."
        ),
        config=config,
    )
    tier1_agent.tool(escalate_to_tier2)
    tier1_agent.tool(resolve)

    tier2_agent = Agent(
        "tier2",
        prompt=(
            "You are tier-2 support — a domain specialist. You handle "
            "most technical questions in your area. For genuinely edge "
            "cases (rare hardware faults, legal compliance questions, "
            "anything requiring senior judgement) call "
            "`escalate_to_senior(reason)`. Otherwise call "
            "`resolve(answer)` with a 1-2 sentence answer.\n"
            "\n"
            "For demo purposes, if the question mentions 'compressor', "
            "'GDPR', or 'compliance', escalate. Otherwise resolve.\n"
            "\n"
            "Call exactly ONE tool. Don't write a separate body."
        ),
        config=config,
    )
    tier2_agent.tool(escalate_to_senior)
    tier2_agent.tool(resolve)

    senior_agent = Agent(
        "senior",
        prompt=(
            "You are senior support — the buck stops here. Whatever the "
            "question, you provide a confident, concise answer. Call "
            "`resolve(answer)` with a 1-2 sentence answer. Never "
            "escalate (you have no escalate tool).\n"
            "\n"
            "Call exactly ONE tool. Don't write a separate body."
        ),
        config=config,
    )
    senior_agent.tool(resolve)

    user = await user_hc.register(user_agent, Passport(name="user"), Resume())
    tier1 = await tier1_hc.register(tier1_agent, Passport(name="tier1"), Resume())
    tier2 = await tier2_hc.register(tier2_agent, Passport(name="tier2"), Resume())
    senior = await senior_hc.register(senior_agent, Passport(name="senior"), Resume())

    graph = TransitionGraph(
        initial_speaker=user.agent_id,
        transitions=[
            # resolve terminates the workflow.
            Transition(when=ToolCalled("resolve"), then=TerminateTarget("resolved")),
            # User's question → tier1. Escalation FROM tier1 / tier2
            # is via Handoff returns from escalate_to_* tools — the
            # framework reads target from the Handoff and stamps it
            # onto the packet, so no graph rules are needed for those
            # edges.
            Transition(when=FromSpeaker(user.agent_id), then=AgentTarget(tier1.agent_id)),
        ],
        default_target=TerminateTarget("fall_through"),
        max_turns=10,
    )

    session = await user.open(
        type=WORKFLOW_TYPE,
        target=[tier1.agent_id, tier2.agent_id, senior.agent_id],
        knobs={"graph": graph.to_dict()},
    )
    print(f"session: {session.session_id}\n")

    name_by_id = {
        user.agent_id: "user",
        tier1.agent_id: "tier1",
        tier2.agent_id: "tier2",
        senior.agent_id: "senior",
    }

    await session.send(
        "My refrigerator's compressor is humming louder than usual and "
        "occasionally clicks. What's wrong and how do I fix it?"
    )

    # Wait for the workflow to terminate (any of the five close routes
    # documented in /docs/beta/network/termination — this demo uses
    # ToolCalled("resolve") → TerminateTarget("resolved")).
    close_env = await user_hc.wait_for_session_event(
        session_id=session.session_id,
        predicate=lambda e: e.event_type == EV_SESSION_CLOSED,
        timeout=180.0,
    )

    # Print the transcript from the WAL after close.
    for env in await hub_obj.read_wal(session.session_id):
        speaker = name_by_id.get(env.sender_id, env.sender_id[:8])
        if env.event_type == EV_TEXT:
            print(f"{speaker:>14}: {env.event_data['text']}")
        elif env.event_type == EV_PACKET:
            routing = env.event_data.get("routing", {}) or {}
            if routing.get("kind") == "handoff":
                line = f"[Handed off via {routing.get('tool', '')}] {routing.get('reason', '')}"
                print(f"{speaker:>14}: {line.rstrip()}")
            body = env.event_data.get("body", "")
            if body:
                print(f"{speaker:>14}: {body}")

    print(f"\nclosed: reason={close_env.event_data.get('reason')!r}")

    print("\n--- final resolution ---")
    state = hub_obj._adapter_states[session.session_id]
    print(state.context_vars.get("resolution", "(no resolution)"))

    await user_hc.close()
    await tier1_hc.close()
    await tier2_hc.close()
    await senior_hc.close()
    await hub_obj.close()

if __name__ == "__main__":
    asyncio.run(main())