scanner.py: универсальный LLM-промпт (LLM сама определяет типы) + fallback-генератор

drhider/scanner.py

@@ -2,8 +2,11 @@
 Проход 1: обнаружение сущностей в тексте документов.
 
 Два метода сканирования:
-1. _scan_regex — быстрое regex-обнаружение (телефоны, email, ИНН, счета, компании, ФИО)
-2. _scan_llm_ner — LLM-обнаружение (имена, адреса, паспорта — то что regex не ловит)
+1. scan_regex — быстрое regex-обнаружение (телефоны, email, ИНН, счета, компании, ФИО)
+2. scan_llm_ner — универсальное LLM-обнаружение (любые приватные данные)
+
+LLM САМА определяет какие типы сущностей присутствуют в документе.
+Никакого фиксированного списка типов.
 """
 
 import re
@@ -12,15 +15,25 @@ import logging
 from typing import Dict
 
 from .config import ENTITY_PATTERNS, COMPANY_PATTERN, PERSON_PATTERN
-from .generators import ENTITY_GENERATORS
+from .generators import ENTITY_GENERATORS, FALLBACK_GENERATOR
 from .generators.company import generate_company
 from .generators.person import generate_person
-from .generators.address import generate_address
-from .generators.passport import generate_passport
 
 log = logging.getLogger("drhider")
 
 
+def _get_generator(entity_type: str):
+    """Получить генератор по типу сущности. Если тип неизвестен — fallback.
+
+    Args:
+        entity_type: Строковый идентификатор типа (например, "phone", "person_name")
+
+    Returns:
+        Функция-генератор: (str) -> str
+    """
+    return ENTITY_GENERATORS.get(entity_type, FALLBACK_GENERATOR)
+
+
 def scan_regex(text: str, mapping: Dict[str, str]) -> None:
     """Сканировать текст regex-паттернами, заполнить словарь замен.
 
@@ -43,8 +56,8 @@ def scan_regex(text: str, mapping: Dict[str, str]) -> None:
             if not original or original in mapping:
                 continue
 
-            # Выбираем генератор по типу сущности
-            generator = ENTITY_GENERATORS.get(entity_type, lambda x: "XXX")
+            # Генератор: специализированный (если есть) или fallback
+            generator = _get_generator(entity_type)
             mapping[original] = generator(original)
 
     # ── Названия компаний ──
@@ -86,17 +99,21 @@ def scan_llm_ner(all_texts: Dict[str, str], mapping: Dict[str, str], llm_client)
         f"FILE: {fname}\n{t[:3000]}" for fname, t in all_texts.items()
     )
 
-    # Промпт для LLM
+    # ── Универсальный промпт — без фиксированного списка типов ──
     prompt = (
-        "Ты — система обнаружения персональных данных в документах. "
-        "Найди ВСЕ следующие сущности в тексте ниже:\n\n"
-        "1. ФИО (полные и сокращённые — 'Иванов И.И.', 'Петров А.С.')\n"
-        "2. Названия компаний-контрагентов (не 'НУБЕС')\n"
-        "3. Почтовые адреса\n"
-        "4. Паспортные данные\n\n"
-        "Формат ответа — JSON-массив:\n"
-        '[{"type": "person"|"company"|"address"|"passport", "value": "найденный текст"}]\n\n'
-        f"Текст:\n{combined[:8000]}"
+        "You are a PII (Personally Identifiable Information) detector. "
+        "Find ALL sensitive or private information in the text below.\n\n"
+        "Return a JSON array. Each item must have:\n"
+        '  - "type": short snake_case label describing the entity\n'
+        '    (e.g. person_name, phone, email, address, inn, company, passport,\n'
+        '     contract_number, employee_id, bank_account — WHATEVER you see)\n'
+        '  - "value": the EXACT string from the text (copy verbatim)\n\n'
+        "Rules:\n"
+        "- Copy value VERBATIM — same spaces, punctuation, case as in text\n"
+        "- If same value appears multiple times — include only once\n"
+        "- Invent the type label yourself based on what you see\n"
+        "- Return ONLY the JSON array, no other text, no markdown wrapping\n\n"
+        f"Text:\n{combined[:8000]}"
     )
 
     try:
@@ -118,15 +135,12 @@ def scan_llm_ner(all_texts: Dict[str, str], mapping: Dict[str, str], llm_client)
             if not val or val in mapping:
                 continue
 
-            ent_type = ent.get("type", "")
-            if ent_type == "person":
-                mapping[val] = generate_person(val)
-            elif ent_type == "company":
-                mapping[val] = generate_company(val)
-            elif ent_type == "address":
-                mapping[val] = generate_address(val)
-            elif ent_type == "passport":
-                mapping[val] = generate_passport(val)
+            # Тип — любой (LLM сама придумала)
+            ent_type = ent.get("type", "").strip().lower().replace(" ", "_")
+
+            # Генератор: специализированный (если тип совпал) или fallback
+            generator = _get_generator(ent_type)
+            mapping[val] = generator(val)
 
     except Exception as e:
         log.warning("LLM NER failed: %s", e)