[Ch 5] Tools Integration, Guardrails & Safety Patterns

# ❌ Bad: pollutes the conversation history
user_message = f"[USER_ID: {user_id}] {actual_query}"

# agent/context.py
from pydantic import BaseModel

class AgentContext(BaseModel):
    """Typed context passed through every agent call without touching messages."""
    user_id: str
    session_id: str
    permissions: list[str] = []   # e.g. ["read", "write", "delete"]
    org_id: str = ""

# agent/tools.py — accessing context inside a tool
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from .context import AgentContext

@tool
def delete_task(task_id: str, config: RunnableConfig) -> str:
    """Permanently delete a task. Requires 'delete' permission."""
    # LangGraph automatically injects `config` — just add it as a parameter
    ctx_data = config.get("configurable", {}).get("agent_context")
    if ctx_data is None:
        return "Error: No agent context provided."

    ctx = AgentContext(**ctx_data) if isinstance(ctx_data, dict) else ctx_data

    # Permission check using context — no need to pass user_id as a tool arg
    if "delete" not in ctx.permissions:
        return f"Error: User {ctx.user_id} does not have delete permission."

    if task_id not in _tasks:
        return f"Error: Task '{task_id}' not found."

    task_title = _tasks.pop(task_id)["title"]
    return f"Deleted task '{task_title}' (ID: {task_id})."

# agent/main.py — how to pass context at call time
from .context import AgentContext

ctx = AgentContext(
    user_id="user-42",
    session_id="sess-abc",
    permissions=["read", "write"],   # no "delete" for this user
    org_id="acme-corp",
)

config = {
    "configurable": {
        "thread_id": "user-42-session-1",
        "agent_context": ctx.model_dump(),   # serialize for LangGraph
    }
}

app.invoke({"messages": [HumanMessage(content="Delete task task-001")]}, config=config)
# → Tool returns: "Error: User user-42 does not have delete permission."

pip install nemoguardrails

guardrails/
├── config.yml
└── main.co         # Colang flow definitions

# guardrails/config.yml
models:
  - type: main
    engine: openai
    model: gpt-4o-mini

rails:
  input:
    flows:
      - check harmful input
  output:
    flows:
      - check harmful output

# guardrails/main.co

# ── Input rail: block harmful requests ───────────────────────────────────────
define flow check harmful input
  $is_harmful = execute check_input_for_harm

  if $is_harmful
    bot refuse harmful input
    stop

define bot refuse harmful input
  "I'm sorry, I can't help with that request."

# ── Output rail: block harmful responses ─────────────────────────────────────
define flow check harmful output
  $is_harmful = execute check_output_for_harm

  if $is_harmful
    bot provide safe response
    stop

define bot provide safe response
  "I'm not able to provide that information."

# agent/guardrail.py
import os
from nemoguardrails import RailsConfig, LLMRails

def load_rails(config_path: str = "guardrails/") -> LLMRails:
    """Load NeMo guardrails from config directory."""
    config = RailsConfig.from_path(config_path)
    return LLMRails(config)

rails = load_rails()

# agent/main.py — wrapped agent call with guardrails
async def safe_chat(app, thread_id: str, user_input: str) -> str:
    """Run user input through guardrails before and after the agent."""

    # ── 1. Input guardrail ────────────────────────────────────────────────────
    input_check = await rails.generate_async(
        messages=[{"role": "user", "content": user_input}]
    )
    # NeMo returns the guardrail response; if it's a refusal, return immediately
    if _is_refusal(input_check):
        return input_check["content"]

    # ── 2. Run the agent ──────────────────────────────────────────────────────
    config = {"configurable": {"thread_id": thread_id}}
    final_state = app.invoke(
        {"messages": [HumanMessage(content=user_input)]},
        config=config,
    )
    agent_response = final_state["messages"][-1].content

    # ── 3. Output guardrail ───────────────────────────────────────────────────
    output_check = await rails.generate_async(
        messages=[
            {"role": "user", "content": user_input},
            {"role": "assistant", "content": agent_response},
        ]
    )
    if _is_refusal(output_check):
        return "I can't provide that response."

    return agent_response


def _is_refusal(response: dict) -> bool:
    """Detect if NeMo triggered a refusal rail."""
    content = response.get("content", "")
    refusal_phrases = ["I'm sorry, I can't", "I'm not able to"]
    return any(phrase in content for phrase in refusal_phrases)

# agent/tools.py — a tool that interrupts for confirmation
from langgraph.types import interrupt, Command

class DeleteCompletedInput(BaseModel):
    confirm_message: str = Field(
        description="Message to show the user when asking for confirmation"
    )

@tool("delete_completed_tasks", args_schema=DeleteCompletedInput)
def delete_completed_tasks(confirm_message: str) -> str:
    """Delete all tasks marked as 'done'. Requires human confirmation."""
    done_tasks = [t for t in _tasks.values() if t["status"] == "done"]

    if not done_tasks:
        return "No completed tasks to delete."

    # ── INTERRUPT: pause and ask the human ───────────────────────────────────
    approval = interrupt({
        "question": f"About to permanently delete {len(done_tasks)} completed task(s). Continue?",
        "tasks_to_delete": [t["title"] for t in done_tasks],
        "action": "delete_completed_tasks",
    })

    # Execution resumes here after the human responds
    if not approval.get("approved", False):
        return "Deletion cancelled by user."

    # Execute the deletion
    for task in done_tasks:
        del _tasks[task["id"]]

    return f"Successfully deleted {len(done_tasks)} completed task(s)."

# agent/main.py — handling the interrupt in your API/CLI layer
from langgraph.types import Command

def chat_with_hitl(app, thread_id: str, user_input: str):
    config = {"configurable": {"thread_id": thread_id}}

    # First invoke — may pause at an interrupt
    for event in app.stream(
        {"messages": [HumanMessage(content=user_input)]},
        config=config,
        stream_mode="values",
    ):
        if "__interrupt__" in event:
            interrupt_data = event["__interrupt__"][0].value
            print(f"\n⚠️  CONFIRMATION REQUIRED")
            print(f"   {interrupt_data['question']}")
            print(f"   Tasks: {', '.join(interrupt_data['tasks_to_delete'])}")

            answer = input("   Approve? (yes/no): ").strip().lower()
            approved = answer in ("yes", "y")

            # Resume the graph with the human's decision
            for resume_event in app.stream(
                Command(resume={"approved": approved}),
                config=config,
                stream_mode="values",
            ):
                last = resume_event.get("messages", [])
                if last:
                    print(f"\nAssistant: {last[-1].content}")
            return

    # No interrupt — normal completion
    state = app.get_state(config)
    if state.values["messages"]:
        print(f"\nAssistant: {state.values['messages'][-1].content}")

# .env.example
OPENAI_API_KEY=your-api-key-here