Update app.py

app.py

@@ -318,6 +318,44 @@ def attack_sus(message, max_tokens, temperature, top_p, available_pic_str):
         top_p=top_p,
     )
 
+    pic_list = re.findall(r'\[([^\]]+)\]', AVAILABLE_PICTURES)
+    pic_list.append("DEFAULT IMAGE")
+    
+    valid_img = False
+    image_name = re.search(r'\[(.*?)\]', output.choices[0].message.content).group(1)
+    for x in pic_list:
+        if image_name == pic_list:
+            valid_img = True
+
+    if valid_img = False:
+        output2 = llm_client.chat.completions.create(
+        	model="Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2",
+        	messages=[
+                {
+                    "role": "system",
+                    "content": sus_copy
+                },
+        		{
+        			"role": "user",
+        			"content": message
+        		},
+                {
+                    "role": "assistant",
+                    "content": output.choices[0].message.content
+                },
+                {
+                    "role": "user",
+                    "content": "This is a system message. The image you selected to reply does not exists in the list. Please try again with a valid image name."
+                }
+        	],
+            max_tokens=max_tokens,
+            stream=False,
+            temperature=temperature,
+            top_p=top_p,
+        )
+
+        return output2.choices[0].message.content
+    
     print("final sus prompt: ", sus_copy)
     return output.choices[0].message.content