chg: [images_engines] improve prompts + ad prompt to detect csam domains and images

Author: Terrtia

bin/lib/images_engine.py

@@ -38,7 +38,7 @@ def get_image_obj(obj_gid):
 
 def create_ollama_domain_data(model, descriptions):
     return json.dumps({'model': model,
-                       'prompt': f'From this list of images descriptions, Can you please describe this domain and check if it\'s related to child exploitation?\n\n{descriptions}',
+                       'prompt': f'From this list of images descriptions, Describe this domain\n\n{descriptions}',
                        'stream': False
                        })
 
@@ -49,6 +49,25 @@ def create_ollama_image_data(model, images):
                        'images': images
                        })
 
+def create_ollama_description_csam_classification(model, description):
+    return json.dumps({'model': model,
+                       'prompt': f'Does this description involve CE or CSAM? Answer "Yes" or "No".\nDescription: {description}',
+                       'stream': False
+                       })
+
+def create_ollama_image_csam_classification(model, images):
+    return json.dumps({'model': model,
+                       'prompt': 'Does this image involve CE or CSAM? Answer "Yes" or "No".',
+                       'stream': False,
+                       'images': images
+                       })
+
+def create_ollama_domain_csam_classification(model, descriptions):
+    return json.dumps({'model': model,
+                       'prompt': f'Is this website domain associated with child exploitation? Respond with only "Yes" or "No"\n\n{descriptions}',
+                       'stream': False
+                       })
+
 # screenshot + image
 def api_get_image_description(obj_gid):
     model = get_default_image_description_model()
@@ -72,7 +91,7 @@ def api_get_image_description(obj_gid):
         return {"status": "error", "reason": f"ollama requests error: {e}"}, 400
     if res.status_code != 200:
         # TODO LOG
-        return {"status": "error", "reason": f" llama requests error: {res.status_code}, {res.text}"}, 400
+        return {"status": "error", "reason": f"ollama requests error: {res.status_code}, {res.text}"}, 400
     else:
         r = res.json()
         if r:
@@ -117,7 +136,7 @@ def get_domain_description(domain_id, reprocess=True):
         return {"status": "error", "reason": f"ollama requests error: {e}"}, 400
     if res.status_code != 200:
         # TODO LOG
-        return {"status": "error", "reason": f" llama requests error: {res.status_code}, {res.text}"}, 400
+        return {"status": "error", "reason": f"ollama requests error: {res.status_code}, {res.text}"}, 400
     else:
         r = res.json()
         if r:
@@ -152,6 +171,98 @@ def _create_image_description():
         # print(f'{done}/{total}        {progress}%')
 
 
+def update_domain_description(domain, model):
+    domain.delete_description(model)
+    search_engine.remove_document('desc-dom', domain.get_global_id())
+    get_domain_description(domain.get_id(), reprocess=False)
+
+def update_domains_descriptions():
+    nb_domains = Domains.get_nb_domains_up_by_type('onion') + Domains.get_nb_domains_up_by_type('web')
+    model = get_default_image_description_model()
+    done = 0
+    for domain in Domains.get_domain_up_iterator():
+        update_domain_description(domain, model)
+        done += 1
+        progress = int(done * 100 / nb_domains)
+        print(f'{done}/{nb_domains}        {progress}%')
+    search_engine.delete_index('desc-dom')
+
+
+def check_is_image_csam(obj_gid, image_description=False):
+    model = get_default_image_description_model()
+
+    image = get_image_obj(obj_gid)
+    if not image:
+        return {"status": "error", "reason": "Unknown image"}, 404
+
+    headers = {"Connection": "close", 'Content-Type': 'application/json', 'Accept': 'application/json'}
+    is_csam = None
+
+    # Check if image description is CSAM related
+    if image_description:
+        description = api_get_image_description(obj_gid)
+        if description[1] == 200:
+            description = description[0]
+        data = create_ollama_description_csam_classification(model, description)
+
+    # Check If image content is CSAM
+    else:
+        b64 = image.get_base64()
+        if not b64:
+            return {"status": "error", "reason": "No Content"}, 404
+        data = create_ollama_image_csam_classification(model, [b64])
+
+    if data:
+        try:
+            res = requests.post(f'{OLLAMA_URL}/api/generate', data=data, headers=headers)
+        except Exception as e:
+            return {"status": "error", "reason": f"ollama requests error: {e}"}, 400
+        if res.status_code != 200:
+            # TODO LOG
+            return {"status": "error", "reason": f"ollama requests error: {res.status_code}, {res.text}"}, 400
+        else:
+            r = res.json()
+            if r:
+                res = image.add_description_model(model, r['response'])
+                if res == 'YES':
+                    is_csam = True
+                elif res == 'NO':
+                    is_csam = False
+
+    # TODO LOG NONE result
+    if is_csam:
+        image.add_tag('dark-web:topic="pornography-child-exploitation"')
+        print(obj_gid, is_csam)
+
+    return is_csam, 200
+
+def check_images_csam(image_description=False):
+    for image in Images.get_all_images_objects():
+        check_is_image_csam(image.get_global_id(), image_description=image_description)
+
+
+def check_if_domain_csam(domain_id):
+    model = get_default_image_description_model()
+    domain = Domains.Domain(domain_id)
+    description = get_domain_description(domain_id)
+
+    headers = {"Connection": "close", 'Content-Type': 'application/json', 'Accept': 'application/json'}
+    try:
+        res = requests.post(f'{OLLAMA_URL}/api/generate', data=create_ollama_domain_csam_classification(model, description), headers=headers)
+    except Exception as e:  # TODO LOG
+        return {"status": "error", "reason": f"ollama requests error: {e}"}, 400
+    if res.status_code != 200:
+        # TODO LOG
+        return {"status": "error", "reason": f"ollama requests error: {res.status_code}, {res.text}"}, 400
+    else:
+        r = res.json()
+        if r:
+            if r['response'] == 'yes':
+                print('yes')
+                domain.add_tag('dark-web:topic="pornography-child-exploitation"')
+
+
 if __name__ == '__main__':
-    # _create_domains_up_description()
-    _create_image_description()
+    update_domains_descriptions()
+    # check_images_csam()
+    # _create_image_description()

bin/lib/objects/Domains.py

@@ -356,6 +356,11 @@ def get_description(self, model=None):
             description = description.replace("`", ' ')
         return description
 
+    def delete_description(self, model=None):
+        if model is None:
+            model = get_default_image_description_model()
+        self._delete_field(f'desc:{model}')
+
     ## -Descriptions- ##
 
     ## Search ##

bin/lib/search_engine.py

@@ -6,6 +6,7 @@
 import logging.config
 import sys
 import time
+import uuid
 
 import meilisearch
 
@@ -31,6 +32,11 @@
 M_KEY = config_loader.get_config_str('Indexer', 'meilisearch_key')
 config_loader = None
 
+
+def get_obj_uuid5(obj_gid):
+    return str(uuid.uuid5(uuid.NAMESPACE_URL, obj_gid))
+
+
 def is_meilisearch_enabled():
     return IS_MEILISEARCH_ENABLED
 
@@ -175,6 +181,14 @@ def index_domains_descriptions():
         index_domain_description(dom_id)
 
 
+def remove_document(index_name, obj_gid):
+    Engine.remove(index_name, get_obj_uuid5(obj_gid))
+
+
+def delete_index(index_name):
+    Engine._delete(index_name)
+    Engine.client.create_index(index_name, {'primaryKey': 'uuid'})
+
 def log(user_id, index, to_search):
     logger.warning(f'{user_id} search: {index} - {to_search}')