180 lines
8.1 KiB
Python
180 lines
8.1 KiB
Python
"""
|
||
Проход 1: обнаружение сущностей в тексте документов.
|
||
|
||
Два метода сканирования:
|
||
1. scan_regex — быстрое regex-обнаружение (телефоны, email, ИНН, счета, компании, ФИО)
|
||
2. scan_llm_ner — универсальное LLM-обнаружение (любые приватные данные)
|
||
|
||
Замена: гибридные токены — читаемый шаблон + уникальный номер.
|
||
Пример: Иванов_0001, ООО_Технология_0034, ул_Ленина_0015
|
||
"""
|
||
|
||
import re
|
||
import json
|
||
import random
|
||
import logging
|
||
from typing import Dict
|
||
|
||
from .config import (
|
||
ENTITY_PATTERNS, COMPANY_PATTERN, PERSON_PATTERN,
|
||
RU_SURNAMES, COMPANY_NOUNS, CITY_POOL, STREET_POOL,
|
||
)
|
||
|
||
log = logging.getLogger("drhider")
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
# Гибридные токены: тип → (пул_шаблонов, префикс_для_счётчика)
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
TYPE_POOLS = {
|
||
"person_name": (RU_SURNAMES, "person"),
|
||
"person": (RU_SURNAMES, "person"),
|
||
"company": (COMPANY_NOUNS, "company"),
|
||
"phone": (["+7_000_000"], "phone"),
|
||
"email": (["email"], "email"),
|
||
"address": (CITY_POOL, "address"),
|
||
"inn": (["ИНН"], "inn"),
|
||
"inn_ul": (["ИНН"], "inn"),
|
||
"inn_fl": (["ИНН"], "inn"),
|
||
"kpp": (["КПП"], "kpp"),
|
||
"ogrn": (["ОГРН"], "ogrn"),
|
||
"bik": (["БИК"], "bik"),
|
||
"passport": (["Паспорт"], "passport"),
|
||
"bank_account": (["Счёт"], "bank"),
|
||
"rs": (["РС"], "bank"),
|
||
"ks": (["КС"], "bank"),
|
||
"contract_number": (["Договор"], "contract"),
|
||
"other_id": (["ID"], "other"),
|
||
}
|
||
_FALLBACK_POOL = (["Данные"], "data")
|
||
|
||
|
||
def _next_token(entity_type: str, counters: Dict[str, int]) -> str:
|
||
"""Сгенерировать токен замены: случайный шаблон из пула + номер.
|
||
|
||
Одна и та же сущность в mapping получает один и тот же токен
|
||
(mapping dict гарантирует согласованность).
|
||
Разные сущности одного типа получают РАЗНЫЕ шаблоны из пула.
|
||
|
||
Args:
|
||
entity_type: Тип сущности (person_name, company, phone...)
|
||
counters: Глобальные счётчики {prefix_key: count}
|
||
|
||
Returns:
|
||
Строка вида "Иванов_0001", "ООО_Технология_0034", "[Адрес_0015]"
|
||
"""
|
||
pool, prefix_key = TYPE_POOLS.get(entity_type, _FALLBACK_POOL)
|
||
template = random.choice(pool)
|
||
key = f"{prefix_key}_{template}"
|
||
counters[key] = counters.get(key, 0) + 1
|
||
return f"{template}_{counters[key]:04d}"
|
||
|
||
log = logging.getLogger("drhider")
|
||
|
||
|
||
def scan_regex(text: str, mapping: Dict[str, str], counters: Dict[str, int]) -> None:
|
||
"""Сканировать текст regex-паттернами, заполнить словарь замен.
|
||
|
||
Замена: гибридные токены (шаблон + номер).
|
||
|
||
Args:
|
||
text: Текст документа для сканирования
|
||
mapping: Словарь замен (мутабельный, пополняется)
|
||
counters: Глобальные счётчики токенов (мутабельный)
|
||
"""
|
||
# ── Сущности по regex-паттернам ──
|
||
for entity_type, pattern in ENTITY_PATTERNS.items():
|
||
for match in re.finditer(pattern, text, re.IGNORECASE | re.MULTILINE):
|
||
original = match.group(0).strip()
|
||
if not original or original in mapping:
|
||
continue
|
||
mapping[original] = _next_token(entity_type, counters)
|
||
|
||
# ── Названия компаний ──
|
||
for match in COMPANY_PATTERN.finditer(text):
|
||
original = match.group(0).strip()
|
||
if not original or original in mapping:
|
||
continue
|
||
mapping[original] = _next_token("company", counters)
|
||
|
||
# ── ФИО ──
|
||
for match in PERSON_PATTERN.finditer(text):
|
||
original = match.group(0).strip()
|
||
if not original or original in mapping:
|
||
continue
|
||
mapping[original] = _next_token("person_name", counters)
|
||
|
||
|
||
def scan_llm_ner(all_texts: Dict[str, str], mapping: Dict[str, str],
|
||
llm_client, counters: Dict[str, int]) -> None:
|
||
"""Универсальное LLM-обнаружение приватных данных.
|
||
|
||
LLM САМА определяет типы. Замена — гибридные токены.
|
||
|
||
Args:
|
||
all_texts: {filename: text_content} — тексты всех файлов
|
||
mapping: Словарь замен (мутабельный, пополняется)
|
||
llm_client: Объект с методом .complete(prompt) -> str
|
||
counters: Глобальные счётчики токенов (мутабельный)
|
||
"""
|
||
# Формируем комбинированный текст: имя файла + первые 3000 символов
|
||
combined = "\n\n---FILE---\n\n".join(
|
||
f"FILE: {fname}\n{t[:3000]}" for fname, t in all_texts.items()
|
||
)
|
||
|
||
# ── Универсальный промпт: STEP 1 (Исполнитель) + STEP 2 (ПДн) ──
|
||
already_found = "\n".join(f"- {k}" for k in mapping.keys()) if mapping else "(none)"
|
||
|
||
prompt = (
|
||
"You are a PII (Personally Identifiable Information) detector for Russian legal documents.\n\n"
|
||
"Find ALL sensitive or private information that is NOT already captured below.\n\n"
|
||
|
||
"Already captured by regex (skip these exact strings):\n"
|
||
f"{already_found}\n\n"
|
||
|
||
"Return a JSON array. Each item must have:\n"
|
||
' - "type": MUST be one of: person_name, company, phone, email, address,\n'
|
||
' inn, kpp, ogrn, bik, passport, bank_account, contract_number, other_id\n'
|
||
' Use "other_id" only if nothing else fits. If unsure — do NOT include.\n'
|
||
' - "value": the EXACT string from the text (copy verbatim)\n\n'
|
||
|
||
"NOT PII — do NOT include these (legal terms, not private data):\n"
|
||
"- Roles: Исполнитель, Заказчик, Стороны, Сторона, Покупатель, Поставщик\n"
|
||
"- Titles: Генеральный директор, Директор, действующего на основании\n"
|
||
"- Legal terms: Услуги, Договор, Приложение, Спецификация, НДС, Устава\n\n"
|
||
|
||
"Rules:\n"
|
||
"- Copy value VERBATIM — same spaces, punctuation, case as in text\n"
|
||
"- Same value → include only once\n"
|
||
"- Return ONLY the JSON array, no other text, no markdown wrapping\n\n"
|
||
|
||
f"Text:\n{combined[:8000]}"
|
||
)
|
||
|
||
try:
|
||
# Отправляем запрос в LLM
|
||
raw = llm_client.complete(prompt)
|
||
|
||
# Чистим markdown-обёртку (если LLM вернула ```json ... ```)
|
||
raw = raw.strip()
|
||
if raw.startswith("```"):
|
||
raw = raw.split("\n", 1)[1]
|
||
if raw.endswith("```"):
|
||
raw = raw[:-3]
|
||
|
||
entities = json.loads(raw)
|
||
|
||
# Добавляем найденные сущности в mapping
|
||
for ent in entities:
|
||
val = ent.get("value", "").strip()
|
||
if not val or val in mapping:
|
||
continue
|
||
|
||
# Тип — любой (LLM сама придумала)
|
||
ent_type = ent.get("type", "").strip().lower().replace(" ", "_")
|
||
mapping[val] = _next_token(ent_type, counters)
|
||
|
||
except Exception as e:
|
||
log.warning("LLM NER failed: %s", e)
|