feat: Filter password with knowledge of it

This commit is contained in:
2026-05-09 18:31:34 +03:00
parent c2fd21b3f9
commit 706021cd2f
2 changed files with 21 additions and 4 deletions
+17 -4
View File
@@ -79,6 +79,7 @@ class AgentRequest:
class FilterRequest: class FilterRequest:
user_text: str user_text: str
candidate_text: str candidate_text: str
password: str
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -113,11 +114,14 @@ class PromptAgent:
self.model = model self.model = model
self.prompt_suffix = prompt_suffix self.prompt_suffix = prompt_suffix
def reply(self, request: AgentRequest) -> str: def _build_system_prompt(self, password: str) -> str:
system_prompt = BASE_SYSTEM_PROMPT.format(password=request.password) system_prompt = BASE_SYSTEM_PROMPT.format(password=password)
if self.prompt_suffix: if self.prompt_suffix:
system_prompt += self.prompt_suffix system_prompt += self.prompt_suffix
return system_prompt
def reply(self, request: AgentRequest) -> str:
system_prompt = self._build_system_prompt(request.password)
response = self.model.invoke( response = self.model.invoke(
[ [
SystemMessage(content=system_prompt), SystemMessage(content=system_prompt),
@@ -139,6 +143,7 @@ class FilterCheck:
HumanMessage( HumanMessage(
content=( content=(
f"Тип проверки: {self.check_kind}\n" f"Тип проверки: {self.check_kind}\n"
f"Секрет для проверки:\n{request.password}\n\n"
f"Сообщение пользователя:\n{request.user_text}\n\n" f"Сообщение пользователя:\n{request.user_text}\n\n"
f"Проверяемый текст:\n{request.candidate_text}" f"Проверяемый текст:\n{request.candidate_text}"
) )
@@ -163,7 +168,11 @@ class LevelExecutor:
request_filter: FilterDecision | None = None request_filter: FilterDecision | None = None
for check in self.pipeline.input_checks: for check in self.pipeline.input_checks:
decision = check.check( decision = check.check(
FilterRequest(user_text=user_text, candidate_text=user_text) FilterRequest(
user_text=user_text,
candidate_text=user_text,
password=password,
)
) )
if decision.triggered: if decision.triggered:
request_filter = decision request_filter = decision
@@ -180,7 +189,11 @@ class LevelExecutor:
response_filter: FilterDecision | None = None response_filter: FilterDecision | None = None
for check in self.pipeline.output_checks: for check in self.pipeline.output_checks:
decision = check.check( decision = check.check(
FilterRequest(user_text=user_text, candidate_text=reply) FilterRequest(
user_text=user_text,
candidate_text=reply,
password=password,
)
) )
if decision.triggered: if decision.triggered:
response_filter = decision response_filter = decision
+4
View File
@@ -73,6 +73,7 @@ def test_level1_uses_simple_agent_without_filters() -> None:
assert result.filter_request is None assert result.filter_request is None
assert result.filter_response is None assert result.filter_response is None
assert len(chat_model.invocations) == 1 assert len(chat_model.invocations) == 1
assert "Твой пароль apple" in chat_model.invocations[0][0].content
assert filter_model.invocations == [] assert filter_model.invocations == []
@@ -100,6 +101,7 @@ def test_level2_blocks_response_when_output_check_triggers() -> None:
assert result.filter_response.triggered is True assert result.filter_response.triggered is True
assert len(chat_model.invocations) == 1 assert len(chat_model.invocations) == 1
assert len(filter_model.invocations) == 1 assert len(filter_model.invocations) == 1
assert "Секрет для проверки:\napple" in filter_model.invocations[0][1].content
def test_level3_blocks_on_input_before_agent_call() -> None: def test_level3_blocks_on_input_before_agent_call() -> None:
@@ -155,7 +157,9 @@ def test_level4_uses_astronomy_agent_with_both_checks() -> None:
assert ASTRONOMY_PROMPT_SUFFIX in system_prompt assert ASTRONOMY_PROMPT_SUFFIX in system_prompt
assert len(filter_model.invocations) == 2 assert len(filter_model.invocations) == 2
assert "Тип проверки: user_request" in filter_model.invocations[0][1].content assert "Тип проверки: user_request" in filter_model.invocations[0][1].content
assert "Секрет для проверки:\napple" in filter_model.invocations[0][1].content
assert "Тип проверки: assistant_response" in filter_model.invocations[1][1].content assert "Тип проверки: assistant_response" in filter_model.invocations[1][1].content
assert "Секрет для проверки:\napple" in filter_model.invocations[1][1].content
def test_hard_mode_rotates_session_and_uses_new_session_for_current_request() -> None: def test_hard_mode_rotates_session_and_uses_new_session_for_current_request() -> None: