Lab 03 · The reliability kit

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 03: a reliability kit for the Northwind support agent.

PII redaction, an audit trail of our own, a tool budget, a fallback model and retries,
composed in the order that works.

Run:  python reliability_kit.py
"""

import operator
from typing import Annotated, Any, Awaitable, Callable, NotRequired

from langchain.agents import create_agent
from langchain.agents.middleware import (
    AgentMiddleware,
    AgentState,
    ModelFallbackMiddleware,
    ModelRetryMiddleware,
    PIIMiddleware,
    ToolCallLimitMiddleware,
    ToolCallRequest,
)
from langchain.messages import ToolMessage
from langchain.tools import tool
from langchain_core.language_models import BaseChatModel
from langgraph.types import Command

from model_setup import get_model

@tool
def get_order_status(order_id: str) -> str:
    """Look up the shipping status of an order, such as A-100."""
    return f"Order {order_id.upper()} shipped via DHL, arriving 2026-09-15."


@tool
def get_invoice_total(order_id: str) -> str:
    """Return the amount charged for an order. Use for questions about price or total."""
    return f"Order {order_id.upper()} total: 249.00 USD."

class AuditState(AgentState):
    """AgentState plus an audit trail. `operator.add` so parallel tool calls both append."""

    audit_log: NotRequired[Annotated[list[dict[str, Any]], operator.add]]

def audit_entry(request: ToolCallRequest, result: ToolMessage) -> dict[str, Any]:
    """One line of the audit trail for an executed tool call."""
    name, args = request.tool_call["name"], request.tool_call["args"]
    return {"tool": name, "args": args, "status": result.status}

class AuditMiddleware(AgentMiddleware):
    """Records every executed tool call in agent state.

    A `Command` is how a wrapper hook makes a persistent update; the ToolMessage has to
    travel inside it, or the model never sees the tool result.
    """

    state_schema = AuditState

    def wrap_tool_call(
        self, request: ToolCallRequest, handler: Callable[[ToolCallRequest], ToolMessage]
    ) -> Command:
        """Run the tool, then record it (sync path: invoke)."""
        result = handler(request)
        return Command(update={"audit_log": [audit_entry(request, result)], "messages": [result]})

    async def awrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage]],
    ) -> Command:
        """Run the tool, then record it (async path: ainvoke)."""
        result = await handler(request)
        return Command(update={"audit_log": [audit_entry(request, result)], "messages": [result]})

SYSTEM_PROMPT = (
    "You are the Northwind support assistant. Answer briefly and only about orders. "
    "Never guess a status or an amount: always call the matching tool."
)

def build_agent(model: BaseChatModel | None = None, fallback: BaseChatModel | None = None) -> Any:
    """Build the agent with the reliability middleware stack."""
    return create_agent(
        model=model or get_model(),
        tools=[get_order_status, get_invoice_total],
        system_prompt=SYSTEM_PROMPT,
        middleware=[
            # Outermost: the raw email never reaches the model, the tools or the trace.
            PIIMiddleware("email", strategy="redact"),
            AuditMiddleware(),
            ToolCallLimitMiddleware(run_limit=4, exit_behavior="end"),
            # Fallback sits OUTSIDE retry, so the fallback is tried only once the retries
            # are spent, and retry must re-raise, or the fallback never fires.
            ModelFallbackMiddleware(fallback or get_model()),
            ModelRetryMiddleware(max_retries=2, on_failure="error"),
        ],
    )

def ask(agent: Any, question: str) -> None:
    """Run one question and print what the middleware did."""
    result = agent.invoke({"messages": [{"role": "user", "content": question}]})
    print("stored question:", result["messages"][0].content)
    for entry in result.get("audit_log", []):
        print("audit log entry:", entry)
    print("answer:", result["messages"][-1].text)


if __name__ == "__main__":
    question = "I am ada@northwind.example - where is A-100 and how much was it?"

    print("=== Normal run: the primary model answers")
    ask(build_agent(), question)

    # A deployment name that does not exist makes every primary call fail.
    print("\n=== Broken primary: retries run out, then the fallback answers")
    ask(build_agent(model=get_model(deployment="no-such-deployment")), question)

=== Normal run: the primary model answers
stored question: I am [REDACTED_EMAIL] - where is A-100 and how much was it?
audit log entry: {'tool': 'get_order_status', 'args': {'order_id': 'A-100'}, 'status': 'success'}
audit log entry: {'tool': 'get_invoice_total', 'args': {'order_id': 'A-100'}, 'status': 'success'}
answer: Order A-100 shipped via DHL and is scheduled to arrive 2026-09-15. Total charged: 249.00 USD.

=== Broken primary: retries run out, then the fallback answers
stored question: I am [REDACTED_EMAIL] - where is A-100 and how much was it?
audit log entry: {'tool': 'get_order_status', 'args': {'order_id': 'A-100'}, 'status': 'success'}
audit log entry: {'tool': 'get_invoice_total', 'args': {'order_id': 'A-100'}, 'status': 'success'}
answer: Order A-100 shipped via DHL — arriving 2026-09-15. Total charged: 249.00 USD.