Improve LLM and CSV processing

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

drhider/builder.py

@@ -8,13 +8,13 @@
 
 import io
 import csv
-import os
 import zipfile
 import re
 import time
 from typing import Dict, List, Tuple
 
 from .config import ENTITY_PATTERNS, COMPANY_PATTERN
+from .names import resolve_name
 
 
 def build_zip(files: List[Tuple[str, bytes]], mapping_csv: str = "") -> bytes:
@@ -39,24 +39,10 @@ def build_zip(files: List[Tuple[str, bytes]], mapping_csv: str = "") -> bytes:
     # Дедупликация имён: имя+одинаковый контент → пропуск; имя+разный контент → суффикс
     seen: Dict[str, bytes] = {}
 
-    def _resolve_name(name: str, content: bytes) -> str:
-        if name not in seen:
-            seen[name] = content
-            return name
-        if seen[name] == content:
-            return None  # точный дубль — пропускаем
-        base, ext = os.path.splitext(name)
-        n = 2
-        while f"{base}_{n}{ext}" in seen:
-            n += 1
-        resolved = f"{base}_{n}{ext}"
-        seen[resolved] = content
-        return resolved
-
     with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
         # Добавляем обфусцированные файлы
         for fname, content in files:
-            out_name = _resolve_name(fname, content)
+            out_name = resolve_name(fname, content, seen)
             if out_name is None:
                 continue
             info = zipfile.ZipInfo(out_name)

drhider/extractor.py

@@ -1,10 +1,11 @@
 """
-Извлечение текста из документов — все форматы → Markdown.
+Извлечение текста из документов и подготовка CSV к обфускации.
 
 Поддерживает:
 - .docx → MD (python-docx: стили, жирный/курсив, таблицы)
 - .pdf → MD (pdfplumber: текст + таблицы, без форматирования)
-- .doc → MD (LibreOffice headless: .doc → .docx → штатный docx_to_markdown)
+- .doc → HTTP-сервис конвертации (CONVERT_SERVICE_URL) → .docx → Markdown
+- .csv → текст ячеек для сканирования и обезличенный CSV на выходе
 - .txt и прочие → как есть
 
 Также содержит:
@@ -12,13 +13,16 @@
 """
 
 import io
+import csv
 import os
 import zipfile
 import tempfile
 import shutil
 import subprocess
 import logging
-from typing import List, Tuple
+from typing import Dict, List, Tuple
+
+from . import replacer
 
 log = logging.getLogger("drhider")
 
@@ -199,6 +203,77 @@ def doc_to_markdown(content: bytes) -> str:
 # Единая точка входа
 # ═══════════════════════════════════════════════════════════════════════════
 
