diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 98c839acf9..0c6df98d7c 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -55,4 +55,4 @@ jobs: output: both thresholds: '60 80' - name: "Run pre-commit hooks" - run: git ls-files -- . | xargs poetry run pre-commit run --files + run: poetry run pre-commit run --all-files diff --git a/.gitignore b/.gitignore index 59da102671..2a11dc1ba1 100644 --- a/.gitignore +++ b/.gitignore @@ -163,3 +163,6 @@ cython_debug/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ + +# PyRIT secrets file +.env diff --git a/pyrit/agent/red_teaming_bot.py b/pyrit/agent/red_teaming_bot.py index 0d31aa4a3c..9588d7a8b8 100644 --- a/pyrit/agent/red_teaming_bot.py +++ b/pyrit/agent/red_teaming_bot.py @@ -62,15 +62,21 @@ def get_session_chat_messages(self) -> list[ChatMessage]: return self._conversation_memory.get_chat_messages_with_session_id(session_id=self.session_id) def complete_chat_user(self, message: str, labels: list[str] = []) -> str: - session_messages: list[ChatMessage] = self.get_session_chat_messages() - if not session_messages: + message_list: list[ChatMessage] = [] + if not self.get_session_chat_messages(): # If there are no messages, then this is the first message of the conversation - self._add_chat_message_to_memory(ChatMessage(role="system", content=self._system_prompt), labels) + message_list.append(ChatMessage(role="system", content=self._system_prompt)) - self._add_chat_message_to_memory(ChatMessage(role="user", content=message), labels) + message_list.append(ChatMessage(role="user", content=message)) - response_msg = self._chat_engine.complete_chat(messages=session_messages) - self._add_chat_message_to_memory(ChatMessage(role="assistant", content=response_msg), labels) + response_msg = self._chat_engine.complete_chat(messages=message_list) + message_list.append(ChatMessage(role="assistant", content=response_msg)) + + self._conversation_memory.add_chat_messages_to_memory( + conversations=message_list, + session=self.session_id, + labels=self._global_memory_labels + labels, + ) return response_msg @@ -91,10 +97,3 @@ def is_conversation_complete(self) -> bool: # If the last message contains the conversation end token, then the conversation is complete return True return False - - def _add_chat_message_to_memory(self, message: ChatMessage, labels: list[str]): - self._conversation_memory.add_chat_message_to_memory( - conversation=message, - session=self.session_id, - labels=self._global_memory_labels + labels, - ) diff --git a/tests/test_red_teaming_bot.py b/tests/test_red_teaming_bot.py index 1b713659fd..35c0cc277b 100644 --- a/tests/test_red_teaming_bot.py +++ b/tests/test_red_teaming_bot.py @@ -1,7 +1,6 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -import os import pathlib from unittest.mock import patch @@ -13,6 +12,7 @@ from pyrit.chat import AzureOpenAIChat from pyrit.models import PromptTemplate from pyrit.memory import FileMemory +from pyrit.common.path import HOME_PATH @pytest.fixture @@ -41,7 +41,7 @@ def chat_completion_engine() -> AzureOpenAIChat: @pytest.fixture def red_teaming_bot(chat_completion_engine: AzureOpenAIChat, tmp_path: pathlib.Path): attack_strategy = PromptTemplate.from_yaml_file( - pathlib.Path(os.getcwd()) / "datasets" / "attack_strategies" / "red_team_chatbot_with_objective.yaml" + pathlib.Path(HOME_PATH) / "datasets" / "attack_strategies" / "red_team_chatbot_with_objective.yaml" ) file_memory = FileMemory(filepath=tmp_path / "test.json.memory") @@ -71,6 +71,17 @@ def test_complete_chat_user(red_teaming_bot: RedTeamingBot): assert "Instructions" in chats[0].content +def test_complete_chat_user_calls_complete_chat(red_teaming_bot: RedTeamingBot): + with patch.object(red_teaming_bot._chat_engine, "complete_chat") as mock: + mock.return_value = "Hello, this is a message sent by the assistant. How can i help you?" + red_teaming_bot.complete_chat_user("new chat") + + args, kwargs = mock.call_args + assert kwargs["messages"] is not None + assert kwargs["messages"][0].role == "system" + assert kwargs["messages"][1].content == "new chat" + + def test_is_conversation_complete_false(red_teaming_bot: RedTeamingBot): with patch.object(red_teaming_bot._chat_engine, "complete_chat") as mock: mock.return_value = "Hello, this is a message sent by the assistant. How can i help you?"