Архив для RAG разросся при распаковке: лимиты до загрузки документов

Пользователь загружает архив документов для базы знаний. Сам ZIP невелик, но после распаковки задача занимает всю память и надолго блокирует обработчик. Размер входного файла не определяет объём результата: одинаковые данные хорошо сжимаются. Поэтому ограничение загрузки в несколько мегабайт само по себе не ограничивает расход распаковщика и последующего PDF-парсера.

Границу стоит поставить до индексации. Для архива нужны отдельные бюджеты: число записей, размер одного распакованного файла и суммарный размер прочитанных данных. Затем каждый поддерживаемый документ проходит собственный разбор с лимитами времени и ресурсов. Подготовим небольшой безопасный пример, который показывает обе проверки размера и не создаёт большой архив.

Сначала прочитайте каталог

ZIP хранит метаданные записей, включая ожидаемый распакованный размер. До чтения содержимого можно отклонить слишком много файлов или сумму заявленных размеров выше бюджета. Это позволяет не тратить ресурсы на очевидно неподходящий ввод. Однако метаданные не должны быть единственным основанием доверия: при чтении нужно считать фактически полученные байты.

В нашем примере разрешено не больше трёх файлов и десяти байтов суммарного содержимого. Это намеренно маленькие учебные значения. Архив создаётся в памяти без сжатия и содержит два коротких текста. Никакие файлы на диске не распаковываются, поэтому пример сосредоточен на объёме чтения. Защита путей и изоляция парсеров понадобятся при добавлении сохранения документов.

from io import BytesIO
from zipfile import ZipFile, ZIP_STORED
def read_bounded(stream, limit):
    result = bytearray()
    while True:
        chunk = stream.read(min(4, limit - len(result) + 1))
        if not chunk:
            return bytes(result)
        result.extend(chunk)
        if len(result) > limit:
            raise ValueError("actual bytes exceeded")
def load_zip(data, max_files, total_limit):
    with ZipFile(BytesIO(data)) as archive:
        infos = archive.infolist()
        if len(infos) > max_files:
            raise ValueError("too many entries")
        if sum(item.file_size for item in infos) > total_limit:
            raise ValueError("declared bytes exceeded")
        names, consumed = [], 0
        for item in infos:
            if item.is_dir():
                continue
            with archive.open(item) as stream:
                payload = read_bounded(stream, total_limit - consumed)
            consumed += len(payload)
            names.append(item.filename)
        return names, consumed
buffer = BytesIO()
with ZipFile(buffer, "w", compression=ZIP_STORED) as archive:
    archive.writestr("a.txt", "abcd")
    archive.writestr("b.txt", "efgh")
print(load_zip(buffer.getvalue(), 3, 10))
for label, action in [
    ("metadata rejected", lambda: load_zip(buffer.getvalue(), 3, 7)),
    ("actual bytes rejected", lambda: read_bounded(BytesIO(b"abcde"), 4)),
]:
    try:
        action()
    except ValueError:
        print(label)
    else:
        raise AssertionError("limit not enforced")

Первый запуск возвращает два имени и восемь прочитанных байтов. Во втором бюджет снижается до семи: каталог архива отклоняется ещё до чтения содержимого. Третья проверка подаёт обычный поток из пяти байтов в функцию с бюджетом четыре. Она читает максимум на один байт больше разрешения и останавливается с ошибкой. Так отдельно проверяется фактический счётчик, даже если предварительная проверка уже отсекла проблемный архив.

Почему считать нужно во время чтения

Если сначала вызвать read без размера, весь файл уже окажется в памяти к моменту сравнения длины. Ограничение должно управлять самим чтением. В примере функция запрашивает не больше оставшегося бюджета плюс один байт. Дополнительный байт нужен, чтобы отличить точное попадание в лимит от превышения. Результат на границе допустим, результат на байт длиннее отклоняется.

В рабочем загрузчике также задайте отдельный предел для одной записи. Иначе один файл использует весь общий бюджет и лишит места остальные. Счётчик общего объёма переносится между записями; его нельзя обнулять при каждом новом имени. Каталог с большим числом пустых записей тоже требует ограничения, поскольку их обработка расходует время и память независимо от суммарного текста.

Вложенные архивы лучше отклонять, если продукт не требует их поддержки. Если они нужны, передавайте один общий бюджет и ограничивайте глубину вложенности. Не выдавайте каждому вложенному ZIP новый полный лимит. Распаковку и разбор сложных форматов полезно выполнять в отдельном процессе с ограничением времени: счётчик выходных байтов не ограничивает всю вычислительную работу декомпрессора.

Как вернуть понятную ошибку пользователю

Сообщение должно объяснять допустимое действие: «Суммарный объём документов после распаковки превышает лимит. Разделите архив на несколько загрузок». Не стоит предлагать сильнее сжать тот же набор: это уменьшит транспортный файл, но сохранит причину отказа. В журнале оставьте этап отказа, число записей и использованный бюджет, без содержимого закрытых документов.

Только после успешной проверки архива отправляйте документы в обычную цепочку извлечения текста. Для RAG-запроса «Найди условия продления договора» модель должна получать уже принятые фрагменты с источниками. Распаковщик не должен зависеть от просьбы модели открыть ещё один вложенный файл. Методы ZipFile.infolist и ZipFile.open позволяют разделить чтение каталога и содержимого.

В платном курсе «LLM Security PRO: Prompt Injection, утечки, tool-abuse» вы построите модель угроз RAG и агента, ограничите полномочия инструментов и стоимость обработки, соберёте проверки безопасности для CI/CD. Следующий шаг здесь — включить загрузку документов в эту модель угроз и проверять бюджеты до обращения к модели. Откройте программу курса и спроектируйте защиту всей цепочки приёма данных.