Fixing bug with red_teaming_bot

.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

.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

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,
-        )

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?"