Lab 10 · The support org chart

langchain==1.4.2
langchain-core==1.6.4
langchain-openai==1.6.2
python-dotenv==1.2.3
langgraph==1.2.11

# Azure OpenAI settings. Your instructor gives you these three values.
AZURE_OPENAI_ENDPOINT=https://<your-resource>.openai.azure.com/openai/v1
AZURE_OPENAI_API_KEY=<paste-your-key-here>
AZURE_OPENAI_DEPLOYMENT=<your-chat-deployment-name>

# Only used when your endpoint does NOT end in /openai/v1.
OPENAI_API_VERSION=2025-04-01-preview

"""Azure OpenAI model setup shared by every lab.

Import it in lab code as:  from model_setup import get_model, get_embeddings
"""

import asyncio
import os
import sys

from dotenv import load_dotenv
from langchain_core.embeddings import Embeddings
from langchain_core.language_models import BaseChatModel

# Read AZURE_OPENAI_* values from the .env file next to your lab code.
load_dotenv()

# Windows consoles default to cp1252 and crash on the curly quotes models love to print.
for _stream in (sys.stdout, sys.stderr):
    if hasattr(_stream, "reconfigure"):
        _stream.reconfigure(encoding="utf-8", errors="replace")

def _setting(name: str) -> str:
    """Return a required setting from .env, or stop with a readable message."""
    value = os.getenv(name, "").strip()
    if not value:
        sys.exit(f"Missing {name} in your .env file. Open .env and fill it in.")
    return value


def _uses_v1_api(endpoint: str) -> bool:
    """True for the OpenAI-compatible Azure endpoint that ends in /openai/v1."""
    return endpoint.rstrip("/").endswith("/openai/v1")

def get_model(deployment: str | None = None, **kwargs) -> BaseChatModel:
    """Return a chat model for your Azure OpenAI deployment."""
    endpoint = _setting("AZURE_OPENAI_ENDPOINT")
    api_key = _setting("AZURE_OPENAI_API_KEY")
    name = deployment or _setting("AZURE_OPENAI_DEPLOYMENT")

    if _uses_v1_api(endpoint):
        from langchain_openai import ChatOpenAI

        return ChatOpenAI(model=name, base_url=endpoint, api_key=api_key, **kwargs)

    from langchain_openai import AzureChatOpenAI

    return AzureChatOpenAI(
        azure_deployment=name,
        azure_endpoint=endpoint,
        api_key=api_key,
        api_version=os.getenv("OPENAI_API_VERSION", "2025-04-01-preview"),
        **kwargs,
    )

def get_embeddings(deployment: str | None = None) -> Embeddings:
    """Return an embeddings model for your Azure OpenAI embeddings deployment."""
    endpoint = _setting("AZURE_OPENAI_ENDPOINT")
    api_key = _setting("AZURE_OPENAI_API_KEY")
    name = deployment or _setting("AZURE_OPENAI_EMBEDDING_DEPLOYMENT")

    if _uses_v1_api(endpoint):
        from langchain_openai import OpenAIEmbeddings

        return OpenAIEmbeddings(model=name, base_url=endpoint, api_key=api_key)

    from langchain_openai import AzureOpenAIEmbeddings

    return AzureOpenAIEmbeddings(
        azure_deployment=name,
        azure_endpoint=endpoint,
        api_key=api_key,
        api_version=os.getenv("OPENAI_API_VERSION", "2025-04-01-preview"),
    )

def use_stable_event_loop() -> None:
    """On Windows, use the Selector event loop so async MCP code exits cleanly."""
    if sys.platform == "win32":
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

"""Lab 10: the Northwind support org chart, built twice.

`build_org(model, pattern=...)` returns the same organisation as subagents or as handoffs, so
the two patterns can be swapped, and costed, with one argument.

Run:  python support_org.py
"""

from typing import Any, Literal, NotRequired

from langchain.agents import AgentState, create_agent
from langchain.messages import AIMessage, HumanMessage, ToolMessage
from langchain.tools import BaseTool, ToolRuntime, tool
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.language_models import BaseChatModel
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import END, START, StateGraph
from langgraph.types import Command

from model_setup import get_model

ORDERS = {"A-100": "shipped, total 249.00 USD", "A-101": "processing, total 1299.50 USD"}
HUMAN_QUEUE: list[str] = []


