Agent Guardrails & Safety

import json
import logging
from pathlib import Path

from langchain.agents import create_agent
from langchain.agents.middleware import (
    ModelCallLimitMiddleware,
    PIIMiddleware,
    ToolCallLimitMiddleware,
    wrap_tool_call,
)
from langchain_core.messages import ToolMessage
from langchain_ollama import ChatOllama

audit = logging.getLogger("docsassistant.audit")

WRITE_TOOLS = {"github_open_issue"}
POLICY = {                                   # role -> tools that role may cause the agent to call
    "anonymous": {"docs_search", "docs_get_section"},
    "reader": {"docs_search", "docs_get_section", "github_open_issue"},
}


KILL_SWITCH = Path("/etc/docsassistant/writes-enabled")   # e.g. a mounted ConfigMap key


def writes_enabled() -> bool:                # the kill switch: read on every call, so no restart
    try:
        return KILL_SWITCH.read_text().strip() == "true"
    except OSError:                          # missing or unreadable file: writes stay off
        return False


@wrap_tool_call
def policy_and_audit(request, handler):
    call, role = request.tool_call, request.runtime.context.role
    allowed = call["name"] in POLICY.get(role, set())
    if call["name"] in WRITE_TOOLS and not writes_enabled():
        allowed = False
    audit.info(json.dumps({"tool": call["name"], "args": call["args"], "role": role,
                           "decision": "allow" if allowed else "deny"}))
    if not allowed:
        return ToolMessage(f"Denied by policy: '{call['name']}' is not available here. "
                           "Tell the user what you would have done instead.",
                           tool_call_id=call["id"], status="error")
    result = handler(request)
    audit.info(json.dumps({"tool": call["name"], "outcome": str(result.content)[:200]}))
    return result


agent = create_agent(
    ChatOllama(model="llama3.1:8b", temperature=0),
    tools=[docs_search, docs_get_section, github_open_issue],
    context_schema=UserContext,                                      # carries the caller's role
    middleware=[
        policy_and_audit,                                            # policy layer + audit trail
        ToolCallLimitMiddleware(tool_name="github_open_issue", run_limit=1),   # one issue per run
        ToolCallLimitMiddleware(run_limit=15),                       # overall tool budget
        ModelCallLimitMiddleware(run_limit=20),                      # model-call budget
        PIIMiddleware("email", strategy="redact", apply_to_input=True),
        # ...plus HumanInTheLoopMiddleware for github_open_issue (see Human-in-the-Loop)
    ],
)

agent.invoke({"messages": [("user", "The memory page has the wrong default window size.")]},
             context=UserContext(user_id="alice", role="reader"))   # role comes from your auth layer

import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.guardrail.InputGuardrail;
import dev.langchain4j.guardrail.InputGuardrailResult;
import dev.langchain4j.service.guardrail.InputGuardrails;

public class OversizeAndInjectionGuard implements InputGuardrail {

    private static final Pattern OVERRIDE = Pattern.compile(
            "(?i)ignore(\\s+(all|your|previous|prior|the))+\\s+(rules|instructions)");

    @Override
    public InputGuardrailResult validate(UserMessage userMessage) {
        String text = userMessage.singleText();
        if (text.length() > 4000) {
            return fatal("Question too long; please shorten it.");
        }
        if (OVERRIDE.matcher(text).find()) {
            return fatal("This request cannot be processed.");
        }
        return success();
    }
}

interface DocsAssistant {
    @InputGuardrails(OversizeAndInjectionGuard.class)
    String chat(@MemoryId String userId, @dev.langchain4j.service.UserMessage String message);
}

DocsAssistant assistant = AiServices.builder(DocsAssistant.class)
        .chatModel(chatModel)
        .tools(new DocsTools())
        .maxSequentialToolsInvocations(15)          // tool-call budget per request
        .build();

import org.springframework.ai.chat.client.advisor.SafeGuardAdvisor;
import org.springframework.ai.model.tool.ToolCallLimitBehavior;
import org.springframework.ai.model.tool.ToolCallingManager;

@Bean
ToolCallingManager toolCallingManager() {
    return ToolCallingManager.builder()
            .maxCallsPerTool(10)
            .maxCallsPerTool("github_open_issue", 1)     // at most one issue per turn
            .maxTotalToolCalls(20)
            .onLimitExceeded(ToolCallLimitBehavior.RETURN_ERROR_RESPONSE)
            .build();
}

@Bean
ChatClient docsAssistant(ChatClient.Builder builder, DocsTools docsTools) {
    return builder
            .defaultAdvisors(SafeGuardAdvisor.builder()
                    .sensitiveWords(List.of("internal use only", "confidential"))
                    .failureResponse("I can't help with that request.")
                    .build())
            .defaultTools(docsTools)
            .build();
}