За пределами RAG: task-aware knowledge compression для корпоративного ИИ на AWS

Если вы используете Retrieval-Augmented Generation (RAG) для сложных аналитических задач, которые охватывают сотни документов, например финансовый due diligence или проверку на соответствие требованиям регуляторов, вы, вероятно, уперлись в его предел. Поиск по сходству находит релевантные фрагменты, но часто не улавливает междокументные связи. В этой статье показано, как закрыть этот пробел с помощью task-aware knowledge compression (TAKC) — техники, которая заранее сжимает целые базы знаний в представления, зависящие от задачи, и разворачивается на AWS. Полностью готовую open-source реализацию можно задеплоить в своем аккаунте.

Task-aware knowledge compression

Представьте private equity-фирму, которая оценивает сделку по приобретению производственной компании за $500 million. Команде due diligence нужно проанализировать финансовую отчетность 12 дочерних компаний за 5 лет. При этом в работе есть 200+ контрактов с поставщиками, отчеты по экологическому соответствию для 8 объектов и 50+ судебных дел. Когда аналитик спрашивает о консолидированных финансовых рисках с учетом текущих условий поставщиков и незавершенных судебных процессов, поиск по сходству в RAG не способен показать такой ответ. Нужная информация распределена по сотням документов, а связи между ними не имеют лексического сходства.

TAKC решает такие задачи, используя LLM для создания более коротких сводок документов, ориентированных на конкретную задачу, причем для разных задач создаются разные сводки.

Разные задачи требуют разной информации из одного и того же документа. Годовой отчет, сжатый для финансового анализа, должен сохранять выручку, маржу и данные о денежном потоке. Тот же отчет, сжатый для compliance-проверки, должен сохранять нормативные ссылки и историю нарушений. Обычное суммаризирование пытается охватить все сразу, из-за чего плотность полезной информации для конкретного сценария снижается. TAKC сжимает документы через призму конкретной задачи, сохраняя важное и отбрасывая остальное. В разделе Ingestion pipeline показано, как промпт на сжатие прямо задает, какую информацию нужно сохранить. Для production-развертываний храните task-type prompts в версионируемой конфигурации, например в AWS Systems Manager Parameter Store или в выделенном префиксе Amazon Simple Storage Service (Amazon S3), чтобы изменения промптов можно было аудитировать, а повторное сжатие запускать при обновлении промптов.

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

TAKC дает доступ ко всей базе знаний в сжатом виде, а не только к top-k фрагментам, которые возвращает поиск по сходству. Система сохраняет связи между документами, потому что сжатие видит документы вместе. Она также формирует разные сжатые представления для разных задач из одного и того же исходного материала. Финансовое сжатие формы 10-K выглядит совершенно иначе, чем сжатие того же отчета для оценки юридических рисков. Сжатие уменьшает число токенов в 8x–64x, сохраняя именно тот материал, который важен для задачи.

Многоуровневое сжатие

Разные запросы требуют разной степени полноты. Вопрос вроде «Какой была выручка в Q3?» требует намного меньше контекста, чем запрос на анализ связей между условиями оплаты поставщиков и квартальным денежным потоком по дочерним компаниям.

TAKC решает это, поддерживая четыре уровня сжатия для каждого типа задачи. При самом мягком сжатии (8x) система сохраняет примерно на 87,5% меньше контекста. Этого достаточно для многошагового рассуждения и синтеза по нескольким документам. На среднем уровне (16x) сокращение контекста достигает примерно 93,8%. Этот уровень подходит для аналитических запросов средней сложности. Высокий уровень (32x) уменьшает контекст примерно на 96,9% и служит для фактических поисков и четко сформулированных вопросов. На ультрауровне (64x) сокращение достигает примерно 98,4%. Этот уровень подходит для задач классификации и поиска по ключевым словам.

Большинство корпоративных запросов — это поисковые обращения, которые можно обслужить на более высоких уровнях сжатия с минимальными затратами. Редкие сложные запросы потребляют больший бюджет контекста только тогда, когда это действительно нужно. Такая маршрутизация по уровням дополняет существующие оптимизации RAG, например фильтрацию по метаданным и переформулирование запроса, которые могут сузить набор документов до применения сжатия. Чтобы проверить качество сжатия, сравнивайте ответы LLM на каждом уровне с ответами, полученными на полном, несжатом документе. В reference implementation есть тестовые скрипты, которые выполняют такое сравнение для ваших конкретных типов задач и документов.

Архитектура на AWS

Реализация работает на AWS как два раздельных serverless-пайплайна: один для ingestion и один для запросов. На рисунке 1 показаны оба пайплайна.

Архитектура TAKC. Поток пайплайна обрабатывает ingestion данных и сжатие. Пользовательский поток обрабатывает аутентифицированные запросы

