package ch.phoenixtechnologies.lc4j.workshop.guardrails;

import dev.langchain4j.data.message.UserMessage;
import io.quarkiverse.langchain4j.guardrails.InputGuardrail;
import io.quarkiverse.langchain4j.guardrails.InputGuardrailResult;
import jakarta.enterprise.context.ApplicationScoped;

@ApplicationScoped
public class PromptInjectionGuard implements InputGuardrail {
    private final PromptInjectionDetectionService service;

    public PromptInjectionGuard(PromptInjectionDetectionService service) {
        this.service = service;
    }

    @Override
    public InputGuardrailResult validate(UserMessage userMessage) {
        double result = service.isInjection(userMessage.singleText());
        if (result > 0.7) {
            return failure("Prompt injection detected");
        }
        return success();
    }
}
