Guardrails & Prompt Injection

import html
import re
import unicodedata

INVISIBLE = dict.fromkeys(map(ord, "​‌‍⁠‪‫‬‭‮⁦⁧⁨⁩"))
HTML_COMMENT = re.compile(r"<!--.*?-->", re.S)
HIDDEN_BLOCK = re.compile(r"<(script|style)\b.*?</\1>|<[^>]+style=\"[^\"]*(display:\s*none|visibility:\s*hidden)[^\"]*\"[^>]*>.*?</[^>]+>",
                          re.S | re.I)
SUSPICIOUS = re.compile(r"\b(ignore (all |any )?(previous|prior|above) instructions|you are now|system prompt|"
                        r"do not (tell|reveal to) the user|developer mode)\b", re.I)


def sanitise(raw_html_or_text: str) -> str:
    text = HIDDEN_BLOCK.sub(" ", raw_html_or_text)
    text = HTML_COMMENT.sub(" ", text)
    text = re.sub(r"<[^>]+>", " ", text)                  # drop remaining tags after the hidden ones are gone
    text = html.unescape(text)
    text = unicodedata.normalize("NFKC", text).translate(INVISIBLE)
    return re.sub(r"[ \t]+", " ", text).strip()


def injection_flags(text: str) -> list[str]:
    return [m.group(0) for m in SUSPICIOUS.finditer(text)]

import os

from langchain.agents import create_agent
from langchain.agents.middleware import AgentMiddleware, PIIMiddleware, hook_config
from langchain.chat_models import init_chat_model
from langchain_core.messages import AIMessage, RemoveMessage

REFUSAL = "I can only help with questions about the Irurueta documentation."
WITHHELD = "I cannot share that answer. The request has been logged for review."


class SafetyClassifierMiddleware(AgentMiddleware):
    """Screens the user turn before the agent runs and the final answer after it."""

    def __init__(self, classifier):
        super().__init__()
        self.classifier = classifier          # a safety model that answers "safe" or "unsafe <categories>"

    def _unsafe(self, role: str, text: str) -> bool:
        verdict = self.classifier.invoke([{"role": role, "content": text}]).text
        return verdict.strip().lower().startswith("unsafe")

    @hook_config(can_jump_to=["end"])
    def before_agent(self, state, runtime):
        if self._unsafe("user", state["messages"][-1].text):
            return {"messages": [AIMessage(content=REFUSAL)], "jump_to": "end"}
        return None

    def after_agent(self, state, runtime):
        last = state["messages"][-1]
        if self._unsafe("assistant", last.text):
            # replace the answer in the persisted state, not just in the response
            return {"messages": [RemoveMessage(id=last.id), AIMessage(content=WITHHELD)]}
        return None


def build_guarded_agent(tools, checkpointer):
    return create_agent(
        init_chat_model(os.environ["CHAT_MODEL"]),
        tools=tools,
        middleware=[
            SafetyClassifierMiddleware(init_chat_model(os.environ["GUARD_MODEL"], temperature=0)),
            PIIMiddleware("email", strategy="redact", apply_to_input=True, apply_to_output=True),
            PIIMiddleware("credit_card", strategy="block", apply_to_input=True),
        ],
        checkpointer=checkpointer,
    )

package com.example.docsassistant.rag.guardrails;

import java.util.List;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.SafeGuardAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.core.Ordered;

public class SafetyClassifierAdvisor implements CallAdvisor {

    static final String REFUSAL = "I can only help with questions about the Irurueta documentation.";

    private final ChatClient guard;

    public SafetyClassifierAdvisor(ChatClient.Builder guardModelBuilder) {
        this.guard = guardModelBuilder.build();
    }

    @Override
    public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
        String question = request.prompt().getUserMessage().getText();
        String verdict = guard.prompt().user(question).call().content();
        if (verdict != null && verdict.strip().toLowerCase().startsWith("unsafe")) {
            ChatResponse refusal = new ChatResponse(List.of(new Generation(new AssistantMessage(REFUSAL))));
            return ChatClientResponse.builder().chatResponse(refusal).context(request.context()).build();
        }
        return chain.nextCall(request);
    }

    @Override
    public String getName() {
        return "safetyClassifier";
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE + 100;       // before memory and RAG advisors
    }

    public static SafeGuardAdvisor wordFilter() {
        return SafeGuardAdvisor.builder()
                .sensitiveWords(List.of("password dump", "credit card numbers"))
                .failureResponse(REFUSAL)
                .order(Ordered.HIGHEST_PRECEDENCE)
                .build();
    }
}

package com.example.docsassistant.l4j.guardrails;

import java.util.regex.Pattern;

import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.guardrail.InputGuardrail;
import dev.langchain4j.guardrail.InputGuardrailResult;
import dev.langchain4j.guardrail.OutputGuardrail;
import dev.langchain4j.guardrail.OutputGuardrailResult;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.service.guardrail.InputGuardrails;
import dev.langchain4j.service.guardrail.OutputGuardrails;

public final class GuardedAssistant {

    public interface DocsAssistant {
        @InputGuardrails(SafetyInputGuardrail.class)
        @OutputGuardrails({NoSecretsOutputGuardrail.class, CanaryOutputGuardrail.class})
        String ask(String question);
    }

    /** Classifies the question with a safety model; needs a public no-arg constructor or a DI framework. */
    public static class SafetyInputGuardrail implements InputGuardrail {
        static ChatModel guardModel;                 // wired at start-up in this example

        @Override
        public InputGuardrailResult validate(UserMessage userMessage) {
            String verdict = guardModel.chat(userMessage.singleText());
            return verdict.strip().toLowerCase().startsWith("unsafe")
                    ? fatal("The question violates the usage policy")
                    : success();
        }
    }

    /** Rejects answers that look like they contain credentials. */
    public static class NoSecretsOutputGuardrail implements OutputGuardrail {
        private static final Pattern SECRET = Pattern.compile("(?i)(api[_-]?key|secret|password)\\s*[:=]\\s*\\S{8,}");

        @Override
        public OutputGuardrailResult validate(AiMessage responseFromLLM) {
            return SECRET.matcher(responseFromLLM.text()).find()
                    ? reprompt("Answer contained a credential", "Answer again without any credentials or secrets.")
                    : success();
        }
    }

    /** Detects system-prompt leakage through a canary string placed in the system prompt. */
    public static class CanaryOutputGuardrail implements OutputGuardrail {
        static final String CANARY = "canary-7f3a";

        @Override
        public OutputGuardrailResult validate(AiMessage responseFromLLM) {
            return responseFromLLM.text().contains(CANARY) ? fatal("System prompt leaked") : success();
        }
    }

    private GuardedAssistant() {
    }
}