Мы выбрали AWS Lambda для вычислений, потому что каждый вызов функции кратковременный и событийно-ориентированный. Рабочая нагрузка обрабатывает данные рывками во время ingestion и принимает переменную нагрузку запросов между такими пиками, поэтому serverless здесь естественно подходит.

Для публикации интерфейса запросов в виде REST endpoint мы выбрали Amazon API Gateway. Для кэширования мы выбрали Amazon ElastiCache Serverless для чтения по составным ключам (takc:{task}:{rate}) без управления shard-ами. Amazon Cognito отвечает за выдачу JWT и обновление токенов без собственного auth-кода, уменьшая объем реализации.

Пайплайн ingestion

Когда документ попадает в Amazon S3 под префикс типа задачи, например raw-data/financial/, уведомление о событии S3 запускает функцию AWS Lambda. Эта функция разбивает документ на сегменты по 256 токенов с перекрытием 50 токенов, чтобы не терять информацию на границах. Затем она асинхронно вызывает функцию сжатия для каждого чанка, что позволяет распараллелить обработку. Для ingestion в больших объемах настройте reserved concurrency для функции сжатия и поставьте очередь Amazon Simple Queue Service (Amazon SQS) между шагами chunking и compression, чтобы мягко обрабатывать throttling.

Вторая функция вызывает Amazon Bedrock для сжатия чанков на всех четырех уровнях. Каждый вызов сжатия включает task-aware prompt, который говорит модели, какую информацию нужно сохранить:

TASK: Financial analysis. Preserve revenue metrics, margins, cash flow, debt obligations, and financial risk indicators.
COMPRESSION TARGET: Reduce to approximately 1/16 of original length.
INSTRUCTIONS:
- Focus on facts and relationships relevant to the task
- Preserve numerical data and metrics
- Maintain entities and their attributes
- Keep causal relationships and dependencies
- Remove redundant or irrelevant information

Модель знает, что нужно сохранить, потому что prompt явно указывает, что важно для задачи. Именно это отличает task-aware compression от generic compression. Система хранит сжатые результаты в Amazon ElastiCache Serverless с ключами вроде takc:financial:medium и дополнительно сохраняет их в S3 для надежности. Модель данных Redis OSS поддерживает иерархическую структуру ключей, нужную для многоуровневых cache lookup. Элементы кэша живут 24 часа и резервно сохраняются в S3. Если запись кэша удалена или истекла, функция запроса возвращается к резервной копии в S3 и заполняет кэш заново при чтении.

Пайплайн запросов

Пользователь проходит аутентификацию через Amazon Cognito, получает JWT и отправляет запрос через Amazon API Gateway. AWS WAF стоит перед API для rate limiting и защиты от угроз. Функция Lambda анализирует сложность запроса с помощью эвристик (сигналы по ключевым словам и длина запроса), получает подходящий сжатый кэш из Amazon ElastiCache Serverless и отправляет сжатый контекст вместе с запросом в Amazon Bedrock для inference. Когда уверенность маршрутизации низкая, система по умолчанию использует средний уровень сжатия как безопасный fallback.

Более дорогие вызовы сжатия в Bedrock происходят один раз во время ingestion. Путь запроса — это lookup в кэше плюс inference на сжатом контексте.

В стеке используется Amazon Bedrock (Anthropic Claude 3 Haiku, Claude 3 Sonnet и Amazon Titan Text) для сжатия и inference. Выбор модели настраивается через значения CDK context без изменений кода. AWS Lambda (Python 3.12+) обрабатывает данные и логику запросов. Amazon ElastiCache Serverless хранит сжатый кэш, а Amazon S3 хранит сырые данные, чанки и резервные копии кэша (KMS-encrypted). Amazon API Gateway публикует REST endpoints, а Amazon Cognito обеспечивает JWT-based authentication. AWS WAF защищает API с помощью rate limiting и управляемых правил безопасности. Amazon CloudWatch обеспечивает мониторинг и метрики, а AWS Key Management Service (AWS KMS) управляет ключами шифрования.

Инфраструктура задается как один стек AWS Cloud Development Kit (AWS CDK) и разворачивается одной командой. AWS CDK обеспечивает повторяемые развертывания и позволяет настраивать параметры, такие как размер чанка для сжатия, объем памяти Lambda и лимиты хранения ElastiCache через context values. Для production-развертываний стоит выделить stateful-ресурсы (S3, ElastiCache, Cognito) в отдельный стек, чтобы уменьшить blast radius и обеспечить независимое управление жизненным циклом вычислительного и storage-слоев.

Сравнение затрат

Использование входных токенов для базы знаний на 100,000 токенов, которую запрашивают 1,000 раз в день (выходные токены не учитываются, потому что длина ответа не зависит от размера контекста):

Экономия токенов за счет сжатия напрямую выводится из коэффициентов сжатия. Фактическая экономия зависит от цен конкретной модели Bedrock и от шаблона запросов.