@tool
def lookup_order(order_id: str) -> str:
    """Look up a Northwind order/invoice by id, such as A-100."""
    return f"Order {order_id.upper()}: {ORDERS.get(order_id.upper(), 'not found')}."


@tool
def search_docs(query: str) -> str:
    """Search the Northwind product documentation (SSO, webhooks, the API)."""
    return "SSO: re-upload the IdP metadata in Admin -> Security -> SAML after a cert roll."


@tool
def search_faq(query: str) -> str:
    """Search the Northwind policy FAQ (refunds, trials, data residency)."""
    return "Refunds run for 30 days from the invoice date, pro-rated."


@tool
def escalate_to_human(summary: str) -> str:
    """Hand the request to a human support engineer."""
    HUMAN_QUEUE.append(summary)
    return f"Escalated as ESC-{len(HUMAN_QUEUE):03d}. A human will reply within a business day."

class SupportState(AgentState):
    """Agent state plus the name of the agent that currently owns the conversation."""

    active_agent: NotRequired[str]


SPECIALISTS = {
    "billing": ("invoices, refunds, order status and anything about money", [lookup_order]),
    "technical": ("product errors, API questions, SSO and webhooks", [search_docs]),
    "knowledge": ("policy: refund windows, trials, data residency", [search_faq]),
}

SUPERVISOR_PROMPT = (
    "You are the Northwind support supervisor. Never answer from your own knowledge: delegate "
    "one topic to one specialist with a short self-contained brief (the specialist cannot see "
    "this conversation). Escalate to a human when the customer asks for one.\n"
    + "\n".join(f"- {name}: {scope}" for name, (scope, _) in SPECIALISTS.items())
)

def build_specialist(name: str, model: BaseChatModel, extra_tools: tuple[BaseTool, ...] = ()):
    """Compile one specialist. `name=` is the label you read in a LangSmith trace."""
    scope, tools = SPECIALISTS[name]
    return create_agent(
        model=model,
        tools=[*tools, *extra_tools],
        system_prompt=f"You are the Northwind {name} specialist. You handle {scope}. Answer "
        "from your tools only, in two sentences.",
        state_schema=SupportState,
        name=f"{name}-specialist",
    )

def specialist_as_tool(name: str, model: BaseChatModel) -> BaseTool:
    """Expose a specialist as a tool: a task brief in, its final text out, nothing else."""
    agent = build_specialist(name, model)

    @tool(f"ask_{name}", description=f"Ask the {name} specialist about {SPECIALISTS[name][0]}.")
    def ask(task: str) -> str:
        """Run the specialist on the brief alone and return its final answer."""
        return agent.invoke({"messages": [HumanMessage(task)]})["messages"][-1].text

    return ask

def handoff_tool(target: str) -> BaseTool:
    """Build a tool that transfers control to `target` in the PARENT graph."""

    @tool(f"transfer_to_{target}", description=f"Transfer the conversation to {target}.")
    def transfer(runtime: ToolRuntime) -> Command:
        """Hand the conversation, and ownership of the thread, to the target."""
        requesting_ai = runtime.state["messages"][-1]  # the AIMessage holding this tool call
        ack = ToolMessage(f"Transferred to {target}.", name=f"transfer_to_{target}",
                          tool_call_id=runtime.tool_call_id)
        # Forward both: a tool result whose tool call was left behind is a broken history.
        return Command(
            goto=target,
            update={"active_agent": target, "messages": [requesting_ai, ack]},
            graph=Command.PARENT,
        )

    return transfer

_Targets = Literal["supervisor", "billing", "technical", "knowledge"]


def entry(state: SupportState) -> Command[_Targets]:
    """Whoever owns the thread keeps it: this is what makes turn 2 cheaper than turn 1."""
    return Command(goto=state.get("active_agent", "supervisor"))


def human_queue(state: SupportState) -> dict[str, Any]:
    """Escalation is a deterministic node: no model call, so this path cannot loop."""
    HUMAN_QUEUE.append(state["messages"][0].text)
    ticket = f"ESC-{len(HUMAN_QUEUE):03d}"
    return {"messages": [AIMessage(f"I have escalated this to a human engineer ({ticket}).")]}

