133 lines
6.0 KiB
Python
133 lines
6.0 KiB
Python
"""
|
||
Проход 1: обнаружение сущностей в тексте документов.
|
||
|
||
Два метода сканирования:
|
||
1. _scan_regex — быстрое regex-обнаружение (телефоны, email, ИНН, счета, компании, ФИО)
|
||
2. _scan_llm_ner — LLM-обнаружение (имена, адреса, паспорта — то что regex не ловит)
|
||
"""
|
||
|
||
import re
|
||
import json
|
||
import logging
|
||
from typing import Dict
|
||
|
||
from .config import ENTITY_PATTERNS, COMPANY_PATTERN, PERSON_PATTERN
|
||
from .generators import ENTITY_GENERATORS
|
||
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 scan_regex(text: str, mapping: Dict[str, str]) -> None:
|
||
"""Сканировать текст regex-паттернами, заполнить словарь замен.
|
||
|
||
Ищет в тексте:
|
||
- Сущности из ENTITY_PATTERNS (телефоны, email, ИНН, ОГРН, КПП, БИК, счета, паспорта)
|
||
- Названия компаний (кроме «НУБЕС» — это Исполнитель, не заменяем)
|
||
- ФИО в формате «Фамилия И.О.»
|
||
|
||
Найденные значения добавляются в mapping: оригинал → фиктивное значение.
|
||
Если сущность уже есть в mapping — не перезаписываем (согласованность).
|
||
|
||
Args:
|
||
text: Текст документа для сканирования
|
||
mapping: Словарь замен (мутабельный, пополняется)
|
||
"""
|
||
# ── Сущности по regex-паттернам (телефоны, email, реквизиты) ──
|
||
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
|
||
|
||
# Выбираем генератор по типу сущности
|
||
generator = ENTITY_GENERATORS.get(entity_type, lambda x: "XXX")
|
||
mapping[original] = generator(original)
|
||
|
||
# ── Названия компаний ──
|
||
for match in COMPANY_PATTERN.finditer(text):
|
||
original = match.group(0).strip()
|
||
if not original or original in mapping:
|
||
continue
|
||
|
||
# «НУБЕС» — Исполнитель, НЕ заменяем
|
||
if re.search(r'НУБЕС|NUBES', original, re.IGNORECASE):
|
||
continue
|
||
|
||
mapping[original] = generate_company(original)
|
||
|
||
# ── ФИО (Фамилия И.О.) ──
|
||
for match in PERSON_PATTERN.finditer(text):
|
||
original = match.group(0).strip()
|
||
if not original or original in mapping:
|
||
continue
|
||
mapping[original] = generate_person(original)
|
||
|
||
|
||
def scan_llm_ner(all_texts: Dict[str, str], mapping: Dict[str, str], llm_client) -> None:
|
||
"""LLM NER для обнаружения имён, адресов, паспортных данных.
|
||
|
||
Отправляет объединённый текст всех файлов в LLM, получает JSON-список
|
||
найденных сущностей, добавляет их в словарь замен.
|
||
|
||
Отправляются первые 3000 символов каждого файла (экономия токенов).
|
||
Общий размер промпта ограничен 8000 символами.
|
||
|
||
Args:
|
||
all_texts: {filename: text_content} — тексты всех файлов
|
||
mapping: Словарь замен (мутабельный, пополняется)
|
||
llm_client: Объект с методом .complete(prompt) -> str
|
||
"""
|
||
# Формируем комбинированный текст: имя файла + первые 3000 символов
|
||
combined = "\n\n---FILE---\n\n".join(
|
||
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]}"
|
||
)
|
||
|
||
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
|
||
|
||
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)
|
||
|
||
except Exception as e:
|
||
log.warning("LLM NER failed: %s", e)
|