TAKC требует первоначальных затрат на сжатие через одноразовые вызовы Bedrock во время ingestion. Эти затраты окупаются для баз знаний, которые меняются редко и запрашиваются многократно. Для баз знаний, которые меняются каждый час, per-query retrieval model RAG может быть практичнее.

Когда выбирать TAKC, RAG или оба подхода

Эта таблица показывает, когда каждый из подходов лучше подходит с учетом характеристик вашей нагрузки:

На практике production-система часто выигрывает от сочетания обоих подходов. RAG эффективно обрабатывает быстрые lookup-запросы. TAKC обслуживает аналитические запросы, в которых retrieval-based подходы не видят связи. Анализатор сложности запроса может маршрутизировать между ними. Используйте RAG, когда пользователю нужно отследить ответ до конкретных исходных документов. Для регулируемых сценариев, где одновременно нужны междокументное рассуждение и аудируемость, комбинируйте TAKC с RAG: TAKC используйте для аналитического ответа, а RAG — для получения поддерживающих исходных документов под audit trail.

Как начать

Предварительные требования

Чтобы развернуть эту reference implementation, проверьте, что у вас есть следующее:

• Аккаунт AWS с доступом к моделям Amazon Bedrock.
• AWS CDK CLI.
• Python 3.12 или новее.
• Node.js 18 или новее.

Развертывание и тестирование

Клонируйте репозиторий aws-samples/sample-bedrock-takc-compression и разверните CDK stack:

git clone https://github.com/aws-samples/sample-bedrock-takc-compression
cd sample-bedrock-takc-compression/cdk
python3 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
cdk deploy

После завершения развертывания проверьте пайплайн:

1. Загрузите документ в S3 bucket:
   Формулы и расчет

   aws s3 cp your-document.pdf s3://$(aws cloudformation describe-stacks --stack-name TakcStack \
    --query 'Stacks[0].Outputs[?OutputKey==`DataBucketName`].OutputValue' \
    --output text)/raw-data/financial/
   

2. Подождите 2-3 минуты, пока ingestion pipeline разобьет и сожмет документ на всех четырех уровнях сжатия.
3. Запросите API endpoint:
   Формулы и расчет

   curl -X POST $(aws cloudformation describe-stacks --stack-name TakcStack \
    --query 'Stacks[0].Outputs[?OutputKey==`ApiEndpoint`].OutputValue' \
    --output text)/query \
    -H "Authorization: Bearer $TOKEN" \
    -H "Content-Type: application/json" \
    -d '{"question": "What are the key financial risks?"}'
   

Система обрабатывает chunking, multi-rate compression, caching и маршрутизацию запросов без дополнительной настройки.

Очистка ресурсов

Чтобы избежать постоянных расходов, удалите CDK stack, когда закончите тестирование. Сначала очистите S3 bucket, потому что CDK не может удалить bucket, если в нем есть объекты:

aws s3 rm s3://$(aws cloudformation describe-stacks --stack-name TakcStack \
 --query 'Stacks[0].Outputs[?OutputKey==`DataBucketName`].OutputValue' \
 --output text) --recursive

Затем удалите stack:

cd sample-bedrock-takc-compression/cdk
source .venv/bin/activate
cdk destroy

Это удалит функции Lambda, API Gateway, кэш Amazon ElastiCache Serverless, пользовательский пул Amazon Cognito, WAF web ACL и предупреждения Amazon CloudWatch. KMS key сохраняется с окном pending deletion на 30 дней. Чтобы назначить удаление немедленно, выполните следующую команду:

aws kms schedule-key-deletion --key-id <key-id> --pending-window-in-days 7

Заключение

Сложные аналитические задачи, охватывающие сотни документов, требуют большего, чем извлечение фрагментов. TAKC предлагает подход, рассчитанный именно на это. Сжимайте всю базу знаний офлайн через призму конкретных задач, кэшируйте эти представления на нескольких уровнях детализации и сопоставляйте каждый запрос с подходящим уровнем сжатия по его сложности.

Реализация на AWS использует Amazon Bedrock для сжатия и inference, Amazon ElastiCache Serverless для кэширования и полностью serverless-архитектуру, которая масштабируется вместе со спросом. Доступ к reference implementation, CDK infrastructure и deployment scripts можно получить в aws-samples/sample-bedrock-takc-compression.

Чтобы начать, разверните CDK stack со своими документами и посмотрите, как система отвечает на разные типы запросов. Если ваша нагрузка включает междокументное рассуждение по стабильным базам знаний, TAKC может снизить затраты на токены и одновременно улучшить качество ответов.

Источники

• Документация Amazon Bedrock
• Документация Amazon ElastiCache Serverless
• Руководство по prompt engineering для Amazon Bedrock

Материал — перевод статьи с английского.

Оригинал: Beyond RAG: Task-aware knowledge compression for enterprise AI on AWS