+def decode_text(content: bytes) -> str:
+    """Декодировать текст как UTF-8, затем CP1251, с Latin-1 последним fallback."""
+    try:
+        return content.decode("utf-8")
+    except UnicodeDecodeError:
+        # Блокнот и выгрузки 1С часто сохраняют кириллицу в CP1251.
+        try:
+            return content.decode("cp1251")
+        except UnicodeDecodeError:
+            return content.decode("latin-1", errors="replace")
+
+
+def detect_delimiter(text: str) -> str:
+    """Определить CSV-разделитель, используя Sniffer и детерминированный fallback."""
+    delimiters = [",", ";", "\t"]
+    nonempty_lines = [line for line in text.splitlines() if line.strip()]
+    if nonempty_lines:
+        sample = "\n".join(nonempty_lines[:10])
+        try:
+            return csv.Sniffer().sniff(sample, delimiters="".join(delimiters)).delimiter
+        except csv.Error:
+            first_line = nonempty_lines[0]
+            return max(delimiters, key=first_line.count)
+    return ","
+
+
+def csv_to_text(content: bytes) -> str:
+    """Распарсить CSV и склеить значения ячеек построчно для сканирования."""
+    text = decode_text(content).removeprefix("\ufeff")
+    delimiter = detect_delimiter(text)
+    reader = csv.reader(
+        io.StringIO(text, newline=""),
+        delimiter=delimiter,
+        quotechar='"',
+        strict=True,
+    )
+    # Каждая ячейка отделена переводом строки, чтобы сущности из соседних ячеек
+    # не объединялись в строку, по которой работают regex и LLM.
+    return "\n".join(cell for row in reader for cell in row)
+
+
+def obfuscate_csv(content: bytes, mapping: Dict[str, str],
+                  sorted_keys: List[str], compiled_re) -> bytes:
+    """Заменить сущности в CSV-ячейках, сохранив порядок строк и столбцов."""
+    text = decode_text(content).removeprefix("\ufeff")
+    delimiter = detect_delimiter(text)
+    reader = csv.reader(
+        io.StringIO(text, newline=""),
+        delimiter=delimiter,
+        quotechar='"',
+        strict=True,
+    )
+    output = io.StringIO(newline="")
+    writer = csv.writer(
+        output,
+        delimiter=delimiter,
+        quotechar='"',
+        lineterminator="\n",
+        strict=True,
+    )
+    for row in reader:
+        writer.writerow([
+            replacer.apply_replacements(cell, mapping, sorted_keys, compiled_re)
+            for cell in row
+        ])
+
+    # csv.writer нормализует кавычки/переводы строк; побайтовое сохранение
+    # исходного формата не является целью. BOM помогает Excel открыть кириллицу.
+    return ("\ufeff" + output.getvalue()).encode("utf-8")
+
+
 def extract_text(fname: str, content: bytes, ctype: str = "") -> str:
     """Извлечь текст из одного файла → Markdown.
 
@@ -217,22 +292,26 @@ def extract_text(fname: str, content: bytes, ctype: str = "") -> str:
         try:
             return docx_to_markdown(content)
         except ImportError:
-            return content.decode('utf-8', errors='replace')
+            return decode_text(content)
 
     # ── .pdf: текст + таблицы ──
     elif ext == '.pdf':
         try:
             return pdf_to_markdown(content)
         except ImportError:
-            return content.decode('utf-8', errors='replace')
+            return decode_text(content)
 
-    # ── .doc: бинарный → LibreOffice → DOCX → Markdown ──
+    # ── .doc: бинарный → HTTP-сервис конвертации → DOCX → Markdown ──
     elif ext == '.doc':
         return doc_to_markdown(content)
 
+    # ── .csv: ячейки для обнаружения сущностей ──
+    elif ext == '.csv':
+        return csv_to_text(content)
+
     # ── .txt и прочие ──
     else:
-        return content.decode('utf-8', errors='replace')
+        return decode_text(content)
 
 
 # ═══════════════════════════════════════════════════════════════════════════
@@ -338,7 +417,7 @@ def expand_zips(files: List[Tuple[str, bytes, str]]) -> List[Tuple[str, bytes, s
                         continue
                     # Синхронизация с фронтом (listZipFiles): из zip берём ТОЛЬКО документы
                     low = name.lower()
-                    if not (low.endswith('.zip') or low.endswith(('.pdf', '.doc', '.docx', '.txt', '.md'))):
+                    if not (low.endswith('.zip') or low.endswith(('.pdf', '.doc', '.docx', '.txt', '.md', '.csv'))):
                         continue
 
                     inner_data = zf.read(info)

drhider/llm_client.py

@@ -11,6 +11,8 @@ LLM-клиент для Obfuscator.
 
 import os
 import time
+from typing import Optional
+
 import httpx
 
 
@@ -21,9 +23,12 @@ class LLMClient:
         LLM_API_KEY или LLM_KEY — ключ API
         LLM_URL — URL эндпоинта (по умолчанию https://api.aillm.ru/v1/chat/completions)
         LLM_MODEL — модель (по умолчанию gpt-oss-120b)
+        LLM_RETRIES — число повторов после первой попытки (по умолчанию 3)
+        LLM_RETRY_DELAY — базовая задержка backoff в секундах (по умолчанию 2)
     """
 
