Merge pull request #88 from RomiconEZ/fix-bugs

Add Logo and `ChatSession` Exception Handling + Fix Docs

Author: nizamovtimur

README.md

@@ -1,6 +1,6 @@
-# LLAMATOR
+# ![LLAMATOR](assets/LLAMATOR.svg)
 
-Red Teaming python-framework for testing chatbots and LLM-systems
+Red Teaming python-framework for testing chatbots and LLM systems
 
 [![License: CC BY-NC-SA 4.0](https://img.shields.io/badge/License-CC_BY--NC--SA_4.0-lightgrey.svg)](https://creativecommons.org/licenses/by-nc-sa/4.0/)
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/llamator)](https://pypi.org/project/llamator)

assets/LLAMATOR.svg

@@ -35,7 +35,7 @@ tested_model = llamator.ClientOpenAI(
     base_url="http://localhost:1234/v1",
     model="model-identifier",
     temperature=0.1,
-    system_prompts=["You are a defensive model."],
+    model_description="Model description",
 )
 
 # List of tests with the number of attempts
@@ -74,8 +74,8 @@ config = {
 
 # Start testing
 llamator.start_testing(
-    attack_model,
-    tested_model,
+    attack_model=attack_model,
+    tested_model=tested_model,
     config=config,
     tests_with_attempts=tests_with_attempts,
     multistage_depth=20,

docker/instruction.md

@@ -1,6 +1,6 @@
 # Project Overview
 
-LLAMATOR - Red Teaming python-framework for testing chatbots and LLM-systems
+LLAMATOR - Red Teaming python-framework for testing chatbots and LLM systems
 
 [![License: CC BY-NC-SA 4.0](https://img.shields.io/badge/License-CC_BY--NC--SA_4.0-lightgrey.svg)](https://creativecommons.org/licenses/by-nc-sa/4.0/)
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/llamator)](https://pypi.org/project/llamator)

docs/howtos.md

@@ -142,7 +142,6 @@
     "        self.model_description = model_description\n",
     "\n",
     "    def interact(self, history: List[Dict[str, str]], messages: List[Dict[str, str]]) -> Dict[str, str]:\n",
-    "        history += messages\n",
     "        try:\n",
     "            r = requests.post(self.api_url, json={\"question\": messages[-1][\"content\"]})\n",
     "            if r.status_code == 200:\n",
@@ -152,7 +151,6 @@
     "        except Exception as e:\n",
     "            print(e)\n",
     "            raise\n",
-    "        history.append(response_message)\n",
     "        return response_message"
    ]
   },
@@ -217,7 +215,7 @@
     "\n",
     "### Available Attacks\n",
     "\n",
-    "Check out the [documentation](https://romiconez.github.io/llamator/attacks_description.html) for an overview of available attacks."
+    "Check out the [documentation](https://romiconez.github.io/llamator/attack_descriptions.html) for an overview of available attacks."
    ]
   },
   {

docs/project_overview.md

@@ -11,7 +11,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {
     "id": "JuO12HZQQEnx"
    },
@@ -35,7 +35,7 @@
      "output_type": "stream",
      "text": [
       "Name: llamator\n",
-      "Version: 1.1.1\n",
+      "Version: 2.2.0\n",
       "Summary: Framework for testing vulnerabilities of large language models (LLM).\n",
       "Home-page: \n",
       "Author: \n",
@@ -239,11 +239,6 @@
     "\n",
     "        # Generating a Prompt by an Attack Model\n",
     "        for attack_prompt_index in range(self.num_attempts):\n",
-    "            chat = ChatSession(self.attack_config.attack_client.get_target_client())\n",
-    "            attack_prompt = chat.say(\n",
-    "                ORDINARY_PROMPT_TEMPLATE.format(description=self.client_config.get_model_description())\n",
-    "            )\n",
-    "            attack_prompts.append(attack_prompt)\n",
     "            yield StatusUpdate(\n",
     "                self.client_config,\n",
     "                self.test_name,\n",
@@ -252,6 +247,15 @@
     "                attack_prompt_index,\n",
     "                self.num_attempts,\n",
     "            )\n",
+    "            chat = ChatSession(self.attack_config.attack_client.get_target_client())\n",
+    "            attack_prompt = chat.say(\n",
+    "                ORDINARY_PROMPT_TEMPLATE.format(description=self.client_config.get_model_description())\n",
+    "            )\n",
+    "            if not attack_prompt:\n",
+    "                logger.warning(\"Error while attacking against target LLM (didn't receive response) ...\")\n",
+    "                self.status.report_error(str(attack_prompt_index))\n",
+    "            else:\n",
+    "                attack_prompts.append(attack_prompt)\n",
     "\n",
     "        # Executing attack prompts and checking the output of the attacked model\n",
     "        yield StatusUpdate(\n",
@@ -291,6 +295,9 @@
     "                    self.status.report_breach(attack_prompt, response)\n",
     "                else:\n",
     "                    self.status.report_resilient(attack_prompt, response)\n",
+    "            # Adding a response and status to a report\n",
+    "            responses.append(response)\n",
+    "            statuses.append(current_status)\n",
     "            yield StatusUpdate(\n",
     "                self.client_config,\n",
     "                self.test_name,\n",
@@ -299,9 +306,6 @@
     "                attack_prompt_index,\n",
     "                self.num_attempts,\n",
     "            )\n",
-    "            # Adding a response and status to a report\n",
-    "            responses.append(response)\n",
-    "            statuses.append(current_status)\n",
     "\n",
     "        # Adding data to the report\n",
     "        self._prepare_attack_data(attack_prompts, responses, statuses)\n",
@@ -332,7 +336,7 @@
      "data": {
       "text/plain": [
        "{'role': 'assistant',\n",
-       " 'content': 'Hello! How are you today? How can I assist you?'}"
+       " 'content': \"Hello! Nice to meet you. How are you today? Is there something you would like to talk about or ask me? I'm here to help. 😊\"}"
       ]
      },
      "execution_count": 9,
@@ -373,7 +377,7 @@
     "\n",
     "### Available Attacks\n",
     "\n",
-    "Check out the [documentation](https://romiconez.github.io/llamator/attacks_description.html) for an overview of available attacks."
+    "Check out the [documentation](https://romiconez.github.io/llamator/attack_descriptions.html) for an overview of available attacks."
    ]
   },
   {
@@ -399,7 +403,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Test progress ........................................:   0%|          | 0/1 [00:00<?, ?it/s]:   0%|          | 0/1 [00:00<?, ?it/s]: 100%|██████████| 1/1 [00:36<00:00, 36.06s/it]: 100%|██████████| 1/1 [00:36<00:00, 36.06s/it]\n"
+      "Test progress ........................................:   0%|          | 0/1 [00:00<?, ?it/s]:   0%|          | 0/1 [00:00<?, ?it/s]: 100%|██████████| 1/1 [01:14<00:00, 74.54s/it]: 100%|██████████| 1/1 [01:14<00:00, 74.55s/it]\n"
      ]
     },
     {
@@ -420,21 +424,22 @@
       "repeat_10_times:\n",
       "    Tests LLM for 10 times response vulnerability\n",
       "\n",
-      "Excel report created: ./artifacts\\LLAMATOR_run_2025-01-12_17-02-05\\attacks_report.xlsx\n"
+      "DISCLAIMER: Report may contain HARMFUL and OFFENSIVE language, reader discretion is recommended.\n",
+      "Excel report created: ./artifacts\\LLAMATOR_run_2025-02-17_11-30-47\\attacks_report.xlsx\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2025-01-12 17:02:46,865 [WARNING] [word_report_generator.py:352]: Test with in_code_name 'repeat_10_times' not found in JSON. Skipping file repeat_10_times.csv.\n"
+      "2025-02-17 11:32:05,188 [WARNING] [word_report_generator.py:352]: Test with in_code_name 'repeat_10_times' not found in JSON. Skipping file repeat_10_times.csv.\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Word report created: ./artifacts\\LLAMATOR_run_2025-01-12_17-02-05\\attacks_report.docx\n",
+      "Word report created: ./artifacts\\LLAMATOR_run_2025-02-17_11-30-47\\attacks_report.docx\n",
       "Thank you for using LLAMATOR!\n"
      ]
     }

examples/llamator-api.ipynb

@@ -292,7 +292,7 @@
     "\n",
     "### Available Attacks\n",
     "\n",
-    "Check out the [documentation](https://romiconez.github.io/llamator/attacks_description.html) for an overview of available attacks."
+    "Check out the [documentation](https://romiconez.github.io/llamator/attack_descriptions.html) for an overview of available attacks."
    ]
   },
   {