65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from agent_framework.guardrails.framework_llm_client import classify_with_framework_llm
|
|
from agent_framework.guardrails.calibrated.llm_client import _BINARY_BLOCK_DIGIT, _BINARY_TASKS
|
|
from agent_framework.guardrails.calibrated.prompts.pinj import build_pinj_prompt
|
|
from agent_framework.guardrails.calibrated.prompts.revprec import build_revprec_prompt
|
|
|
|
|
|
class FakeLLM:
|
|
def __init__(self, content: str):
|
|
self.content = content
|
|
self.profile_resolver = None
|
|
|
|
async def ainvoke(self, messages, **kwargs):
|
|
return self.content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("task", "content", "allowed"),
|
|
[
|
|
("PINJ", "1", True),
|
|
("PINJ", "0", False),
|
|
("PINJ", "prosa 1 e veredito 0", False),
|
|
("PINJ", "sem digito", True),
|
|
("PINJ", "", True),
|
|
("REVPREC", "1", False),
|
|
("REVPREC", "0", True),
|
|
("REVPREC", "Resposta: 1", False),
|
|
("REVPREC", "não sei", True),
|
|
("REVPREC", "", True),
|
|
],
|
|
)
|
|
async def test_binary_guardrail_semantics(task: str, content: str, allowed: bool, monkeypatch):
|
|
monkeypatch.setenv("USE_MOCK_LLM", "false")
|
|
out = await classify_with_framework_llm(
|
|
FakeLLM(content), task, {"text": "texto", "context": {}}, profile_name="guardrail"
|
|
)
|
|
assert bool(out["allowed"]) is allowed
|
|
|
|
|
|
def test_binary_polarity_contract():
|
|
assert _BINARY_BLOCK_DIGIT["REVPREC"] == "1"
|
|
assert _BINARY_BLOCK_DIGIT.get("PINJ", "0") == "0"
|
|
assert "PINJ" in _BINARY_TASKS and "REVPREC" in _BINARY_TASKS
|
|
|
|
|
|
def test_pinj_prompt_is_binary_and_keeps_examples():
|
|
prompt = build_pinj_prompt("texto qualquer")
|
|
assert "Responda APENAS um caractere" in prompt
|
|
assert '"allowed"' not in prompt
|
|
assert prompt.count(" Texto: ") == 17
|
|
assert prompt.count(" Saída: 0") == 11
|
|
assert prompt.count(" Saída: 1") == 6
|
|
|
|
|
|
def test_revprec_prompt_preserves_boundaries():
|
|
prompt = build_revprec_prompt("FALA_NOVA", "\nHistorico da conversa:\n[assistant] CANCELEI_ANTES\n")
|
|
assert "APENAS 1 ou 0" in prompt
|
|
assert 'somente a fala do bloco "Resposta:"' in prompt
|
|
assert "PROTOCOLO" in prompt
|
|
assert "DESCRIÇÃO DA FATURA" in prompt
|