-    def __init__(self):
+    def __init__(self, retries: Optional[int] = None,
+                 retry_delay: Optional[float] = None):
         """Инициализировать клиент. httpx импортируется лениво.
 
         Ведёт кумулятивные счётчики за всё время жизни клиента:
@@ -36,6 +41,13 @@ class LLMClient:
             llm_started          — time.time() начала текущего LLM-вызова
         """
         self._httpx = httpx
+        self.retries = int(os.environ.get("LLM_RETRIES", "3")) if retries is None else retries
+        self.retry_delay = (
+            float(os.environ.get("LLM_RETRY_DELAY", "2.0"))
+            if retry_delay is None else retry_delay
+        )
+        if self.retries < 0 or self.retry_delay < 0:
+            raise ValueError("LLM retries and retry delay must be non-negative")
         self.tokens_prompt = 0
         self.tokens_completion = 0
         self.llm_sec = 0.0
@@ -52,6 +64,31 @@ class LLMClient:
             return time.time() - self.llm_started
         return 0.0
 
+    def _post_with_retries(self, url: str, json: dict, headers: dict,
+                           timeout: int) -> httpx.Response:
+        """Повторить только временные транспортные ошибки и выбранные HTTP-статусы."""
+        for attempt in range(self.retries + 1):
+            try:
+                response = self._httpx.post(
+                    url, json=json, headers=headers, timeout=timeout
+                )
+                self.llm_sec += response.elapsed.total_seconds()
+                response.raise_for_status()
+                return response
+            except (self._httpx.TimeoutException, self._httpx.TransportError):
+                if attempt >= self.retries:
+                    raise
+            except self._httpx.HTTPStatusError as e:
+                # Retry only throttling and server failures; 400/401/403 require
+                # a corrected request or credentials and will not improve on retry.
+                if e.response.status_code not in (429, 500, 502, 503, 504):
+                    raise
+                if attempt >= self.retries:
+                    raise
+
+            # 120 секунд на попытку и backoff означают около 8 минут в худшем случае.
+            time.sleep(self.retry_delay * (attempt + 1))
+
     def complete(self, prompt: str) -> str:
         """Отправить промпт в LLM и вернуть текст ответа.
 
@@ -76,7 +113,7 @@ class LLMClient:
         self.llm_active = True
         self.llm_started = time.time()
         try:
-            r = self._httpx.post(
+            r = self._post_with_retries(
                 url,
                 json={
                     "model": model,
@@ -90,10 +127,6 @@ class LLMClient:
                 },
                 timeout=120,
             )
-            r.raise_for_status()
-
-            # Время запроса (реальное время HTTP-вызова LLM)
-            self.llm_sec += r.elapsed.total_seconds()
 
             data = r.json()
             # Счётчики токенов из usage (могут отсутствовать у некоторых прокси)

drhider/names.py

@@ -0,0 +1,25 @@
+"""Helpers for assigning stable, collision-free output filenames."""
+
+import os
+from typing import Dict, Optional
+
+
+def resolve_name(name: str, content: bytes, seen: Dict[str, bytes]) -> Optional[str]:
+    """Return a unique name, or None when the same name and content were seen.
+
+    Name collisions with different content receive numeric suffixes before
+    the extension, while exact duplicates are omitted.
+    """
+    if name not in seen:
+        seen[name] = content
+        return name
+    if seen[name] == content:
+        return None
+
+    base, ext = os.path.splitext(name)
+    index = 2
+    while f"{base}_{index}{ext}" in seen:
+        index += 1
+    resolved = f"{base}_{index}{ext}"
+    seen[resolved] = content
+    return resolved

drhider/obfuscator.py

@@ -17,10 +17,19 @@ from . import extractor
 from . import scanner
 from . import replacer
 from . import builder
+from .names import resolve_name
 
 log = logging.getLogger("drhider")
 
 
+class LLMIncompleteError(Exception):
+    """Processing cannot safely return files when one or more LLM chunks failed."""
+
+    def __init__(self, llm_errors: int):
+        self.llm_errors = llm_errors
+        super().__init__("LLM processing incomplete")
+
+
 def _dedupe_file_names(
     files: List[Tuple[str, bytes, str]]
 ) -> List[Tuple[str, bytes, str]]:
@@ -39,17 +48,9 @@ def _dedupe_file_names(
     out: List[Tuple[str, bytes, str]] = []
 
     for fname, content, ctype in files:
-        candidate = fname
-        if candidate in seen:
-            if seen[candidate] == content:
-                continue  # точный дубль — пропускаем
-            base, ext = os.path.splitext(fname)
-            n = 2
-            while f"{base}_{n}{ext}" in seen:
-                n += 1
-            candidate = f"{base}_{n}{ext}"
-        seen[candidate] = content
-        out.append((candidate, content, ctype))
+        candidate = resolve_name(fname, content, seen)
+        if candidate is not None:
+            out.append((candidate, content, ctype))
 
     return out
 
@@ -177,13 +178,18 @@ class TwoPassObfuscator:
             cancelled_llm = cancelled_extract
             if not cancelled_llm and self._llm_client:
                 try:
-                    scanner.scan_llm_ner(
+                    llm_failures = scanner.scan_llm_ner(
                         all_texts, self._mapping, self._llm_client, self._counters,
                         cancel_event=cancel_event, file_progress=_fp,
                     )
-                except scanner.CancelRequested:
+                except scanner.CancelRequested as e:
+                    if e.llm_failures:
+                        raise LLMIncompleteError(e.llm_failures)
                     log.info("obfuscate: LLM cancelled, llm_done=%d", len(llm_done))
                     cancelled_llm = True
+                else:
+                    if llm_failures:
+                        raise LLMIncompleteError(llm_failures)
 
             # Предсортировать ключи один раз (по убыванию длины)
             self._sorted_keys = sorted(
@@ -220,7 +226,18 @@ class TwoPassObfuscator:
                 t0 = time.time()
                 obf_content = content  # По умолчанию — без изменений
 
-                if fname.endswith('.doc'):
+                if fname.lower().endswith('.csv'):
+                    try:
+                        obf_content = extractor.obfuscate_csv(
+                            content, self._mapping, self._sorted_keys, self._compiled_re
+                        )
+                    except Exception as e:
+                        log.warning("obfuscate: skip CSV %r during replacement: %r", fname, e)
+                        skipped.add(fname)
+                        if progress_cb:
+                            progress_cb("done", i, total, display_name, 0.0)
+                        continue
+                elif fname.endswith('.doc'):
                     # .doc — конвертируется через liberta, обрабатывается как .docx
                     md_text = all_texts.get(fname, '')
                     replaced = replacer.apply_replacements(

drhider/scanner.py

@@ -115,6 +115,9 @@ def scan_regex(text: str, mapping: Dict[str, str], counters: Dict[str, int]) ->
 class CancelRequested(Exception):
     """Обработка прервана пользователем (мягкая остановка между файлами LLM)."""
 
+    def __init__(self, llm_failures: int = 0):
+        self.llm_failures = llm_failures
+
 # Приоритет границ для чанков (от предпочтительных к жёстким)
 _CHUNK_BOUNDARIES = ['\n\n', '\n', '. ', '? ', '! ', '; ', ', ', ' ']
 _CHUNK_SIZE = 6000          # символов на чанк (≈1500-2000 токенов RU — NER-качество)
@@ -202,26 +205,53 @@ def _build_llm_prompt(text: str, already_found: List[str]) -> str:
 
 
 def _call_llm(text: str, llm_client) -> List[dict]:
-    """Один вызов LLM по фрагменту → список сущностей (пустой при ошибке)."""
+    """Один вызов LLM по фрагменту → проверенный список сущностей."""
+    raw = llm_client.complete(text).strip()
+    if raw.startswith("```"):
+        lines = raw.splitlines()
+        if lines and lines[0].startswith("```"):
+            lines = lines[1:]
+        if lines and lines[-1].strip() == "```":
+            lines = lines[:-1]
+        raw = "\n".join(lines).strip()
+
     try:
-        raw = llm_client.complete(text)
-        raw = raw.strip()
-        if raw.startswith("```"):
-            raw = raw.split("\n", 1)[1]
-            if raw.endswith("```"):
-                raw = raw[:-3]
         entities = json.loads(raw)
-        if not isinstance(entities, list):
-            return []
-        return entities
-    except Exception as e:
-        log.warning("LLM NER failed: %s", e)
-        return []
+    except json.JSONDecodeError as original_error:
+        # Модель иногда добавляет текст вокруг массива; raw_decode корректно
+        # учитывает вложенные массивы и квадратные скобки внутри JSON-строк.
+        decoder = json.JSONDecoder()
+        entities = None
+        for index, char in enumerate(raw):
+            if char != "[":
+                continue
+            try:
+                candidate, _ = decoder.raw_decode(raw, index)
+            except json.JSONDecodeError:
+                continue
+            if isinstance(candidate, list):
+                entities = candidate
+                break
+        if entities is None:
+            raise ValueError("LLM response does not contain a JSON array") from original_error
+
+    if not isinstance(entities, list):
+        raise ValueError("LLM response must be a JSON array")
+    for entity in entities:
+        if not isinstance(entity, dict):
+            raise ValueError("LLM entity must be a JSON object")
+        value = entity.get("value")
+        if not isinstance(value, str) or not value.strip():
+            raise ValueError("LLM entity value must be a non-empty string")
+        entity_type = entity.get("type")
+        if entity_type is not None and not isinstance(entity_type, str):
+            raise ValueError("LLM entity type must be a string")
+    return entities
 
 
 def scan_llm_ner(all_texts: Dict[str, str], mapping: Dict[str, str],
                   llm_client, counters: Dict[str, int],
-                  cancel_event=None, file_progress=None) -> None:
+                  cancel_event=None, file_progress=None) -> int:
     """Универсальное LLM-обнаружение приватных данных (пофайлово, целиком).
 
     Каждый файл обрабатывается ПОЛНОСТЬЮ: текст разбивается на чанки