def build_org(model: BaseChatModel | None = None,
              pattern: Literal["subagents", "handoffs"] = "subagents",
              checkpointer: BaseCheckpointSaver | None = None):
    """Build the Northwind support org with one of the two multi-agent patterns."""
    model = model or get_model()
    if pattern == "subagents":
        return create_agent(
            model=model,
            tools=[*(specialist_as_tool(name, model) for name in SPECIALISTS), escalate_to_human],
            system_prompt=SUPERVISOR_PROMPT,
            state_schema=SupportState,
            checkpointer=checkpointer,
            name="support-supervisor",
        )

    supervisor = create_agent(
        model=model,
        tools=[*(handoff_tool(name) for name in SPECIALISTS), handoff_tool("human_queue")],
        system_prompt=SUPERVISOR_PROMPT + "\nTransfer control; never answer yourself.",
        state_schema=SupportState,
        name="support-supervisor",
    )
    builder = StateGraph(SupportState)
    builder.add_node("entry", entry)
    builder.add_node("supervisor", supervisor)
    builder.add_node("human_queue", human_queue)
    for name in SPECIALISTS:
        peers = [handoff_tool(other) for other in SPECIALISTS if other != name]  # never itself
        extra = (*peers, handoff_tool("human_queue"))
        builder.add_node(name, build_specialist(name, model, extra))
        builder.add_edge(name, END)
    builder.add_edge(START, "entry")
    builder.add_edge("human_queue", END)
    return builder.compile(checkpointer=checkpointer)

class CallCounter(BaseCallbackHandler):
    """Counts every model call in a run, the ones inside subagents included."""

    def __init__(self) -> None:
        """Start at zero."""
        self.model_calls = 0

    def on_chat_model_start(self, serialized: Any, messages: Any, **kwargs: Any) -> None:
        """Called by LangChain each time any chat model starts."""
        self.model_calls += 1

TURNS = ["How much was order A-100?", "And what about A-101?"]

if __name__ == "__main__":
    model = get_model()
    for pattern in ("subagents", "handoffs"):
        print(f"== {pattern}")
        org = build_org(model, pattern=pattern, checkpointer=InMemorySaver())
        total = 0
        for turn, question in enumerate(TURNS, 1):
            counter = CallCounter()
            config = {"configurable": {"thread_id": pattern}, "callbacks": [counter],
                      "recursion_limit": 12}
            seen = len(org.get_state(config).values.get("messages", []))
            result = org.invoke({"messages": [HumanMessage(question)]}, config=config)
            new_messages = result["messages"][seen:]
            route = [call["name"] for m in new_messages for call in getattr(m, "tool_calls", [])]
            total += counter.model_calls
            print(f"turn {turn}: {question}")
            print(f"  tools:       {', '.join(route) or 'none'}")
            print(f"  owner:       {result.get('active_agent', 'supervisor')}")
            print(f"  model calls: {counter.model_calls}")
            print(f"  answer:      {result['messages'][-1].text}")
        print(f"  total model calls for {pattern}: {total}\n")

== subagents
turn 1: How much was order A-100?
  tools:       ask_billing
  owner:       supervisor
  model calls: 4
  answer:      I'm forwarding your question to the billing specialist.

Billing brief (for the billing specialist):
- Customer asked: "How much was order A-100?"
- Please provide: total amount charged, currency, invoice number (if available), current payment/status (paid/refunded/pending), and a short breakdown of line items/discounts if available.
- Keep reply succinct and customer-facing.

Billing response:
- Total charged: 249.00 USD
- Current status: shipped
- No invoice number or line-item/discount breakdown available in records

If you'd like, I can ask billing to (a) resend an invoice if one exists, (b) provide a line-item breakdown, or (c) escalate to a human for deeper investigation. Which would you prefer?
turn 2: And what about A-101?
  tools:       ask_billing
  owner:       supervisor
  model calls: 4
  answer:      Billing update for order A-101:

- Total charged: 1,299.50 USD  
- Current status/payment: processing  
- No invoice number or line-item/discount breakdown available

Would you like me to ask billing to resend an invoice, provide a line-item breakdown, or escalate this to a human for further investigation?
  total model calls for subagents: 8

== handoffs
turn 1: How much was order A-100?
  tools:       transfer_to_billing, lookup_order
  owner:       billing
  model calls: 3
  answer:      Order A-100 total was 249.00 USD. The order status is shipped.
turn 2: And what about A-101?
  tools:       lookup_order
  owner:       billing
  model calls: 2
  answer:      Order A-101 total is 1,299.50 USD. The order status is processing.
  total model calls for handoffs: 5