@@ -243,11 +273,12 @@ def scan_llm_ner(all_texts: Dict[str, str], mapping: Dict[str, str],
     file_items = sorted(
         all_texts.items(), key=lambda kv: len(kv[1]), reverse=True
     )
+    llm_failures = 0
 
     for fname, full_text in file_items:
         # Мягкая остановка ТОЛЬКО между файлами — текущий файл добирается до конца
         if cancel_event is not None and cancel_event.is_set():
-            raise CancelRequested()
+            raise CancelRequested(llm_failures)
 
         if not full_text or full_text.startswith("[DOC binary"):
             continue
@@ -265,11 +296,19 @@ def scan_llm_ner(all_texts: Dict[str, str], mapping: Dict[str, str],
                            for c in chunks]
                 all_ent = []
                 for fut in futures:
-                    all_ent.extend(fut.result())
+                    try:
+                        all_ent.extend(fut.result())
+                    except Exception as e:
+                        llm_failures += 1
+                        log.warning("LLM NER chunk failed in %r: %s", fname, e)
         else:
             all_ent = []
             for k, c in enumerate(chunks):
-                all_ent.extend(_call_llm(_build_llm_prompt(c, mapping.keys()), llm_client))
+                try:
+                    all_ent.extend(_call_llm(_build_llm_prompt(c, mapping.keys()), llm_client))
+                except Exception as e:
+                    llm_failures += 1
+                    log.warning("LLM NER chunk failed in %r: %s", fname, e)
                 if file_progress:
                     file_progress("file_chunk", fname,
                                   chunks_done=k + 1, chunks_total=len(chunks))
@@ -287,8 +326,12 @@ def scan_llm_ner(all_texts: Dict[str, str], mapping: Dict[str, str],
             if not _verify_entity(val, full_text):
                 log.debug("LLM hallucination dropped: %r", val)
                 continue
-            ent_type = ent.get("type", "").strip().lower().replace(" ", "_")
+            raw_type = ent.get("type")
+            ent_type = raw_type.strip().lower().replace(" ", "_") if raw_type else ""
+            ent_type = ent_type or "other_id"
             mapping[val] = _next_token(ent_type, counters)
 
         if file_progress:
             file_progress("file_done", fname, elapsed=round(time.time() - t0, 2))
+
+    return llm_failures